docs: explain extension layer boundaries
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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 `<kernel>.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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user