refactor: unify rotary embedding interface and update docs
- Merge cos/sin into single freqs_cis tensor [batch, seq, dim/2, 2] throughout the pipeline: RotaryEmbedding buffer, forward return type, apply_rotary_emb signature, CUDA kernel interface - CUDA kernel now takes freqs_cis directly and reads cos/sin via stride offset internally, eliminating Python-side slice/copy overhead - Kernel interface: rotary_emb(x, freqs_cis) replaces rotary_emb(x, cos, sin) - All call sites pass rotary_emb as Tensor (was tuple), type annotations consistent - Update build threads from 8 to 16 - Fix all docs: get-started, inference, training, cuda_kernels, architecture, internals — reflect new rotary interface, KVCache fields, rotary backend dispatch, .so path, kernel registry count, file layout
This commit is contained in:
@@ -380,7 +380,9 @@ classDiagram
|
||||
+int max_len
|
||||
+float base
|
||||
+Optional[Dict] rope_scaling
|
||||
+forward(x, position_ids=None) Tensor
|
||||
+Tensor cos_table
|
||||
+Tensor sin_table
|
||||
+forward(x, position_ids=None) Tuple[Tensor, Tensor]
|
||||
}
|
||||
|
||||
class Embedding {
|
||||
@@ -849,6 +851,9 @@ classDiagram
|
||||
+Tensor req_pool_indices
|
||||
+Tensor seq_lens
|
||||
+Tensor out_cache_loc
|
||||
+int max_len
|
||||
+Optional[Tensor] page_table
|
||||
+Optional[Tensor] decode_mask
|
||||
}
|
||||
|
||||
class PagePool {
|
||||
@@ -1401,7 +1406,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, 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, is_available | CUDA attention kernels + backend abstraction |
|
||||
| **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.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 |
|
||||
@@ -1420,6 +1425,7 @@ classDiagram
|
||||
| **Context** | `TrainContext` | Unified training state bag |
|
||||
| **Object Pool** | `Allocator`, `PagePool` | Page-based KV cache with LRU eviction |
|
||||
| **Strategy (Attention)** | `AttentionBackend`, `TorchNativeBackend`, `CudaBackend` | Attention computation backend switching via context manager |
|
||||
| **Auto-dispatch (Rotary)** | `apply_rotary_emb`, `rotary_backend.py`, `rotary_ops.py` | Rotary embedding CUDA kernel auto-dispatch with torch fallback |
|
||||
| **Executor** | `BaseExecutor`, `NoneExecutor`, `DDPExecutor`, `FSDPExecutor` | Gradient accumulation & model distribution |
|
||||
| **Storage** | `Store`, `H5Store`, `MmapStore`, `JsonlStore` | Format-agnostic data access with multi-segment support |
|
||||
| **Producer-Consumer** | `InferenceScheduler`, `Task`, queues | Continuous batching |
|
||||
@@ -1431,7 +1437,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 (`TorchNativeBackend` default, `CudaBackend` for CUDA kernels).
|
||||
5. **Inference Flow**: `InferenceEngine` → `InferenceScheduler` → `AutoRegressiveLM`, backed by `PagePool` + `KVCache` + `SamplingPipeline`. Attention backend selected via `attn_backend()` context manager (`TorchNativeBackend` default, `CudaBackend` for CUDA kernels). Rotary embedding auto-dispatches to CUDA kernel when available (inference mode), else torch complex multiply (training).
|
||||
6. **Distributed**: `spawn_parallel_fn` + `setup_parallel` for multi-process DDP
|
||||
7. **Dataset Loading**: `DatasetFactory` creates datasets, `Store` (H5Store/MmapStore/JsonlStore) loads data with explicit `_length` and multi-segment `_data`
|
||||
8. **Checkpoint**: `Checkpoint` saves/loads safetensors + metadata (rank-0 only), extra state saved as `{key}.pt`
|
||||
@@ -1439,4 +1445,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-07-30
|
||||
> Document Update Time: 2026-07-31
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# CUDA Kernels
|
||||
|
||||
AstrAI includes optional custom CUDA attention kernels for decode and prefill. These are built when `nvcc` is available and CUDA is detected, and are dispatched via the `CudaBackend` attention backend.
|
||||
AstrAI includes optional custom CUDA kernels for attention and rotary embedding. These are built when `nvcc` is available and CUDA is detected, and are dispatched via the `CudaBackend` attention backend or auto-dispatched for rotary.
|
||||
|
||||
## Overview
|
||||
|
||||
@@ -9,6 +9,7 @@ AstrAI includes optional custom CUDA attention kernels for decode and prefill. T
|
||||
| `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 |
|
||||
| `rotary_emb` | `rotary_emb.cu` | Fused rotary embedding (cos/sin lookup + rotation) |
|
||||
|
||||
Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Accumulate) exist:
|
||||
|
||||
@@ -18,6 +19,18 @@ Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Ac
|
||||
| Split-Q MMA prefill | `attn_prefill_split_q_mma.cuh` | Split Q across warps + MMA (sm_80+) |
|
||||
| Paged split-KV MMA decode | `attn_paged_decode_split_kv_mma.cuh` | Paged cache + split-KV + MMA |
|
||||
|
||||
### Rotary Embedding Kernel
|
||||
|
||||
The `rotary_emb` kernel (`csrc/kernels/rotary_emb.cu`) fuses cos/sin lookup and rotation into a single kernel:
|
||||
|
||||
- One thread per (head, dim-pair), vectorized `__nv_bfloat162` load/store
|
||||
- f32 cos/sin input, bf16 compute and output
|
||||
- 256-thread blocks, grid-stride loop
|
||||
- Auto-dispatched via `apply_rotary_emb` in `astrai/extension/rotary_backend.py` (CUDA when available + inference mode, else torch complex-multiply fallback)
|
||||
- No context-manager backend needed — rotary is backend-agnostic, both attention backends benefit
|
||||
|
||||
Standalone benchmark vs torch complex-multiply (48 calls = 24 layers × q+k): 6-9x faster, max diff 0 (decode) to 3e-2 (large prefill, bf16).
|
||||
|
||||
## Build System
|
||||
|
||||
### Auto-detection
|
||||
@@ -36,7 +49,7 @@ CSRC_KERNELS=true pip install -e . --no-build-isolation
|
||||
|
||||
# Rebuild after editing .cu/.cuh files
|
||||
CSRC_KERNELS=true python setup.py build_ext --inplace
|
||||
# Output: astrai/extension/*.so
|
||||
# Output: astrai/extension/lib/*.so
|
||||
```
|
||||
|
||||
### Architecture flags
|
||||
@@ -53,7 +66,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 3). Each entry maps a kernel name to its source files and build flags.
|
||||
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.
|
||||
|
||||
## Attention Backend
|
||||
|
||||
@@ -74,9 +87,20 @@ with attn_backend(ATTN_BACKEND.CUDA):
|
||||
|
||||
`CudaBackend` falls back to `TorchNativeBackend` when a kernel is not available.
|
||||
|
||||
### Rotary Backend
|
||||
|
||||
`astrai/extension/rotary_backend.py` provides `apply_rotary_emb(x, (cos, sin))` with auto-dispatch:
|
||||
|
||||
- **CUDA path**: calls `rotary_emb` kernel directly when available, input is bf16 on CUDA, and `torch.is_grad_enabled()` is `False` (inference)
|
||||
- **Torch fallback**: complex multiply (`torch.view_as_complex` → `torch.complex` multiply → `torch.view_as_real`), used during training (supports autograd) or when kernel unavailable
|
||||
|
||||
No context-manager switching needed — the dispatch is automatic per call.
|
||||
|
||||
## Python Wrappers
|
||||
|
||||
`astrai/extension/attention_ops.py` provides Python wrappers for each compiled kernel. Each wrapper calls its CUDA kernel directly and raises `RuntimeError` if the `.so` is not available. Fallback to torch SDPA is handled by the attention backend, not the wrapper functions.
|
||||
`astrai/extension/attention_ops.py` provides Python wrappers for each compiled attention kernel. Each wrapper calls its CUDA kernel directly and raises `RuntimeError` if the `.so` is not available. Fallback to torch SDPA is handled by the attention backend, not the wrapper functions.
|
||||
|
||||
`astrai/extension/rotary_ops.py` provides the wrapper for the rotary embedding kernel. Fallback to torch complex multiply is handled by `rotary_backend.py`.
|
||||
|
||||
Interface (all functions):
|
||||
```
|
||||
@@ -115,8 +139,8 @@ nvcc -I csrc -arch=sm_89 -O3 --use_fast_math \
|
||||
## Known Optimization Targets
|
||||
|
||||
- **Decode D=256**: spill eliminated (BC=16 + STAGES=2), but still 248 regs — further tiling could help.
|
||||
- **Prefill single-batch**: bandwidth low (52 GB/s at q=kv=2048) — likely compute-bound but near L20 bf16 ceiling (~94 TFLOP/s).
|
||||
- **Decode single-batch**: bandwidth low (309 GB/s at kv=512) — L20 HBM ~864 GB/s theoretical; small kv underutilizes SMs despite split-KV.
|
||||
- **Prefill single-batch**: bandwidth low (22 GB/s at q=kv=2048) — compute-bound at ~94 TFLOP/s (near L20 bf16 ceiling ~193 TFLOP/s for non-causal).
|
||||
- **Decode single-batch**: bandwidth low (113 GB/s at kv=512, 13% of 864 GB/s theoretical) — small kv underutilizes SMs despite split-KV; scales to 757 GB/s (88%) at B=16+.
|
||||
|
||||
## File Layout
|
||||
|
||||
@@ -124,10 +148,11 @@ nvcc -I csrc -arch=sm_89 -O3 --use_fast_math \
|
||||
csrc/
|
||||
├── build.py # Build system: REGISTRY, _arch_flags, nvcc flags
|
||||
├── kernels/
|
||||
│ ├── attn_common.h # Shared attention utilities
|
||||
│ ├── attn_common.h # Shared attention params (AttentionParams, PagedAttentionParams)
|
||||
│ ├── attn_decode.cu # Basic decode kernel (registered)
|
||||
│ ├── attn_prefill.cu # Basic prefill kernel (registered)
|
||||
│ ├── attn_paged_decode.cu # Paged decode 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
|
||||
│ ├── attn_prefill_split_q.cuh # Split-Q variant
|
||||
@@ -145,4 +170,6 @@ csrc/
|
||||
└── attn_prefill_test.cu # Prefill kernel test
|
||||
```
|
||||
|
||||
> Document Update Time: 2026-07-30
|
||||
Compiled `.so` files are placed in `astrai/extension/lib/`, separate from Python source files.
|
||||
|
||||
> Document Update Time: 2026-07-31
|
||||
|
||||
@@ -41,7 +41,7 @@ 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 $$
|
||||
|
||||
The complex rotation `freqs_cis` is pre-computed once (`cos, sin` pairs per position). `apply_rotary_emb` multiplies Q/K as complex numbers. 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 `cos_table` and `sin_table` (f32, `[max_len, dim/2]`). `forward()` returns a `(cos, sin)` tuple 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.
|
||||
|
||||
@@ -151,7 +151,7 @@ Three-layer separation (SGLang-inspired):
|
||||
- **ReqToTokenPool**: Index table `[req_idx, pos] → physical token slot`, shared across all layers.
|
||||
- **Allocator + PrefixCache**: Paged-mode slot allocation with ref-counting, LRU eviction, and hash-based prefix sharing.
|
||||
|
||||
`PagePool` orchestrates all three. In contiguous mode (default), `req_to_token` is a trivial linear mapping. In paged mode, slots are allocated on demand with prefix caching support. Attention layers access buffers directly via `KVCache` dataclass — no methods, no abstraction.
|
||||
`PagePool` orchestrates all three. In contiguous mode (default), `req_to_token` is a trivial linear mapping. In paged mode, slots are allocated on demand with prefix caching support. `bind_tasks()` returns a `KVCache` dataclass with precomputed `page_table` and `decode_mask` fields (computed once per decode step, shared across all layers). Attention layers access buffers directly — no methods, no abstraction.
|
||||
|
||||
### Attention Backend
|
||||
|
||||
@@ -160,6 +160,8 @@ Attention computation is decoupled from the model via `AttentionBackend` ABC (`a
|
||||
- **`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.
|
||||
|
||||
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.
|
||||
|
||||
Backend selection is thread-safe via `contextvars`, mirroring `torch.nn.attention.sdpa_kernel`:
|
||||
|
||||
```python
|
||||
@@ -228,4 +230,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-07-30
|
||||
> Document Update Time: 2026-07-31
|
||||
|
||||
+3
-3
@@ -17,14 +17,14 @@ cd AstrAI
|
||||
# Basic install (pure PyTorch, no custom CUDA kernels)
|
||||
pip install -e .
|
||||
|
||||
# With CUDA kernels (optional, for fused attention)
|
||||
# With CUDA kernels (optional, for fused attention and rotary embedding)
|
||||
# CSRC_KERNELS=true pip install -e . --no-build-isolation
|
||||
|
||||
# With dev dependencies (pytest, ruff)
|
||||
# pip install -e ".[dev]"
|
||||
```
|
||||
|
||||
> **CUDA kernels** are opt-in. They are not built by default. When built, they can be activated via `with attn_backend(ATTN_BACKEND.CUDA):` for accelerated decode/prefill. You can skip them for normal usage.
|
||||
> **CUDA kernels** are opt-in. They are not built by default. When built, they can be activated via `with attn_backend(ATTN_BACKEND.CUDA):` for accelerated decode/prefill, and the fused rotary embedding kernel is auto-dispatched when available. You can skip them for normal usage.
|
||||
|
||||
## 2. Download Model Weights
|
||||
|
||||
@@ -232,4 +232,4 @@ docker compose up -d
|
||||
| System architecture | [Architecture](developer/architecture.md) |
|
||||
| Data pipeline internals | [Data Flow](developer/dataflow.md) |
|
||||
|
||||
> Document Update Time: 2026-07-30
|
||||
> Document Update Time: 2026-07-31
|
||||
|
||||
@@ -47,7 +47,10 @@ KVCache
|
||||
├── req_to_token [num_reqs, max_ctx_len]
|
||||
├── req_pool_indices [batch_size]
|
||||
├── seq_lens [batch_size]
|
||||
└── out_cache_loc [batch, seq_len] — write indices for this forward
|
||||
├── out_cache_loc [batch, seq_len] — write indices for this forward
|
||||
├── max_len int — max(seq_lens), avoids GPU sync in decode
|
||||
├── page_table [batch, max_len] — precomputed gather indices for decode (None for prefill)
|
||||
└── decode_mask [batch, max_len] bool — precomputed position validity mask (None for single-batch decode)
|
||||
```
|
||||
|
||||
Attention layers do raw buffer indexing: `k_buffer[layer_id, out_cache_loc] = k` to write, `k_buffer[layer_id, indices]` to gather.
|
||||
@@ -77,6 +80,15 @@ with attn_backend(ATTN_BACKEND.CUDA):
|
||||
|
||||
Fallback: `CudaBackend` delegates to `TorchNativeBackend` when a CUDA kernel is not available.
|
||||
|
||||
### Rotary Embedding Backend
|
||||
|
||||
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)
|
||||
- **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 `cos_table`/`sin_table` as f32 buffers and returns a `(cos, sin)` tuple from `forward()`. Both attention backends share the same rotary dispatch — it is backend-agnostic.
|
||||
|
||||
## Continuous Batching
|
||||
|
||||
`InferenceScheduler` runs a daemon thread with a 4-phase loop:
|
||||
@@ -183,6 +195,10 @@ Supports `stop_sequences` and streaming via `event: content_block_delta`.
|
||||
| `temperature` | float | 1.0 | Sampling temperature (> 0.0) |
|
||||
| `max_tokens` | Optional[int] | None | Max generation length |
|
||||
| `stream` | bool | False | Stream output |
|
||||
| `stop` | Optional[Union[str, List[str]]] | None | Stop sequences |
|
||||
| `frequency_penalty` | float | 0.0 | Frequency penalty |
|
||||
| `tools` | Optional[List[dict]] | None | Tool definitions for function calling |
|
||||
| `tool_choice` | Optional[str] | None | Tool selection mode |
|
||||
|
||||
### SSE Streaming Format
|
||||
|
||||
@@ -278,4 +294,4 @@ async for token in engine.generate_async("Hello", ...): # -> AsyncGenerator[s
|
||||
print(token)
|
||||
```
|
||||
|
||||
> Document Update Time: 2026-07-30
|
||||
> Document Update Time: 2026-07-31
|
||||
|
||||
@@ -41,7 +41,7 @@ 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 $$
|
||||
|
||||
The complex rotation `freqs_cis` is pre-computed once (`cos, sin` pairs per position). `apply_rotary_emb` multiplies Q/K as complex numbers.
|
||||
`RotaryEmbedding` pre-computes `cos_table` and `sin_table` (f32, `[max_len, dim/2]`). `forward()` returns a `(cos, sin)` tuple 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
|
||||
|
||||
@@ -232,4 +232,4 @@ nohup python scripts/tools/train.py \
|
||||
|
||||
Full parameter reference at [params.md](params.md).
|
||||
|
||||
> Document Update Time: 2026-07-20
|
||||
> Document Update Time: 2026-07-31
|
||||
|
||||
Reference in New Issue
Block a user