From 0654b4b916cb8e3d05dd964deab3255e0a752250 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Wed, 15 Jul 2026 21:43:29 +0800 Subject: [PATCH] refactor: template combine kernel, fix mask bug, unify dispatch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Template combine kernel, share macros, extract entry_utils helpers - Fix mask indexing (pass stride not pre-multiplied base) - Remove !p.use_mask — MMA handles mask --- csrc/kernels/attn_decode.cu | 4 +- csrc/kernels/attn_decode_split_kv_mma.cuh | 3 +- csrc/kernels/attn_entry_utils.cuh | 161 +++++++----------- csrc/kernels/attn_paged_decode.cu | 4 +- .../attn_paged_decode_split_kv_mma.cuh | 3 +- csrc/kernels/attn_prefill.cu | 2 +- csrc/kernels/attn_prefill_split_q_mma.cuh | 3 +- 7 files changed, 70 insertions(+), 110 deletions(-) diff --git a/csrc/kernels/attn_decode.cu b/csrc/kernels/attn_decode.cu index 8fd3fb0..9d7fcc8 100644 --- a/csrc/kernels/attn_decode.cu +++ b/csrc/kernels/attn_decode.cu @@ -38,7 +38,7 @@ template static void dispatch_decode(AttentionParams& p) { #ifndef ASTRAI_NO_MMA int G = p.q_head / p.kv_head; - if (!p.use_mask && G >= 1 && G <= 16) { + if (G >= 1 && G <= 16) { launch_mma_decode(p); return; } @@ -65,7 +65,7 @@ torch::Tensor attn_decode( auto O_view = (layout == 1) ? O.transpose(1, 2) : O; p.o = (bf16*)O_view.data_ptr(); - dispatch_head_dim(p.head_dim, [&]() { dispatch_decode(p); }); + DISPATCH_HEAD_DIM(p.head_dim, dispatch_decode, p); return O; } diff --git a/csrc/kernels/attn_decode_split_kv_mma.cuh b/csrc/kernels/attn_decode_split_kv_mma.cuh index 30be5bb..a78a3e1 100644 --- a/csrc/kernels/attn_decode_split_kv_mma.cuh +++ b/csrc/kernels/attn_decode_split_kv_mma.cuh @@ -64,7 +64,6 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams p) { // KV: stride-based base — [batch, kv_head, kv_len, head_dim] const int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h; - const int mask_batch_base = batch * p.mask_b_stride; const int tiles_total = (p.kv_len + BC - 1) / BC; const int tiles_per_split = (tiles_total + p.num_splits - 1) / p.num_splits; const int ti_begin = split * tiles_per_split; @@ -125,7 +124,7 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams p) { int maxc = (p.causal_offset >= 0) ? min(p.kv_len, p.causal_offset + 1) : p.kv_len; mma_softmax_tile(kv0, maxc, maxc, 0, 0, - mask_batch_base, 0, + p.mask_b_stride, 0, batch, p.mask, has_mask, Sacc, Oacc, m0, m1, l0, l1, lane); diff --git a/csrc/kernels/attn_entry_utils.cuh b/csrc/kernels/attn_entry_utils.cuh index edcc286..0833bfe 100644 --- a/csrc/kernels/attn_entry_utils.cuh +++ b/csrc/kernels/attn_entry_utils.cuh @@ -5,10 +5,6 @@ using bf16 = __nv_bfloat16; -// --------------------------------------------------------------------------- -// Shared dispatch helpers — eliminates duplication across .cu entry files. -// --------------------------------------------------------------------------- - inline int compute_num_splits(int base_blocks, int tiles_total) { int sm_count = 0; cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0); @@ -16,21 +12,20 @@ inline int compute_num_splits(int base_blocks, int tiles_total) { return std::max(1, std::min(n, std::min(tiles_total, 32))); } -// Dispatch head_dim to a generic lambda FN (callable as FN.operator()()). -template -inline void dispatch_head_dim(int hd, Fn&& fn) { - switch (hd) { - case 32: fn.template operator()<32>(); break; - case 64: fn.template operator()<64>(); break; - case 128: fn.template operator()<128>(); break; - case 256: fn.template operator()<256>(); break; - default: - TORCH_CHECK(false, "unsupported head_dim ", hd, - " (supported: 32, 64, 128, 256)"); +// Dispatch head_dim: shared macro — avoids C++20 lambda template syntax. +// Usage: DISPATCH_HEAD_DIM(hd, fn, arg) +// Expands to: fn<32>(arg); fn<64>(arg); etc. +#define DISPATCH_HEAD_DIM(hd, fn, arg) \ + switch (hd) { \ + case 32: fn<32>(arg); break; \ + case 64: fn<64>(arg); break; \ + case 128: fn<128>(arg); break; \ + case 256: fn<256>(arg); break; \ + default: \ + TORCH_CHECK(false, "unsupported head_dim ", hd, \ + " (supported: 32, 64, 128, 256)"); \ } -} -// Allocate split-KV partial buffers and wire into params. template inline void alloc_split_partials(P& p) { auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); @@ -40,10 +35,48 @@ inline void alloc_split_partials(P& p) { p.ml_part = (float*)ml_part.data_ptr(); } -// --------------------------------------------------------------------------- -// Param packing — fills AttentionParams from torch tensors with validation. -// --------------------------------------------------------------------------- +// ---- Shared Q-dims + strides extraction ---- +template +inline void extract_q_dims_and_strides(torch::Tensor& q, int64_t layout, P& p) { + if (layout == 1) q = q.transpose(1, 2); + p.batch = (int)q.size(0); + p.q_head = (int)q.size(1); + p.q_len = (int)q.size(2); + p.head_dim = (int)q.size(3); + p.q_stride_b = (int)q.stride(0); + p.q_stride_h = (int)q.stride(1); + p.q_stride_l = (int)q.stride(2); + p.q_stride_d = (int)q.stride(3); +} +// ---- Shared mask packing ---- +template +inline void pack_mask(const c10::optional& mask, P& p) { + if (p.use_mask) { + auto m = mask.value(); + TORCH_CHECK(m.is_cuda(), "mask must be on CUDA"); + TORCH_CHECK(m.dtype() == torch::kBool, "mask must be bool"); + TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch"); + TORCH_CHECK(m.size(m.dim() - 1) == p.kv_len, "mask kv_len mismatch"); + if (m.dim() == 2) { + p.mask_b_stride = (int)m.stride(0); + p.mask_q_stride = 0; + } else if (m.dim() == 3) { + TORCH_CHECK(m.size(1) == p.q_len, "mask q_len mismatch"); + p.mask_b_stride = (int)m.stride(0); + p.mask_q_stride = (int)m.stride(1); + } else { + TORCH_CHECK(false, "mask must be 2D [batch, kv_len] or 3D [batch, q_len, kv_len]"); + } + p.mask = m.data_ptr(); + } else { + p.mask = nullptr; + p.mask_b_stride = 0; + p.mask_q_stride = 0; + } +} + +// ---- attn_pack_params (contiguous KV) ---- template inline void attn_pack_params( torch::Tensor q, @@ -52,7 +85,7 @@ inline void attn_pack_params( c10::optional mask, int64_t causal_offset, double scale, - int64_t layout, // 0 = b h l d, 1 = b l h d + int64_t layout, AttentionParams& p ) { const at::cuda::OptionalCUDAGuard device_guard(device_of(q)); @@ -64,26 +97,14 @@ inline void attn_pack_params( TORCH_CHECK(k.sizes() == v.sizes(), "K and V must have identical shapes"); TORCH_CHECK(q.dim() == 4 && k.dim() == 4, "Q/K/V must be 4D"); - // Normalize to b h l d view (zero-copy transpose if user passed b l h d) - if (layout == 1) { - q = q.transpose(1, 2); - k = k.transpose(1, 2); - v = v.transpose(1, 2); - } + extract_q_dims_and_strides(q, layout, p); + + if (layout == 1) k = k.transpose(1, 2), v = v.transpose(1, 2); - p.batch = (int)q.size(0); - p.q_head = (int)q.size(1); - p.q_len = (int)q.size(2); - p.head_dim = (int)q.size(3); p.kv_head = (int)k.size(1); p.kv_len = (int)k.size(2); TORCH_CHECK(k.size(3) == p.head_dim, "K/V head_dim must match Q"); - // Strides (layout-agnostic: works for b h l d and b l h d) - p.q_stride_b = (int)q.stride(0); - p.q_stride_h = (int)q.stride(1); - p.q_stride_l = (int)q.stride(2); - p.q_stride_d = (int)q.stride(3); p.kv_stride_b = (int)k.stride(0); p.kv_stride_h = (int)k.stride(1); p.kv_stride_l = (int)k.stride(2); @@ -100,34 +121,10 @@ inline void attn_pack_params( p.o_part = nullptr; p.ml_part = nullptr; - if (p.use_mask) { - auto m = mask.value(); - TORCH_CHECK(m.is_cuda(), "mask must be on CUDA"); - TORCH_CHECK(m.dtype() == torch::kBool, "mask must be bool"); - TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch"); - TORCH_CHECK(m.size(m.dim() - 1) == p.kv_len, "mask kv_len mismatch"); - if (m.dim() == 2) { - p.mask_b_stride = (int)m.stride(0); - p.mask_q_stride = 0; - } else if (m.dim() == 3) { - TORCH_CHECK(m.size(1) == p.q_len, "mask q_len mismatch"); - p.mask_b_stride = (int)m.stride(0); - p.mask_q_stride = (int)m.stride(1); - } else { - TORCH_CHECK(false, "mask must be 2D [batch, kv_len] or 3D [batch, q_len, kv_len]"); - } - p.mask = m.data_ptr(); - } else { - p.mask = nullptr; - p.mask_b_stride = 0; - p.mask_q_stride = 0; - } + pack_mask(mask, p); } -// --------------------------------------------------------------------------- -// Param packing for paged attention. -// --------------------------------------------------------------------------- - +// ---- attn_pack_paged_params ---- template inline void attn_pack_paged_params( torch::Tensor q, @@ -139,7 +136,7 @@ inline void attn_pack_paged_params( c10::optional mask, int64_t causal_offset, double scale, - int64_t layout, // 0 = b h l d, 1 = b l h d + int64_t layout, PagedAttentionParams& p ) { const at::cuda::OptionalCUDAGuard device_guard(device_of(q)); @@ -151,15 +148,8 @@ inline void attn_pack_paged_params( TORCH_CHECK(page_table.dtype() == torch::kLong, "page_table must be int64"); TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must have identical shapes"); - // Normalize Q to b h l d view if user passed b l h d - if (layout == 1) { - q = q.transpose(1, 2); - } + extract_q_dims_and_strides(q, layout, p); - p.batch = (int)q.size(0); - p.q_head = (int)q.size(1); - p.q_len = (int)q.size(2); - p.head_dim = (int)q.size(3); p.kv_head = (int)k_cache.size(2); p.kv_len = (int)kv_len; p.page_size = (int)page_size; @@ -171,12 +161,6 @@ inline void attn_pack_paged_params( "k_cache dim 1 must equal page_size, got ", k_cache.size(1), " vs ", page_size); - // Q strides - p.q_stride_b = (int)q.stride(0); - p.q_stride_h = (int)q.stride(1); - p.q_stride_l = (int)q.stride(2); - p.q_stride_d = (int)q.stride(3); - 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); @@ -189,26 +173,5 @@ inline void attn_pack_paged_params( p.o_part = nullptr; p.ml_part = nullptr; - if (p.use_mask) { - auto m = mask.value(); - TORCH_CHECK(m.is_cuda(), "mask must be on CUDA"); - TORCH_CHECK(m.dtype() == torch::kBool, "mask must be bool"); - TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch"); - TORCH_CHECK(m.size(m.dim() - 1) == p.kv_len, "mask kv_len mismatch"); - if (m.dim() == 2) { - p.mask_b_stride = (int)m.stride(0); - p.mask_q_stride = 0; - } else if (m.dim() == 3) { - TORCH_CHECK(m.size(1) == p.q_len, "mask q_len mismatch"); - p.mask_b_stride = (int)m.stride(0); - p.mask_q_stride = (int)m.stride(1); - } else { - TORCH_CHECK(false, "mask must be 2D [batch, kv_len] or 3D [batch, q_len, kv_len]"); - } - p.mask = m.data_ptr(); - } else { - p.mask = nullptr; - p.mask_b_stride = 0; - p.mask_q_stride = 0; - } + pack_mask(mask, p); } diff --git a/csrc/kernels/attn_paged_decode.cu b/csrc/kernels/attn_paged_decode.cu index 0af4c8c..71ebe62 100644 --- a/csrc/kernels/attn_paged_decode.cu +++ b/csrc/kernels/attn_paged_decode.cu @@ -34,7 +34,7 @@ template static void dispatch_paged_decode(PagedAttentionParams& p) { #ifndef ASTRAI_NO_MMA int G = p.q_head / p.kv_head; - if (!p.use_mask && G >= 1 && G <= 16 && p.page_size >= 32) { + if (G >= 1 && G <= 16 && p.page_size >= 32) { launch_paged_mma_decode(p); return; } @@ -62,7 +62,7 @@ torch::Tensor attn_paged_decode( auto O_view = (layout == 1) ? O.transpose(1, 2) : O; p.o = (bf16*)O_view.data_ptr(); - dispatch_head_dim(p.head_dim, [&]() { dispatch_paged_decode(p); }); + DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_decode, p); return O; } diff --git a/csrc/kernels/attn_paged_decode_split_kv_mma.cuh b/csrc/kernels/attn_paged_decode_split_kv_mma.cuh index dca6b7e..99c0d21 100644 --- a/csrc/kernels/attn_paged_decode_split_kv_mma.cuh +++ b/csrc/kernels/attn_paged_decode_split_kv_mma.cuh @@ -51,7 +51,6 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f; float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f; - const int mask_batch_base = batch * p.mask_b_stride; const int tiles_total = (p.kv_len + BC - 1) / BC; const int tiles_per_split = (tiles_total + p.num_splits - 1) / p.num_splits; const int ti_begin = split * tiles_per_split; @@ -119,7 +118,7 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams int maxc = (p.causal_offset >= 0) ? min(p.kv_len, p.causal_offset + 1) : p.kv_len; mma_softmax_tile(kv0, maxc, maxc, 0, 0, - mask_batch_base, 0, + p.mask_b_stride, 0, batch, p.mask, has_mask, Sacc, Oacc, m0, m1, l0, l1, lane); diff --git a/csrc/kernels/attn_prefill.cu b/csrc/kernels/attn_prefill.cu index 7d621e6..d3c979f 100644 --- a/csrc/kernels/attn_prefill.cu +++ b/csrc/kernels/attn_prefill.cu @@ -47,7 +47,7 @@ torch::Tensor attn_prefill( auto O_view = (layout == 1) ? O.transpose(1, 2) : O; p.o = (bf16*)O_view.data_ptr(); - dispatch_head_dim(p.head_dim, [&]() { dispatch_prefill(p); }); + DISPATCH_HEAD_DIM(p.head_dim, dispatch_prefill, p); return O; } diff --git a/csrc/kernels/attn_prefill_split_q_mma.cuh b/csrc/kernels/attn_prefill_split_q_mma.cuh index 6374162..edf3890 100644 --- a/csrc/kernels/attn_prefill_split_q_mma.cuh +++ b/csrc/kernels/attn_prefill_split_q_mma.cuh @@ -94,7 +94,6 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams p) { const int block_max_kv = blockIdx.x * WARPS * BR + WARPS * BR - 1 + p.causal_offset; const int has_mask = p.use_mask && p.mask; - const int mask_batch_base = batch * p.mask_b_stride; // Last active tile: block-level causal bound (all warps in the block share // the K/V load, so the prefetch range is the block max, not per-warp). @@ -165,7 +164,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams p) { : p.kv_len; mma_softmax_tile(kv0, maxc0, maxc1, qr0, qr1, - mask_batch_base, p.mask_q_stride, + p.mask_b_stride, p.mask_q_stride, batch, p.mask, has_mask, Sacc, Oacc, m0, m1, l0, l1, lane);