diff --git a/docs/developer/architecture.md b/docs/developer/architecture.md index 18f23ce..4b609e9 100644 --- a/docs/developer/architecture.md +++ b/docs/developer/architecture.md @@ -1437,7 +1437,7 @@ 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** | 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.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.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 | @@ -1468,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 (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. +5. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `AutoRegressiveLM`, backed by `PagePool` + `KVCache` + `SamplingPipeline`. `astrai.extension.backend` owns attention/rotary dispatch, fallback, and KV cache policy; it calls the stateless compiled-kernel wrappers in `astrai.extension.ops`. Attention uses cuda > flash > torch priority unless explicitly selected by `ASTR_BACKEND` or `attn_backend()`. Rotary embedding auto-dispatches to the CUDA op when supported, 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` @@ -1476,4 +1476,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-02 +> Document Update Time: 2026-08-16 diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index 93846c2..70047bb 100644 --- a/docs/developer/cuda_kernels.md +++ b/docs/developer/cuda_kernels.md @@ -81,6 +81,113 @@ NVCC_FLAGS = -O3 --expt-relaxed-constexpr --use_fast_math 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`. +## Python Extension Architecture + +The Python extension package separates low-level kernel bindings from execution +policy: + +```text +astrai/extension/ +├── __init__.py # Stable public API +├── loader.py # Optional compiled-module discovery and loading +├── ops/ +│ ├── attention.py # Stateless attention kernel wrappers +│ └── rotary.py # Stateless rotary kernel wrapper +└── backend/ + ├── attention.py # Backend selection, KV cache I/O, and fallback + └── rotary.py # Per-call CUDA/torch rotary dispatch +``` + +The dependency direction is one-way: + +```text +model / inference + | + v +extension public API + | + v +backend policy ---> ops wrappers ---> loader ---> compiled .so + | + +-----------> torch / flash-attn fallback +``` + +`ops` must not import `backend`. This keeps direct kernel bindings independent +of model, cache, fallback, and backend-selection policy. + +### Ops Layer + +`astrai.extension.ops` is the low-level boundary around compiled extensions: + +- Wrappers are stateless and map Python arguments to pybind or + `torch.library.custom_op` calls. +- Wrappers validate kernel availability and raise `RuntimeError` when a + requested extension was not built. +- Wrappers do not choose another implementation, gather KV cache entries, or + decide whether an input is supported by a backend. +- Tests that specifically exercise a compiled kernel may import from + `astrai.extension.ops`. + +For example, `attn_prefill(...)` means "run this CUDA kernel" rather than "run +attention using the best available implementation": + +```python +from astrai.extension.ops import attn_prefill + +output = attn_prefill(q, k, v, mask=mask, is_causal=True) +``` + +If the kernel is unavailable, this call fails. Callers that need fallback and +capability dispatch must use the public `attention(...)` entry point instead. + +### Backend Layer + +`astrai.extension.backend` owns execution policy: + +- It selects CUDA, FlashAttention, or torch-native attention. +- It checks per-call constraints such as dtype, shape, head dimension, cache + availability, and installed optional dependencies. +- It owns KV cache writes and reads because those operations differ by backend. +- It provides torch fallbacks and raises when an explicitly requested backend + cannot handle a call. +- Rotary dispatch follows the same boundary without a backend class: the + policy layer chooses the fused op for supported inference calls and otherwise + uses the autograd-compatible torch implementation. + +Normal model and inference code should import the stable API from +`astrai.extension`: + +```python +from astrai.extension import ATTN_BACKEND, attention, attn_backend + +output = attention(q, k, v, kv_cache=cache, layer_id=layer_id, fwd="decode") + +with attn_backend(ATTN_BACKEND.TORCH_NATIVE): + output = attention(q, k, v) +``` + +The package root re-exports the supported high-level API and selected direct +kernel wrappers. Internal code should use `astrai.extension.backend` only when +it needs a backend type or policy implementation, and `astrai.extension.ops` +only when it deliberately requires one exact kernel. + +### Placement Rules + +When extending this package: + +| Change | Location | +|--------|----------| +| Add a pybind call for a compiled kernel | `astrai/extension/ops/` | +| Add argument translation required by the compiled ABI | `astrai/extension/ops/` | +| Add capability checks or implementation selection | `astrai/extension/backend/` | +| Add a torch or third-party fallback | `astrai/extension/backend/` | +| Add attention KV cache behavior | `astrai/extension/backend/attention.py` | +| Expose a supported user-facing symbol | `astrai/extension/__init__.py` | + +Imports belong at module scope. Optional dependencies such as `flash_attn` may +use a module-level guarded import. Type-only imports that would create a runtime +cycle belong under `TYPE_CHECKING`. + ## Attention Backend `astrai/extension/backend/attention.py` provides the backend abstraction: @@ -102,7 +209,11 @@ with attn_backend(ATTN_BACKEND.CUDA): engine.generate("hello") ``` -`CudaBackend` falls back to `FlashAttnBackend` (when flash-attn is installed and supports the input dtype) or `TorchNativeBackend` otherwise. +The `attention(...)` policy entry point falls back to `FlashAttnBackend` (when +flash-attn is installed and supports the call) or `TorchNativeBackend` when the +automatically selected CUDA backend cannot handle an input. An explicit +`ASTR_BACKEND` or `attn_backend(...)` selection is strict and raises instead of +silently switching implementations. ### Rotary Backend @@ -187,4 +298,4 @@ csrc/ Compiled `.so` files are placed in `astrai/extension/lib/`, separate from Python source files. -> Document Update Time: 2026-07-31 +> Document Update Time: 2026-08-16 diff --git a/docs/developer/internals.md b/docs/developer/internals.md index bb28f6c..2a8ff74 100644 --- a/docs/developer/internals.md +++ b/docs/developer/internals.md @@ -176,12 +176,19 @@ Three-layer separation (SGLang-inspired): ### Attention Backend +The extension package separates mechanism from policy: + +- `astrai/extension/ops/` contains stateless wrappers that invoke one exact compiled kernel and fail when it is unavailable. +- `astrai/extension/backend/` owns capability checks, implementation selection, fallback, and KV cache I/O. +- Model and inference code use the stable `astrai.extension` API instead of selecting ops directly. + Attention computation is decoupled from the model via `AttentionBackend` ABC (`astrai/extension/backend/attention.py`): -- **`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`** (default when supported): 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). - **`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. +- The `attention(...)` entry point uses cuda > flash > torch priority and chooses another compatible backend when an automatically selected backend cannot handle a call. +- `ASTR_BACKEND=cuda|torch_native|flash` and `attn_backend(...)` are explicit selections; incompatible calls raise instead of silently changing backend. Rotary embedding is applied via `apply_rotary_emb` in `astrai/extension/backend/rotary.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. @@ -196,6 +203,8 @@ with attn_backend(ATTN_BACKEND.CUDA): Layout convention: all q/k/v are `[batch, seq_len, n_heads, head_dim]` (blhd). Scale is always `1/sqrt(head_dim)`. +Direct imports from `astrai.extension.ops` are reserved for low-level kernel tests and code that intentionally requires a specific compiled implementation. They do not provide fallback. + ## Mask Algorithm Internals ### Template mode (`template: true`) @@ -253,4 +262,4 @@ total_steps = (batches_per_replica // grad_accum_steps) * n_epoch This accounts for data-parallel sharding — each rank processes `1/nprocs` of the dataset. -> Document Update Time: 2026-08-02 +> Document Update Time: 2026-08-16 diff --git a/docs/guides/inference.md b/docs/guides/inference.md index 7544755..8f9ab4d 100644 --- a/docs/guides/inference.md +++ b/docs/guides/inference.md @@ -61,6 +61,14 @@ Attention layers do raw buffer indexing: `k_buffer[layer_id, out_cache_loc] = k` ## Attention Backend +Inference code calls the policy API exported by `astrai.extension`. The +extension implementation is split into two layers: + +- `astrai.extension.backend` owns capability checks, backend selection, + fallback, and KV cache I/O. +- `astrai.extension.ops` contains direct wrappers around compiled CUDA kernels; + these wrappers raise if a kernel is unavailable and do not fall back. + Attention computation (cache I/O + SDPA/kernel dispatch) is decoupled from the model via `AttentionBackend` ABC: ``` @@ -70,8 +78,9 @@ AttentionBackend (ABC) └── TorchNativeBackend SDPA + indirect KV cache gather (always-available fallback) ``` -Default priority: cuda > flash > torch. Set ``ASTR_BACKEND=cuda|torch_native|flash`` -to override. +Default priority is cuda > flash > torch. Automatic selection may choose a +compatible fallback for a particular call. Set +`ASTR_BACKEND=cuda|torch_native|flash` to require one backend process-wide. Select via context manager (mirrors `torch.nn.attention.sdpa_kernel`): @@ -82,12 +91,20 @@ with attn_backend(ATTN_BACKEND.CUDA): engine.generate("hello") ``` +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` 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: when `CudaBackend` cannot handle an input (wrong dtype or head_dim), `FlashAttnBackend` is tried next (if installed), then `TorchNativeBackend`. +This fallback is performed by the public `attention(...)` policy entry point +only when no backend was explicitly selected. Import from +`astrai.extension.ops` only for direct kernel tests or when failure on a missing +kernel is the intended behavior. + ### Rotary Embedding Backend Rotary embedding is applied via `apply_rotary_emb` in `astrai/extension/backend/rotary.py`, which auto-dispatches: @@ -329,4 +346,4 @@ async for token in engine.generate_async("Hello", ...): # -> AsyncGenerator[s print(token) ``` -> Document Update Time: 2026-07-31 +> Document Update Time: 2026-08-16