fix: stabilize paged decode attention kernels

- zero-fill split partials so combine skips unwritten splits deterministically
- skip loading masked KV in paged decode kernels to avoid 0*NaN output poisoning
- zero-fill shared memory tile buffers to prevent stale NaN leaking into softmax

astrai/extension/rotary_backend.py

@@ -1,17 +1,16 @@
 """Rotary embedding with auto-dispatch to CUDA kernel.
 
-Single entry point ``apply_rotary_emb(x, cos, sin)`` — uses the fused
+Single entry point ``apply_rotary_emb(x, freqs_cis)`` — uses the fused
 CUDA kernel when available, falls back to torch complex multiply otherwise.
 
 Layout: x is [batch, seq_len, n_heads, head_dim] (bf16).
-cos/sin are [batch, seq_len, head_dim/2] (f32).
+freqs_cis is [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs.
 """
 
 import torch
 from torch import Tensor
 
 from astrai.extension.loader import is_available
-from astrai.extension.rotary_ops import rotary_emb as _cuda_rotary
 
 _cache = {"available": None}
 
@@ -22,32 +21,34 @@ def _cuda_available() -> bool:
     return _cache["available"]
 
 
-def _torch_apply(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
+def _torch_apply(x: Tensor, freqs_cis: Tensor) -> Tensor:
+    cos, sin = freqs_cis[..., 0], freqs_cis[..., 1]
     dtype = x.dtype
     x_ = x.float().reshape(*x.shape[:-1], -1, 2)
     x_complex = torch.view_as_complex(x_)
-    freqs_cis = torch.complex(cos, sin).unsqueeze(2)
-    x_rotated = x_complex * freqs_cis
+    freqs_cis_complex = torch.complex(cos, sin).unsqueeze(2)
+    x_rotated = x_complex * freqs_cis_complex
     x_out = torch.view_as_real(x_rotated).flatten(-2)
     return x_out.to(dtype)
 
 
-def apply_rotary_emb(x: Tensor, rotary_emb: tuple[Tensor, Tensor]) -> Tensor:
+def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
     """Apply rotary embedding to x.
 
     Args:
         x: [batch, seq_len, n_heads, head_dim] (bf16)
-        rotary_emb: (cos, sin) tuple, each [batch, seq_len, head_dim/2] (f32)
+        freqs_cis: [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs
 
     Returns:
         [batch, seq_len, n_heads, head_dim] (bf16)
     """
-    cos, sin = rotary_emb
     if (
         _cuda_available()
         and not torch.is_grad_enabled()
         and x.is_cuda
         and x.dtype == torch.bfloat16
     ):
-        return _cuda_rotary(x, cos, sin)
-    return _torch_apply(x, cos, sin)
+        from astrai.extension.rotary_ops import rotary_emb as _cuda_rotary
+
+        return _cuda_rotary(x, freqs_cis)
+    return _torch_apply(x, freqs_cis)

astrai/extension/rotary_ops.py

@@ -2,10 +2,10 @@
 
 Calls the compiled CUDA kernel directly. If the kernel is not available,
 raises ``RuntimeError``. Fallback to torch complex multiply is the
-responsibility of ``astrai.model.components.rope.apply_rotary_emb``.
+responsibility of ``astrai.extension.rotary_backend.apply_rotary_emb``.
 
-Layout convention: x is ``[batch, seq_len, n_heads, head_dim]`` (blhd, bf16).
-cos/sin are ``[batch, seq_len, head_dim/2]`` (f32).
+Layout: x is [batch, seq_len, n_heads, head_dim] (bf16, contiguous).
+freqs_cis is [batch, seq_len, head_dim/2, 2] (f32, contiguous) — [cos, sin] pairs.
 """
 
 import torch
@@ -21,21 +21,12 @@ def _check_available():
         )
 
 
-def rotary_emb(
-    x: torch.Tensor,
-    cos: torch.Tensor,
-    sin: torch.Tensor,
-) -> torch.Tensor:
+def rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
     """Fused rotary embedding kernel.
 
-    Applies rotation: for each pair (x_even, x_odd):
-        out_even = x_even * cos - x_odd * sin
-        out_odd  = x_even * sin + x_odd * cos
-
     Args:
         x: [batch, seq_len, n_heads, head_dim] (bf16, contiguous)
-        cos: [batch, seq_len, head_dim/2] (f32)
-        sin: [batch, seq_len, head_dim/2] (f32)
+        freqs_cis: [batch, seq_len, head_dim/2, 2] (f32, contiguous) — [cos, sin] pairs
 
     Returns:
         [batch, seq_len, n_heads, head_dim] (bf16)
@@ -43,4 +34,6 @@ def rotary_emb(
     _check_available()
     if not x.is_contiguous():
         x = x.contiguous()
-    return _modules["rotary_emb"].rotary_emb(x, cos, sin)
+    if not freqs_cis.is_contiguous():
+        freqs_cis = freqs_cis.contiguous()
+    return _modules["rotary_emb"].rotary_emb(x, freqs_cis)

astrai/model/components/rope.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional, Tuple
+from typing import Dict, Optional
 
 import torch
 import torch.nn as nn
@@ -10,16 +10,18 @@ def get_rotary_emb(
     max_len: int,
     base: float = 10000,
     device: Optional[torch.device] = None,
-) -> Tuple[Tensor, Tensor]:
+) -> Tensor:
     """Precompute cos/sin tables for rotary embedding.
 
     Returns:
-        (cos, sin) each of shape [max_len, dim/2] (f32)
+        [max_len, dim/2, 2] (f32) — [cos, sin] pairs.
     """
     theta = base ** (-torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim)
     t = torch.arange(0, max_len, dtype=torch.float64, device=device)
     freqs = torch.outer(t, theta).float()
-    return torch.cos(freqs), torch.sin(freqs)
+    cos = torch.cos(freqs)
+    sin = torch.sin(freqs)
+    return torch.stack([cos, sin], dim=-1)
 
 
 def ntk_base(base: float, dim: int, factor: float) -> float:
@@ -49,13 +51,10 @@ class RotaryEmbedding(nn.Module):
         self._set_rotary_buffer(self.max_len)
 
     def _set_rotary_buffer(self, max_len: int):
-        cos, sin = get_rotary_emb(self.dim, max_len, self.base)
-        self.register_buffer("cos_table", cos, persistent=False)
-        self.register_buffer("sin_table", sin, persistent=False)
+        freqs_cis = get_rotary_emb(self.dim, max_len, self.base)
+        self.register_buffer("freqs_cis", freqs_cis, persistent=False)
 
-    def forward(
-        self, x: Tensor, position_ids: Optional[Tensor] = None
-    ) -> Tuple[Tensor, Tensor]:
+    def forward(self, x: Tensor, position_ids: Optional[Tensor] = None) -> Tensor:
         """Lookup cos/sin for the given positions.
 
         Args:
@@ -63,7 +62,7 @@ class RotaryEmbedding(nn.Module):
             position_ids: [batch, seq_len] optional position indices.
 
         Returns:
-            (cos, sin) each of shape [batch, seq_len, dim/2] (f32)
+            [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs.
         """
         if position_ids is None:
             position_ids = (
@@ -71,6 +70,4 @@ class RotaryEmbedding(nn.Module):
                 .unsqueeze(0)
                 .expand(x.size(0), -1)
             )
-        cos = self.cos_table[position_ids].float()
-        sin = self.sin_table[position_ids].float()
-        return cos, sin
+        return self.freqs_cis[position_ids].float()

csrc/build.py

@@ -53,7 +53,7 @@ NVCC_FLAGS = [
     "--use_fast_math",
     "--ptxas-options=-O3,-v",
     "--extra-device-vectorization",
-    "--threads=8",
+    "--threads=16",
 ]
 
 

csrc/kernels/attn_entry_utils.cuh

@@ -1,4 +1,5 @@
 #pragma once
+#include <float.h>
 #include <torch/extension.h>
 #include <c10/cuda/CUDAGuard.h>
 #include "attn_common.h"
@@ -23,8 +24,8 @@ using bf16 = __nv_bfloat16;
 template<typename P>
 inline void alloc_split_partials(P& p) {
     auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
-    auto o_part = torch::empty(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt);
-    auto ml_part = torch::empty(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, 2}, fopt);
+    auto o_part = torch::zeros(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt);
+    auto ml_part = torch::full(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, 2}, -FLT_MAX, fopt);
     p.o_part = (float*)o_part.data_ptr();
     p.ml_part = (float*)ml_part.data_ptr();
 }

