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

Benchmark: NVIDIA L20, BF16, 1B model, paged KV cache, CUDA Graph, prompt 512, generation 128 (median of 3 alternating runs)
- batch 1: 234.5 -> 242.6 tok/s (1.034x, +3.4%)
- batch 8: 1243.1 -> 1286.6 tok/s (1.035x, +3.5%)

astrai/extension/backend/attention.py

@@ -253,28 +253,6 @@ def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
     )
 
 
-def _write_and_gather_kv(
-    kv_cache: "KVCache",
-    k: Tensor,
-    v: Tensor,
-    layer_id: int,
-    q: Tensor,
-    attn_mask: Optional[Tensor],
-) -> tuple[Tensor, Tensor]:
-    kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
-    kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
-    max_len = kv_cache.max_len
-    indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
-    if q.size(1) == 1 and attn_mask is not None and attn_mask.dim() == 4:
-        pos_mask = attn_mask[:, 0, 0]
-    else:
-        pos_mask = (
-            torch.arange(max_len, device=q.device)[None, :] < kv_cache.seq_lens[:, None]
-        )
-    indices = torch.where(pos_mask, indices, torch.zeros_like(indices))
-    return kv_cache.k_buffer[layer_id, indices], kv_cache.v_buffer[layer_id, indices]
-
-
 def attention(
     q: Tensor,
     k: Tensor,
@@ -565,10 +543,6 @@ class CudaBackend(AttentionBackend):
         if kv_cache is None:
             raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
 
-        loc = kv_cache.out_cache_loc
-        kv_cache.k_buffer[layer_id, loc] = k
-        kv_cache.v_buffer[layer_id, loc] = v
-
         kv_indptr = kv_cache.kv_indptr
 
         out = attn_paged_decode(
@@ -578,6 +552,8 @@ class CudaBackend(AttentionBackend):
             kv_cache.req_to_token,
             kv_cache.req_pool_indices,
             kv_indptr,
+            new_k=k,
+            new_v=v,
             is_causal=True,
             o_part_buf=kv_cache.decode_o_part,
             ml_part_buf=kv_cache.decode_ml_part,

astrai/extension/ops/attention.py

@@ -97,6 +97,8 @@ def attn_paged_decode(
     req_to_token: torch.Tensor,
     req_pool_indices: torch.Tensor,
     kv_indptr: torch.Tensor,
+    new_k: Optional[torch.Tensor] = None,
+    new_v: Optional[torch.Tensor] = None,
     mask: Optional[torch.Tensor] = None,
     is_causal: bool = False,
     o_part_buf: Optional[torch.Tensor] = None,
@@ -116,6 +118,8 @@ def attn_paged_decode(
         req_to_token: [num_reqs, max_context_len] (int32) — token -> slot
         req_pool_indices: [batch] (int32) — rows into req_to_token
         kv_indptr: [batch+1] (int32) — prefix sum of per-request seq_lens
+        new_k: current-token K to append, [batch, n_kv_heads, head_dim]
+        new_v: current-token V to append, same shape as new_k
         mask: 2D [batch, max_context_len] (bool, True=keep) or None
         is_causal: apply causal mask
         o_part_buf: pre-allocated split-KV o partial buffer (workflow bypass)
@@ -134,6 +138,8 @@ def attn_paged_decode(
         req_to_token,
         req_pool_indices,
         kv_indptr,
+        new_k=new_k,
+        new_v=new_v,
         mask=mask,
         causal_offset=causal_offset,
         o_part_buf=o_part_buf,