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,

csrc/kernels/attn_common.h

@@ -36,6 +36,8 @@ struct AttentionParams {
     const T* __restrict__ q_ptr;
     const T* __restrict__ k_ptr;
     const T* __restrict__ v_ptr;
+    const T* __restrict__ new_k_ptr;
+    const T* __restrict__ new_v_ptr;
     T* __restrict__ o_ptr;
     const bool* __restrict__ mask;
 
@@ -50,6 +52,9 @@ struct AttentionParams {
     int kv_l_stride;
     int kv_d_stride;
 
+    int new_kv_b_stride;
+    int new_kv_h_stride;
+
     int mask_b_stride;
     int mask_h_stride;
     int mask_l_stride;

csrc/kernels/attn_decode_split_kv.cuh

@@ -57,8 +57,8 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
             int s = i / p.head_dim;
             int d_dim = i % p.head_dim;
             int kc = chunk_start + s;
-            int token = KV::resolve_token(p, kctx, kc, true);
+            KVAddr a = KV::template decode_addr<1>(
-            KVAddr a = KV::kv_addr_from_token(p, kctx, token, d_dim);
+                p, kctx, batch, kv_head, kc, d_dim, true, true);
             k_smem[i] = a.valid ? *reinterpret_cast<const bf16*>(a.k) : (bf16)0.f;
             v_smem[i] = a.valid ? *reinterpret_cast<const bf16*>(a.v) : (bf16)0.f;
         }

csrc/kernels/attn_decode_split_kv_mma.cuh

@@ -73,8 +73,10 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
             int r = i / Traits::HEAD_DIM, d = i % Traits::HEAD_DIM;
             int kc = kv0 + r;
             bool valid = kc < seq_len;
-            int token = KV::resolve_token(p, kctx, kc, valid);
+            // All GQA passes consume new K/V directly. Only the first pass
-            KVAddr a = KV::kv_addr_from_token(p, kctx, token, d);
+            // persists it, so no cross-block synchronization is required.
+            KVAddr a = KV::template decode_addr<Traits::VEC>(
+                p, kctx, batch, kv_head, kc, d, valid, pass == 0);
             int off = r * Traits::LD + swiz_col(d, r, Traits::SWIZ_MASK);
             cp_async_16_pred(&dK[off], a.k, a.valid);
             cp_async_16_pred(&dV[off], a.v, a.valid);

csrc/kernels/attn_entry_utils.cuh

@@ -130,6 +130,8 @@ inline void attn_pack_params(
     p.q_ptr = (const T*)q.data_ptr();
     p.k_ptr = (const T*)k.data_ptr();
     p.v_ptr = (const T*)v.data_ptr();
+    p.new_k_ptr = nullptr;
+    p.new_v_ptr = nullptr;
     p.o_ptr = nullptr;
     p.o_part = nullptr;
     p.ml_part = nullptr;
@@ -148,6 +150,8 @@ inline void attn_pack_paged_decode_params(
     torch::Tensor req_to_token,
     torch::Tensor req_pool_indices,
     torch::Tensor kv_indptr,
+    const c10::optional<torch::Tensor>& new_k,
+    const c10::optional<torch::Tensor>& new_v,
     c10::optional<torch::Tensor> mask,
     int64_t causal_offset,
     double scale,
@@ -191,6 +195,33 @@ inline void attn_pack_paged_decode_params(
     p.qo_indptr = nullptr;
     p.max_context_len = (int)req_to_token.size(1);
 
+    TORCH_CHECK(new_k.has_value() == new_v.has_value(),
+                "new_k and new_v must be provided together");
+    if (new_k.has_value()) {
+        auto nk = new_k.value();
+        auto nv = new_v.value();
+        TORCH_CHECK(nk.is_cuda() && nv.is_cuda(), "new K/V must be CUDA tensors");
+        TORCH_CHECK(nk.dtype() == torch::kBFloat16 && nv.dtype() == torch::kBFloat16,
+                    "new K/V must be bf16");
+        TORCH_CHECK(nk.dim() == 3 && nv.dim() == 3,
+                    "new K/V must be 3D [batch, kv_head, head_dim]");
+        TORCH_CHECK(nk.sizes() == nv.sizes(), "new K and V must have identical shapes");
+        TORCH_CHECK(nk.strides() == nv.strides(),
+                    "new K and V must have identical strides");
+        TORCH_CHECK(nk.size(0) == p.batch && nk.size(1) == p.kv_head
+                    && nk.size(2) == p.head_dim, "new K/V shape mismatch");
+        TORCH_CHECK(nk.stride(2) == 1 && nv.stride(2) == 1,
+                    "new K/V head_dim must be contiguous");
+        p.new_k_ptr = (const T*)nk.data_ptr();
+        p.new_v_ptr = (const T*)nv.data_ptr();
+        p.new_kv_b_stride = (int)nk.stride(0);
+        p.new_kv_h_stride = (int)nk.stride(1);
+    } else {
+        p.new_k_ptr = nullptr;
+        p.new_v_ptr = nullptr;
+        p.new_kv_b_stride = p.new_kv_h_stride = 0;
+    }
+
     p.causal_offset = (int)causal_offset;
     p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
     p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim);
@@ -279,6 +310,8 @@ inline void attn_pack_paged_prefill_params(
 
     p.k_ptr = (const T*)k_cache.data_ptr();
     p.v_ptr = (const T*)v_cache.data_ptr();
+    p.new_k_ptr = nullptr;
+    p.new_v_ptr = nullptr;
     p.q_ptr = (const T*)q.data_ptr();
     p.req_to_token = req_to_token.data_ptr<int>();
     p.req_pool_indices = req_pool_indices.data_ptr<int>();

csrc/kernels/attn_layout_policies.cuh

@@ -160,6 +160,14 @@ struct ContigKV {
             + (int64_t)d * p.kv_d_stride;
         return {&p.k_ptr[gmem_off], &p.v_ptr[gmem_off], valid};
     }
+
+    template <int VEC>
+    DEVICE_FORCEINLINE KVAddr decode_addr(
+        const AttentionParams<bf16>& p, const KVContext& c,
+        int, int, int kc, int d, bool valid, bool) {
+        int token = resolve_token(p, c, kc, valid);
+        return kv_addr_from_token(p, c, token, d);
+    }
 };
 
 // ---- Paged (SGLang-style flat pool) K/V ----
@@ -209,4 +217,39 @@ struct PagedKV {
         const int64_t gmem_off = (int64_t)safe_slot * c.pool_stride + c.head_off + d;
         return {&p.k_ptr[gmem_off], &p.v_ptr[gmem_off], valid};
     }
+
+    DEVICE_FORCEINLINE KVAddr new_kv_addr(
+        const AttentionParams<bf16>& p, int batch, int kv_head, int d) {
+        const int64_t off = (int64_t)batch * p.new_kv_b_stride
+            + (int64_t)kv_head * p.new_kv_h_stride + d;
+        return {&p.new_k_ptr[off], &p.new_v_ptr[off], true};
+    }
+
+    DEVICE_FORCEINLINE void store_new_kv(
+        const AttentionParams<bf16>& p, const KVContext& c,
+        int seq_len, int d, const KVAddr& src) {
+        int slot = resolve_token(p, c, seq_len - 1, true);
+        const int64_t off = (int64_t)slot * c.pool_stride + c.head_off + d;
+        const_cast<bf16*>(p.k_ptr)[off] = *reinterpret_cast<const bf16*>(src.k);
+        const_cast<bf16*>(p.v_ptr)[off] = *reinterpret_cast<const bf16*>(src.v);
+    }
+
+    template <int VEC>
+    DEVICE_FORCEINLINE KVAddr decode_addr(
+        const AttentionParams<bf16>& p, const KVContext& c,
+        int batch, int kv_head, int kc, int d, bool valid, bool persist) {
+        if (p.new_k_ptr && valid && kc == kv_len(p, batch) - 1) {
+            KVAddr src = new_kv_addr(p, batch, kv_head, d);
+            if (persist) {
+                #pragma unroll
+                for (int j = 0; j < VEC; j++) {
+                    KVAddr value = new_kv_addr(p, batch, kv_head, d + j);
+                    store_new_kv(p, c, kc + 1, d + j, value);
+                }
+            }
+            return src;
+        }
+        int token = resolve_token(p, c, kc, valid);
+        return kv_addr_from_token(p, c, token, d);
+    }
 };

csrc/kernels/attn_paged_decode.cu

@@ -8,6 +8,8 @@ torch::Tensor attn_paged_decode(
     torch::Tensor req_to_token,
     torch::Tensor req_pool_indices,
     torch::Tensor kv_indptr,
+    c10::optional<torch::Tensor> new_k,
+    c10::optional<torch::Tensor> new_v,
     c10::optional<torch::Tensor> mask,
     int64_t causal_offset,
     double scale,
@@ -21,6 +23,7 @@ torch::Tensor attn_paged_decode(
     AttentionParams<bf16> p;
     attn_pack_paged_decode_params(q, k_cache, v_cache,
                                    req_to_token, req_pool_indices, kv_indptr,
+                                   new_k, new_v,
                                    mask, causal_offset, scale, p);
 
     torch::Tensor O;
@@ -71,6 +74,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         py::arg("req_to_token"),
         py::arg("req_pool_indices"),
         py::arg("kv_indptr"),
+        py::arg("new_k") = py::none(),
+        py::arg("new_v") = py::none(),
         py::arg("mask") = py::none(),
         py::arg("causal_offset") = -1,
         py::arg("scale") = 0.0,

tests/extension/test_backend_equivalence.py

@@ -7,6 +7,7 @@ seq_lens with padding mask), and end-to-end scheduler.run_batch.
 import torch
 
 from astrai.extension import ATTN_BACKEND, attn_backend
+from astrai.extension.ops.attention import attn_paged_decode
 from astrai.inference.cache import PagePool, TaskCacheManager
 from astrai.inference.runtime.graph import CudaGraphContext
 from astrai.inference.scheduler import InferenceScheduler
@@ -160,6 +161,56 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
     assert diff < 0.05, f"Decode diff (mixed seq_lens): {diff}"
 
 
+@skip_no_kernel
+def test_paged_decode_appends_new_kv_in_kernel():
+    """Fused decode writes current-token K/V to each request's paged slot."""
+    pool = PagePool(
+        n_layers=1,
+        n_kv_heads=1,
+        head_dim=D,
+        max_batch_size=2,
+        max_seq_len=64,
+        device="cuda",
+        dtype=torch.bfloat16,
+        page_size=8,
+        n_tokens=128,
+    )
+    task_cache = _mk_task_cache(pool)
+    ws = _ws(pool)
+    task_cache.task_alloc("t1", list(range(8)))
+    task_cache.task_alloc("t2", list(range(6)))
+    task_cache.task_extend("t1", 8)
+    task_cache.task_extend("t2", 6)
+    kv_cache = task_cache.bind(["t1", "t2"], ws)
+
+    q = torch.randn(2, 2, D, device="cuda", dtype=torch.bfloat16)
+    new_k = torch.randn(2, 1, D, device="cuda", dtype=torch.bfloat16)
+    new_v = torch.randn(2, 1, D, device="cuda", dtype=torch.bfloat16)
+    out = attn_paged_decode(
+        q,
+        kv_cache.k_buffer[0],
+        kv_cache.v_buffer[0],
+        kv_cache.req_to_token,
+        kv_cache.req_pool_indices,
+        kv_cache.kv_indptr,
+        new_k=new_k,
+        new_v=new_v,
+        is_causal=True,
+        o_part_buf=kv_cache.decode_o_part,
+        ml_part_buf=kv_cache.decode_ml_part,
+        out_buf=kv_cache.decode_out,
+    )
+    torch.cuda.synchronize()
+
+    torch.testing.assert_close(
+        kv_cache.k_buffer[0, kv_cache.out_cache_loc], new_k, rtol=0, atol=0
+    )
+    torch.testing.assert_close(
+        kv_cache.v_buffer[0, kv_cache.out_cache_loc], new_v, rtol=0, atol=0
+    )
+    assert torch.isfinite(out).all()
+
+
 @skip_no_kernel
 def test_decode_cuda_graph_replay_is_exact(cuda_model):
     """INT32 cache indices must remain graph-capturable and replay exactly."""