diff --git a/csrc/kernels/gqa_decode_attn.cu b/csrc/kernels/gqa_decode_attn.cu index e9387e6..f7261d7 100644 --- a/csrc/kernels/gqa_decode_attn.cu +++ b/csrc/kernels/gqa_decode_attn.cu @@ -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); +} + +#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 +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, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + gqa_decode_attn_mma_kernel + <<>>(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, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + gqa_decode_attn_mma_splitk_kernel + <<>>( + p, o_part.data_ptr(), ml_part.data_ptr(), num_splits); + gqa_decode_combine_kernel<<>>( + o_part.data_ptr(), ml_part.data_ptr(), p.o, + num_splits, p.head_dim); +} +#endif + template 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, - cudaFuncAttributeMaxDynamicSharedMemorySize, smem); - gqa_decode_attn_mma_kernel<<>>(p); + if (decode_use_mma(p)) { + launch_mma_decode(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<<>>(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<<>>(p); #endif + launch_scalar_decode(p); } torch::Tensor gqa_decode_attn( diff --git a/csrc/kernels/gqa_decode_attn_mma.cuh b/csrc/kernels/gqa_decode_attn_mma.cuh index 7d73fe0..41a24b0 100644 --- a/csrc/kernels/gqa_decode_attn_mma.cuh +++ b/csrc/kernels/gqa_decode_attn_mma.cuh @@ -216,4 +216,254 @@ __global__ void gqa_decode_attn_mma_kernel(GQAParams p) { p.o[o_off + 1] = __float2bfloat16(Oacc[dn8][3] * rl1); } } +} + +// --------------------------------------------------------------------------- +// 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 +__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); } \ No newline at end of file