perf: add split-K (FlashDecoding) to decode MMA kernel
Decode has only batch*kv_head independent tasks, so the grid was tiny (e.g. 16 blocks) leaving most SMs idle (ncu: 0.04 waves/SM, 11% DRAM). - Partition KV across gridDim.z blocks emitting unnormalised (O, m, l) partials, reduced by a new combine kernel - Choose split count to fill the device (~2 blocks/SM), capped by tile count and 32; fall back to single-pass direct-write when batch*kv_head already saturates the SMs - Refactor decode dispatch into named helpers, de-duplicate scalar fallback Result: now DRAM-bound at 63% (99->543 GB/s), 2.1-2.5x over torch SDPA in the low-parallelism regime, on par at high parallelism
This commit is contained in:
parent
dea59f7e1d
commit
7ba43a7c6f
|
|
@ -5,37 +5,73 @@
|
|||
#include "gqa_decode_attn_mma.cuh"
|
||||
#endif
|
||||
|
||||
// Scalar fallback: one warp per query head, per (batch, kv_head) block.
|
||||
static void launch_scalar_decode(const GQAParams& p) {
|
||||
int group_size = p.q_head / p.kv_head;
|
||||
size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
|
||||
gqa_decode_attn_kernel<<<p.batch * p.kv_head, dim3(32, group_size), smem>>>(p);
|
||||
}
|
||||
|
||||
#ifndef ASTRAI_NO_MMA
|
||||
// Tensor-core head-packing requires 1 < G <= 16 (the MMA M dim) and no mask.
|
||||
static bool decode_use_mma(const GQAParams& p) {
|
||||
int G = p.q_head / p.kv_head;
|
||||
return !p.use_mask && G > 1 && G <= 16;
|
||||
}
|
||||
|
||||
// Decode has only batch*kv_head independent tasks; without split-K the grid is
|
||||
// tiny (e.g. 16 blocks) and leaves most SMs idle. Pick the smallest split count
|
||||
// that fills the device (~2 blocks/SM), capped by the tile count and 32.
|
||||
static int decode_num_splits(const GQAParams& p, int tiles_total) {
|
||||
int sm_count = 0;
|
||||
cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
|
||||
int base_blocks = p.kv_head * p.batch;
|
||||
int desired = 2 * (sm_count > 0 ? sm_count : 64);
|
||||
int n = (desired + base_blocks - 1) / base_blocks;
|
||||
return std::max(1, std::min(n, std::min(tiles_total, 32)));
|
||||
}
|
||||
|
||||
template <int HEAD_DIM, int BC>
|
||||
static void launch_mma_decode(GQAParams& p) {
|
||||
constexpr int BR = 16, LD = HEAD_DIM; // XOR swizzle → no padding
|
||||
int smem = (2 * BC * LD + BR * LD) * (int)sizeof(bf16);
|
||||
int tiles_total = (p.kv_len + BC - 1) / BC;
|
||||
int num_splits = decode_num_splits(p, tiles_total);
|
||||
|
||||
// Enough (batch, kv_head) work to fill the SMs → single pass, direct write.
|
||||
if (num_splits <= 1) {
|
||||
cudaFuncSetAttribute(gqa_decode_attn_mma_kernel<HEAD_DIM, BC>,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
||||
gqa_decode_attn_mma_kernel<HEAD_DIM, BC>
|
||||
<<<dim3(p.kv_head, p.batch), 32, smem>>>(p);
|
||||
return;
|
||||
}
|
||||
|
||||
// Split-K (FlashDecoding): partition kv across blocks, then reduce.
|
||||
auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
|
||||
auto o_part = torch::empty({p.batch, p.q_head, num_splits, p.head_dim}, fopt);
|
||||
auto ml_part = torch::empty({p.batch, p.q_head, num_splits, 2}, fopt);
|
||||
|
||||
cudaFuncSetAttribute(gqa_decode_attn_mma_splitk_kernel<HEAD_DIM, BC>,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
||||
gqa_decode_attn_mma_splitk_kernel<HEAD_DIM, BC>
|
||||
<<<dim3(p.kv_head, p.batch, num_splits), 32, smem>>>(
|
||||
p, o_part.data_ptr<float>(), ml_part.data_ptr<float>(), num_splits);
|
||||
gqa_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(
|
||||
o_part.data_ptr<float>(), ml_part.data_ptr<float>(), p.o,
|
||||
num_splits, p.head_dim);
|
||||
}
|
||||
#endif
|
||||
|
||||
template <int HEAD_DIM>
|
||||
static void dispatch_decode(GQAParams& p) {
|
||||
#ifndef ASTRAI_NO_MMA
|
||||
constexpr int BC = 32, BR = 16, LD = HEAD_DIM; // XOR swizzle → no padding
|
||||
int G = p.q_head / p.kv_head;
|
||||
// head-packing tensor-core path needs 1 < G <= 16 (MMA M dim) and no mask;
|
||||
// everything else uses the scalar kernel
|
||||
if (!p.use_mask && G > 1 && G <= 16) {
|
||||
dim3 grid(p.kv_head, p.batch, 1);
|
||||
dim3 block(32, 1, 1);
|
||||
// sK + sV + sQ, each BC/BR * LD (single buffer for high occupancy)
|
||||
int smem = (2 * BC * LD + BR * LD) * (int)sizeof(bf16);
|
||||
cudaFuncSetAttribute(gqa_decode_attn_mma_kernel<HEAD_DIM, BC>,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
|
||||
gqa_decode_attn_mma_kernel<HEAD_DIM, BC><<<grid, block, smem>>>(p);
|
||||
if (decode_use_mma(p)) {
|
||||
launch_mma_decode<HEAD_DIM, 32>(p);
|
||||
return;
|
||||
}
|
||||
// scalar fallback (per-KV-head, one warp per query head)
|
||||
int group_size = p.q_head / p.kv_head;
|
||||
size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
|
||||
dim3 block(32, group_size);
|
||||
dim3 grid(p.batch * p.kv_head);
|
||||
gqa_decode_attn_kernel<<<grid, block, smem>>>(p);
|
||||
#else
|
||||
// scalar fallback (per-KV-head, one warp per query head)
|
||||
int group_size = p.q_head / p.kv_head;
|
||||
size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
|
||||
dim3 block(32, group_size);
|
||||
dim3 grid(p.batch * p.kv_head);
|
||||
gqa_decode_attn_kernel<<<grid, block, smem>>>(p);
|
||||
#endif
|
||||
launch_scalar_decode(p);
|
||||
}
|
||||
|
||||
torch::Tensor gqa_decode_attn(
|
||||
|
|
|
|||
|
|
@ -217,3 +217,253 @@ __global__ void gqa_decode_attn_mma_kernel(GQAParams p) {
|
|||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Split-K (FlashDecoding) decode: identical math to the kernel above, but the
|
||||
// KV sequence is partitioned across gridDim.z blocks so that a decode with only
|
||||
// batch*kv_head independent tasks can fill all SMs. Each (batch, kv_head, split)
|
||||
// block computes an UN-normalised partial (Oacc, m, l) over its KV slice; the
|
||||
// combine kernel below reduces across splits. Fixes the "grid too small"
|
||||
// bottleneck (0.04 waves/SM → many blocks) for long-context, small-batch decode.
|
||||
//
|
||||
// Partial layout (float, contiguous):
|
||||
// o_part : [batch, q_head, num_splits, HEAD_DIM]
|
||||
// ml_part: [batch, q_head, num_splits, 2] (m, l)
|
||||
template <int HEAD_DIM, int BC>
|
||||
__global__ void gqa_decode_attn_mma_splitk_kernel(GQAParams p,
|
||||
float* __restrict__ o_part,
|
||||
float* __restrict__ ml_part,
|
||||
int num_splits) {
|
||||
constexpr int BR = 16;
|
||||
constexpr int KD = HEAD_DIM / 16;
|
||||
constexpr int NC8 = BC / 8;
|
||||
constexpr int KT2 = BC / 16;
|
||||
constexpr int DN8 = HEAD_DIM / 8;
|
||||
constexpr int LD = HEAD_DIM;
|
||||
constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1);
|
||||
|
||||
const int lane = threadIdx.x;
|
||||
const int gid = lane >> 2;
|
||||
const int tid4 = lane & 3;
|
||||
|
||||
const int kv_head = blockIdx.x;
|
||||
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;
|
||||
|
||||
extern __shared__ __align__(16) bf16 smem[];
|
||||
bf16* sK = smem;
|
||||
bf16* sV = sK + BC * LD;
|
||||
bf16* sQ = sV + BC * LD;
|
||||
|
||||
bf16 scale_bf16 = __float2bfloat16(p.scale);
|
||||
for (int i = lane; i < BR * HEAD_DIM; i += 32) {
|
||||
int r = i / HEAD_DIM, d = i % HEAD_DIM;
|
||||
bf16 val = __float2bfloat16(0.0f);
|
||||
if (r < G) {
|
||||
int qh = q_head0 + r;
|
||||
val = p.q[(batch * p.q_head + qh) * HEAD_DIM + d];
|
||||
}
|
||||
sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = __hmul(val, scale_bf16);
|
||||
}
|
||||
__syncwarp();
|
||||
|
||||
unsigned Qa[KD][4];
|
||||
int qrow_l = (lane & 7) + (lane & 8);
|
||||
int qcol_l = (lane & 16) ? 8 : 0;
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < KD; kt++)
|
||||
ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + swiz_col(kt * 16 + qcol_l, qrow_l, SWIZ_MASK)]);
|
||||
|
||||
float Oacc[DN8][4];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < DN8; j++)
|
||||
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 kv_base = (batch * p.kv_head + kv_head) * p.kv_len * HEAD_DIM;
|
||||
const int mask_base = batch * p.kv_len;
|
||||
const int tiles_total = (p.kv_len + BC - 1) / BC;
|
||||
const int tiles_per_split = (tiles_total + num_splits - 1) / num_splits;
|
||||
const int ti_begin = split * tiles_per_split;
|
||||
const int ti_end = min(tiles_total, ti_begin + tiles_per_split);
|
||||
const int has_mask = p.use_mask && p.mask;
|
||||
|
||||
for (int ti = ti_begin; ti < ti_end; ti++) {
|
||||
int kv0 = ti * BC;
|
||||
|
||||
bool full_tile = (kv0 + BC <= p.kv_len);
|
||||
if (full_tile) {
|
||||
constexpr int VEC = 8;
|
||||
int total = BC * HEAD_DIM;
|
||||
#pragma unroll
|
||||
for (int i = lane * VEC; i < total; i += 32 * VEC) {
|
||||
int r = i / HEAD_DIM, d = i % HEAD_DIM;
|
||||
int kc = kv0 + r;
|
||||
cp_async_16(&sK[r * LD + swiz_col(d, r, SWIZ_MASK)],
|
||||
&p.k[kv_base + kc * HEAD_DIM + d]);
|
||||
cp_async_16(&sV[r * LD + swiz_col(d, r, SWIZ_MASK)],
|
||||
&p.v[kv_base + kc * HEAD_DIM + d]);
|
||||
}
|
||||
cp_async_commit();
|
||||
cp_async_wait_all();
|
||||
} else {
|
||||
for (int i = lane; i < BC * HEAD_DIM; i += 32) {
|
||||
int r = i / HEAD_DIM, d = i % HEAD_DIM;
|
||||
int kc = kv0 + r;
|
||||
bf16 z = __float2bfloat16(0.0f);
|
||||
sK[r * LD + swiz_col(d, r, SWIZ_MASK)] =
|
||||
(kc < p.kv_len) ? p.k[kv_base + kc * HEAD_DIM + d] : z;
|
||||
sV[r * LD + swiz_col(d, r, SWIZ_MASK)] =
|
||||
(kc < p.kv_len) ? p.v[kv_base + kc * HEAD_DIM + d] : z;
|
||||
}
|
||||
}
|
||||
__syncwarp();
|
||||
|
||||
float Sacc[NC8][4];
|
||||
#pragma unroll
|
||||
for (int n8 = 0; n8 < NC8; n8++) {
|
||||
Sacc[n8][0] = Sacc[n8][1] = Sacc[n8][2] = Sacc[n8][3] = 0.0f;
|
||||
int krow_l = n8 * 8 + (lane & 7);
|
||||
int kcol_h = (lane & 8) ? 8 : 0;
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < KD; kt++) {
|
||||
unsigned b[2];
|
||||
ldmatrix_x2(b, &sK[krow_l * LD + swiz_col(kt * 16 + kcol_h, krow_l, SWIZ_MASK)]);
|
||||
mma16816(Sacc[n8], Qa[kt], b, Sacc[n8]);
|
||||
}
|
||||
}
|
||||
|
||||
float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX;
|
||||
#pragma unroll
|
||||
for (int n8 = 0; n8 < NC8; n8++) {
|
||||
int cc = kv0 + n8 * 8 + 2 * tid4;
|
||||
bool bc0 = (cc >= p.kv_len) || (has_mask && !p.mask[mask_base + cc]);
|
||||
bool bc1 = (cc + 1 >= p.kv_len) || (has_mask && !p.mask[mask_base + cc + 1]);
|
||||
bool cz = p.is_causal;
|
||||
int off = p.causal_offset;
|
||||
bool bad0 = bc0 || (cz && cc > off);
|
||||
bool bad1 = bc1 || (cz && (cc + 1) > off);
|
||||
float s0 = bad0 ? -FLT_MAX : Sacc[n8][0];
|
||||
float s1 = bad1 ? -FLT_MAX : Sacc[n8][1];
|
||||
float s2 = bad0 ? -FLT_MAX : Sacc[n8][2];
|
||||
float s3 = bad1 ? -FLT_MAX : Sacc[n8][3];
|
||||
Sacc[n8][0] = s0; Sacc[n8][1] = s1; Sacc[n8][2] = s2; Sacc[n8][3] = s3;
|
||||
rmax0 = fmaxf(rmax0, fmaxf(s0, s1));
|
||||
rmax1 = fmaxf(rmax1, fmaxf(s2, s3));
|
||||
}
|
||||
rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 1));
|
||||
rmax0 = fmaxf(rmax0, __shfl_xor_sync(0xFFFFFFFF, rmax0, 2));
|
||||
rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 1));
|
||||
rmax1 = fmaxf(rmax1, __shfl_xor_sync(0xFFFFFFFF, rmax1, 2));
|
||||
|
||||
float nm0 = fmaxf(m0, rmax0), nm1 = fmaxf(m1, rmax1);
|
||||
float corr0 = (nm0 == -FLT_MAX) ? 1.0f : __expf(m0 - nm0);
|
||||
float corr1 = (nm1 == -FLT_MAX) ? 1.0f : __expf(m1 - nm1);
|
||||
|
||||
float rsum0 = 0.0f, rsum1 = 0.0f;
|
||||
#pragma unroll
|
||||
for (int n8 = 0; n8 < NC8; n8++) {
|
||||
float p0 = (Sacc[n8][0] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][0] - nm0);
|
||||
float p1 = (Sacc[n8][1] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][1] - nm0);
|
||||
float p2 = (Sacc[n8][2] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][2] - nm1);
|
||||
float p3 = (Sacc[n8][3] == -FLT_MAX) ? 0.0f : __expf(Sacc[n8][3] - nm1);
|
||||
Sacc[n8][0] = p0; Sacc[n8][1] = p1; Sacc[n8][2] = p2; Sacc[n8][3] = p3;
|
||||
rsum0 += p0 + p1;
|
||||
rsum1 += p2 + p3;
|
||||
}
|
||||
rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 1);
|
||||
rsum0 += __shfl_xor_sync(0xFFFFFFFF, rsum0, 2);
|
||||
rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 1);
|
||||
rsum1 += __shfl_xor_sync(0xFFFFFFFF, rsum1, 2);
|
||||
l0 = l0 * corr0 + rsum0;
|
||||
l1 = l1 * corr1 + rsum1;
|
||||
m0 = nm0; m1 = nm1;
|
||||
|
||||
#pragma unroll
|
||||
for (int j = 0; j < DN8; j++) {
|
||||
Oacc[j][0] *= corr0; Oacc[j][1] *= corr0;
|
||||
Oacc[j][2] *= corr1; Oacc[j][3] *= corr1;
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int kt2 = 0; kt2 < KT2; kt2++) {
|
||||
unsigned Pa[4];
|
||||
Pa[0] = pk2(Sacc[kt2 * 2][0], Sacc[kt2 * 2][1]);
|
||||
Pa[1] = pk2(Sacc[kt2 * 2][2], Sacc[kt2 * 2][3]);
|
||||
Pa[2] = pk2(Sacc[kt2 * 2 + 1][0], Sacc[kt2 * 2 + 1][1]);
|
||||
Pa[3] = pk2(Sacc[kt2 * 2 + 1][2], Sacc[kt2 * 2 + 1][3]);
|
||||
int vrow_l = kt2 * 16 + (lane & 15);
|
||||
#pragma unroll
|
||||
for (int dn8 = 0; dn8 < DN8; dn8++) {
|
||||
unsigned b[2];
|
||||
ldmatrix_x2_trans(b, &sV[vrow_l * LD + swiz_col(dn8 * 8, vrow_l, SWIZ_MASK)]);
|
||||
mma16816(Oacc[dn8], Pa, b, Oacc[dn8]);
|
||||
}
|
||||
}
|
||||
__syncwarp();
|
||||
}
|
||||
|
||||
// ---- write UN-normalised partials for this split ----
|
||||
#pragma unroll
|
||||
for (int dn8 = 0; dn8 < DN8; dn8++) {
|
||||
int d = dn8 * 8 + 2 * tid4;
|
||||
int r0 = gid, r1 = gid + 8;
|
||||
if (r0 < G) {
|
||||
int hh = q_head0 + r0;
|
||||
float* op = o_part + ((size_t)(batch * p.q_head + hh) * num_splits + split) * HEAD_DIM;
|
||||
op[d] = Oacc[dn8][0];
|
||||
op[d + 1] = Oacc[dn8][1];
|
||||
}
|
||||
if (r1 < G) {
|
||||
int hh = q_head0 + r1;
|
||||
float* op = o_part + ((size_t)(batch * p.q_head + hh) * num_splits + split) * HEAD_DIM;
|
||||
op[d] = Oacc[dn8][2];
|
||||
op[d + 1] = Oacc[dn8][3];
|
||||
}
|
||||
}
|
||||
if (tid4 == 0) {
|
||||
int r0 = gid, r1 = gid + 8;
|
||||
if (r0 < G) {
|
||||
float* mp = ml_part + ((size_t)(batch * p.q_head + q_head0 + r0) * num_splits + split) * 2;
|
||||
mp[0] = m0; mp[1] = l0;
|
||||
}
|
||||
if (r1 < G) {
|
||||
float* mp = ml_part + ((size_t)(batch * p.q_head + q_head0 + r1) * num_splits + split) * 2;
|
||||
mp[0] = m1; mp[1] = l1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Reduce split-K partials into the final bf16 output. One block per (batch,
|
||||
// q_head); each thread owns one head_dim element and folds across all splits
|
||||
// with a numerically-stable online rescale.
|
||||
__global__ void gqa_decode_combine_kernel(const float* __restrict__ o_part,
|
||||
const float* __restrict__ ml_part,
|
||||
bf16* __restrict__ out,
|
||||
int num_splits, int head_dim) {
|
||||
int bh = blockIdx.x; // batch * q_head + head
|
||||
int d = threadIdx.x;
|
||||
if (d >= head_dim) return;
|
||||
|
||||
const float* mlp = ml_part + (size_t)bh * num_splits * 2;
|
||||
float mstar = -FLT_MAX;
|
||||
for (int s = 0; s < num_splits; s++)
|
||||
mstar = fmaxf(mstar, mlp[s * 2]);
|
||||
|
||||
float lstar = 0.0f;
|
||||
for (int s = 0; s < num_splits; s++) {
|
||||
float mi = mlp[s * 2];
|
||||
if (mi > -FLT_MAX) lstar += mlp[s * 2 + 1] * __expf(mi - mstar);
|
||||
}
|
||||
|
||||
const float* op = o_part + (size_t)bh * num_splits * head_dim;
|
||||
float acc = 0.0f;
|
||||
for (int s = 0; s < num_splits; s++) {
|
||||
float mi = mlp[s * 2];
|
||||
if (mi > -FLT_MAX) acc += op[s * head_dim + d] * __expf(mi - mstar);
|
||||
}
|
||||
float inv = (lstar > 1e-20f) ? (1.0f / lstar) : 0.0f;
|
||||
out[(size_t)bh * head_dim + d] = __float2bfloat16(acc * inv);
|
||||
}
|
||||
Loading…
Reference in New Issue