docs: update for attention backend and extension API

- Remove stale 'not yet wired' references
- Add AttentionBackend/CudaBackend sections to cuda_kernels.md, internals.md, inference.md
- Add astrai.extension to architecture.md module table and design patterns
- Update get-started.md: CUDA kernels activatable via attn_backend()

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

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_decode` | `attn_decode.cu` | GQA decode attention (split-KV) |
-| `attn_prefill` | `attn_prefill.cu` | Basic GQA prefill attention |
+| `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-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 waraps + 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
+is_causal: True = causal mask; False = non-causal
-mask:          2D [batch, kv_len] or 3D [batch, q_len, kv_len] (bool)
+mask:      2D [batch, kv_len] or 3D [batch, q_len, kv_len] (bool, True=keep)
-scale:         0.0 = auto (1/sqrt(head_dim)); >0 = explicit
-layout:        "bhld" (default) or "blhd"
 ```
 
-> **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
 

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`)

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
 

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