diff --git a/astrai/extension/backend/attention.py b/astrai/extension/backend/attention.py index 82ef03f..bca3794 100644 --- a/astrai/extension/backend/attention.py +++ b/astrai/extension/backend/attention.py @@ -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 diff --git a/docs/developer/architecture.md b/docs/developer/architecture.md index 1e9fd0f..d5a7186 100644 --- a/docs/developer/architecture.md +++ b/docs/developer/architecture.md @@ -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 { + <> + +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 diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index 2b878d4..c8b600d 100644 --- a/docs/developer/cuda_kernels.md +++ b/docs/developer/cuda_kernels.md @@ -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 diff --git a/docs/developer/internals.md b/docs/developer/internals.md index 584f183..262c89f 100644 --- a/docs/developer/internals.md +++ b/docs/developer/internals.md @@ -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.