From e9b03f4fcae85a7eca64aa2fa220bfe639e32472 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Wed, 8 Jul 2026 16:03:40 +0800 Subject: [PATCH] perf: apply cp.async, XOR swizzle, pre-scaled Q to decode MMA kernel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Decode MMA kernel previously used scalar global→shared loads with LD=HEAD_DIM+8 padding and per-tile scale multiply. This commit brings it in line with the prefill MMA kernel (which already had these optimizations): - cp.async K/V loads (bypasses registers, halves load instructions) - XOR swizzle: LD=HEAD_DIM instead of HEAD_DIM+8 (zero waste smem) - Pre-scale Q during load (removes per-tile scale multiply in softmax) - Clean up prefill MMA kernel comments (no code change) ~2x speedup on decode (0.47ms→0.24ms at seq_len=512) --- csrc/kernels/gqa_decode_attn.cu | 5 +- csrc/kernels/gqa_decode_attn_mma.cuh | 74 +++++++++++++++++++-------- csrc/kernels/gqa_prefill_attn.cu | 2 +- csrc/kernels/gqa_prefill_attn_mma.cuh | 16 +++--- 4 files changed, 62 insertions(+), 35 deletions(-) diff --git a/csrc/kernels/gqa_decode_attn.cu b/csrc/kernels/gqa_decode_attn.cu index 72286a1..e9387e6 100644 --- a/csrc/kernels/gqa_decode_attn.cu +++ b/csrc/kernels/gqa_decode_attn.cu @@ -8,14 +8,15 @@ template static void dispatch_decode(GQAParams& p) { #ifndef ASTRAI_NO_MMA - constexpr int BC = 32, LD = HEAD_DIM + 8; + 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); - int smem = (2 * BC * LD + 16 * LD) * (int)sizeof(bf16); + // 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); diff --git a/csrc/kernels/gqa_decode_attn_mma.cuh b/csrc/kernels/gqa_decode_attn_mma.cuh index 8ecdc9b..7d73fe0 100644 --- a/csrc/kernels/gqa_decode_attn_mma.cuh +++ b/csrc/kernels/gqa_decode_attn_mma.cuh @@ -2,7 +2,7 @@ #include "gqa_common.cuh" #include "gqa_mma_utils.cuh" -// Tensor-core decode via GQA head-packing. +// Tensor-core decode via GQA head-packing with cp.async loads. // // Decode has q_len == 1, so S = q @ K^T is a GEMV per head — no tensor-core work // on its own. But GQA gives us G = q_head / kv_head query heads that all share @@ -14,6 +14,12 @@ // instead of different sequence positions of one head, and (2) causal masking is // a single scalar bound shared by every row. One warp owns one (batch, kv_head); // requires G <= 16. +// +// Optimizations: +// - cp.async global→shared for K/V (bypasses registers, cuts instruction count) +// - XOR swizzle (swiz_col): LD=HEAD_DIM, zero waste, no bank conflicts +// - pre-scaled Q: Q scaled during load, softmax skips per-tile multiply +// - single-buffer: keeps smem small for high occupancy template __global__ void gqa_decode_attn_mma_kernel(GQAParams p) { @@ -22,7 +28,8 @@ __global__ void gqa_decode_attn_mma_kernel(GQAParams p) { constexpr int NC8 = BC / 8; // S n-tiles (N=8 each) constexpr int KT2 = BC / 16; // P k-tiles (K=16 each) constexpr int DN8 = HEAD_DIM / 8; // O n-tiles (N=8 each) - constexpr int LD = HEAD_DIM + 8; // padded smem row stride + constexpr int LD = HEAD_DIM; // XOR swizzle handles bank conflicts, zero waste + constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1); const int lane = threadIdx.x; // single warp const int gid = lane >> 2; // 0..7 → rows gid, gid+8 @@ -38,7 +45,8 @@ __global__ void gqa_decode_attn_mma_kernel(GQAParams p) { bf16* sV = sK + BC * LD; // [BC][LD] bf16* sQ = sV + BC * LD; // [BR][LD] - // stage Q: row r == query head (q_head0 + r) at position 0 (zero-padded past G) + // ---- stage Q into shared (pre-scaled, swizzled) ---- + 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); @@ -46,7 +54,7 @@ __global__ void gqa_decode_attn_mma_kernel(GQAParams p) { int qh = q_head0 + r; val = p.q[(batch * p.q_head + qh) * HEAD_DIM + d]; // q_len == 1 } - sQ[r * LD + d] = val; + sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = __hmul(val, scale_bf16); } __syncwarp(); @@ -56,7 +64,7 @@ __global__ void gqa_decode_attn_mma_kernel(GQAParams p) { int qcol_l = (lane & 16) ? 8 : 0; #pragma unroll for (int kt = 0; kt < KD; kt++) - ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + kt * 16 + qcol_l]); + ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + swiz_col(kt * 16 + qcol_l, qrow_l, SWIZ_MASK)]); float Oacc[DN8][4]; #pragma unroll @@ -65,21 +73,43 @@ __global__ void gqa_decode_attn_mma_kernel(GQAParams p) { 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 = (p.kv_len + BC - 1) / BC; + const int has_mask = p.use_mask && p.mask; for (int ti = 0; ti < tiles; ti++) { int kv0 = ti * BC; - 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 + d] = (kc < p.kv_len) ? p.k[kv_base + kc * HEAD_DIM + d] : z; - sV[r * LD + d] = (kc < p.kv_len) ? p.v[kv_base + kc * HEAD_DIM + d] : z; + // ---- load K/V tile to shared (cp.async on full tiles) ---- + bool full_tile = (kv0 + BC <= p.kv_len); + if (full_tile) { + constexpr int VEC = 8; // 8 bf16 = 16 bytes per cp.async + 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(); - // S = Q @ K^T + // S = Q @ K^T (Q already pre-scaled, so Sacc includes scale) float Sacc[NC8][4]; #pragma unroll for (int n8 = 0; n8 < NC8; n8++) { @@ -89,28 +119,28 @@ __global__ void gqa_decode_attn_mma_kernel(GQAParams p) { #pragma unroll for (int kt = 0; kt < KD; kt++) { unsigned b[2]; - ldmatrix_x2(b, &sK[krow_l * LD + kt * 16 + kcol_h]); + 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]); } } - // ---- online softmax (single query position → one causal bound per col) ---- + // ---- online softmax (Q pre-scaled → no per-tile scale multiply) ---- 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) || - (p.use_mask && p.mask && !p.mask[batch * p.kv_len + cc]); + (has_mask && !p.mask[mask_base + cc]); bool bc1 = (cc + 1 >= p.kv_len) || - (p.use_mask && p.mask && !p.mask[batch * p.kv_len + cc + 1]); + (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] * p.scale; - float s1 = bad1 ? -FLT_MAX : Sacc[n8][1] * p.scale; - float s2 = bad0 ? -FLT_MAX : Sacc[n8][2] * p.scale; - float s3 = bad1 ? -FLT_MAX : Sacc[n8][3] * p.scale; + 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)); @@ -161,7 +191,7 @@ __global__ void gqa_decode_attn_mma_kernel(GQAParams p) { #pragma unroll for (int dn8 = 0; dn8 < DN8; dn8++) { unsigned b[2]; - ldmatrix_x2_trans(b, &sV[vrow_l * LD + dn8 * 8]); + ldmatrix_x2_trans(b, &sV[vrow_l * LD + swiz_col(dn8 * 8, vrow_l, SWIZ_MASK)]); mma16816(Oacc[dn8], Pa, b, Oacc[dn8]); } } @@ -186,4 +216,4 @@ __global__ void gqa_decode_attn_mma_kernel(GQAParams p) { p.o[o_off + 1] = __float2bfloat16(Oacc[dn8][3] * rl1); } } -} +} \ No newline at end of file diff --git a/csrc/kernels/gqa_prefill_attn.cu b/csrc/kernels/gqa_prefill_attn.cu index 0936573..5aa1f9c 100644 --- a/csrc/kernels/gqa_prefill_attn.cu +++ b/csrc/kernels/gqa_prefill_attn.cu @@ -11,7 +11,7 @@ static void dispatch_prefill(GQAParams& p) { constexpr int WARPS = 4, BC = 32, BR = 16, LD = HEAD_DIM; dim3 grid((p.q_len + BR * WARPS - 1) / (BR * WARPS), p.q_head, p.batch); dim3 block(WARPS * 32, 1, 1); - // shared sQ: single staging area (BR*LD), not per-warp + // sK + sV (each BC*LD) + shared sQ staging (BR*LD) int smem = (2 * BC * LD + BR * LD) * (int)sizeof(bf16); cudaFuncSetAttribute(gqa_prefill_attn_mma_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); diff --git a/csrc/kernels/gqa_prefill_attn_mma.cuh b/csrc/kernels/gqa_prefill_attn_mma.cuh index ccced71..474fe9a 100644 --- a/csrc/kernels/gqa_prefill_attn_mma.cuh +++ b/csrc/kernels/gqa_prefill_attn_mma.cuh @@ -12,15 +12,11 @@ // repack; row reductions fold across the 4-lane thread group. Templated on // with BC a multiple of 16. // -// Optimizations vs v6 baseline: shared sQ staging (single area, serialized -// per-warp load) → cuts smem from (2*BC + WARPS*BR)*LD to (2*BC + BR)*LD bf16, -// raising occupancy; pre-scale Q by attention scale during Q load → removes -// per-tile scale multiply in the softmax loop; cp.async global→shared for K/V -// full-tile loads → eliminates shared-store bank conflicts and register staging, -// scalar fallback only for the last partial tile; causal tile skipping (block- -// level early break + warp-level skip); XOR swizzle (swiz_col) at 8-bf16 chunk -// granularity → eliminates ldmatrix bank conflicts without LD padding, setting -// LD=HEAD_DIM (zero waste). +// Optimizations: shared sQ staging (single area, serialized per-warp load) +// → cuts smem; pre-scale Q by attention scale during Q load; cp.async global→ +// shared for K/V; scalar fallback only for the last partial tile; causal tile +// skipping (block-level early break + warp-level skip); XOR swizzle (swiz_col) +// → eliminates ldmatrix bank conflicts without LD padding (LD=HEAD_DIM). template __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) { @@ -247,4 +243,4 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) { __float2bfloat16(Oacc[dn8][3] * rl1); } } -} +} \ No newline at end of file