diff --git a/docs/README-zh-CN.md b/docs/README-zh-CN.md index a3ebbd0..b10f447 100644 --- a/docs/README-zh-CN.md +++ b/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 diff --git a/docs/developer/architecture.md b/docs/developer/architecture.md index 989b27f..446e819 100644 --- a/docs/developer/architecture.md +++ b/docs/developer/architecture.md @@ -315,6 +315,7 @@ classDiagram <> +Tensor hidden_states +Optional[Tensor] aux_loss + +Optional[RouterStats] router_stats } class GQA { @@ -361,6 +362,7 @@ classDiagram <> +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 { <> +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` diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index 9fcd08f..88a6f6b 100644 --- a/docs/developer/cuda_kernels.md +++ b/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`) +- **`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. +- **`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 diff --git a/docs/developer/dataflow.md b/docs/developer/dataflow.md index 4ba066d..2b3d3e7 100644 --- a/docs/developer/dataflow.md +++ b/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. diff --git a/docs/developer/internals.md b/docs/developer/internals.md index 49b2a48..329a774 100644 --- a/docs/developer/internals.md +++ b/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`**: 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. +- **`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. +- **`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. diff --git a/docs/get-started.md b/docs/get-started.md index c60e386..c09cd55 100644 --- a/docs/get-started.md +++ b/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 diff --git a/docs/guides/evaluation.md b/docs/guides/evaluation.md index 2e13491..d8e3571 100644 --- a/docs/guides/evaluation.md +++ b/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). diff --git a/docs/guides/inference.md b/docs/guides/inference.md index 2e5c7d7..b996417 100644 --- a/docs/guides/inference.md +++ b/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 (attn_paged_decode, attn_paged_prefill) + ├── CudaBackend CUDA kernel dispatch (default on GPU) + ├── 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`):