perf: extend MMA decode to arbitrary GQA ratio, add launch bounds, vectorize combine

- Multi-pass MMA: encode pass in grid blockIdx.x, compute q_head0/G in-kernel
- Fixes crash for G>32 (previously block(32,G) exceeded 1024 threads)
- Fixes alloc_split_partials using uninitialized num_splits (MAX_SPLITS=32)
- __launch_bounds__ on all MMA and prefill kernels for better register allocation
- 4x vectorized combine kernel (4 head_dim per thread)
- uint4 vectorized K loads in scalar decode kernels
- cp.async .L2::128B cache hint for K/V tile streaming
- Extract warp_reduce_sum, bf16, MAX_SPLITS to attn_warp_utils.cuh

csrc/kernels/attn_decode_split_kv.cuh

@@ -2,16 +2,9 @@
 #include <cuda_bf16.h>
 #include <float.h>
 #include "attn_common.h"
-
+#include "attn_warp_utils.cuh"
-using bf16 = __nv_bfloat16;
 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 <int HEAD_DIM, bool IsCausal, bool HasMask>
 __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
     int batch = blockIdx.x / p.kv_head;
@@ -93,7 +86,7 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> 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<bf16> 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;
 

csrc/kernels/attn_decode_split_kv_mma.cuh

@@ -3,6 +3,7 @@
 #include <cuda_bf16.h>
 #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<bf16> 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<bf16> 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++) {

csrc/kernels/attn_dispatchers.cuh

@@ -4,6 +4,7 @@
 
 #include <cuda_runtime.h>
 #include <algorithm>
+#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<bf16>& p) {
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
 static inline void launch_decode_mma(AttentionParams<bf16>& p, int group_size) {
     int G = p.q_head / p.kv_head;
-    if (G >= 1 && G <= 16) {
+    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<HEAD_DIM, 32, 1, STAGES>;
-        dim3 grid(p.kv_head, p.batch, p.num_splits);
+    dim3 grid(p.kv_head * num_passes, p.batch, p.num_splits);
     attn_decode_split_kv_mma_kernel<Traits, IsCausal, HasMask><<<grid, 32>>>(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<HEAD_DIM, IsCausal, HasMask><<<grid, block, smem>>>(p);
-    }
 }
 #endif
 
@@ -100,8 +94,9 @@ static inline void launch_decode_scalar(AttentionParams<bf16>& 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<HEAD_DIM, IsCausal, HasMask><<<grid, block, smem>>>(p);
 }
 
@@ -140,13 +135,16 @@ static inline void dispatch_decode(AttentionParams<bf16>& p) {
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
 static inline void launch_paged_decode_mma(PagedAttentionParams<bf16>& 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<HEAD_DIM, 32, 1, STAGES>;
-        dim3 grid(p.kv_head, p.batch, p.num_splits);
+        dim3 grid(p.kv_head * num_passes, p.batch, p.num_splits);
-        paged_attn_decode_split_kv_mma_kernel<Traits, IsCausal, HasMask><<<grid, 32>>>(p);
+        paged_attn_decode_split_kv_mma_kernel<Traits, IsCausal, HasMask> <<<grid, 32>>>(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<bf16>& 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<HEAD_DIM, IsCausal, HasMask><<<grid, block, smem>>>(p);
 }
 

csrc/kernels/attn_entry_utils.cuh

@@ -2,6 +2,7 @@
 #include <torch/extension.h>
 #include <c10/cuda/CUDAGuard.h>
 #include "attn_common.h"
+#include "attn_warp_utils.cuh"
 
 using bf16 = __nv_bfloat16;
 
@@ -22,8 +23,8 @@ using bf16 = __nv_bfloat16;
 template<typename P>
 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 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, p.num_splits, 2}, 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();
 }

csrc/kernels/attn_paged_decode_split_kv.cuh

@@ -2,16 +2,9 @@
 #include <cuda_bf16.h>
 #include <float.h>
 #include "attn_common.h"
-
+#include "attn_warp_utils.cuh"
-using bf16 = __nv_bfloat16;
 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 <int HEAD_DIM, bool IsCausal, bool HasMask>
 __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> p) {
     int batch = blockIdx.x / p.kv_head;
@@ -71,7 +64,7 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> 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<bf16> 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<bf16> 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;
 

csrc/kernels/attn_paged_decode_split_kv_mma.cuh

@@ -3,6 +3,7 @@
 #include <cuda_bf16.h>
 #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<bf16>
     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<bf16>
 
     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++) {

csrc/kernels/attn_warp_utils.cuh

@@ -0,0 +1,13 @@
+#pragma once
+#include <cuda_bf16.h>
+
+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;
+}