fix: skip CUDA rotary kernel when grad is enabled

- apply_rotary_emb now checks torch.is_grad_enabled() before dispatching to CUDA kernel
- Training (grad enabled) uses torch complex multiply path which supports autograd backward
- Inference (inference_mode/no_grad) uses CUDA kernel as before
- Without this fix, training backward would crash — the CUDA kernel has no autograd backward()

.dockerignore

@@ -5,5 +5,7 @@
 !astrai/
 !scripts/
 !docs/
+!csrc/
+!setup.py
 !pyproject.toml
 !README.md

.github/workflows/release.yml

@@ -26,22 +26,30 @@ jobs:
           if-no-files-found: error
 
   build-cuda-linux:
-    name: Build CUDA wheel (Linux)
+    name: Build CUDA wheel (Linux, ${{ matrix.cuda_tag }})
     runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - cuda_tag: "cu128"
+            cuda_ver: "12.8.0"
+          - cuda_tag: "cu130"
+            cuda_ver: "13.0.0"
     steps:
       - uses: actions/checkout@v4
       - uses: actions/setup-python@v5
         with:
           python-version: "3.12"
 
-      - name: Install torch (CUDA 12.8)
+      - name: Install torch (${{ matrix.cuda_tag }})
         run: |
-          pip install torch --index-url https://download.pytorch.org/whl/cu128
+          pip install torch --index-url https://download.pytorch.org/whl/${{ matrix.cuda_tag }}
 
-      - name: Setup CUDA
+      - name: Setup CUDA (${{ matrix.cuda_ver }})
         uses: Jimver/cuda-toolkit@v0.2.35
         with:
-          cuda: "12.8.0"
+          cuda: "${{ matrix.cuda_ver }}"
 
       - name: Build wheel (with CUDA kernels)
         run: |
@@ -49,7 +57,7 @@ jobs:
 
       - uses: actions/upload-artifact@v4
         with:
