docs: sync all documentation with current codebase

- remove GenerationRequest and generate_with_request references (class deleted)
- document cuda>flash>torch default priority and FlashAttnBackend
- add ASTR_BACKEND env var to backend docs, TorchNativeBackend (default) → (fallback)
- fix JsonlStore transform routing → DatasetFactory ownership
- fix CudaBackend fallback chain description (FlashAttn → TorchNative)
- add FlashAttnBackend to architecture strategy table
- add router_stats to DecoderOutput/FFNOutput TypedDict diagrams
- add decode_o_part/ml_part/decode_out to KVCache diagram
- add --append_eos/--no-append_eos to IFD evaluation parameter table
- update get-started CUDA kernel note (no longer requires explicit attn_backend activation)
- fix python -m scripts.tools.server (no __init__.py) → direct script call

docs/README-zh-CN.md

@@ -187,7 +187,7 @@ docker run --gpus all -it astrai:latest
 
 # 运行推理服务
 docker run --gpus all -p 8000:8000 astrai:latest \
-  python -m scripts.tools.server --port 8000 --device cuda
+  python scripts/tools/server.py --port 8000 --device cuda
 
 # 挂载数据卷
 docker run --gpus all -v /path/to/data:/data -it astrai:latest

docs/developer/architecture.md

@@ -315,6 +315,7 @@ classDiagram
             <<TypedDict>>
             +Tensor hidden_states
             +Optional[Tensor] aux_loss
+            +Optional[RouterStats] router_stats
         }
 
         class GQA {
@@ -361,6 +362,7 @@ classDiagram
             <<TypedDict>>
             +Tensor hidden_states
             +Optional[Tensor] aux_loss
+            +Optional[RouterStats] router_stats
         }
 
         class DeepSeekMoE {
@@ -807,7 +809,6 @@ classDiagram
             +AutoTokenizer tokenizer
             +InferenceScheduler scheduler
             +generate(prompt, stream, max_tokens, temperature, top_p, top_k, frequency_penalty, rep_window) Union[Generator, str, List[str]]
-            +generate_with_request(request) Union[Generator, str, List[str]]
             +generate_async(prompt, max_tokens, temperature, top_p, top_k, frequency_penalty, rep_window) AsyncGenerator
             +get_stats() Dict
             +shutdown()
@@ -915,6 +916,9 @@ classDiagram
             +int max_len
             +Optional[Tensor] kv_indptr
             +Optional[Tensor] qo_indptr
+            +Optional[Tensor] decode_o_part
+            +Optional[Tensor] decode_ml_part
+            +Optional[Tensor] decode_out
         }
 
         class PagePool {
@@ -980,17 +984,6 @@ classDiagram
             +get_stats() Dict
         }
 
