diff --git a/docs/developer/architecture.md b/docs/developer/architecture.md index cae26f1..b264729 100644 --- a/docs/developer/architecture.md +++ b/docs/developer/architecture.md @@ -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 | diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index 63a5f7b..6806809 100644 --- a/docs/developer/cuda_kernels.md +++ b/docs/developer/cuda_kernels.md @@ -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 diff --git a/docs/developer/internals.md b/docs/developer/internals.md index 4a0608b..06eab29 100644 --- a/docs/developer/internals.md +++ b/docs/developer/internals.md @@ -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. diff --git a/docs/guides/inference.md b/docs/guides/inference.md index 19785b8..dd1ae08 100644 --- a/docs/guides/inference.md +++ b/docs/guides/inference.md @@ -49,7 +49,8 @@ KVCache ├── seq_lens [batch_size] ├── out_cache_loc [batch, seq_len] — write indices for this forward ├── max_len int — max(seq_lens), avoids GPU sync in decode - └── kv_indptr [batch + 1] int32 — prefix sum of seq_lens, precomputed once per step + ├── kv_indptr [batch + 1] int32 — prefix sum of seq_lens, precomputed once per step + └── qo_indptr [batch + 1] int32 — prefix sum of per-request q_lens (prefill), precomputed once per step ``` Attention layers do raw buffer indexing: `k_buffer[layer_id, out_cache_loc] = k` to write, `k_buffer[layer_id, indices]` to gather. @@ -61,7 +62,7 @@ Attention computation (cache I/O + SDPA/kernel dispatch) is decoupled from the m ``` AttentionBackend (ABC) ├── TorchNativeBackend SDPA + indirect KV cache gather (default) - └── CudaBackend CUDA kernel dispatch (attn_paged_decode, attn_prefill) + └── CudaBackend CUDA kernel dispatch (attn_paged_decode, attn_paged_prefill) ``` Select via context manager (mirrors `torch.nn.attention.sdpa_kernel`): @@ -75,7 +76,7 @@ with attn_backend(ATTN_BACKEND.CUDA): `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, gathers full-sequence K/V via indirect indexing (same as `TorchNativeBackend`), then calls `attn_prefill`. +`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: `CudaBackend` delegates to `TorchNativeBackend` when a CUDA kernel is not available. @@ -83,12 +84,13 @@ Fallback: `CudaBackend` delegates to `TorchNativeBackend` when a CUDA kernel is Rotary embedding is applied via `apply_rotary_emb` in `astrai/extension/rotary_backend.py`, which auto-dispatches: -- **CUDA kernel** (`rotary_emb.cu`): fused cos/sin lookup + rotation in a single kernel, used when the kernel is available, input is on CUDA, and `torch.is_grad_enabled()` is `False` (inference mode) +- **CUDA kernel** (`rotary_emb.cu`): fused cos/sin lookup + rotation in a single kernel, used when the kernel is available, the input is bf16 on CUDA, and `torch.is_grad_enabled()` is `False` (inference mode) - **Torch fallback**: complex multiply path (`torch.view_as_complex` → `torch.complex` multiply → `torch.view_as_real`), used during training (supports autograd backward) or when the CUDA kernel is not available -`RotaryEmbedding` stores a complex `freqs_cis` buffer and returns a tensor -from `forward()`. Both attention backends share the same rotary dispatch — it -is backend-agnostic. +`RotaryEmbedding` stores a cos/sin table `freqs_cis` of shape +`[max_len, dim/2, 2]` (f32 — `[cos, sin]` pairs) and `forward()` returns +a `[batch, seq_len, dim/2, 2]` slice indexed by `position_ids`. Both +attention backends share the same rotary dispatch — it is backend-agnostic. ## Continuous Batching diff --git a/docs/guides/training.md b/docs/guides/training.md index e4ea14a..fd7efad 100644 --- a/docs/guides/training.md +++ b/docs/guides/training.md @@ -41,10 +41,12 @@ 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. +`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. ## Training Loop