csrc/kernels/attn_paged_decode_split_kv.cuh

@@ -67,14 +67,17 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p)
             partial = warp_reduce_sum(partial) * p.scale;
 
             int kv_idx = chunk_start + s;
+            bool masked = false;
             if constexpr (HasMask) {
                 if (!p.mask[mask_base + kv_idx])
-                    partial = -FLT_MAX;
+                    masked = true;
             }
             if constexpr (IsCausal) {
                 if (kv_idx > p.causal_offset)
-                    partial = -FLT_MAX;
+                    masked = true;
             }
+            if (masked)
+                partial = -FLT_MAX;
 
             float new_m = fmaxf(m, partial);
             float alpha = __expf(m - new_m);
@@ -85,7 +88,11 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p)
             int logical_page = pos / p.page_size;
             int page_offset = pos % p.page_size;
             int phys_page = p.page_table[batch * p.max_pages + logical_page];
-            if (phys_page >= 0) {
+            if (masked) {
+                #pragma unroll
+                for (int i = 0; i < hd_per_thread; i++)
+                    acc_reg[i] = fmaf(acc_reg[i], alpha, 0.0f);
+            } else if (phys_page >= 0) {
                 int64_t v_base = (int64_t)phys_page * p.page_size * p.kv_head * p.head_dim
                                + (int64_t)page_offset * p.kv_head * p.head_dim
                                + (int64_t)kv_head * p.head_dim;

csrc/kernels/attn_paged_decode_split_kv_mma.cuh

@@ -31,6 +31,13 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
     __shared__ __align__(16) bf16 sK[Traits::STAGES * Traits::BC * Traits::LD];
     __shared__ __align__(16) bf16 sV[Traits::STAGES * Traits::BC * Traits::LD];
 
+    #pragma unroll
+    for (int i = lane; i < Traits::STAGES * Traits::BC * Traits::LD; i += 32) {
+        sK[i] = __float2bfloat16(0.0f);
+        sV[i] = __float2bfloat16(0.0f);
+    }
+    __syncwarp();
+
     const int q_base = batch * p.q_stride_b + q_head0 * p.q_stride_h;
     const int qra = gid;
     const int qrb = gid + 8;
@@ -68,6 +75,9 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
             int r = i / Traits::HEAD_DIM, d = i % Traits::HEAD_DIM;
             int kc = kv0 + r;
             bool valid = (kc < p.kv_len);
+            if constexpr (HasMask) {
+                valid = valid && p.mask[batch * p.mask_b_stride + kc];
+            }
             int phys_page = valid ? p.page_table[batch * p.max_pages + kc] : 0;
             valid = valid && (phys_page >= 0);
             int page_off = kc % p.page_size;

csrc/kernels/rotary_emb.cu

@@ -3,8 +3,7 @@
 
 __global__ void rotary_emb_kernel(
     const __nv_bfloat16* __restrict__ x,
-    const float* __restrict__ cos,
-    const float* __restrict__ sin,
+    const float* __restrict__ freqs_cis,
     __nv_bfloat16* __restrict__ out,
     int batch,
     int seq_len,
@@ -26,14 +25,14 @@ __global__ void rotary_emb_kernel(
         int b = tmp / seq_len;
 
         int x_offset = ((b * seq_len + seq) * n_heads + head) * head_dim + (pair << 1);
-        int cs_offset = (b * seq_len + seq) * half_dim + pair;
+        int cs_offset = ((b * seq_len + seq) * half_dim + pair) * 2;
 
         __nv_bfloat162 x_pair = *reinterpret_cast<const __nv_bfloat162*>(x + x_offset);
         float x_even = __bfloat162float(__low2bfloat16(x_pair));
         float x_odd  = __bfloat162float(__high2bfloat16(x_pair));
 
-        float c = cos[cs_offset];
-        float s = sin[cs_offset];
+        float c = freqs_cis[cs_offset];
+        float s = freqs_cis[cs_offset + 1];
 
         float out_even = x_even * c - x_odd * s;
         float out_odd  = x_even * s + x_odd * c;
@@ -45,23 +44,21 @@ __global__ void rotary_emb_kernel(
 
 torch::Tensor rotary_emb(
     torch::Tensor x,
-    torch::Tensor cos,
-    torch::Tensor sin
+    torch::Tensor freqs_cis
 ) {
+    TORCH_CHECK(x.is_cuda(), "x must be on CUDA");
+    TORCH_CHECK(freqs_cis.is_cuda(), "freqs_cis must be on CUDA");
+    TORCH_CHECK(x.scalar_type() == torch::kBFloat16, "x must be bf16");
+    TORCH_CHECK(x.dim() == 4, "x must be 4D [batch, seq_len, n_heads, head_dim]");
+    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
+    TORCH_CHECK(freqs_cis.dim() == 4, "freqs_cis must be 4D [batch, seq_len, dim/2, 2]");
+    TORCH_CHECK(freqs_cis.is_contiguous(), "freqs_cis must be contiguous");
 
     int batch = x.size(0);
     int seq_len = x.size(1);
     int n_heads = x.size(2);
     int head_dim = x.size(3);
 
-    TORCH_CHECK(x.is_cuda(), "x must be on CUDA");
-    TORCH_CHECK(cos.is_cuda(), "cos must be on CUDA");
-    TORCH_CHECK(sin.is_cuda(), "sin must be on CUDA");
-    TORCH_CHECK(x.scalar_type() == torch::kBFloat16, "x must be bf16");
-    TORCH_CHECK(x.dim() == 4, "x must be 4D [batch, seq_len, n_heads, head_dim]");
-    TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
-    TORCH_CHECK(cos.dim() == 3, "cos must be 3D [batch, seq_len, head_dim/2]");
-    TORCH_CHECK(sin.dim() == 3, "sin must be 3D [batch, seq_len, head_dim/2]");
     TORCH_CHECK(head_dim % 2 == 0, "head_dim must be even");
 
     auto out = torch::empty_like(x);
@@ -73,8 +70,7 @@ torch::Tensor rotary_emb(
 
     rotary_emb_kernel<<<grid, block>>>(
         reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()),
-        cos.data_ptr<float>(),
-        sin.data_ptr<float>(),
+        freqs_cis.data_ptr<float>(),
         reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
         batch, seq_len, n_heads, head_dim
     );
@@ -85,8 +81,7 @@ torch::Tensor rotary_emb(
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("rotary_emb", &rotary_emb,
         py::arg("x"),
-        py::arg("cos"),
-        py::arg("sin"),
-        "Fused rotary embedding (bf16 x, f32 cos/sin, bf16 out)"
+        py::arg("freqs_cis"),
+        "Fused rotary embedding (bf16 x, f32 freqs_cis [b,s,d/2,2], bf16 out)"
     );
 }

docs/developer/architecture.md

@@ -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

docs/developer/cuda_kernels.md

@@ -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

docs/developer/internals.md

@@ -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

docs/get-started.md

@@ -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

docs/guides/inference.md

@@ -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

docs/guides/training.md

@@ -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