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
This commit is contained in:
@@ -2,16 +2,9 @@
|
||||
#include <cuda_bf16.h>
|
||||
#include <float.h>
|
||||
#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 <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;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user