diff --git a/docs/developer/architecture.md b/docs/developer/architecture.md index 3158307..56db898 100644 --- a/docs/developer/architecture.md +++ b/docs/developer/architecture.md @@ -1401,7 +1401,8 @@ 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.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, ParallelModel, RowParallelLinear, ColumnParallelLinear | Distributed parallel & gradient accumulation | +| **astrai.extension** | AttentionBackend, TorchNativeBackend, CudaBackend, attn_backend, ATTN_BACKEND, attn_decode, attn_prefill, attn_paged_decode, is_available | CUDA attention kernels + backend abstraction | +| **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 | @@ -1418,6 +1419,7 @@ classDiagram | **Observer** | `TrainCallback`, callback implementations | Training process monitoring | | **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 | | **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 | @@ -1429,7 +1431,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 `KVCache` + `SamplingPipeline` +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). 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` @@ -1437,4 +1439,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-20 +> Document Update Time: 2026-07-30 diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index 347f4f3..900440f 100644 --- a/docs/developer/cuda_kernels.md +++ b/docs/developer/cuda_kernels.md @@ -1,21 +1,21 @@ # CUDA Kernels -AstrAI includes optional custom CUDA attention kernels for decode and prefill. These are **not built by default** and are **not yet wired into the model or inference path** — they are standalone kernels with benchmarks and tests. +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. ## Overview | Kernel | File | Description | |--------|------|-------------| -| `attn_decode` | `attn_decode.cu` | Basic GQA decode attention | -| `attn_prefill` | `attn_prefill.cu` | Basic GQA prefill attention | +| `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 | Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Accumulate) exist: | Variant | File | Optimization | |---------|------|--------------| -| Split-KV MMA decode | `attn_decode_split_kv_mma.cuh` | Split KV across waraps + MMA (sm_80+) | -| Split-Q MMA prefill | `attn_prefill_split_q_mma.cuh` | Split Q across waraps + MMA (sm_80+) | +| Split-KV MMA decode | `attn_decode_split_kv_mma.cuh` | Split KV across warps + MMA (sm_80+) | +| 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 | ## Build System @@ -55,19 +55,36 @@ NVCC_FLAGS = -O3 --expt-relaxed-constexpr --use_fast_math 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. +## Attention Backend + +`astrai/extension/attention_backend.py` provides the backend abstraction: + +- **`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` + +Select a backend via context manager (mirrors `torch.nn.attention.sdpa_kernel`): + +```python +from astrai.extension import attn_backend, ATTN_BACKEND + +with attn_backend(ATTN_BACKEND.CUDA): + engine.generate("hello") +``` + +`CudaBackend` falls back to `TorchNativeBackend` when a kernel is not available. + ## Python Wrappers -`astrai/extension/ops.py` provides Python wrappers for each compiled kernel. When the `.so` is not available, wrappers **fall back to `torch.nn.functional.scaled_dot_product_attention`** (SDPA). +`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. -Interface: +Interface (all functions): ``` -causal_offset: -1 = non-causal; >=0 = absolute position of first Q token -mask: 2D [batch, kv_len] or 3D [batch, q_len, kv_len] (bool) -scale: 0.0 = auto (1/sqrt(head_dim)); >0 = explicit -layout: "bhld" (default) or "blhd" +is_causal: True = causal mask; False = non-causal +mask: 2D [batch, kv_len] or 3D [batch, q_len, kv_len] (bool, True=keep) ``` -> **Note**: Wrappers are not yet called from `model/transformer.py` or `inference/`. The model uses PyTorch's built attention. Integration is future work. +Layout convention: all q/k/v are `[batch, seq_len, n_heads, head_dim]` (blhd). Scale is always `1/sqrt(head_dim)`. ## Standalone Testing diff --git a/docs/developer/internals.md b/docs/developer/internals.md index 2d98bc5..f3ed54d 100644 --- a/docs/developer/internals.md +++ b/docs/developer/internals.md @@ -153,6 +153,24 @@ Three-layer separation (SGLang-inspired): `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. +### Attention Backend + +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. + +Backend selection is thread-safe via `contextvars`, mirroring `torch.nn.attention.sdpa_kernel`: + +```python +from astrai.extension import attn_backend, ATTN_BACKEND + +with attn_backend(ATTN_BACKEND.CUDA): + engine.generate("hello") +``` + +Layout convention: all q/k/v are `[batch, seq_len, n_heads, head_dim]` (blhd). Scale is always `1/sqrt(head_dim)`. + ## Mask Algorithm Internals ### Template mode (`template: true`) diff --git a/docs/get-started.md b/docs/get-started.md index d1ac064..1d20715 100644 --- a/docs/get-started.md +++ b/docs/get-started.md @@ -24,7 +24,7 @@ pip install -e . # pip install -e ".[dev]" ``` -> **CUDA kernels** are opt-in. They are not built by default and are not yet wired into the model or inference path. 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. You can skip them for normal usage. ## 2. Download Model Weights diff --git a/docs/guides/inference.md b/docs/guides/inference.md index 98b852d..fd43bc4 100644 --- a/docs/guides/inference.md +++ b/docs/guides/inference.md @@ -4,6 +4,7 @@ - [KV Cache](#kv-cache) - [KVCache System](#kvcache-system) +- [Attention Backend](#attention-backend) - [Continuous Batching](#continuous-batching) - [Sampling](#sampling-strategy-pattern) - [Protocol Handlers](#protocol-handlers-strategy-pattern) @@ -51,6 +52,31 @@ KVCache Attention layers do raw buffer indexing: `k_buffer[layer_id, out_cache_loc] = k` to write, `k_buffer[layer_id, indices]` to gather. +## Attention Backend + +Attention computation (cache I/O + SDPA/kernel dispatch) is decoupled from the model via `AttentionBackend` ABC: + +``` +AttentionBackend (ABC) + ├── TorchNativeBackend SDPA + indirect KV cache gather (default) + └── CudaBackend CUDA kernel dispatch (attn_paged_decode, attn_prefill) +``` + +Select via context manager (mirrors `torch.nn.attention.sdpa_kernel`): + +```python +from astrai.extension import attn_backend, ATTN_BACKEND + +with attn_backend(ATTN_BACKEND.CUDA): + engine.generate("hello") +``` + +`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`. + +Fallback: `CudaBackend` delegates to `TorchNativeBackend` when a CUDA kernel is not available. + ## Continuous Batching `InferenceScheduler` runs a daemon thread with a 4-phase loop: @@ -252,4 +278,4 @@ async for token in engine.generate_async("Hello", ...): # -> AsyncGenerator[s print(token) ``` -> Document Update Time: 2026-07-09 +> Document Update Time: 2026-07-30