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