From 3d3ea47d373a42fefe1d3b7f04de47f8b9cabc80 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Wed, 19 Aug 2026 00:30:14 +0800 Subject: [PATCH] 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 | 28 +---------- astrai/extension/ops/attention.py | 6 +++ csrc/kernels/attn_common.h | 5 ++ csrc/kernels/attn_decode_split_kv.cuh | 4 +- csrc/kernels/attn_decode_split_kv_mma.cuh | 6 ++- csrc/kernels/attn_entry_utils.cuh | 33 +++++++++++++ csrc/kernels/attn_layout_policies.cuh | 43 +++++++++++++++++ csrc/kernels/attn_paged_decode.cu | 5 ++ tests/extension/test_backend_equivalence.py | 51 +++++++++++++++++++++ 9 files changed, 151 insertions(+), 30 deletions(-) diff --git a/astrai/extension/backend/attention.py b/astrai/extension/backend/attention.py index bba16d3..ee1a8f8 100644 --- a/astrai/extension/backend/attention.py +++ b/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, diff --git a/astrai/extension/ops/attention.py b/astrai/extension/ops/attention.py index b3345e7..4dccecd 100644 --- a/astrai/extension/ops/attention.py +++ b/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, diff --git a/csrc/kernels/attn_common.h b/csrc/kernels/attn_common.h index 1929679..92f66fe 100644 --- a/csrc/kernels/attn_common.h +++ b/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; diff --git a/csrc/kernels/attn_decode_split_kv.cuh b/csrc/kernels/attn_decode_split_kv.cuh index a840b6e..337f5f7 100644 --- a/csrc/kernels/attn_decode_split_kv.cuh +++ b/csrc/kernels/attn_decode_split_kv.cuh @@ -57,8 +57,8 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams 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::kv_addr_from_token(p, kctx, token, d_dim); + KVAddr a = KV::template decode_addr<1>( + p, kctx, batch, kv_head, kc, d_dim, true, true); k_smem[i] = a.valid ? *reinterpret_cast(a.k) : (bf16)0.f; v_smem[i] = a.valid ? *reinterpret_cast(a.v) : (bf16)0.f; } diff --git a/csrc/kernels/attn_decode_split_kv_mma.cuh b/csrc/kernels/attn_decode_split_kv_mma.cuh index 7718fa0..6a66da6 100644 --- a/csrc/kernels/attn_decode_split_kv_mma.cuh +++ b/csrc/kernels/attn_decode_split_kv_mma.cuh @@ -73,8 +73,10 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams 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); - KVAddr a = KV::kv_addr_from_token(p, kctx, token, d); + // All GQA passes consume new K/V directly. Only the first pass + // persists it, so no cross-block synchronization is required. + KVAddr a = KV::template decode_addr( + 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); diff --git a/csrc/kernels/attn_entry_utils.cuh b/csrc/kernels/attn_entry_utils.cuh index 5ec9efc..309db1b 100644 --- a/csrc/kernels/attn_entry_utils.cuh +++ b/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& new_k, + const c10::optional& new_v, c10::optional 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(); p.req_pool_indices = req_pool_indices.data_ptr(); diff --git a/csrc/kernels/attn_layout_policies.cuh b/csrc/kernels/attn_layout_policies.cuh index 99646e2..84ebaa3 100644 --- a/csrc/kernels/attn_layout_policies.cuh +++ b/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 + DEVICE_FORCEINLINE KVAddr decode_addr( + const AttentionParams& 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& 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& 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(p.k_ptr)[off] = *reinterpret_cast(src.k); + const_cast(p.v_ptr)[off] = *reinterpret_cast(src.v); + } + + template + DEVICE_FORCEINLINE KVAddr decode_addr( + const AttentionParams& 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); + } }; diff --git a/csrc/kernels/attn_paged_decode.cu b/csrc/kernels/attn_paged_decode.cu index 0cf2c65..4cca860 100644 --- a/csrc/kernels/attn_paged_decode.cu +++ b/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 new_k, + c10::optional new_v, c10::optional mask, int64_t causal_offset, double scale, @@ -21,6 +23,7 @@ torch::Tensor attn_paged_decode( AttentionParams 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, diff --git a/tests/extension/test_backend_equivalence.py b/tests/extension/test_backend_equivalence.py index 2ec240e..baeede1 100644 --- a/tests/extension/test_backend_equivalence.py +++ b/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."""