docs: fix stale docs and align with code
- update cuda_kernels layout, arch flags, and add FP8 section - fix install docs: kernels auto-build when nvcc + CUDA detected - mark ignored OpenAI request params and complete KVCache fields - add docker docs to indexes and astrai.optim to module overview - refresh document update timestamps
This commit is contained in:
+3
-2
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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`） |
|
||||
|
||||
### 贡献
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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<Fmt, BlockM, BlockN, K, Stages>`, `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 `<kernel>.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 `<kernel>.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<InT> + mma_shape<InT> (bf16 m16n8k16 / fp8 m16n8k32) + ldmatrix_x2/x4<T>
|
||||
│ ├── 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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
+13
-4
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user