docs: align CUDA kernel and RoPE docs with code

- Fix rotary docs to describe cos/sin freqs_cis table, not complex buffer
- Replace attn_prefill with attn_paged_prefill for the CudaBackend path
- Register attn_paged_prefill in kernel overview, layout, and module list
- Add qo_indptr and InferenceWorkspace to architecture class diagram
- Add FrequencyPenaltyStrategy to sampling design patterns

docs/guides/inference.md

@@ -49,7 +49,8 @@ KVCache
   ├── seq_lens               [batch_size]
   ├── out_cache_loc          [batch, seq_len] — write indices for this forward
   ├── max_len                int — max(seq_lens), avoids GPU sync in decode
-  └── kv_indptr              [batch + 1] int32 — prefix sum of seq_lens, precomputed once per step
+  ├── kv_indptr              [batch + 1] int32 — prefix sum of seq_lens, precomputed once per step
+  └── qo_indptr              [batch + 1] int32 — prefix sum of per-request q_lens (prefill), precomputed once per step
 ```
 
 Attention layers do raw buffer indexing: `k_buffer[layer_id, out_cache_loc] = k` to write, `k_buffer[layer_id, indices]` to gather.
@@ -61,7 +62,7 @@ Attention computation (cache I/O + SDPA/kernel dispatch) is decoupled from the m
 ```
 AttentionBackend (ABC)
   ├── TorchNativeBackend   SDPA + indirect KV cache gather (default)
-  └── CudaBackend          CUDA kernel dispatch (attn_paged_decode, attn_prefill)
+  └── CudaBackend          CUDA kernel dispatch (attn_paged_decode, attn_paged_prefill)
 ```
 
 Select via context manager (mirrors `torch.nn.attention.sdpa_kernel`):
@@ -75,7 +76,7 @@ with attn_backend(ATTN_BACKEND.CUDA):
 
 `CudaBackend` decode path: writes K/V to cache, then calls `attn_paged_decode` with `page_size=1` — the `req_to_token` table serves directly as the page table, each token slot is a single-token "page". No explicit K/V gather needed.
 
-`CudaBackend` prefill path: writes K/V, gathers full-sequence K/V via indirect indexing (same as `TorchNativeBackend`), then calls `attn_prefill`.
+`CudaBackend` prefill path: writes K/V, then calls `attn_paged_prefill` — a ragged-batch (paged) prefill kernel that reads K/V directly from the flat pool via `req_to_token`, addressing each request's `q_len`/`kv_len` through `qo_indptr` and `kv_indptr`. No explicit K/V gather needed.
 
 Fallback: `CudaBackend` delegates to `TorchNativeBackend` when a CUDA kernel is not available.
 
@@ -83,12 +84,13 @@ Fallback: `CudaBackend` delegates to `TorchNativeBackend` when a CUDA kernel is
 
 Rotary embedding is applied via `apply_rotary_emb` in `astrai/extension/rotary_backend.py`, which auto-dispatches:
 
-- **CUDA kernel** (`rotary_emb.cu`): fused cos/sin lookup + rotation in a single kernel, used when the kernel is available, input is on CUDA, and `torch.is_grad_enabled()` is `False` (inference mode)
+- **CUDA kernel** (`rotary_emb.cu`): fused cos/sin lookup + rotation in a single kernel, used when the kernel is available, the input is bf16 on CUDA, and `torch.is_grad_enabled()` is `False` (inference mode)
 - **Torch fallback**: complex multiply path (`torch.view_as_complex` → `torch.complex` multiply → `torch.view_as_real`), used during training (supports autograd backward) or when the CUDA kernel is not available
 
-`RotaryEmbedding` stores a complex `freqs_cis` buffer and returns a tensor
-from `forward()`. Both attention backends share the same rotary dispatch — it
-is backend-agnostic.
+`RotaryEmbedding` stores a cos/sin table `freqs_cis` of shape
+`[max_len, dim/2, 2]` (f32 — `[cos, sin]` pairs) and `forward()` returns
+a `[batch, seq_len, dim/2, 2]` slice indexed by `position_ids`. Both
+attention backends share the same rotary dispatch — it is backend-agnostic.
 
 ## Continuous Batching
 

docs/guides/training.md

@@ -41,10 +41,12 @@ RoPE embeds position into Q/K vectors via complex rotation:
 
 $$ q_i = R_i W_q x_i, \quad k_j = R_j W_k x_j, \quad q_i^T k_j = x_i^T W_q^T R_{i-j} W_k x_j $$
 
-`RotaryEmbedding` pre-computes a complex `freqs_cis` buffer. `forward()` returns
-a tensor indexed by `position_ids`. `apply_rotary_emb` applies the rotation:
-during training it uses torch complex multiply (autograd-compatible); during
-inference it auto-dispatches to a fused CUDA kernel when available.
+`RotaryEmbedding` pre-computes a cos/sin table `freqs_cis` of shape
+`[max_len, dim/2, 2]` (f32 — `[cos, sin]` pairs). `forward()` returns
+a `[batch, seq_len, dim/2, 2]` slice indexed by `position_ids`.
+`apply_rotary_emb` applies the rotation: during training it uses torch
+complex multiply (autograd-compatible); during inference it auto-dispatches
+to a fused CUDA kernel when available.
 
 ## Training Loop