ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

docs: explain extension layer boundaries

Browse Source
This commit is contained in:
ViperEkura
2026-08-16 21:30:03 +08:00
parent 6ac3b51496
commit 0d0dc64884
4 changed files with 148 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/developer/architecture.md
+3 -3
View File
@@ -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
docs/developer/cuda_kernels.md
+113 -2
View File
@@ -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
docs/developer/internals.md
+12 -3
View File
@@ -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
docs/guides/inference.md
+20 -3
View File
@@ -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