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