docs: fix stale developer documentation claims
- Move task_alloc/task_free/task_extend/task_cached/task_record_hashes and bind from the PagePool card to a new TaskCacheManager card matching pool.py - Drop the nonexistent Executor tokenizer attribute and association, add task_cache instead - Add AllocationStrategy/ContiguousStrategy/PagedStrategy cards and point Allocator/RadixCache composition at PagedStrategy - Add TaskCacheManager and the allocation strategies to the module overview, add _task_cache to InferenceScheduler - Fix the design-pattern count in the table of contents (15 -> 16) - Rewrite the FlashAttnBackend class docstring: packed decode gathers flat K/V via req_to_token and calls flash_attn_varlen_func; dense prefill uses flash_attn_func (no flash_attn_with_kvcache exists) - Apply the same correction to the backend bullets in internals.md and cuda_kernels.md - Rename the stale fp8_mma_test.cu reference to fp8_test.cu in cuda_kernels.md
This commit is contained in:
@@ -694,12 +694,13 @@ class CudaBackend(AttentionBackend):
|
||||
class FlashAttnBackend(AttentionBackend):
|
||||
"""FlashAttention backend via the optional ``flash-attn`` package.
|
||||
|
||||
Decode (q_len=1, contiguous cache): uses ``flash_attn_with_kvcache``,
|
||||
which reads K/V directly from the flat pool via cache_batch_idx +
|
||||
cache_seqlens — no materialized KV gather.
|
||||
Decode (q_len=1, contiguous cache): writes K/V to the pool, gathers
|
||||
flat K/V via the ``req_to_token`` page table, and calls
|
||||
``flash_attn_varlen_func`` over the ragged batch
|
||||
(``qo_indptr``/``kv_indptr``).
|
||||
|
||||
Prefill / non-contiguous decode: falls back to KV gather +
|
||||
``flash_attn_func``.
|
||||
Prefill: packed 3-D calls share the ``flash_attn_varlen_func`` path;
|
||||
dense 4-D calls go through ``flash_attn_func`` (mask-free only).
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
|
||||
@@ -4,7 +4,7 @@
|
||||
|
||||
- [Class Diagram](#class-diagram) — Full Mermaid class diagram across 10+ namespaces
|
||||
- [Module Overview](#module-overview) — Component inventory per module
|
||||
- [Design Patterns](#design-patterns) — 15 documented patterns with classes
|
||||
- [Design Patterns](#design-patterns) — 16 documented patterns with classes
|
||||
- [Core Relationships](#core-relationships) — 11 key inter-component relationships
|
||||
|
||||
## Class Diagram
|
||||
@@ -816,8 +816,8 @@ classDiagram
|
||||
|
||||
class Executor {
|
||||
+AutoModel model
|
||||
+AutoTokenizer tokenizer
|
||||
+PagePool kv_cache
|
||||
+TaskCacheManager task_cache
|
||||
+InferenceWorkspace _workspace
|
||||
+Optional[str] device
|
||||
+Optional[torch.dtype] dtype
|
||||
@@ -845,6 +845,7 @@ classDiagram
|
||||
|
||||
class InferenceScheduler {
|
||||
+PagePool _cache
|
||||
+TaskCacheManager _task_cache
|
||||
+Executor _executor
|
||||
+TaskManager _task_mgr
|
||||
+Event _stop_event
|
||||
@@ -888,6 +889,24 @@ classDiagram
|
||||
+release(pages)
|
||||
}
|
||||
|
||||
class AllocationStrategy {
|
||||
<<abstract>>
|
||||
+alloc(state, prompt_ids) bool
|
||||
+free(state)
|
||||
+extend(state, pos) bool
|
||||
+write_indices(state, prompt_ids)
|
||||
+record_hashes(state, prompt_ids, start_logical_page)
|
||||
}
|
||||
|
||||
class ContiguousStrategy {
|
||||
+write_indices(state, prompt_ids)
|
||||
}
|
||||
|
||||
class PagedStrategy {
|
||||
-Allocator _alloc
|
||||
-RadixCache _prefix
|
||||
}
|
||||
|
||||
class KVStorage {
|
||||
+int size
|
||||
+Tensor k_buffer
|
||||
@@ -926,14 +945,21 @@ classDiagram
|
||||
+bool contiguous
|
||||
-KVStorage _storage
|
||||
-ReqToTokenPool _req_pool
|
||||
-Allocator _alloc
|
||||
-RadixCache _prefix
|
||||
-AllocationStrategy _strategy
|
||||
+strategy AllocationStrategy
|
||||
+req_pool ReqToTokenPool
|
||||
+bind_tasks(req_indices, seq_lens, workspace, device, start_pos, incremental) KVCache
|
||||
}
|
||||
|
||||
class TaskCacheManager {
|
||||
-PagePool _pool
|
||||
-Dict _states
|
||||
+task_alloc(task_id, prompt_ids) bool
|
||||
+task_free(task_id)
|
||||
+task_extend(task_id, pos) bool
|
||||
+task_cached(task_id) int
|
||||
+task_record_hashes(task_id, prompt_ids, start_logical_page)
|
||||
+bind_tasks(task_ids, workspace, device, start_pos) KVCache
|
||||
+bind(task_ids, workspace) KVCache
|
||||
}
|
||||
|
||||
class Task {
|
||||
@@ -1316,17 +1342,22 @@ classDiagram
|
||||
PositionIdStrategy <|-- DocResetPositionId
|
||||
PositionIdStrategy <|-- ContinuousPositionId
|
||||
StoreWriter <|-- BinWriter
|
||||
AllocationStrategy <|-- ContiguousStrategy
|
||||
AllocationStrategy <|-- PagedStrategy
|
||||
RawRollout <|-- RolloutResult
|
||||
LaunchStrategy <|-- TorchrunStrategy
|
||||
LaunchStrategy <|-- LocalStrategy
|
||||
%% --- Composition (strong ownership, part destroyed with whole) ---
|
||||
PagePool *-- KVStorage
|
||||
PagePool *-- ReqToTokenPool
|
||||
PagePool *-- Allocator
|
||||
PagePool *-- RadixCache
|
||||
PagePool *-- AllocationStrategy
|
||||
PagedStrategy *-- Allocator
|
||||
PagedStrategy *-- RadixCache
|
||||
TaskCacheManager o-- PagePool
|
||||
RadixCache *-- RadixNode
|
||||
InferenceEngine *-- InferenceScheduler
|
||||
InferenceScheduler *-- PagePool
|
||||
InferenceScheduler *-- TaskCacheManager
|
||||
InferenceScheduler *-- Executor
|
||||
Executor *-- InferenceWorkspace
|
||||
InferenceScheduler *-- TaskManager
|
||||
@@ -1419,7 +1450,7 @@ classDiagram
|
||||
Task --> TaskStatus
|
||||
InferenceEngine --> AutoModel
|
||||
Executor --> AutoModel
|
||||
Executor --> AutoTokenizer
|
||||
Executor --> TaskCacheManager
|
||||
TaskManager --> AutoTokenizer
|
||||
|
||||
```
|
||||
@@ -1436,7 +1467,7 @@ classDiagram
|
||||
| **astrai.model** | ModelFactory, AutoModel, AutoRegressiveLM, EmbeddingEncoder, DecoderBlock, GQA, MLA, MLP, DeepSeekMoE, AttnFactory, FFNFactory, RMSNorm, Linear, LoRAConfig, LoRALinear, RotaryEmbedding, Embedding | Neural network model |
|
||||
| **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.inference** | InferenceEngine, InferenceScheduler, Executor, InferenceWorkspace, PagePool, TaskCacheManager, KVStorage, ReqToTokenPool, KVCache, Allocator, RadixCache, AllocationStrategy, ContiguousStrategy, PagedStrategy, 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, `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 |
|
||||
@@ -1478,4 +1509,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-22
|
||||
> Document Update Time: 2026-08-29
|
||||
|
||||
@@ -305,7 +305,7 @@ cycle belong under `TYPE_CHECKING`.
|
||||
|
||||
- **`AttentionBackend`** (ABC): `fwd_decode` / `fwd_prefill` abstract methods, `forward` dispatches by q_len
|
||||
- **`CudaBackend`**: CUDA kernel dispatch — decode via `attn_paged_decode` (page_size=1), prefill via `attn_paged_prefill` (ragged batch, `qo_indptr` + `kv_indptr`). Default on GPU.
|
||||
- **`FlashAttnBackend`**: Optional flash-attn dispatch with `flash_attn_with_kvcache` fast path.
|
||||
- **`FlashAttnBackend`**: Optional flash-attn dispatch via `flash_attn_varlen_func` over gathered flat K/V.
|
||||
- **`TorchNativeBackend`**: SDPA with indirect KV cache gather (always-available fallback)
|
||||
|
||||
Default priority: cuda > flash > torch. Set ``ASTR_BACKEND=cuda|torch_native|flash``
|
||||
@@ -415,7 +415,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)
|
||||
- `fp8_test.cu` — single-warp bf16→fp8→mma.sync sanity check + full FP8 GEMM correctness (sm_89)
|
||||
|
||||
## Benchmarks
|
||||
|
||||
@@ -482,4 +482,4 @@ csrc/
|
||||
|
||||
Compiled `.so` files are placed in `astrai/extension/lib/`, separate from Python source files.
|
||||
|
||||
> Document Update Time: 2026-08-22
|
||||
> Document Update Time: 2026-08-29
|
||||
|
||||
@@ -185,7 +185,7 @@ The extension package separates mechanism from policy:
|
||||
Attention computation is decoupled from the model via `AttentionBackend` ABC (`astrai/extension/backend/attention.py`):
|
||||
|
||||
- **`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`.
|
||||
- **`FlashAttnBackend`**: optional flash-attn dispatch; inference paths gather flat K/V from the pool via `req_to_token` and call `flash_attn_varlen_func` over the ragged batch (fp16/bf16 only); dense mask-free training calls use `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`.
|
||||
- The `attention(...)` entry point uses cuda > flash > torch priority and chooses another compatible backend when an automatically selected backend cannot handle a call.
|
||||
- Resolution precedence is: explicit `attn_backend(...)` context > `ASTR_BACKEND` env > default. An explicit `attn_backend(...)` selection is strict (incompatible calls raise); `ASTR_BACKEND` is a default-level override that falls back to a compatible backend when incapable. Training calls (`fwd=None`, no KV cache) resolve by capability: the CUDA cache kernels cannot run without a cache, so they fall back to flash (mask-free/causal calls only) and finally to torch SDPA.
|
||||
|
||||
Reference in New Issue
Block a user