diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 30db257..882881a 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -39,11 +39,12 @@ ruff format . # re-format after fix python -u -m pytest tests/ -v ``` -> Failed tests may leave orphan tempdirs under `%TEMP%`. Clean them manually if needed. +> Failed tests may leave orphan tempdirs under the system temp directory +> (`$TMPDIR` on Linux/macOS, `%TEMP%` on Windows). Clean them manually if needed. ### 4. (Optional) Full pre-commit check script -If you have Git Bash available: +If you have `bash` available (Git Bash on Windows works too): ```bash bash scripts/pre_commit.sh diff --git a/README.md b/README.md index 98bc401..e3fe00a 100644 --- a/README.md +++ b/README.md @@ -51,7 +51,7 @@ AstrAI is an end-to-end Transformer framework for building, training, evaluating | **Data** | Declarative JSON preprocessing, configurable masking and packing, binary/JSONL storage, and streaming datasets | | **Inference** | Continuous batching, paged KV cache, radix prefix caching, streaming generation, and Torch/CUDA/FlashAttention backends | | **Serving** | FastAPI server with OpenAI and Anthropic chat completion protocols, including SSE streaming and tool calls | -| **Evaluation** | Perplexity, MMLU, HumanEval, IFEval, IFD, and ROUGE evaluation tools | +| **Evaluation** | Perplexity, MMLU, HumanEval, IFEval, IFD, ROUGE, and weight-analysis evaluation tools | | **Extensibility** | Factory and registry architecture for models, datasets, training strategies, callbacks, kernels, and protocol components | ### Getting Started @@ -65,8 +65,9 @@ AstrAI requires Python 3.12+ and pins PyTorch exactly to `2.11.0`. Training, `sc ```bash git clone https://github.com/ViperEkura/AstrAI.git cd AstrAI -pip install -e . # pure PyTorch (no CUDA kernels) -# CSRC_KERNELS=true pip install -e . --no-build-isolation # optional: fused CUDA kernels +pip install -e . # kernels auto-build when nvcc + CUDA are detected +# CSRC_KERNELS=false pip install -e . # skip kernels (pure PyTorch) +# CSRC_KERNELS=true pip install -e . --no-build-isolation # force the fused CUDA kernel build # pip install -e ".[dev]" # dev dependencies (pytest, ruff) ``` @@ -239,6 +240,8 @@ See [Inference Guide](docs/guides/inference.md) for SSE streaming format, error | [Data Flow](./docs/developer/dataflow.md) | Data pipeline, storage backends & dataset architecture | | [Internals](./docs/developer/internals.md) | Training internals: loss formulas, callback lifecycle, KV cache | | [CUDA Kernels](./docs/developer/cuda_kernels.md) | Custom CUDA attention kernels & benchmarks | +| [Docker Serving](./docs/developer/docker-serving.md) | YAML-driven containerized serving (`serve.yaml`, `serve.sh`) | +| [Docker Training](./docs/developer/docker-training.md) | YAML-driven containerized training (`train.yaml`, `train.sh`) | ### Contributing diff --git a/docs/README-zh-CN.md b/docs/README-zh-CN.md index b10f447..4056e87 100644 --- a/docs/README-zh-CN.md +++ b/docs/README-zh-CN.md @@ -57,7 +57,7 @@ AstrAI 是一个覆盖模型构建、训练、评测与部署的端到端 Transf | **数据** | 声明式 JSON 预处理、可配置掩码与样本打包、二进制/JSONL 存储和流式数据集 | | **推理** | 连续批处理、分页 KV Cache、Radix 前缀缓存、流式生成，以及 Torch/CUDA/FlashAttention 后端 | | **服务** | 基于 FastAPI 的 OpenAI 与 Anthropic 聊天补全协议，支持 SSE 流式输出和工具调用 | -| **评测** | Perplexity、MMLU、HumanEval、IFEval、IFD 和 ROUGE 评测工具 | +| **评测** | Perplexity、MMLU、HumanEval、IFEval、IFD、ROUGE 和权重分析评测工具 | | **扩展** | 基于工厂与注册表扩展模型、数据集、训练策略、回调、内核和协议组件 | ### 快速上手 @@ -71,8 +71,9 @@ AstrAI 需要 Python 3.12+，并精确固定 PyTorch 版本为 `2.11.0`。训练 ```bash git clone https://github.com/ViperEkura/AstrAI.git cd AstrAI -pip install -e . # 纯 PyTorch（不含 CUDA 内核） -# CSRC_KERNELS=true pip install -e . --no-build-isolation # 可选：融合 CUDA 内核加速 +pip install -e . # 检测到 nvcc + CUDA 时自动构建内核 +# CSRC_KERNELS=false pip install -e . # 跳过内核（纯 PyTorch） +# CSRC_KERNELS=true pip install -e . --no-build-isolation # 强制构建融合 CUDA 内核 # pip install -e ".[dev]" # 可选：开发依赖（pytest, ruff） ``` @@ -242,6 +243,8 @@ SSE 流式格式、错误码和统计端点详见[推理文档](guides/inference | [数据流程](./developer/dataflow.md) | 数据管道、存储后端与数据集架构 | | [内部实现](./developer/internals.md) | 训练原理：损失公式、回调生命周期、KV Cache | | [CUDA 内核](./developer/cuda_kernels.md) | 自定义 CUDA 注意力内核与基准测试 | +| [Docker 服务部署](./developer/docker-serving.md) | YAML 驱动的容器化服务（`serve.yaml`、`serve.sh`） | +| [Docker 训练部署](./developer/docker-training.md) | YAML 驱动的容器化训练（`train.yaml`、`train.sh`） | ### 贡献 diff --git a/docs/developer/architecture.md b/docs/developer/architecture.md index 4b609e9..1e9fd0f 100644 --- a/docs/developer/architecture.md +++ b/docs/developer/architecture.md @@ -1437,7 +1437,8 @@ classDiagram | **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, GenerateResult, BaseSamplingStrategy–SamplingPipeline, FrequencyPenaltyStrategy, ProtocolHandler, ResponseBuilder, OpenAIResponseBuilder, AnthropicResponseBuilder, StopChecker, GenContext, StopInfo, ChatMessage, FunctionDef, ToolDef, ChatCompletionRequest, AnthropicMessage, MessagesRequest, BaseToolParser, ToolParserFactory, SimpleJsonToolParser | Inference service | -| **astrai.extension** | `backend` policy package, `ops` kernel-wrapper package, AttentionBackend, TorchNativeBackend, CudaBackend, FlashAttnBackend, attention, attn_backend, ATTN_BACKEND, apply_rotary_emb, is_available | Stable API over attention/rotary execution policy and optional CUDA kernels | +| **astrai.extension** | `backend` policy package, `ops` kernel-wrapper package, `fp8.py` FP8 strategy layer, AttentionBackend, TorchNativeBackend, CudaBackend, FlashAttnBackend, attention, attn_backend, ATTN_BACKEND, apply_rotary_emb, is_available | Stable API over attention/rotary/FP8 execution policy and optional CUDA kernels | +| **astrai.optim** | OptimizerFactory, MuonAdamW, NoraNadamW, ManoAdamW, composite_step/composite_zero_grad/composite_state_dict, partition_optimizer_parameters | Built-in optimizers (`muon_adamw` / `nora_nadamw` / `mano_adamw`) with shared composite-optimizer helpers | | **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 | | **astrai.protocols** | OptimizerProtocol, SchedulerProtocol | Structural subtyping for optimizer/scheduler wrappers | @@ -1461,6 +1462,7 @@ classDiagram | **Storage** | `Store`, `MmapStore`, `JsonlStore` | Format-agnostic data access with multi-segment support | | **Producer-Consumer** | `InferenceScheduler`, `Task`, queues | Continuous batching | | **Model Registry** | `ModelFactory`, `AutoRegressiveLM`, `EmbeddingEncoder` | Model-type dynamic loading | +| **Optimizer Routing** | `OptimizerFactory`, `MuonAdamW`, `NoraNadamW`, `ManoAdamW` | Route parameter groups (matrices vs. embeddings/heads/norms) through different optimizers | ## Core Relationships @@ -1476,4 +1478,4 @@ classDiagram 10. **AutoModel**: `from_pretrained()` loads `config.json` + `model.safetensors`, `_disable_random_init` replaces `nn.init.*` with no-ops 11. **Protocols**: `OptimizerProtocol` / `SchedulerProtocol` — structural subtyping for `AccumOptimizer` / `AccumScheduler` wrappers -> Document Update Time: 2026-08-16 +> Document Update Time: 2026-08-22 diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index e3833d7..44f0167 100644 --- a/docs/developer/cuda_kernels.md +++ b/docs/developer/cuda_kernels.md @@ -1,23 +1,24 @@ # CUDA Kernels -AstrAI includes optional custom CUDA kernels for attention and rotary embedding. These are built when `nvcc` is available and CUDA is detected, and are dispatched via the `CudaBackend` attention backend or auto-dispatched for rotary. +AstrAI includes optional custom CUDA kernels for attention, rotary embedding, and FP8 GEMM. These are built when `nvcc` is available and CUDA is detected, and are dispatched via the `CudaBackend` attention backend, auto-dispatched for rotary, or invoked through the FP8 linear primitives. ## Overview | Kernel | File | Description | |--------|------|-------------| -| `attn_decode` | `attn_decode.cu` | GQA decode attention (split-KV) | -| `attn_prefill` | `attn_prefill.cu` | GQA prefill attention (split-Q) | -| `attn_paged_decode` | `attn_paged_decode.cu` | Paged KV cache decode attention | -| `attn_paged_prefill` | `attn_paged_prefill.cu` | Paged KV cache prefill attention (ragged batch) | -| `rotary_emb` | `rotary_emb.cu` | Fused rotary embedding (cos/sin lookup + rotation) | +| `attn_decode` | `attention/decode.cu` | GQA decode attention (split-KV) | +| `attn_prefill` | `attention/prefill.cu` | GQA prefill attention (split-Q) | +| `attn_paged_decode` | `attention/paged_decode.cu` | Paged KV cache decode attention | +| `attn_paged_prefill` | `attention/paged_prefill.cu` | Paged KV cache prefill attention (ragged batch) | +| `rotary_emb` | `rotary/rotary_emb.cu` | Fused rotary embedding (cos/sin lookup + rotation) | +| `fp8_mm` | `fp8/mm.cu` | FP8 quantization + tensor-core GEMM (sm_89+) | Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Accumulate) exist: | Variant | File | Optimization | |---------|------|--------------| -| Split-KV MMA decode | `attn_decode_split_kv_mma.cuh` | Split KV across warps + MMA (sm_80+) | -| Split-Q MMA prefill | `attn_prefill_split_q_mma.cuh` | Split Q across warps + MMA (sm_80+) | +| Split-KV MMA decode | `attention/decode_split_kv_mma.cuh` | Split KV across warps + MMA (sm_80+) | +| Split-Q MMA prefill | `attention/prefill_split_q_mma.cuh` | Split Q across warps + MMA (sm_80+) | > The paged and non-paged paths share one kernel body. Prefill is templated on > an independent Q schedule (`DenseQSchedule` / `PackedQSchedule`) and KV @@ -26,7 +27,7 @@ Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Ac ### Rotary Embedding Kernel -The `rotary_emb` kernel (`csrc/kernels/rotary_emb.cu`) fuses cos/sin lookup and rotation into a single kernel: +The `rotary_emb` kernel (`csrc/kernels/rotary/rotary_emb.cu`) fuses cos/sin lookup and rotation into a single kernel: - One thread per (head, dim-pair), vectorized `__nv_bfloat162` load/store - f32 cos/sin input, bf16 compute and output @@ -36,6 +37,30 @@ The `rotary_emb` kernel (`csrc/kernels/rotary_emb.cu`) fuses cos/sin lookup and Standalone benchmark vs torch complex-multiply (48 calls = 24 layers × q+k): 6-9x faster, max diff 0 (decode) to 3e-2 (large prefill, bf16). +### FP8 GEMM / Linear Kernel + +The `fp8_mm` family (`csrc/kernels/fp8/`) accelerates bf16 linear layers by +quantizing to FP8 and running tensor-core GEMMs (**requires sm_89+**; fp8 +`mma.sync.m16n8k32` only exists on Ada/Hopper). It follows the same three-layer +style as attention, but split into **three** files: + +| File | Role | +|------|------| +| `fp8/common.h` | `FP8Format` enum (E4M3/E5M2), `Fp8GemmTraits`, `FP8Params` POD — no torch | +| `fp8/gemm.cuh` | pure-CUDA device code: `fp8_quantize_kernel` (BF16→FP8 + amax), `fp8_pq_gemm_kernel` (pre-quantized GEMM, 128×64 CTA / 64×16 warp / 3-stage cp.async) — no torch | +| `fp8/mm.cu` | binding only: `check_fp8_device` (sm_89+), param packing, launch dispatch, pybind → module `fp8_mm` | + +Scale semantics follow `torch._scaled_mm` (quantization step size: divide by +`scale`; the kernel computes the reciprocal internally — the interface never +takes `*_inv`). `amax` is always returned in the original bf16 domain. + +Python layer (two levels): `astrai/extension/ops/fp8.py` provides stateless +primitives (`quantize_bf16` / `mm_fp8` / `linear_forward_fp8` / +`linear_backward_fp8`) via `torch.library.custom_op`, and +`astrai/extension/fp8.py` is the strategy layer (`fp8_autocast`, delayed / +dynamic scaling recipes, `fp8_linear_forward/backward` wiring `aten::linear` +on CUDA). See the FP8 section in `AGENTS.md` for full detail. + ## Build System ### Auto-detection @@ -66,10 +91,17 @@ cmake --build build/cmake -j 16 ### Architecture flags -`setup.py` passes the GPU compute capability to CMake via `ASTRAI_CUDA_ARCH` (default `89`, i.e. sm_89 / L20): +`setup.py` passes the GPU compute capability to CMake via `ASTRAI_CUDA_ARCH`. When +unset, `setup.py` auto-detects the real GPU capability through +`torch.cuda.get_device_capability()`; the CMake fallback default is `80` (sm_80): -- **sm_80+** (Ampere and later): enables tensor-core MMA path (`mma.sync.m16n8k16.bf16`) -- **Below sm_80**: adds `-DASTRAI_NO_MMA` to disable the MMA path at compile time +- **sm_80+** (Ampere and later): enables the tensor-core MMA path + (`mma.sync.m16n8k16.bf16` for bf16 attention, `mma.sync.m16n8k32` for FP8). +- **sm_89+**: required for the FP8 family (`fp8_mm`) — FP8 tensor-core + instructions only exist on Ada/Hopper and newer. +- **`-DASTRAI_NO_MMA`** is a manual escape hatch only — the build never defines + it automatically. To disable the MMA path, add it to `NVCC_FLAGS` yourself; + all supported build targets are sm_80+. ### Build configuration @@ -80,7 +112,7 @@ NVCC_FLAGS = -O3 --expt-relaxed-constexpr --use_fast_math --ptxas-options=-O3,-v --extra-device-vectorization --threads=16 ``` -Each kernel in `astrai/extension/lib` is compiled as an independent pybind11 module (one `.so` per kernel, named `.cpython-*-x86_64-linux-gnu.so`). CMake builds all five kernel targets in parallel via `cmake --build -j N`. +Each kernel in `astrai/extension/lib` is compiled as an independent pybind11 module (one `.so` per kernel, named `.cpython-*-x86_64-linux-gnu.so`). CMake builds all six kernel targets in parallel via `cmake --build -j N`. The target list is the **single source of truth**: `KERNEL_NAMES` and the parallel `KERNEL_SRCS` list in `csrc/CMakeLists.txt`; `astrai/extension/loader.py` auto-discovers the compiled `.so` files. ## Python Extension Architecture @@ -93,7 +125,9 @@ astrai/extension/ ├── loader.py # Optional compiled-module discovery and loading ├── ops/ │ ├── attention.py # Stateless attention kernel wrappers -│ └── rotary.py # Stateless rotary kernel wrapper +│ ├── rotary.py # Stateless rotary kernel wrapper +│ └── fp8.py # Stateless FP8 primitives (custom_op) +├── fp8.py # FP8 strategy layer (fp8_autocast, recipes) └── backend/ ├── attention.py # Backend selection, KV cache I/O, and fallback └── rotary.py # Per-call CUDA/torch rotary dispatch @@ -293,6 +327,7 @@ nvcc -I csrc -arch=sm_89 -O3 --use_fast_math \ Test files: - `attn_test.cu` — decode + prefill kernels (correctness tables + benchmarks) - `attn_paged_test.cu` — paged decode/prefill kernels +- `fp8_mma_test.cu` — BF16→FP8→BF16 MMA demo (sm_89) ## Benchmarks @@ -315,29 +350,39 @@ nvcc -I csrc -arch=sm_89 -O3 --use_fast_math \ ``` csrc/ -├── CMakeLists.txt # CMake build: 5 kernel targets, torch/pybind11 linking +├── CMakeLists.txt # CMake build: kernel registry (KERNEL_NAMES / KERNEL_SRCS), torch/pybind11 linking ├── kernels/ -│ ├── attn_common.h # Unified attention params (contig + paged modes) -│ ├── attn_decode.cu # Basic decode kernel (registered) -│ ├── attn_prefill.cu # Basic prefill kernel (registered) -│ ├── attn_paged_decode.cu # Paged decode kernel (registered) -│ ├── attn_paged_prefill.cu # Paged prefill kernel (registered) -│ ├── rotary_emb.cu # Fused rotary embedding kernel (registered) -│ ├── attn_decode_split_kv.cuh # Split-KV variant (contig + paged via KVSource) -│ ├── attn_decode_split_kv_mma.cuh # Split-KV + MMA variant (contig + paged) -│ ├── attn_prefill_split_q.cuh # Split-Q variant (contig + paged via KVSource) -│ ├── attn_prefill_split_q_mma.cuh # Split-Q + MMA variant (contig + paged) -│ ├── attn_layout_policies.cuh # Q schedules and KVSource policies -│ ├── attn_dispatchers.cuh # Kernel dispatch macros + KV-templated launchers -│ ├── attn_entry_utils.cuh # Entry point helpers -│ ├── attn_mma_utils.cuh # MMA utilities -│ └── attn_warp_utils.cuh # Warp-level utilities +│ ├── common/ # cross-family pure-CUDA helpers (no torch) +│ │ ├── device.cuh # sm_at_least(), kMinSmForFp8* constants +│ │ └── mma.cuh # shared mma_sync + mma_shape (bf16 m16n8k16 / fp8 m16n8k32) + ldmatrix_x2/x4 +│ ├── attention/ # attention family (module names keep the attn_* prefix) +│ │ ├── common.h # AttentionParams POD, TensorLayout enum (BHLD/BLHD) +│ │ ├── warp_utils.cuh # warp reduction helpers +│ │ ├── layout_policies.cuh # KV addressing policies: DenseQSchedule/PackedQSchedule, ContigKV/PagedKV +│ │ ├── mma_utils.cuh # ldmatrix/pack helpers + online-softmax (bf16 mma via common/mma.cuh) +│ │ ├── entry_utils.cuh # torch binding helpers: DISPATCH_HEAD_DIM, pack_*_params +│ │ ├── dispatchers.cuh # pure-CUDA launchers: dispatch_decode/prefill (+paged), split-K math +│ │ ├── decode_split_kv.cuh # decode kernel, scalar (split-KV) +│ │ ├── decode_split_kv_mma.cuh # decode kernel, MMA + split-K +│ │ ├── prefill_split_q.cuh # prefill kernel, scalar (split-Q) +│ │ ├── prefill_split_q_mma.cuh # prefill kernel, MMA (split-Q, packed/ragged Q schedule) +│ │ ├── decode.cu # → module attn_decode +│ │ ├── prefill.cu # → module attn_prefill +│ │ ├── paged_decode.cu # → module attn_paged_decode +│ │ └── paged_prefill.cu # → module attn_paged_prefill +│ ├── rotary/ +│ │ └── rotary_emb.cu # rotary embedding (kernel + binding in one file) → module rotary_emb +│ └── fp8/ # FP8 family (module name fp8_mm) +│ ├── common.h # FP8Format enum, Fp8GemmTraits, FP8Params POD (no torch) +│ ├── gemm.cuh # FP8 device code: quantize + pre-quantized GEMM kernels (no torch) +│ └── mm.cu # binding only: validation, param packing, launch dispatch, pybind └── tests/ - ├── test_utils.cuh # Shared test utilities - ├── attn_test.cu # Decode + prefill kernels - └── attn_paged_test.cu # Paged decode/prefill kernels + ├── test_utils.cuh # Shared test utilities (now_ms, f2bf, bf2f, randf) + ├── attn_test.cu # Decode + prefill kernels + ├── attn_paged_test.cu # Paged decode/prefill kernels + └── fp8_mma_test.cu # BF16→FP8→BF16 MMA demo ``` Compiled `.so` files are placed in `astrai/extension/lib/`, separate from Python source files. -> Document Update Time: 2026-08-16 +> Document Update Time: 2026-08-22 diff --git a/docs/developer/docker-serving.md b/docs/developer/docker-serving.md index 6e0c917..698854c 100644 --- a/docs/developer/docker-serving.md +++ b/docs/developer/docker-serving.md @@ -106,3 +106,5 @@ working with defaults (port 8000, `./params`). `8000`); change `runtime.port` to publish on a different host port. 5. The image user is built with the host UID/GID so the mounted model directory stays readable. + +> Document Update Time: 2026-08-22 diff --git a/docs/developer/docker-training.md b/docs/developer/docker-training.md index b5e4f93..3b6266b 100644 --- a/docs/developer/docker-training.md +++ b/docs/developer/docker-training.md @@ -14,7 +14,7 @@ train.yaml │ scripts/train.sh preflight, Compose wrapper, lifecycle, timer └── docker-compose.yml GPU passthrough, mounts, image, container limits - └── train-entrypoint.sh process count, parallel mode, auto-resume + └── scripts/docker/train-entrypoint.sh process count, parallel mode, auto-resume └── train.py --config /run/astrai/train.yaml ``` @@ -123,3 +123,5 @@ the Docker timeout expires. 3. Do not force DDP for a model that requires FSDP; declare the mode explicitly. 4. Do not use `kill -9` for routine shutdown; use `scripts/train.sh stop CONFIG`. 5. The image user is built with the host UID/GID so mounted checkpoints retain usable ownership. + +> Document Update Time: 2026-08-22 diff --git a/docs/get-started.md b/docs/get-started.md index 8e23196..fc73f11 100644 --- a/docs/get-started.md +++ b/docs/get-started.md @@ -26,17 +26,24 @@ This guide walks you through installing AstrAI, downloading a model, running inf git clone https://github.com/ViperEkura/AstrAI.git cd AstrAI -# Basic install (pure PyTorch, no custom CUDA kernels) +# Kernels auto-build when nvcc + CUDA are detected; skip with CSRC_KERNELS=false pip install -e . -# With CUDA kernels (optional, for fused attention and rotary embedding) +# Force the CUDA kernel build (fused attention, rotary embedding, FP8 GEMM) # CSRC_KERNELS=true pip install -e . --no-build-isolation # With dev dependencies (pytest, ruff) # pip install -e ".[dev]" ``` -> **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. +> **CUDA kernels** build automatically when `nvcc` is on `PATH` and +> `torch.cuda.is_available()` returns `True`; set `CSRC_KERNELS=false` to skip +> them, or `CSRC_KERNELS=true` to force them (required when building in an +> isolated environment with `--no-build-isolation`). 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 @@ -258,5 +265,7 @@ docker compose up -d | Multi-GPU DDP / FSDP | [Distributed Guide](guides/distributed.md) | | System architecture | [Architecture](developer/architecture.md) | | Data pipeline internals | [Data Flow](developer/dataflow.md) | +| YAML-driven containerized serving | [Docker Serving](developer/docker-serving.md) | +| YAML-driven containerized training | [Docker Training](developer/docker-training.md) | -> Document Update Time: 2026-07-31 +> Document Update Time: 2026-08-22 diff --git a/docs/guides/inference.md b/docs/guides/inference.md index bfc1a87..ea335b3 100644 --- a/docs/guides/inference.md +++ b/docs/guides/inference.md @@ -54,7 +54,12 @@ KVCache ├── out_cache_loc [batch, seq_len] — write indices for this forward ├── max_len int — max(seq_lens), avoids GPU sync in decode ├── kv_indptr [batch + 1] int32 — prefix sum of seq_lens, precomputed once per step - └── qo_indptr [batch + 1] int32 — prefix sum of per-request q_lens (prefill), precomputed once per step + ├── qo_indptr [batch + 1] int32 — prefix sum of per-request q_lens (prefill), precomputed once per step + ├── q_tile_to_batch [num_q_tiles] int32 — prefill: Q tile → request (precomputed once per step) + ├── q_tile_to_index [num_q_tiles] int32 — prefill: Q tile → request-local tile index + ├── decode_o_part [batch, n_heads, head_dim] — decode split-K partial output buffer + ├── decode_ml_part [batch, n_heads] — decode split-K partial max/logsum buffer + └── decode_out [batch, n_heads, head_dim] — decode output accumulator ``` Attention layers do raw buffer indexing: `k_buffer[layer_id, out_cache_loc] = k` to write, `k_buffer[layer_id, indices]` to gather. @@ -94,7 +99,7 @@ with attn_backend(ATTN_BACKEND.CUDA): Environment and context selections are strict: if the selected backend cannot handle the call, inference raises an error rather than silently switching. -`CudaBackend` decode path: writes K/V to cache, then calls `attn_paged_decode` with `page_size=1` — the `req_to_token` table serves directly as the page table, each token slot is a single-token "page". No explicit K/V gather needed. +`CudaBackend` decode path: writes K/V via `new_k`/`new_v` while calling `attn_paged_decode` — the `req_to_token` table serves directly as the page table (conceptually a single-token "page" per slot, i.e. `page_size=1`; the op itself takes no `page_size` argument). No explicit K/V gather needed. `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. @@ -253,14 +258,18 @@ The HTTP protocols and direct engine API have distinct request models and defaul | `max_tokens` | Optional[int] | 2048 | Max generation length | | `stream` | Optional[bool] | False | Stream output | | `stop` | Optional[Union[str, List[str]]] | None | Stop sequences | -| `n` | Optional[int] | 1 | Number of choices requested | -| `presence_penalty` | Optional[float] | 0.0 | Presence penalty (-2.0 to 2.0) | +| `n` | Optional[int] | 1 | Accepted for API compatibility, **ignored** (always returns a single choice) | +| `presence_penalty` | Optional[float] | 0.0 | Accepted for API compatibility, **ignored** | | `frequency_penalty` | Optional[float] | 0.0 | Frequency penalty (-2.0 to 2.0) | -| `logit_bias` | Optional[Dict[int, float]] | None | Per-token logit bias | -| `user` | Optional[str] | None | End-user identifier | +| `logit_bias` | Optional[Dict[int, float]] | None | Accepted for API compatibility, **ignored** | +| `user` | Optional[str] | None | Accepted for API compatibility, **ignored** | | `tools` | Optional[List[ToolDef]] | None | Tool definitions for function calling | | `tool_choice` | Optional[Union[str, Dict[str, Any]]] | `"auto"` | Tool selection mode or explicit tool choice | +> `n`, `presence_penalty`, `logit_bias`, and `user` are validated by the request +> model but ignored by the server (a warning is logged when a non-default value +> is supplied). + **Anthropic** (`MessagesRequest`): | Param | Type | Default | Description | @@ -371,4 +380,4 @@ async for token in engine.generate_async("Hello", ...): # -> AsyncGenerator[s print(token) ``` -> Document Update Time: 2026-08-16 +> Document Update Time: 2026-08-22 diff --git a/docs/guides/params.md b/docs/guides/params.md index 0d7261a..9b23ebd 100644 --- a/docs/guides/params.md +++ b/docs/guides/params.md @@ -28,7 +28,7 @@ |-----------|-------------|---------| | `--warmup_ratio` | Fraction of total steps used for LR warmup | 0.05 | | `--max_lr` | Maximum learning rate (cosine decay after warmup) | 3e-4 | -| `--max_grad_norm` | Maximum gradient norm for clipping; the current CLI requires a positive number | 1.0 | +| `--max_grad_norm` | Maximum gradient norm for clipping; `TrainConfig` validates it as positive (or `None`) | 1.0 | ### Optimizer @@ -281,4 +281,4 @@ See [Preprocessing Guide](preprocessing.md) for config file format and examples. --- -> Document Update Time: 2026-07-20 +> Document Update Time: 2026-08-22