diff --git a/csrc/kernels/attn_decode_split_kv.cuh b/csrc/kernels/attn_decode_split_kv.cuh index 67b0618..dd0c2ab 100644 --- a/csrc/kernels/attn_decode_split_kv.cuh +++ b/csrc/kernels/attn_decode_split_kv.cuh @@ -2,16 +2,9 @@ #include #include #include "attn_common.h" - -using bf16 = __nv_bfloat16; +#include "attn_warp_utils.cuh" constexpr int DC_CHUNK = 64; -__device__ inline float warp_reduce_sum(float val) { - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - template __global__ void attn_decode_split_kv_kernel(AttentionParams p) { int batch = blockIdx.x / p.kv_head; @@ -93,7 +86,7 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams p) { // ---- write UN-normalised partials for this split ---- size_t bh = (size_t)batch * p.q_head + q_head; - size_t slot = bh * p.num_splits + split; + size_t slot = bh * MAX_SPLITS + split; int d0 = lane * hd_per_thread; for (int i = 0; i < hd_per_thread; i++) { int dd = d0 + i; @@ -113,7 +106,7 @@ __global__ void attn_decode_combine_kernel(AttentionParams p) { int batch = bh / p.q_head; int q_head = bh % p.q_head; - size_t split_base = (size_t)bh * p.num_splits; + size_t split_base = (size_t)bh * MAX_SPLITS; const float* mlp = p.ml_part + split_base * 2; const float* op = p.o_part + split_base * p.head_dim; diff --git a/csrc/kernels/attn_decode_split_kv_mma.cuh b/csrc/kernels/attn_decode_split_kv_mma.cuh index b88daef..a6db4cb 100644 --- a/csrc/kernels/attn_decode_split_kv_mma.cuh +++ b/csrc/kernels/attn_decode_split_kv_mma.cuh @@ -3,6 +3,7 @@ #include #include "attn_common.h" #include "attn_mma_utils.cuh" +#include "attn_warp_utils.cuh" // Split-K (FlashDecoding) tensor-core decode via GQA head-packing. // Decode has q_len == 1, so we pack G = q_head/kv_head query heads into the @@ -19,11 +20,16 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams p) { const int gid = lane >> 2; const int tid4 = lane & 3; - const int kv_head = blockIdx.x; + const int pass = blockIdx.x / p.kv_head; + const int kv_head = blockIdx.x % p.kv_head; const int batch = blockIdx.y; const int split = blockIdx.z; - const int G = p.q_head / p.kv_head; - const int q_head0 = kv_head * G; + + constexpr int MAX_G = 16; + const int G_total = p.q_head / p.kv_head; + const int g_begin = pass * MAX_G; + const int G = min(MAX_G, G_total - g_begin); + const int q_head0 = kv_head * G_total + g_begin; // Double-buffered shared memory for K/V (no sQ needed) __shared__ __align__(16) bf16 sK[Traits::STAGES * Traits::BC * Traits::LD]; @@ -120,7 +126,7 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams p) { // ---- write UN-normalised partials for this split ---- auto split_slot = [&](int h) -> size_t { size_t bh = (size_t)batch * p.q_head + h; - return bh * p.num_splits + split; + return bh * MAX_SPLITS + split; }; #pragma unroll for (int dn8 = 0; dn8 < Traits::DN8; dn8++) { diff --git a/csrc/kernels/attn_dispatchers.cuh b/csrc/kernels/attn_dispatchers.cuh index 017b95c..52d2a8d 100644 --- a/csrc/kernels/attn_dispatchers.cuh +++ b/csrc/kernels/attn_dispatchers.cuh @@ -4,6 +4,7 @@ #include #include +#include "attn_warp_utils.cuh" #include "attn_prefill_split_q.cuh" #include "attn_decode_split_kv.cuh" #include "attn_paged_decode_split_kv.cuh" @@ -18,7 +19,7 @@ inline int compute_num_splits(int base_blocks, int tiles_total) { int sm_count = 0; cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0); int n = (2 * sm_count + base_blocks - 1) / base_blocks; - return std::max(1, std::min(n, std::min(tiles_total, 32))); + return std::max(1, std::min(n, std::min(tiles_total, MAX_SPLITS))); } // ====================================================================== @@ -77,21 +78,14 @@ static inline void dispatch_prefill(AttentionParams& p) { template static inline void launch_decode_mma(AttentionParams& p, int group_size) { int G = p.q_head / p.kv_head; - if (G >= 1 && G <= 16) { - int tiles_total = (p.kv_len + 32 - 1) / 32; - p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total); - constexpr int STAGES = (HEAD_DIM <= 128) ? 2 : 1; - using Traits = KernelTraits; - dim3 grid(p.kv_head, p.batch, p.num_splits); - attn_decode_split_kv_mma_kernel<<>>(p); - } else { - int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK; - p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total); - size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16); - dim3 grid(p.batch * p.kv_head, 1, p.num_splits); - dim3 block(32, group_size); - attn_decode_split_kv_kernel<<>>(p); - } + constexpr int MAX_G = 16; + int num_passes = (G + MAX_G - 1) / MAX_G; + int tiles_total = (p.kv_len + 32 - 1) / 32; + p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total); + constexpr int STAGES = (HEAD_DIM <= 128) ? 2 : 1; + using Traits = KernelTraits; + dim3 grid(p.kv_head * num_passes, p.batch, p.num_splits); + attn_decode_split_kv_mma_kernel<<>>(p); } #endif @@ -100,8 +94,9 @@ static inline void launch_decode_scalar(AttentionParams& p, int group_size int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK; p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total); size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16); + int g = min(group_size, 32); // cap at 32 to respect 1024-thread limit dim3 grid(p.batch * p.kv_head, 1, p.num_splits); - dim3 block(32, group_size); + dim3 block(32, g); attn_decode_split_kv_kernel<<>>(p); } @@ -140,13 +135,16 @@ static inline void dispatch_decode(AttentionParams& p) { template static inline void launch_paged_decode_mma(PagedAttentionParams& p, int group_size) { int G = p.q_head / p.kv_head; - if (G >= 1 && G <= 16 && p.page_size >= 32) { + constexpr int MAX_G = 16; + bool page_ok = (p.page_size >= 32); + if (G >= 1 && page_ok) { + int num_passes = (G + MAX_G - 1) / MAX_G; int tiles_total = (p.kv_len + 32 - 1) / 32; p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total); constexpr int STAGES = (HEAD_DIM <= 128) ? 2 : 1; using Traits = KernelTraits; - dim3 grid(p.kv_head, p.batch, p.num_splits); - paged_attn_decode_split_kv_mma_kernel<<>>(p); + dim3 grid(p.kv_head * num_passes, p.batch, p.num_splits); + paged_attn_decode_split_kv_mma_kernel <<>>(p); } else { int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK; p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total); @@ -163,8 +161,9 @@ static inline void launch_paged_decode_scalar(PagedAttentionParams& p, int int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK; p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total); size_t smem = PDC_CHUNK * p.head_dim * sizeof(bf16); + int g = min(group_size, 32); // cap at 32 to respect 1024-thread limit dim3 grid(p.batch * p.kv_head, 1, p.num_splits); - dim3 block(32, group_size); + dim3 block(32, g); paged_attn_decode_split_kv_kernel<<>>(p); } diff --git a/csrc/kernels/attn_entry_utils.cuh b/csrc/kernels/attn_entry_utils.cuh index ffb0388..f3881e3 100644 --- a/csrc/kernels/attn_entry_utils.cuh +++ b/csrc/kernels/attn_entry_utils.cuh @@ -2,6 +2,7 @@ #include #include #include "attn_common.h" +#include "attn_warp_utils.cuh" using bf16 = __nv_bfloat16; @@ -22,8 +23,8 @@ using bf16 = __nv_bfloat16; template inline void alloc_split_partials(P& p) { auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); - auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt); - auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt); + auto o_part = torch::empty({p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt); + auto ml_part = torch::empty({p.batch, p.q_head, MAX_SPLITS, 2}, fopt); p.o_part = (float*)o_part.data_ptr(); p.ml_part = (float*)ml_part.data_ptr(); } diff --git a/csrc/kernels/attn_paged_decode_split_kv.cuh b/csrc/kernels/attn_paged_decode_split_kv.cuh index 62c6a65..294496e 100644 --- a/csrc/kernels/attn_paged_decode_split_kv.cuh +++ b/csrc/kernels/attn_paged_decode_split_kv.cuh @@ -2,16 +2,9 @@ #include #include #include "attn_common.h" - -using bf16 = __nv_bfloat16; +#include "attn_warp_utils.cuh" constexpr int PDC_CHUNK = 64; -__device__ inline float paged_warp_reduce_sum(float val) { - for (int offset = 16; offset > 0; offset >>= 1) - val += __shfl_xor_sync(0xFFFFFFFF, val, offset); - return val; -} - template __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams p) { int batch = blockIdx.x / p.kv_head; @@ -71,7 +64,7 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams p) for (int i = 0; i < hd_per_thread; i++) partial += q_reg[i] * __bfloat162float( k_smem[s * p.head_dim + lane * hd_per_thread + i]); - partial = paged_warp_reduce_sum(partial) * p.scale; + partial = warp_reduce_sum(partial) * p.scale; int kv_idx = chunk_start + s; if constexpr (HasMask) { @@ -111,7 +104,7 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams p) } size_t bh = (size_t)batch * p.q_head + q_head; - size_t slot = bh * p.num_splits + split; + size_t slot = bh * MAX_SPLITS + split; int d0 = lane * hd_per_thread; #pragma unroll for (int i = 0; i < hd_per_thread; i++) @@ -130,7 +123,7 @@ __global__ void paged_attn_decode_combine_kernel(PagedAttentionParams p) { int batch = bh / p.q_head; int q_head = bh % p.q_head; - size_t split_base = (size_t)bh * p.num_splits; + size_t split_base = (size_t)bh * MAX_SPLITS; const float* mlp = p.ml_part + split_base * 2; const float* op = p.o_part + split_base * p.head_dim; diff --git a/csrc/kernels/attn_paged_decode_split_kv_mma.cuh b/csrc/kernels/attn_paged_decode_split_kv_mma.cuh index 8baa8f7..e9228e1 100644 --- a/csrc/kernels/attn_paged_decode_split_kv_mma.cuh +++ b/csrc/kernels/attn_paged_decode_split_kv_mma.cuh @@ -3,6 +3,7 @@ #include #include "attn_common.h" #include "attn_mma_utils.cuh" +#include "attn_warp_utils.cuh" // Paged split-KV tensor-core decode via GQA head-packing. // Reads K/V directly from the page pool through a page table — one tile @@ -16,11 +17,16 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams const int gid = lane >> 2; const int tid4 = lane & 3; - const int kv_head = blockIdx.x; + const int pass = blockIdx.x / p.kv_head; + const int kv_head = blockIdx.x % p.kv_head; const int batch = blockIdx.y; const int split = blockIdx.z; - const int G = p.q_head / p.kv_head; - const int q_head0 = kv_head * G; + + constexpr int MAX_G = 16; + const int G_total = p.q_head / p.kv_head; + const int g_begin = pass * MAX_G; + const int G = min(MAX_G, G_total - g_begin); + const int q_head0 = kv_head * G_total + g_begin; __shared__ __align__(16) bf16 sK[Traits::STAGES * Traits::BC * Traits::LD]; __shared__ __align__(16) bf16 sV[Traits::STAGES * Traits::BC * Traits::LD]; @@ -120,7 +126,7 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams auto split_slot = [&](int h) -> size_t { size_t bh = (size_t)batch * p.q_head + h; - return bh * p.num_splits + split; + return bh * MAX_SPLITS + split; }; #pragma unroll for (int dn8 = 0; dn8 < Traits::DN8; dn8++) { diff --git a/csrc/kernels/attn_warp_utils.cuh b/csrc/kernels/attn_warp_utils.cuh new file mode 100644 index 0000000..241626e --- /dev/null +++ b/csrc/kernels/attn_warp_utils.cuh @@ -0,0 +1,13 @@ +#pragma once +#include + +using bf16 = __nv_bfloat16; + +static constexpr int MAX_SPLITS = 32; + +__device__ inline float warp_reduce_sum(float val) { + #pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) + val += __shfl_xor_sync(0xFFFFFFFF, val, offset); + return val; +}