From c5fba9c2384cb7d6c3cc137a483d9f12db9f0aa4 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sun, 9 Aug 2026 23:12:01 +0800 Subject: [PATCH] perf: flatten paged prefill tile dispatch - remove the host-provided max_q_len argument - dispatch only the ragged prefill tile upper bound - validate the rebuilt CUDA backend end to end --- astrai/extension/attention_backend.py | 1 - astrai/extension/attention_ops.py | 3 -- csrc/kernels/attn_common.h | 4 +-- csrc/kernels/attn_dispatchers.cuh | 10 +++---- csrc/kernels/attn_entry_utils.cuh | 6 ++-- csrc/kernels/attn_kv_source.cuh | 34 ++++++++++++++++++++--- csrc/kernels/attn_paged_prefill.cu | 4 +-- csrc/kernels/attn_prefill_split_q.cuh | 15 ++++++++-- csrc/kernels/attn_prefill_split_q_mma.cuh | 17 ++++++++++-- csrc/tests/attn_paged_test.cu | 8 ++---- 10 files changed, 69 insertions(+), 33 deletions(-) diff --git a/astrai/extension/attention_backend.py b/astrai/extension/attention_backend.py index 81ae586..8a11ca4 100644 --- a/astrai/extension/attention_backend.py +++ b/astrai/extension/attention_backend.py @@ -590,7 +590,6 @@ class CudaBackend(AttentionBackend): kv_indptr, qo_indptr, attn_mask, - q_len, is_causal=is_causal, ) return out.reshape(b, q_len, q.size(2), q.size(3)).flatten(2) diff --git a/astrai/extension/attention_ops.py b/astrai/extension/attention_ops.py index 5363e58..adca269 100644 --- a/astrai/extension/attention_ops.py +++ b/astrai/extension/attention_ops.py @@ -151,7 +151,6 @@ def attn_paged_prefill( kv_indptr: torch.Tensor, qo_indptr: torch.Tensor, mask: Optional[torch.Tensor] = None, - max_q_len: int = 0, is_causal: bool = False, ) -> torch.Tensor: """SGLang-style paged prefill (ragged batch, flat KV pool). @@ -169,7 +168,6 @@ def attn_paged_prefill( kv_indptr: [batch+1] (int32) — prefix sum of per-request kv_lens qo_indptr: [batch+1] (int32) — prefix sum of per-request q_lens mask: 4D [batch, 1, q_len, kv_len] (bool, True=keep) or None - max_q_len: max per-request q_len (Python int, for grid computation) is_causal: apply causal mask Returns: @@ -186,6 +184,5 @@ def attn_paged_prefill( kv_indptr, qo_indptr, mask, - max_q_len, causal_offset=causal_offset, ) diff --git a/csrc/kernels/attn_common.h b/csrc/kernels/attn_common.h index f88923a..1d2d78a 100644 --- a/csrc/kernels/attn_common.h +++ b/csrc/kernels/attn_common.h @@ -23,7 +23,7 @@ struct AttentionParams { int q_head; int kv_head; int head_dim; - int q_len; // Contiguous mode; paged mode uses qo_indptr. + int q_len; // Per-request in contiguous mode; total_q in paged mode. int kv_len; // Contiguous mode; paged mode uses kv_indptr. // Attention behavior @@ -66,6 +66,4 @@ struct AttentionParams { AT* __restrict__ o_part; AT* __restrict__ ml_part; - // Host-provided paged prefill grid bound - int max_q_len; }; diff --git a/csrc/kernels/attn_dispatchers.cuh b/csrc/kernels/attn_dispatchers.cuh index 1e120c2..4ab84a9 100644 --- a/csrc/kernels/attn_dispatchers.cuh +++ b/csrc/kernels/attn_dispatchers.cuh @@ -66,9 +66,9 @@ struct PrefillLauncherMMA { constexpr int WARPS = 4; constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16; using Traits = KernelTraits; - int q_len = KV::host_q_len(p); - dim3 grid((q_len + Traits::BR * WARPS - 1) / (Traits::BR * WARPS), - p.q_head, p.batch); + constexpr int ROWS = Traits::BR * WARPS; + dim3 grid(KV::host_q_blocks(p, ROWS), p.q_head, + KV::kPaged ? 1 : p.batch); dim3 block(Traits::NUM_THREADS); attn_prefill_split_q_mma_kernel <<>>(p); @@ -81,8 +81,8 @@ struct PrefillLauncherScalar { template static void launch(AttentionParams& p, cudaStream_t stream) { constexpr int G = (HEAD_DIM == 32) ? 4 : 8, ROWS = 32, P_BC = 32; - int q_len = KV::host_q_len(p); - dim3 grid((q_len + ROWS - 1) / ROWS, p.q_head, p.batch); + dim3 grid(KV::host_q_blocks(p, ROWS), p.q_head, + KV::kPaged ? 1 : p.batch); dim3 block(G, ROWS); attn_prefill_split_q_kernel_t <<>>(p); diff --git a/csrc/kernels/attn_entry_utils.cuh b/csrc/kernels/attn_entry_utils.cuh index 1b722f8..53d228a 100644 --- a/csrc/kernels/attn_entry_utils.cuh +++ b/csrc/kernels/attn_entry_utils.cuh @@ -189,7 +189,6 @@ inline void attn_pack_paged_decode_params( p.kv_indptr = kv_indptr.data_ptr(); p.qo_indptr = nullptr; p.max_context_len = (int)req_to_token.size(1); - p.max_q_len = 1; p.causal_offset = (int)causal_offset; p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0; @@ -228,7 +227,6 @@ inline void attn_pack_paged_prefill_params( torch::Tensor kv_indptr, torch::Tensor qo_indptr, c10::optional mask, - int64_t max_q_len, int64_t causal_offset, double scale, AttentionParams& p @@ -251,6 +249,7 @@ inline void attn_pack_paged_prefill_params( p.q_head = (int)q.size(1); p.head_dim = (int)q.size(2); + p.q_len = (int)q.size(0); p.kv_head = (int)k_cache.size(1); p.batch = (int)req_pool_indices.size(0); TORCH_CHECK(k_cache.size(2) == p.head_dim, "k_cache head_dim mismatch"); @@ -273,7 +272,6 @@ inline void attn_pack_paged_prefill_params( p.kv_indptr = kv_indptr.data_ptr(); p.qo_indptr = qo_indptr.data_ptr(); p.max_context_len = (int)req_to_token.size(1); - p.max_q_len = (int)max_q_len; p.causal_offset = (int)causal_offset; p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0; @@ -288,7 +286,7 @@ inline void attn_pack_paged_prefill_params( p.mask_l_stride = 0; } else if (m.dim() == 4) { TORCH_CHECK(m.size(1) == 1 || m.size(1) == p.q_head, "mask head mismatch"); - TORCH_CHECK(m.size(2) == 1 || m.size(2) == p.max_q_len, "mask q_len mismatch"); + TORCH_CHECK(m.size(2) > 0 && m.size(2) <= p.q_len, "mask q_len mismatch"); TORCH_CHECK(m.size(3) <= p.max_context_len, "mask kv_len mismatch"); p.mask_b_stride = (int)m.stride(0); p.mask_h_stride = (m.size(1) == 1) ? 0 : (int)m.stride(1); diff --git a/csrc/kernels/attn_kv_source.cuh b/csrc/kernels/attn_kv_source.cuh index e820e56..168faed 100644 --- a/csrc/kernels/attn_kv_source.cuh +++ b/csrc/kernels/attn_kv_source.cuh @@ -57,8 +57,16 @@ struct ContigKV { static constexpr bool kPaged = false; // host-side length hooks (grid + split computation in the launchers) - HOST_DEV_FORCEINLINE int host_q_len(const AttentionParams& p) { - return p.q_len; + HOST_DEV_FORCEINLINE int host_q_blocks(const AttentionParams& p, int rows) { + return (p.q_len + rows - 1) / rows; + } + template + HOST_DEV_FORCEINLINE bool map_q_tile(const AttentionParams&, + int flat_tile, int grid_batch, + int& batch, int& q_tile) { + batch = grid_batch; + q_tile = flat_tile; + return true; } HOST_DEV_FORCEINLINE int host_kv_len(const AttentionParams& p) { return p.kv_len; @@ -107,8 +115,26 @@ struct ContigKV { struct PagedKV { static constexpr bool kPaged = true; - HOST_DEV_FORCEINLINE int host_q_len(const AttentionParams& p) { - return p.max_q_len; + HOST_DEV_FORCEINLINE int host_q_blocks(const AttentionParams& p, int rows) { + // sum(ceil(q_len[b] / rows)) <= ceil(total_q / rows) + batch - 1. + return (p.q_len + rows - 1) / rows + p.batch - 1; + } + template + HOST_DEV_FORCEINLINE bool map_q_tile(const AttentionParams& p, + int flat_tile, int, + int& batch, int& q_tile) { + int tile_base = 0; + for (int b = 0; b < p.batch; ++b) { + int len = p.qo_indptr[b + 1] - p.qo_indptr[b]; + int tiles = (len + ROWS - 1) / ROWS; + if (flat_tile < tile_base + tiles) { + batch = b; + q_tile = flat_tile - tile_base; + return true; + } + tile_base += tiles; + } + return false; } HOST_DEV_FORCEINLINE int host_kv_len(const AttentionParams& p) { return p.max_context_len; diff --git a/csrc/kernels/attn_paged_prefill.cu b/csrc/kernels/attn_paged_prefill.cu index a490d22..99d40d7 100644 --- a/csrc/kernels/attn_paged_prefill.cu +++ b/csrc/kernels/attn_paged_prefill.cu @@ -10,7 +10,6 @@ torch::Tensor attn_paged_prefill( torch::Tensor kv_indptr, torch::Tensor qo_indptr, c10::optional mask, - int64_t max_q_len, int64_t causal_offset, double scale ) { @@ -21,7 +20,7 @@ torch::Tensor attn_paged_prefill( attn_pack_paged_prefill_params(q, k_cache, v_cache, req_to_token, req_pool_indices, kv_indptr, qo_indptr, mask, - max_q_len, causal_offset, scale, p); + causal_offset, scale, p); auto O = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options()); p.o_ptr = (bf16*)O.data_ptr(); @@ -41,7 +40,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("kv_indptr"), py::arg("qo_indptr"), py::arg("mask") = py::none(), - py::arg("max_q_len"), py::arg("causal_offset") = -1, py::arg("scale") = 0.0, "SGLang-style paged prefill: flat KV pool + ragged batch."); diff --git a/csrc/kernels/attn_prefill_split_q.cuh b/csrc/kernels/attn_prefill_split_q.cuh index b0850f0..945b6b6 100644 --- a/csrc/kernels/attn_prefill_split_q.cuh +++ b/csrc/kernels/attn_prefill_split_q.cuh @@ -36,9 +36,20 @@ template p) { constexpr int DPT = HEAD_DIM / G; - int q_tile = blockIdx.x; + __shared__ int mapped_batch; + __shared__ int mapped_q_tile; + if (threadIdx.x == 0 && threadIdx.y == 0) { + mapped_batch = -1; + KV::template map_q_tile( + p, blockIdx.x, blockIdx.z, mapped_batch, mapped_q_tile); + } + __syncthreads(); + if (mapped_batch < 0) + return; + + int q_tile = mapped_q_tile; int q_head = blockIdx.y; - int batch = blockIdx.z; + int batch = mapped_batch; int gpos = threadIdx.x; // 0..G-1 (which d-chunk) int row = threadIdx.y; // 0..ROWS-1 int q_row = q_tile * ROWS + row; diff --git a/csrc/kernels/attn_prefill_split_q_mma.cuh b/csrc/kernels/attn_prefill_split_q_mma.cuh index 7bb2a1e..f9140da 100644 --- a/csrc/kernels/attn_prefill_split_q_mma.cuh +++ b/csrc/kernels/attn_prefill_split_q_mma.cuh @@ -24,9 +24,20 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams p) { const int tid4 = lane & 3; // 0..3 const int q_head = blockIdx.y; - const int batch = blockIdx.z; + __shared__ int mapped_batch; + __shared__ int mapped_q_tile; + if (threadIdx.x == 0) { + mapped_batch = -1; + KV::template map_q_tile( + p, blockIdx.x, blockIdx.z, mapped_batch, mapped_q_tile); + } + __syncthreads(); + if (mapped_batch < 0) + return; + const int batch = mapped_batch; + const int q_tile = mapped_q_tile; const int kv_head = q_head / (p.q_head / p.kv_head); - const int qrow0 = (blockIdx.x * Traits::WARPS + warp) * Traits::BR; + const int qrow0 = (q_tile * Traits::WARPS + warp) * Traits::BR; // Per-request dims (from KV policy — paged reads kv_indptr/qo_indptr). const int seq_len = KV::kv_len(p, batch); @@ -61,7 +72,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams p) { // Causal tile-skip bounds (dead code when IsCausal == false) const int max_kv = qrow0 + Traits::BR - 1 + causal_off; const int block_max_kv = - blockIdx.x * Traits::WARPS * Traits::BR + Traits::WARPS * Traits::BR - 1 + q_tile * Traits::WARPS * Traits::BR + Traits::WARPS * Traits::BR - 1 + causal_off; int t_end = tiles - 1; diff --git a/csrc/tests/attn_paged_test.cu b/csrc/tests/attn_paged_test.cu index bc120df..d6438f1 100644 --- a/csrc/tests/attn_paged_test.cu +++ b/csrc/tests/attn_paged_test.cu @@ -489,9 +489,7 @@ static int run_prefill_test(int B, int Hq, int Hkv, p.head_dim = HEAD_DIM; p.q_l_stride = Hq * HEAD_DIM; p.q_h_stride = HEAD_DIM; p.q_d_stride = 1; p.max_context_len = max_ctx; - int max_ql = 0; - for (int b = 0; b < B; b++) max_ql = max(max_ql, q_lens[b]); - p.max_q_len = max_ql; + p.q_len = total_q; p.causal_offset = causal ? 0 : -1; p.use_mask = 0; p.mask = nullptr; p.mask_b_stride = 0; p.mask_h_stride = 0; p.mask_l_stride = 0; @@ -626,7 +624,7 @@ static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) { p.head_dim = HEAD_DIM; p.q_l_stride = Hq * HEAD_DIM; p.q_h_stride = HEAD_DIM; p.q_d_stride = 1; p.max_context_len = max_ctx; - p.max_q_len = q_len; + p.q_len = B * q_len; p.causal_offset = -1; p.use_mask = 1; p.mask = d_mask; p.mask_b_stride = q_len * q_len; p.mask_h_stride = 0; p.mask_l_stride = q_len; @@ -793,7 +791,7 @@ static void bench_prefill(int B, int Hq, int Hkv, int q_len, int kv_len, int cau p.head_dim = HEAD_DIM; p.q_l_stride = Hq * HEAD_DIM; p.q_h_stride = HEAD_DIM; p.q_d_stride = 1; p.max_context_len = max_ctx; - p.max_q_len = q_len; + p.q_len = B * q_len; p.causal_offset = causal ? 0 : -1; p.use_mask = 0; p.mask = nullptr; p.mask_b_stride = 0; p.scale = 1.0f / sqrtf((float)HEAD_DIM);