-          name: cuda-wheel-linux
+          name: cuda-wheel-linux-${{ matrix.cuda_tag }}
           path: dist/*.whl
           if-no-files-found: error
 
@@ -66,10 +74,11 @@ jobs:
           name: pure-wheel
           path: release-assets/pure
 
-      - name: Download CUDA wheel
+      - name: Download CUDA wheels (all variants)
         uses: actions/download-artifact@v4
         with:
-          name: cuda-wheel-linux
+          pattern: cuda-wheel-linux-*
+          merge-multiple: true
           path: release-assets/cuda
 
       - name: Verify release assets
@@ -79,8 +88,7 @@ jobs:
           pure_wheels=(release-assets/pure/*.whl)
           cuda_wheels=(release-assets/cuda/*.whl)
           test "${#pure_wheels[@]}" -eq 1
-          test "${#cuda_wheels[@]}" -eq 1
-          test "$(basename "${pure_wheels[0]}")" != "$(basename "${cuda_wheels[0]}")"
+          test "${#cuda_wheels[@]}" -ge 1
 
       - name: Create release & upload assets
         uses: softprops/action-gh-release@v2

Dockerfile

@@ -1,8 +1,16 @@
 # AstrAI Dockerfile - Multi-stage Build (Optimized)
+#
+# CUDA version selection:
+#   docker build -t astrai .
+#   docker build -t astrai --build-arg CUDA_TAG=cu128 .
+#   docker build -t astrai --build-arg CUDA_TAG=cu130 .
+# Default: cu128
 
 # Build stage - use base image with minimal build tools
 FROM ubuntu:24.04 AS builder
 
+ARG CUDA_TAG=cu128
+
 WORKDIR /app
 
 # Install Python 3.12 and minimal build dependencies
@@ -20,10 +28,12 @@ ENV PATH="/opt/venv/bin:$PATH"
 
 # Copy source code and install (deps read from pyproject.toml)
 COPY astrai/ ./astrai/
+COPY csrc/ ./csrc/
+COPY setup.py .
 COPY pyproject.toml .
 RUN pip install --no-cache-dir --upgrade pip \
     && pip install --no-cache-dir . \
-    --extra-index-url https://download.pytorch.org/whl/cu128
+    --extra-index-url "https://download.pytorch.org/whl/${CUDA_TAG}"
 
 # Production stage
 FROM ubuntu:24.04 AS production

astrai/extension/__init__.py

@@ -30,6 +30,7 @@ from astrai.extension.attention_ops import (
     attn_prefill,
 )
 from astrai.extension.loader import KERNEL_NAMES, is_available
+from astrai.extension.rotary_backend import apply_rotary_emb
 
 __all__ = [
     "ATTN_BACKEND",
@@ -44,4 +45,5 @@ __all__ = [
     "attn_prefill",
     "is_available",
     "KERNEL_NAMES",
+    "apply_rotary_emb",
 ]

astrai/extension/attention_backend.py

@@ -272,12 +272,18 @@ class TorchNativeBackend(AttentionBackend):
             kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
             kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
 
-            max_len = kv_cache.seq_lens.max()
-            indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
-            pos_mask = (
-                torch.arange(max_len, device=q.device)[None, :]
-                < kv_cache.seq_lens[:, None]
-            )
+            max_len = kv_cache.max_len
+            if kv_cache.page_table is not None:
+                indices = kv_cache.page_table
+            else:
+                indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
+            if kv_cache.decode_mask is not None:
+                pos_mask = kv_cache.decode_mask
+            else:
+                pos_mask = (
+                    torch.arange(max_len, device=q.device)[None, :]
+                    < kv_cache.seq_lens[:, None]
+                )
             indices = torch.where(pos_mask, indices, torch.zeros_like(indices))
             k = kv_cache.k_buffer[layer_id, indices]
             v = kv_cache.v_buffer[layer_id, indices]
@@ -338,18 +344,25 @@ class CudaBackend(AttentionBackend):
         kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
         kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
 
-        seq_lens = kv_cache.seq_lens
         max_len = kv_cache.max_len
 
-        page_table = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
+        if kv_cache.page_table is not None:
+            page_table = kv_cache.page_table
+        else:
+            page_table = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
 
         k_cache = kv_cache.k_buffer[layer_id].unsqueeze(1)
         v_cache = kv_cache.v_buffer[layer_id].unsqueeze(1)
 
         if q.size(0) == 1:
             mask = None
+        elif kv_cache.decode_mask is not None:
+            mask = kv_cache.decode_mask
         else:
-            mask = torch.arange(max_len, device=q.device)[None, :] < seq_lens[:, None]
+            mask = (
+                torch.arange(max_len, device=q.device)[None, :]
+                < kv_cache.seq_lens[:, None]
+            )
 
         out = attn_paged_decode(
             q,

astrai/extension/lib/__init__.py

@@ -0,0 +1 @@
+"""Compiled CUDA kernel modules (``*.so``) live here, kept separate from Python source."""

astrai/extension/loader.py

@@ -11,14 +11,14 @@ import logging
 
 logger = logging.getLogger(__name__)
 
-KERNEL_NAMES = ["attn_decode", "attn_prefill", "attn_paged_decode"]
+KERNEL_NAMES = ["attn_decode", "attn_prefill", "attn_paged_decode", "rotary_emb"]
 
 _available: dict[str, bool] = {}
 _modules: dict[str, object] = {}
 
 for _name in KERNEL_NAMES:
     try:
-        _mod = importlib.import_module(f".{_name}", package=__package__)
+        _mod = importlib.import_module(f".lib.{_name}", package=__package__)
         _available[_name] = True
         _modules[_name] = _mod
     except ImportError:

astrai/extension/rotary_backend.py

@@ -0,0 +1,53 @@
+"""Rotary embedding with auto-dispatch to CUDA kernel.
+
+Single entry point ``apply_rotary_emb(x, cos, sin)`` — 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).
+"""
+
+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}
+
+
+def _cuda_available() -> bool:
+    if _cache["available"] is None:
+        _cache["available"] = is_available("rotary_emb")
+    return _cache["available"]
+
+
+def _torch_apply(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
+    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
+    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:
+    """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)
+
+    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)

astrai/extension/rotary_ops.py

@@ -0,0 +1,46 @@
+"""Rotary embedding CUDA kernel wrapper.
+
+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``.
+
+Layout convention: x is ``[batch, seq_len, n_heads, head_dim]`` (blhd, bf16).
+cos/sin are ``[batch, seq_len, head_dim/2]`` (f32).
+"""
+
+import torch
+
+from astrai.extension.loader import _available, _modules
+
+
+def _check_available():
+    if not _available.get("rotary_emb"):
+        raise RuntimeError(
+            "CUDA kernel 'rotary_emb' is not available. "
+            "Build with CSRC_KERNELS=true or use the torch fallback."
+        )
+
+
+def rotary_emb(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: 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)
+
+    Returns:
+        [batch, seq_len, n_heads, head_dim] (bf16)
+    """
+    _check_available()
+    if not x.is_contiguous():
+        x = x.contiguous()
+    return _modules["rotary_emb"].rotary_emb(x, cos, sin)

astrai/inference/core/cache.py

@@ -203,6 +203,10 @@ class KVCache:
         seq_lens: [batch_size] — per-request total sequence lengths
         out_cache_loc: [batch, new_seq_len] or [batch, 1] — write indices
         max_len: max(seq_lens) as Python int — avoids GPU sync in decode
+        page_table: [batch, max_len] — precomputed gather indices for decode;
+            None for prefill or when not yet computed.
+        decode_mask: [batch, max_len] bool — precomputed position validity
+            mask for decode; None for prefill or single-batch decode.
     """
 
     k_buffer: Tensor
@@ -212,6 +216,8 @@ class KVCache:
     seq_lens: Tensor
     out_cache_loc: Tensor
     max_len: int = 0
+    page_table: Optional[Tensor] = None
+    decode_mask: Optional[Tensor] = None
 
 
 class PagePool:
@@ -428,11 +434,21 @@ class PagePool:
             out_cache_loc = self._req_pool.req_to_token[
                 req_pool_indices, start_pos:seq_len
             ]
+            page_table = None
+            decode_mask = None
         else:
             write_pos = seq_lens_t - 1
             out_cache_loc = self._req_pool.req_to_token[
                 req_pool_indices, write_pos
             ].unsqueeze(-1)
+            ml = max(seq_lens)
+            page_table = self._req_pool.req_to_token[req_pool_indices, :ml]
+            if len(task_ids) > 1:
+                decode_mask = (
+                    torch.arange(ml, device=device)[None, :] < seq_lens_t[:, None]
+                )
+            else:
+                decode_mask = None
 
         return KVCache(
             k_buffer=self._storage.k_buffer,
@@ -442,6 +458,8 @@ class PagePool:
             seq_lens=seq_lens_t,
             out_cache_loc=out_cache_loc,
             max_len=max(seq_lens),
+            page_table=page_table,
+            decode_mask=decode_mask,
         )
 
     # ---- internals ----

astrai/inference/core/executor.py

@@ -105,26 +105,33 @@ class Executor:
             [t.frequency_penalty for t in tasks], device=self.device
         )
 
-        history_lists = []
-        history_lens = []
-        for t in tasks:
-            window = t.rep_window
-            prompt_part = t.prompt_ids[-window:]
-            ids = prompt_part + t.output_ids
-            history_lists.append(ids)
-            history_lens.append(len(ids))
+        has_freq = bool((freq_penalties != 0).any())
+        if has_freq:
+            history_lists = []
+            history_lens = []
+            for t in tasks:
+                window = t.rep_window
+                prompt_part = t.prompt_ids[-window:]
+                ids = prompt_part + t.output_ids
+                history_lists.append(ids)
+                history_lens.append(len(ids))
 
-        max_len = max(history_lens) if history_lens else 0
-        padded_ids = torch.zeros(
-            len(tasks), max_len, dtype=torch.long, device=self.device
-        )
-        padded_mask = torch.zeros(
-            len(tasks), max_len, dtype=torch.bool, device=self.device
-        )
-        for i, h in enumerate(history_lists):
-            L = history_lens[i]
-            padded_ids[i, :L] = torch.as_tensor(h, dtype=torch.long, device=self.device)
-            padded_mask[i, :L] = True
+            max_len = max(history_lens) if history_lens else 0
+            padded_ids = torch.zeros(
+                len(tasks), max_len, dtype=torch.long, device=self.device
+            )
+            padded_mask = torch.zeros(
+                len(tasks), max_len, dtype=torch.bool, device=self.device
+            )
+            for i, h in enumerate(history_lists):
+                L = history_lens[i]
+                padded_ids[i, :L] = torch.as_tensor(
+                    h, dtype=torch.long, device=self.device
+                )
+                padded_mask[i, :L] = True
+        else:
+            padded_ids = None
+            padded_mask = None
 
         with torch.inference_mode():
             outputs = self.model(

astrai/inference/core/scheduler.py

@@ -109,9 +109,11 @@ class InferenceScheduler:
                     self._task_mgr.wait_for_tasks(timeout=1.0)
                     continue
 
+                active = self._task_mgr.get_active_tasks()
+
                 to_prefill = [
                     t
-                    for t in self._task_mgr.get_active_tasks()
+                    for t in active
                     if t.output_tokens == 0
                     and cache.task_cached(t.task_id) < len(t.prompt_ids)
                 ]
@@ -137,10 +139,10 @@ class InferenceScheduler:
                                 t.task_id, t.prompt_ids, start_logical_page
                             )
 
-                decode_tasks = self._task_mgr.get_active_tasks()
+                decode_tasks = active
 
                 valid: List[Task] = []
-                for t in sorted(decode_tasks, key=lambda t: t.task_id):
+                for t in decode_tasks:
                     if cache.task_extend(t.task_id, t.next_pos):
                         valid.append(t)
                     else:

astrai/inference/sample.py

@@ -343,6 +343,10 @@ def sample(
     When **temperature** is exactly 0 (scalar or single-element tensor)
     the function short-circuits to ``argmax`` for deterministic decode.
 
+    When **frequency_penalty** is 0 (the common decode case), the entire
+    frequency penalty computation — including the O(batch * vocab) count
+    tensor allocation — is skipped.
+
     Args:
         logits: Raw logits ``[batch, vocab_size]``.
         frequency_penalty: Penalty per occurrence for repeated tokens
@@ -359,14 +363,39 @@ def sample(
         ``True`` — a ``(token_ids, chosen_logprobs)`` tuple where
         ``chosen_logprobs`` has shape ``[batch]``.
     """
-    return SamplingPipeline(
-        [
-            TemperatureStrategy(temperature),
-            TopKStrategy(top_k),
-            TopPStrategy(top_p),
-            FrequencyPenaltyStrategy(frequency_penalty),
-        ]
-    ).sample(
+    greedy = (
+        (
+            isinstance(temperature, Tensor)
+            and temperature.numel() == 1
+            and temperature.item() == 0
+        )
+        if isinstance(temperature, Tensor)
+        else temperature == 0
+    )
+
+    if greedy:
+        tokens = logits.argmax(dim=-1)
+        if not return_logprobs:
+            return tokens
+        log_probs = torch.log_softmax(logits.float(), dim=-1)
+        chosen = torch.gather(log_probs, -1, tokens.unsqueeze(-1)).squeeze(-1)
+        return tokens, chosen
+
+    has_freq = (
+        (isinstance(frequency_penalty, Tensor) and (frequency_penalty != 0).any())
+        if isinstance(frequency_penalty, Tensor)
+        else frequency_penalty != 0
+    )
+
+    strategies: List[BaseSamplingStrategy] = [
+        TemperatureStrategy(temperature),
+        TopKStrategy(top_k),
+        TopPStrategy(top_p),
+    ]
+    if has_freq:
+        strategies.append(FrequencyPenaltyStrategy(frequency_penalty))
+
+    return SamplingPipeline(strategies).sample(
         logits,
         filter_value=filter_value,
         input_ids=input_ids,

astrai/model/components/__init__.py

@@ -1,3 +1,4 @@
+from astrai.extension.rotary_backend import apply_rotary_emb
 from astrai.model.components.attention import GQA, MLA
 from astrai.model.components.decoder_block import DecoderBlock
 from astrai.model.components.embedding import Embedding
@@ -6,7 +7,6 @@ from astrai.model.components.mlp import MLP
 from astrai.model.components.norm import RMSNorm
 from astrai.model.components.rope import (
     RotaryEmbedding,
-    apply_rotary_emb,
     get_rotary_emb,
 )
 

astrai/model/components/attention.py

@@ -6,11 +6,11 @@ import torch.nn.functional as F
 from torch import Tensor
 
 from astrai.extension import attention
+from astrai.extension.rotary_backend import apply_rotary_emb
 from astrai.factory import BaseFactory
 from astrai.inference.core.cache import KVCache
 from astrai.model.components.linear import Linear
 from astrai.model.components.norm import RMSNorm
-from astrai.model.components.rope import apply_rotary_emb
 
 
 class AttnFactory(BaseFactory[nn.Module]):

astrai/model/components/rope.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional
+from typing import Dict, Optional, Tuple
 
 import torch
 import torch.nn as nn
@@ -10,29 +10,22 @@ def get_rotary_emb(
     max_len: int,
     base: float = 10000,
     device: Optional[torch.device] = None,
-) -> Tensor:
+) -> Tuple[Tensor, Tensor]:
+    """Precompute cos/sin tables for rotary embedding.
+
+    Returns:
+        (cos, sin) each of shape [max_len, dim/2] (f32)
+    """
     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()
-    cos = torch.cos(freqs)
-    sin = torch.sin(freqs)
-    return torch.complex(cos, sin)
+    return torch.cos(freqs), torch.sin(freqs)
 
 
 def ntk_base(base: float, dim: int, factor: float) -> float:
     return base * (factor ** (dim / (dim - 2)))
 
 
-def apply_rotary_emb(x: torch.Tensor, freqs_cis: Tensor) -> Tensor:
-    dtype = x.dtype
-    x_ = x.float().reshape(*x.shape[:-1], -1, 2)
-    x_complex = torch.view_as_complex(x_)
-    freqs_cis = freqs_cis.unsqueeze(2)
-    x_rotated = x_complex * freqs_cis
-    x_out = torch.view_as_real(x_rotated).flatten(-2)
-    return x_out.to(dtype)
-
-
 class RotaryEmbedding(nn.Module):
     def __init__(
         self,
@@ -56,16 +49,28 @@ class RotaryEmbedding(nn.Module):
         self._set_rotary_buffer(self.max_len)
 
     def _set_rotary_buffer(self, max_len: int):
-        rotary_emb = get_rotary_emb(self.dim, max_len, self.base)
-        freqs_cis = torch.view_as_real(rotary_emb)
-        self.register_buffer("freqs_cis", freqs_cis, persistent=False)
+        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)
 
-    def forward(self, x: Tensor, position_ids: Optional[Tensor] = None) -> Tensor:
+    def forward(
+        self, x: Tensor, position_ids: Optional[Tensor] = None
+    ) -> Tuple[Tensor, Tensor]:
+        """Lookup cos/sin for the given positions.
+
+        Args:
+            x: [batch, seq_len, ...] — only batch and seq_len are used.
+            position_ids: [batch, seq_len] optional position indices.
+
+        Returns:
+            (cos, sin) each of shape [batch, seq_len, dim/2] (f32)
+        """
         if position_ids is None:
             position_ids = (
                 torch.arange(x.size(1), device=x.device)
                 .unsqueeze(0)
                 .expand(x.size(0), -1)
             )
-        position_freq_cis = self.freqs_cis[position_ids].float()
-        return torch.view_as_complex(position_freq_cis)
+        cos = self.cos_table[position_ids].float()
+        sin = self.sin_table[position_ids].float()
+        return cos, sin

csrc/build.py

@@ -1,6 +1,32 @@
 from pathlib import Path
 
 
+def cuda_toolkit_version() -> tuple[int, int] | None:
+    """Return ``(major, minor)`` of the nvcc on PATH, or ``None``.
+
+    Used by ``setup.py`` to detect nvcc/torch CUDA version mismatches
+    (e.g. nvcc 13.0 with a cu128 torch wheel) which cause cryptic ABI errors.
+    """
+    import shutil
+    import subprocess
+
+    nvcc = shutil.which("nvcc")
+    if nvcc is None:
+        return None
+    try:
+        out = subprocess.check_output(
+            [nvcc, "--version"], stderr=subprocess.STDOUT, text=True
+        )
+        for line in out.splitlines():
+            if "release" in line:
+                ver = line.split("release")[1].split(",")[0].strip()
+                major, minor = ver.split(".")
+                return (int(major), int(minor))
+    except Exception:
+        pass
+    return None
+
+
 def _arch_flags() -> list[str]:
     import torch
 
@@ -46,3 +72,4 @@ def register(name: str, sources: list[str] | None = None, **kwargs):
 register("attn_decode")
 register("attn_prefill")
 register("attn_paged_decode")
+register("rotary_emb")

csrc/kernels/attn_entry_utils.cuh

@@ -23,8 +23,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({p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt);
-    auto ml_part = torch::empty({p.batch, p.q_head, MAX_SPLITS, 2}, fopt);
+    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);
     p.o_part = (float*)o_part.data_ptr();
     p.ml_part = (float*)ml_part.data_ptr();
 }

csrc/kernels/attn_mma_utils.cuh

@@ -3,6 +3,12 @@
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
 
+// Predicated cp.async (4-operand form) requires CUDA 11.2+.
+// bf16 mma.sync requires sm_80+ (guarded at build time by ASTRAI_NO_MMA).
+#if CUDART_VERSION < 11020
+#error "AstrAI CUDA kernels require CUDA 11.2 or later (CUDART_VERSION >= 11020)."
+#endif
+
 // ============================================================================
 // KernelTraits — FlashAttention-v2 style compile-time configuration bundle.
 //

csrc/kernels/rotary_emb.cu

@@ -0,0 +1,92 @@
+#include <torch/extension.h>
+#include <cuda_bf16.h>
+
+__global__ void rotary_emb_kernel(
+    const __nv_bfloat16* __restrict__ x,
+    const float* __restrict__ cos,
+    const float* __restrict__ sin,
+    __nv_bfloat16* __restrict__ out,
+    int batch,
+    int seq_len,
+    int n_heads,
+    int head_dim
+) {
+    const int half_dim = head_dim >> 1;
+    const int total = batch * seq_len * n_heads * half_dim;
+
+    for (int idx = blockIdx.x * blockDim.x + threadIdx.x;
+         idx < total;
+         idx += gridDim.x * blockDim.x) {
+
+        int pair = idx % half_dim;
+        int tmp = idx / half_dim;
+        int head = tmp % n_heads;
+        tmp /= n_heads;
+        int seq = tmp % seq_len;
+        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;
+
+        __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 out_even = x_even * c - x_odd * s;
+        float out_odd  = x_even * s + x_odd * c;
+
+        __nv_bfloat162 out_pair = __floats2bfloat162_rn(out_even, out_odd);
+        *reinterpret_cast<__nv_bfloat162*>(out + x_offset) = out_pair;
+    }
+}
+
+torch::Tensor rotary_emb(
+    torch::Tensor x,
+    torch::Tensor cos,
+    torch::Tensor sin
+) {
+
+    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);
+
+    int half_dim = head_dim / 2;
+    int total = batch * seq_len * n_heads * half_dim;
+    int block = 256;
+    int grid = std::min((total + block - 1) / block, 1024);
+
+    rotary_emb_kernel<<<grid, block>>>(
+        reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()),
+        cos.data_ptr<float>(),
+        sin.data_ptr<float>(),
+        reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
+        batch, seq_len, n_heads, head_dim
+    );
+
+    return out;
+}
+
+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)"
+    );
+}

docker-compose.yml

@@ -3,6 +3,8 @@ services:
     build:
       context: .
       dockerfile: Dockerfile
+      args:
+        CUDA_TAG: ${CUDA_TAG:-cu128}
     user: "${UID:-1000}:${GID:-1000}"
     ports:
       - "8000:8000"
@@ -29,6 +31,8 @@ services:
     build:
       context: .
       dockerfile: Dockerfile
+      args:
+        CUDA_TAG: ${CUDA_TAG:-cu128}
     user: "${UID:-1000}:${GID:-1000}"
     ports:
       - "8000:8000"

pyproject.toml

@@ -36,9 +36,6 @@ dev = ["pytest==9.0.2", "ruff", "httpx2"]
 [tool.setuptools.packages.find]
 where = ["."]
 
-[tool.pip]
-extra-index-url = "https://download.pytorch.org/whl/cu128"
-
 [tool.setuptools.dynamic]
 version = { attr = "astrai.__version__" }
 

setup.py

@@ -1,12 +1,13 @@
 import os
 import sys
+import warnings
 from pathlib import Path
 
 from setuptools import setup
 from setuptools.command.build_ext import build_ext as _build_ext
 
 sys.path.insert(0, str(Path(__file__).parent))
-os.makedirs("astrai/extension", exist_ok=True)
+os.makedirs("astrai/extension/lib", exist_ok=True)
 
 
 def _should_build():
@@ -32,14 +33,31 @@ if _should_build():
     import torch
     from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
-    from csrc.build import REGISTRY
+    from csrc.build import REGISTRY, cuda_toolkit_version
+
+    # Preflight: warn if nvcc major version != torch's bundled CUDA major version.
+    # A mismatch (e.g. nvcc 13.0 + cu128 torch) causes cryptic ABI/header errors.
+    nvcc_ver = cuda_toolkit_version()
+    torch_cuda = torch.version.cuda
+    if nvcc_ver is not None and torch_cuda is not None:
+        torch_major = int(torch_cuda.split(".")[0])
+        if nvcc_ver[0] != torch_major:
+            warnings.warn(
+                f"CUDA version mismatch: nvcc is {nvcc_ver[0]}.{nvcc_ver[1]} "
+                f"but torch was built with CUDA {torch_cuda}. "
+                f"This may cause compilation errors. "
+                f"Install a matching torch wheel: "
+                f"pip install torch --index-url "
+                f"https://download.pytorch.org/whl/cu{nvcc_ver[0]}{nvcc_ver[1]}",
+                stacklevel=2,
+            )
 
     _torch_lib = torch.utils.cpp_extension.library_paths()[0]
 
     for name, info in REGISTRY.items():
         ext_modules.append(
             CUDAExtension(
-                f"astrai.extension.{name}",
+                f"astrai.extension.lib.{name}",
                 info["sources"],
                 extra_compile_args={
                     "cxx": info["cxx_flags"],