refactor: unify rotary embedding interface and update docs

- Merge cos/sin into single freqs_cis tensor [batch, seq, dim/2, 2] throughout the pipeline: RotaryEmbedding buffer, forward return type, apply_rotary_emb signature, CUDA kernel interface
- CUDA kernel now takes freqs_cis directly and reads cos/sin via stride offset internally, eliminating Python-side slice/copy overhead
- Kernel interface: rotary_emb(x, freqs_cis) replaces rotary_emb(x, cos, sin)
- All call sites pass rotary_emb as Tensor (was tuple), type annotations consistent
- Update build threads from 8 to 16
- Fix all docs: get-started, inference, training, cuda_kernels, architecture, internals — reflect new rotary interface, KVCache fields, rotary backend dispatch, .so path, kernel registry count, file layout

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