perf: add fused CUDA rotary embedding kernel

- Single-kernel rotary embedding (cos/sin lookup + rotation) replaces PyTorch complex-multiply path (3 kernel launches + f32 upcast per call)
- RotaryEmbedding now stores cos_table/sin_table and returns (cos, sin) f32 tuple instead of a complex tensor
- apply_rotary_emb in rotary_backend.py auto-dispatches: CUDA kernel if available, else torch complex-multiply fallback; backend-agnostic (both attention backends benefit)
- Kernel: 256-thread blocks, grid-stride loop, vectorized __nv_bfloat162 load/store, f32 compute, bf16 out
- Standalone kernel 6-9x faster than torch across decode/prefill shapes, max diff 0 (decode) to 3e-2 (large prefill, bf16)
- Benchmark (L20, bf16, CUDA backend): B=1 9.48->7.25ms (+31%), B=4 10.73->7.67ms (+40%), B=8 10.77->7.81ms (+38%), B=16 10.79->7.83ms (+38%)

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/loader.py

@@ -11,7 +11,7 @@ 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] = {}

astrai/extension/rotary_backend.py

@@ -0,0 +1,48 @@
+"""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 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)