docs: align CUDA kernel and RoPE docs with code
- Fix rotary docs to describe cos/sin freqs_cis table, not complex buffer - Replace attn_prefill with attn_paged_prefill for the CudaBackend path - Register attn_paged_prefill in kernel overview, layout, and module list - Add qo_indptr and InferenceWorkspace to architecture class diagram - Add FrequencyPenaltyStrategy to sampling design patterns
This commit is contained in:
@@ -817,12 +817,31 @@ classDiagram
|
||||
+AutoModel model
|
||||
+AutoTokenizer tokenizer
|
||||
+PagePool kv_cache
|
||||
+InferenceWorkspace _workspace
|
||||
+Optional[str] device
|
||||
+Optional[torch.dtype] dtype
|
||||
+execute_prefill(tasks, prompt_len, start_pos)
|
||||
+execute_prefill(tasks, prompt_len, start_pos=0)
|
||||
+execute_decode(tasks, return_logprobs=False) Union[List[int], List[Tuple[int, float]]]
|
||||
}
|
||||
|
||||
class InferenceWorkspace {
|
||||
+int max_batch_size
|
||||
+int max_seq_len
|
||||
+torch.device device
|
||||
+torch.dtype dtype
|
||||
+Tensor arange
|
||||
+Tensor input_mask
|
||||
+Tensor input_ids
|
||||
+Tensor req_pool_indices
|
||||
+Tensor seq_lens
|
||||
+Tensor kv_indptr
|
||||
+Tensor qo_indptr
|
||||
+Tensor inc
|
||||
+Tensor out_cache_loc
|
||||
+fill_input_ids(ids) Tensor
|
||||
+decode_mask(position_ids, total_len) Tensor
|
||||
}
|
||||
|
||||
class InferenceScheduler {
|
||||
+PagePool _cache
|
||||
+Executor _executor
|
||||
@@ -886,6 +905,7 @@ classDiagram
|
||||
+Tensor out_cache_loc
|
||||
+int max_len
|
||||
+Optional[Tensor] kv_indptr
|
||||
+Optional[Tensor] qo_indptr
|
||||
}
|
||||
|
||||
class PagePool {
|
||||
@@ -900,7 +920,7 @@ classDiagram
|
||||
+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, seq_lens, device, start_pos) KVCache
|
||||
+bind_tasks(task_ids, workspace, device, start_pos) KVCache
|
||||
}
|
||||
|
||||
class Task {
|
||||
@@ -1305,6 +1325,7 @@ classDiagram
|
||||
InferenceEngine *-- InferenceScheduler
|
||||
InferenceScheduler *-- PagePool
|
||||
InferenceScheduler *-- Executor
|
||||
Executor *-- InferenceWorkspace
|
||||
InferenceScheduler *-- TaskManager
|
||||
AutoRegressiveLM *-- DecoderBlock
|
||||
AutoRegressiveLM *-- RotaryEmbedding
|
||||
@@ -1375,6 +1396,7 @@ classDiagram
|
||||
Checkpoint ..> Checkpoint : serializes
|
||||
CheckpointCallback ..> Checkpoint : creates
|
||||
PagePool ..> KVCache : binds
|
||||
PagePool ..> InferenceWorkspace : fills
|
||||
InferenceEngine ..> GenerationRequest : uses
|
||||
InferenceEngine ..> GenerateResult : creates
|
||||
OpenAIResponseBuilder ..> ChatCompletionRequest : receives
|
||||
@@ -1411,8 +1433,8 @@ 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, PagePool, KVStorage, ReqToTokenPool, KVCache, Allocator, PrefixCache, Task, TaskManager, TaskStatus, StreamDecoder, GenerationRequest, 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, rotary_emb, apply_rotary_emb, rotary_backend, is_available | CUDA attention + rotary kernels, backend abstraction, auto-dispatch |
|
||||
| **astrai.inference** | InferenceEngine, InferenceScheduler, Executor, InferenceWorkspace, PagePool, KVStorage, ReqToTokenPool, KVCache, Allocator, PrefixCache, Task, TaskManager, TaskStatus, StreamDecoder, GenerationRequest, 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.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 |
|
||||
@@ -1424,7 +1446,7 @@ classDiagram
|
||||
| **Factory** | `ModelFactory`, `AttnFactory`, `FFNFactory`, `StrategyFactory`, `DatasetFactory`, `SchedulerFactory`, `CallbackFactory`, `StoreFactory`, `ConfigFactory`, `ExecutorFactory`, `MaskBuilderFactory`, `StoreWriterFactory`, `PackingStrategyFactory`, `PositionIdStrategyFactory`, `ToolParserFactory` | Decorator-based component creation |
|
||||
| **Registry** | `BaseFactory` | Component registration |
|
||||
| **Strategy** | `SEQStrategy`, `SFTStrategy`, `DPOStrategy`, `GRPOStrategy` | Training strategy switching |
|
||||
| **Strategy (Sampling)** | `TemperatureStrategy`, `TopKStrategy`, `TopPStrategy`, `SamplingPipeline` | Composable logit transformations |
|
||||
| **Strategy (Sampling)** | `TemperatureStrategy`, `TopKStrategy`, `TopPStrategy`, `FrequencyPenaltyStrategy`, `SamplingPipeline` | Composable logit transformations |
|
||||
| **Strategy (API)** | `ResponseBuilder`, `OpenAIResponseBuilder`, `AnthropicResponseBuilder` | HTTP API handler with format hooks |
|
||||
| **Builder** | `TrainContextBuilder` | Chain-building training context |
|
||||
| **Observer** | `TrainCallback`, callback implementations | Training process monitoring |
|
||||
|
||||
@@ -9,6 +9,7 @@ AstrAI includes optional custom CUDA kernels for attention and rotary embedding.
|
||||
| `attn_decode` | `attn_decode.cu` | GQA decode attention (split-KV) |
|
||||
| `attn_prefill` | `attn_prefill.cu` | GQA prefill attention (split-Q) |
|
||||
| `attn_paged_decode` | `attn_paged_decode.cu` | Paged KV cache decode attention |
|
||||
| `attn_paged_prefill` | `attn_paged_prefill.cu` | Paged KV cache prefill attention (ragged batch) |
|
||||
| `rotary_emb` | `rotary_emb.cu` | Fused rotary embedding (cos/sin lookup + rotation) |
|
||||
|
||||
Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Accumulate) exist:
|
||||
@@ -66,7 +67,7 @@ NVCC_FLAGS = -O3 --expt-relaxed-constexpr --use_fast_math
|
||||
--ptxas-options=-O3,-v --extra-device-vectorization --threads=8
|
||||
```
|
||||
|
||||
The `REGISTRY` in `csrc/build.py` lists all registered kernels (currently 4). Each entry maps a kernel name to its source files and build flags.
|
||||
The `REGISTRY` in `csrc/build.py` lists all registered kernels (currently 5). Each entry maps a kernel name to its source files and build flags.
|
||||
|
||||
## Attention Backend
|
||||
|
||||
@@ -74,7 +75,7 @@ The `REGISTRY` in `csrc/build.py` lists all registered kernels (currently 4). Ea
|
||||
|
||||
- **`AttentionBackend`** (ABC): `fwd_decode` / `fwd_prefill` abstract methods, `forward` dispatches by q_len
|
||||
- **`TorchNativeBackend`**: SDPA with indirect KV cache gather (default)
|
||||
- **`CudaBackend`**: CUDA kernel dispatch — decode via `attn_paged_decode` (page_size=1), prefill via `attn_prefill`
|
||||
- **`CudaBackend`**: CUDA kernel dispatch — decode via `attn_paged_decode` (page_size=1), prefill via `attn_paged_prefill` (ragged batch, `qo_indptr` + `kv_indptr`)
|
||||
|
||||
Select a backend via context manager (mirrors `torch.nn.attention.sdpa_kernel`):
|
||||
|
||||
@@ -151,6 +152,7 @@ csrc/
|
||||
│ ├── attn_decode.cu # Basic decode kernel (registered)
|
||||
│ ├── attn_prefill.cu # Basic prefill kernel (registered)
|
||||
│ ├── attn_paged_decode.cu # Paged decode kernel (registered)
|
||||
│ ├── attn_paged_prefill.cu # Paged prefill kernel (registered)
|
||||
│ ├── rotary_emb.cu # Fused rotary embedding kernel (registered)
|
||||
│ ├── attn_decode_split_kv.cuh # Split-KV variant
|
||||
│ ├── attn_decode_split_kv_mma.cuh # Split-KV + MMA variant
|
||||
@@ -158,6 +160,8 @@ csrc/
|
||||
│ ├── attn_prefill_split_q_mma.cuh # Split-Q + MMA variant
|
||||
│ ├── attn_paged_decode_split_kv.cuh # Paged + split-KV variant
|
||||
│ ├── attn_paged_decode_split_kv_mma.cuh # Paged + split-KV + MMA variant
|
||||
│ ├── attn_paged_prefill_split_q.cuh # Paged + split-Q variant
|
||||
│ ├── attn_paged_prefill_split_q_mma.cuh # Paged + split-Q + MMA variant
|
||||
│ ├── attn_dispatchers.cuh # Kernel dispatch macros
|
||||
│ ├── attn_entry_utils.cuh # Entry point helpers
|
||||
│ ├── attn_mma_utils.cuh # MMA utilities
|
||||
|
||||
@@ -41,12 +41,14 @@ RoPE embeds position into Q/K vectors via complex rotation:
|
||||
|
||||
$$ q_i = R_i W_q x_i, \quad k_j = R_j W_k x_j, \quad q_i^T k_j = x_i^T W_q^T R_{i-j} W_k x_j $$
|
||||
|
||||
`RotaryEmbedding` pre-computes a complex `freqs_cis` buffer. `forward()` returns
|
||||
a tensor indexed by `position_ids`. `apply_rotary_emb` applies the rotation:
|
||||
during training it uses torch complex multiply (autograd-compatible); during
|
||||
inference it auto-dispatches to a fused CUDA kernel when available. The key
|
||||
property is that the dot product $q_i^T k_j$ depends only on the relative
|
||||
position $i - j$, not the absolute positions.
|
||||
`RotaryEmbedding` pre-computes a cos/sin table `freqs_cis` of shape
|
||||
`[max_len, dim/2, 2]` (f32 — `[cos, sin]` pairs). `forward()` returns
|
||||
a `[batch, seq_len, dim/2, 2]` slice indexed by `position_ids`.
|
||||
`apply_rotary_emb` applies the rotation: during training it uses torch
|
||||
complex multiply (autograd-compatible); during inference it auto-dispatches
|
||||
to a fused CUDA kernel when available. The key property is that the dot
|
||||
product $q_i^T k_j$ depends only on the relative position $i - j$, not the
|
||||
absolute positions.
|
||||
|
||||
**Critical for inference**: RoPE is applied **before** KV cache write, not after. If applied after caching, position encoding drift occurs because cached K/V would have stale rotation factors.
|
||||
|
||||
@@ -175,7 +177,7 @@ Three-layer separation (SGLang-inspired):
|
||||
Attention computation is decoupled from the model via `AttentionBackend` ABC (`astrai/extension/attention_backend.py`):
|
||||
|
||||
- **`TorchNativeBackend`** (default): writes K/V to cache, gathers via `req_to_token` indirect indexing, calls `F.scaled_dot_product_attention`.
|
||||
- **`CudaBackend`**: 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 gathers K/V then calls `attn_prefill`. Falls back to `TorchNativeBackend` when kernel unavailable.
|
||||
- **`CudaBackend`**: 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 `TorchNativeBackend` when kernel unavailable.
|
||||
|
||||
Rotary embedding is applied via `apply_rotary_emb` in `astrai/extension/rotary_backend.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.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user