-        class GenerationRequest {
-            +List[Dict] messages
-            +int top_k
-            +float top_p
-            +float temperature
-            +Optional[int] max_tokens
-            +float frequency_penalty
-            +int rep_window
-            +bool stream
-        }
-
         class BaseSamplingStrategy {
             <<abstract>>
             +apply(logits, filter_value, input_ids, input_mask) Tensor
@@ -1407,7 +1400,7 @@ classDiagram
     CheckpointCallback ..> Checkpoint : creates
     PagePool ..> KVCache : binds
     PagePool ..> InferenceWorkspace : fills
-    InferenceEngine ..> GenerationRequest : uses
+    InferenceEngine ..> GenerateResult : uses
     InferenceEngine ..> GenerateResult : creates
     OpenAIResponseBuilder ..> ChatCompletionRequest : receives
     AnthropicResponseBuilder ..> MessagesRequest : receives
@@ -1443,7 +1436,7 @@ classDiagram
 | **astrai.model** | ModelFactory, AutoModel, AutoRegressiveLM, EmbeddingEncoder, DecoderBlock, GQA, MLA, MLP, DeepSeekMoE, AttnFactory, FFNFactory, RMSNorm, Linear, LoRAConfig, LoRALinear, RotaryEmbedding, Embedding | Neural network model |
 | **astrai.tokenize** | AutoTokenizer, ChatTemplate | Tokenizer and chat template |
 | **astrai.trainer** | Trainer, TrainContext, TrainContextBuilder, BaseStrategy–GRPOStrategy, StrategyFactory, BaseScheduler–WSDScheduler, SchedulerFactory, TrainCallback(Protocol)–MetricCallback, CallbackFactory, RawRollout, RolloutResult, BaseRewardModel, RolloutGenerator, RolloutRunner | Training workflow |
-| **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, InferenceWorkspace, PagePool, KVStorage, ReqToTokenPool, KVCache, Allocator, RadixCache, Task, TaskManager, TaskStatus, StreamDecoder, GenerationRequest, GenerateResult, BaseSamplingStrategy–SamplingPipeline, FrequencyPenaltyStrategy, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, StopInfo, ChatMessage, FunctionDef, ToolDef, ChatCompletionRequest, AnthropicMessage, MessagesRequest, BaseToolParser, ToolParserFactory, SimpleJsonToolParser | Inference service |
+| **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, InferenceWorkspace, PagePool, KVStorage, ReqToTokenPool, KVCache, Allocator, RadixCache, Task, TaskManager, TaskStatus, StreamDecoder, GenerateResult, BaseSamplingStrategy–SamplingPipeline, FrequencyPenaltyStrategy, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, StopInfo, ChatMessage, FunctionDef, ToolDef, ChatCompletionRequest, AnthropicMessage, MessagesRequest, BaseToolParser, ToolParserFactory, SimpleJsonToolParser | Inference service |
 | **astrai.extension** | AttentionBackend, TorchNativeBackend, CudaBackend, attn_backend, ATTN_BACKEND, attn_decode, attn_prefill, attn_paged_decode, attn_paged_prefill, rotary_emb, apply_rotary_emb, rotary_backend, is_available | CUDA attention + rotary kernels, backend abstraction, auto-dispatch |
 | **astrai.parallel** | spawn_parallel_fn, setup_parallel, get_rank/get_world_size/get_current_device, only_on_rank, LaunchStrategy, TorchrunStrategy, LocalStrategy, BaseExecutor, ExecutorFactory, NoneExecutor, DDPExecutor, FSDPExecutor, GradientState, AccumOptimizer, AccumScheduler | Distributed parallel & gradient accumulation |
 | **astrai.factory** | BaseFactory | Component registration |
@@ -1462,7 +1455,7 @@ classDiagram
 | **Observer** | `TrainCallback`, callback implementations | Training process monitoring |
 | **Context** | `TrainContext` | Unified training state bag |
 | **Object Pool** | `Allocator`, `PagePool` | Page-based KV cache with LRU eviction |
-| **Strategy (Attention)** | `AttentionBackend`, `TorchNativeBackend`, `CudaBackend` | Attention computation backend switching via context manager |
+| **Strategy (Attention)** | `AttentionBackend`, `CudaBackend`, `FlashAttnBackend`, `TorchNativeBackend` | Attention computation backend switching via context manager |
 | **Auto-dispatch (Rotary)** | `apply_rotary_emb`, `rotary_backend.py`, `rotary_ops.py` | Rotary embedding CUDA kernel auto-dispatch with torch fallback |
 | **Executor** | `BaseExecutor`, `NoneExecutor`, `DDPExecutor`, `FSDPExecutor` | Gradient accumulation & model distribution |
 | **Storage** | `Store`, `MmapStore`, `JsonlStore` | Format-agnostic data access with multi-segment support |
@@ -1475,7 +1468,7 @@ classDiagram
 2. **Training Flow**: `Trainer` → `TrainContextBuilder` → `TrainContext`, uses `BaseStrategy` for loss, `BaseExecutor` for gradient accumulation + model distribution
 3. **Strategy Selection**: `StrategyFactory` creates strategy by `train_type`
 4. **Executor Selection**: `ExecutorFactory.create(cfg.parallel_mode, grad_accum_steps=cfg.grad_accum_steps, **cfg.executor_kwargs)` → `NoneExecutor` / `DDPExecutor` / `FSDPExecutor`
-5. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `AutoRegressiveLM`, backed by `PagePool` + `KVCache` + `SamplingPipeline`. Attention backend selected via `attn_backend()` context manager (`TorchNativeBackend` default, `CudaBackend` for CUDA kernels). Rotary embedding auto-dispatches to CUDA kernel when available (inference mode), else torch complex multiply (training).
+5. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `AutoRegressiveLM`, backed by `PagePool` + `KVCache` + `SamplingPipeline`. Attention backend selected via `attn_backend()` context manager (cuda > flash > torch priority; `ASTR_BACKEND` env var overrides default; `TorchNativeBackend` fallback). Rotary embedding auto-dispatches to CUDA kernel when available, else torch complex multiply.
 6. **Distributed**: `spawn_parallel_fn` + `setup_parallel` for multi-process DDP
 7. **Dataset Loading**: `DatasetFactory` creates datasets, `Store` (`MmapStore`/`JsonlStore`) loads data with explicit `_length` and multi-segment `_data`
 8. **Checkpoint**: `Checkpoint` saves/loads safetensors + metadata; `CheckpointCallback` performs rank-0 training saves, with extra state saved as `{key}.pt`

docs/developer/cuda_kernels.md

@@ -86,8 +86,12 @@ Each kernel in `astrai/extension/lib` is compiled as an independent pybind11 mod
 `astrai/extension/attention_backend.py` provides the backend abstraction:
 
 - **`AttentionBackend`** (ABC): `fwd_decode` / `fwd_prefill` abstract methods, `forward` dispatches by q_len
-- **`TorchNativeBackend`**: SDPA with indirect KV cache gather (default)
+- **`CudaBackend`**: CUDA kernel dispatch — decode via `attn_paged_decode` (page_size=1), prefill via `attn_paged_prefill` (ragged batch, `qo_indptr` + `kv_indptr`). Default on GPU.
-- **`CudaBackend`**: CUDA kernel dispatch — decode via `attn_paged_decode` (page_size=1), prefill via `attn_paged_prefill` (ragged batch, `qo_indptr` + `kv_indptr`)
+- **`FlashAttnBackend`**: Optional flash-attn dispatch with `flash_attn_with_kvcache` fast path.
+- **`TorchNativeBackend`**: SDPA with indirect KV cache gather (always-available fallback)
+
+Default priority: cuda > flash > torch. Set ``ASTR_BACKEND=cuda|torch_native|flash``
+to override the default.
 
 Select a backend via context manager (mirrors `torch.nn.attention.sdpa_kernel`):
 
@@ -98,7 +102,7 @@ with attn_backend(ATTN_BACKEND.CUDA):
     engine.generate("hello")
 ```
 
-`CudaBackend` falls back to `TorchNativeBackend` when a kernel is not available.
+`CudaBackend` falls back to `FlashAttnBackend` (when flash-attn is installed and supports the input dtype) or `TorchNativeBackend` otherwise.
 
 ### Rotary Backend
 

docs/developer/dataflow.md

@@ -72,7 +72,7 @@ metadata, but the preprocessing `BinWriter` currently does not write offsets.
 
 - If `load_path` is a file: `.jsonl` selects `"jsonl"`; other suffixes raise `ValueError`.
 - If `load_path` is a directory: any recursive `*.bin` plus a `meta.json` selects `"bin"`; otherwise any recursive `*.jsonl` selects `"jsonl"`.
-- Detection does not require `dataset_config.json`; configuration is selected later when `JsonlStore.load()` chooses a transform.
+- Detection does not require `dataset_config.json`; configuration is selected later when `DatasetFactory.load()` constructs a transform via `_build_jsonl_transform()` and passes it to `JsonlStore.load()`.
 
 ### Store Backends
 
@@ -89,12 +89,17 @@ access methods.
 **MmapStore**: Memory-maps `.bin` files. OS page cache sharing is native — no explicit `share_memory_()` needed. Uses `torch.from_numpy(np.memmap(...))`. `segments_are_records=False` — bin segments are contiguous streams; record access is driven by `_offsets` (written when `save_bin(..., record_keys=...)` was used at preprocessing time).
 
 **JsonlStore**: Reads a `.jsonl` file or the sorted top-level `*.jsonl` files in
-a directory. Eager transform selection uses the first available route:
+a directory. Eager transform selection is owned by
+`DatasetFactory._build_jsonl_transform()` (called from `DatasetFactory.load()`) —
+the factory picks the first available route:
 
-1. An explicit `transform=` argument.
+1. An explicit `transform=` argument passed through `store.load()`.
 2. `dataset_config.json` in the JSONL directory. It follows `PipelineConfig` and may add `tokenizer_path`; when omitted, the config directory is used.
 3. The built-in `messages` transform when `tokenizer_path=` is supplied. It masks system/user turns, trains assistant turns, and emits document-reset position IDs.
 
+`JsonlStore.load()` requires `transform=` to be passed explicitly for eager mode
+(raises `ValueError` if missing).
+
 Only DPO gets an automatic lazy route from `DatasetFactory`: raw JSONL plus
 `tokenizer_path` installs `dpo_processor` and tokenizes each record in
 `fetch_record`. GRPO does not currently have an automatic lazy processor.

docs/developer/internals.md

@@ -178,8 +178,10 @@ Three-layer separation (SGLang-inspired):
 
 Attention computation is decoupled from the model via `AttentionBackend` ABC (`astrai/extension/attention_backend.py`):
 
-- **`TorchNativeBackend`** (default): writes K/V to cache, gathers via `req_to_token` indirect indexing, calls `F.scaled_dot_product_attention`.
+- **`CudaBackend`** (default): decode path uses `attn_paged_decode` with `page_size=1` (the `req_to_token` table serves as the page table, each token slot is a single-token "page"); prefill path uses the ragged-batch `attn_paged_prefill` (addresses each request via `qo_indptr` + `kv_indptr` directly against the flat pool). Falls back to `FlashAttnBackend` when dtype unsupported.
-- **`CudaBackend`**: decode path uses `attn_paged_decode` with `page_size=1` (the `req_to_token` table serves as the page table, each token slot is a single-token "page"); prefill path uses the ragged-batch `attn_paged_prefill` (addresses each request via `qo_indptr` + `kv_indptr` directly against the flat pool). Falls back to `TorchNativeBackend` when kernel unavailable.
+- **`FlashAttnBackend`**: optional flash-attn dispatch with `flash_attn_with_kvcache` fast path for contiguous cache; falls back to KV gather + `flash_attn_func`.
+- **`TorchNativeBackend`** (always-available fallback): writes K/V to cache, gathers via `req_to_token` indirect indexing, calls `F.scaled_dot_product_attention`.
+- Default priority: cuda > flash > torch. Set `ASTR_BACKEND=cuda|torch_native|flash` to override.
 
 Rotary embedding is applied via `apply_rotary_emb` in `astrai/extension/rotary_backend.py`, which auto-dispatches to the fused CUDA kernel (`rotary_emb.cu`) during inference or torch complex multiply during training (for autograd compatibility). Both attention backends share the same rotary dispatch.
 

docs/get-started.md

@@ -36,7 +36,7 @@ pip install -e .
 # pip install -e ".[dev]"
 ```
 
-> **CUDA kernels** are opt-in. They are not built by default. When built, they can be activated via `with attn_backend(ATTN_BACKEND.CUDA):` for accelerated decode/prefill, and the fused rotary embedding kernel is auto-dispatched when available. You can skip them for normal usage.
+> **CUDA kernels** are opt-in at build time (`CSRC_KERNELS=true`). Once built, `CudaBackend` is the default attention backend on GPU (cuda > flash > torch priority). Override via `ASTR_BACKEND` env var or `attn_backend()` context manager. Fused rotary embedding kernel is auto-dispatched when available. Skip for CPU-only usage.
 
 ## 2. Download Model Weights
 
@@ -232,7 +232,7 @@ docker build -t astrai:latest .
 
 # Run inference server with GPU
 docker run --gpus all -p 8000:8000 astrai:latest \
-  python -m scripts.tools.server --port 8000 --device cuda
+  python scripts/tools/server.py --port 8000 --device cuda
 
 # Docker Compose (GPU)
 docker compose up -d

docs/guides/evaluation.md

@@ -214,6 +214,7 @@ python scripts/eval/evaluate_ifd.py \
 | `--sentinel_text` | `\n` | Prefix for unconditional pass (`""` → bos/pad fallback) |
 | `--per_token` | False | Include per-token IFD breakdown |
 | `--max_samples` | None | Random subsample per file |
+| `--append_eos` / `--no-append_eos` | `True` | Append (or skip) EOS token to instruction/response |
 
 **How it works**: Two forward passes per batch — (1) conditional: packed BFD sequence with context + response, (2) unconditional: response prefixed with a sentinel. IFD = mean_conditional_loss / mean_unconditional_loss. IFD > 1 means the instruction makes the response harder to predict (higher quality data).
 

docs/guides/inference.md

@@ -65,10 +65,14 @@ Attention computation (cache I/O + SDPA/kernel dispatch) is decoupled from the m
 
 ```
 AttentionBackend (ABC)
-  ├── TorchNativeBackend   SDPA + indirect KV cache gather (default)
+  ├── CudaBackend          CUDA kernel dispatch (default on GPU)
-  └── CudaBackend          CUDA kernel dispatch (attn_paged_decode, attn_paged_prefill)
+  ├── FlashAttnBackend     Optional flash-attn dispatch (fallback)
+  └── TorchNativeBackend   SDPA + indirect KV cache gather (always-available fallback)
 ```
 
+Default priority: cuda > flash > torch.  Set ``ASTR_BACKEND=cuda|torch_native|flash``
+to override.
+
 Select via context manager (mirrors `torch.nn.attention.sdpa_kernel`):
 
 ```python
@@ -82,7 +86,7 @@ with attn_backend(ATTN_BACKEND.CUDA):
 
 `CudaBackend` prefill path: writes K/V, then calls `attn_paged_prefill` — a ragged-batch (paged) prefill kernel that reads K/V directly from the flat pool via `req_to_token`, addressing each request's `q_len`/`kv_len` through `qo_indptr` and `kv_indptr`. No explicit K/V gather needed.
 
-Fallback: `CudaBackend` delegates to `TorchNativeBackend` when a CUDA kernel is not available.
+Fallback: when `CudaBackend` cannot handle an input (wrong dtype or head_dim), `FlashAttnBackend` is tried next (if installed), then `TorchNativeBackend`.
 
 ### Rotary Embedding Backend
 
@@ -143,7 +147,6 @@ Adding a protocol = one builder file, no handler subclassing needed.
 ```
 InferenceEngine
   ├── generate(prompt, stream, ...) → str | List[str] | Generator
-  ├── generate_with_request(req)    → same
   ├── generate_async(prompt, ...)   → AsyncGenerator
   ├── get_stats()                   → Dict
   └── shutdown()
@@ -230,19 +233,6 @@ The HTTP protocols and direct engine API have distinct request models and defaul
 | `stream` | Optional[bool] | False | Stream output |
 | `stop_sequences` | Optional[List[str]] | None | Stop sequences |
 
-**Engine** (`GenerationRequest`):
-
-| Param | Type | Default | Description |
-|-------|------|---------|-------------|
-| `messages` | List[Dict[str, str]] | required | Messages to format before generation |
-| `top_k` | int | 50 | Top-k count; 0 disables filtering |
-| `top_p` | float | 1.0 | Nucleus threshold |
-| `temperature` | float | 1.0 | Sampling temperature; 0 enables greedy decoding |
-| `max_tokens` | Optional[int] | None | Max generation length |
-| `frequency_penalty` | float | 0.0 | Frequency penalty (-2.0 to 2.0) |
-| `rep_window` | int | 64 | Recent-token window used by the frequency penalty |
-| `stream` | bool | False | Stream output |
-
 ### SSE Streaming Format
 
 **OpenAI** (`/v1/chat/completions`, `stream=true`):