refactor: standardize packed 3d inference

- keep training attention on dense 4d tensors
- use packed 3d tensors with KV cache for inference
- extend CUDA rotary embedding to packed 3d inputs
- adapt torch, CUDA and FlashAttention backend dispatch

astrai/extension/rotary_ops.py

@@ -4,8 +4,8 @@ Calls the compiled CUDA kernel directly. If the kernel is not available,
 raises ``RuntimeError``. Fallback to torch complex multiply is the
 responsibility of ``astrai.extension.rotary_backend.apply_rotary_emb``.
 
-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.
+Layout: x is packed [tokens, n_heads, head_dim] or dense
+[batch, seq_len, n_heads, head_dim]. ``freqs_cis`` has matching token axes.
 """
 
 import torch
@@ -25,11 +25,11 @@ def rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
     """Fused rotary embedding kernel.
 
     Args:
-        x: [batch, seq_len, n_heads, head_dim] (bf16, contiguous)
-        freqs_cis: [batch, seq_len, head_dim/2, 2] (f32, contiguous) — [cos, sin] pairs
+        x: packed 3D or dense 4D bf16 tensor.
+        freqs_cis: matching token axes followed by [head_dim/2, 2].
 
     Returns:
-        [batch, seq_len, n_heads, head_dim] (bf16)
+        Tensor with the same shape as ``x``.
     """
     _check_available()
     if not x.is_contiguous():