From ddc4bd1cf6e18684e58e76d9380d7556ddc57468 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Mon, 6 Jul 2026 19:51:55 +0800 Subject: [PATCH] feat: tensor-core mma prefill with build-time dispatch - add register-resident flash-attention kernel using mma.sync.m16n8k16 - dispatch mma vs scalar at build time: pre-sm_80 defines -DASTRAI_NO_MMA, else defaults to mma - scalar path vectorized with float4 smem loads (ld8) --- csrc/build.py | 16 +- csrc/kernels/gqa_prefill_attn.cu | 38 +++-- csrc/kernels/gqa_prefill_attn.cuh | 31 +++- csrc/kernels/gqa_prefill_attn_mma.cuh | 223 ++++++++++++++++++++++++++ csrc/tests/gqa_prefill_test.cu | 4 +- 5 files changed, 291 insertions(+), 21 deletions(-) create mode 100644 csrc/kernels/gqa_prefill_attn_mma.cuh diff --git a/csrc/build.py b/csrc/build.py index 7e86212..c33f955 100644 --- a/csrc/build.py +++ b/csrc/build.py @@ -1,14 +1,20 @@ from pathlib import Path -def _arch_flag(): +def _arch_flags() -> list[str]: import torch if torch.cuda.is_available(): cap = torch.cuda.get_device_capability() - ver = f"{cap[0]}{cap[1]}" - return f"-gencode=arch=compute_{ver},code=sm_{ver}" - return "-gencode=arch=compute_80,code=sm_80" + else: + cap = (8, 0) + ver = f"{cap[0]}{cap[1]}" + flags = [f"-gencode=arch=compute_{ver},code=sm_{ver}"] + # tensor-core mma path (mma.sync.m16n8k16.bf16) requires sm_80+; decide the + # kernel dispatch at build time via this define rather than at runtime. + if cap[0] < 8: + flags.append("-DASTRAI_NO_MMA") + return flags _kernels_dir = Path("csrc/kernels") @@ -20,7 +26,7 @@ def register(name: str, sources: list[str] | None = None, **kwargs): sources = [str(_kernels_dir / f"{name}.cu")] REGISTRY[name] = { "sources": sources, - "nvcc_flags": ["-O3", "--expt-relaxed-constexpr", _arch_flag()], + "nvcc_flags": ["-O3", "--expt-relaxed-constexpr", *_arch_flags()], "extra_link_args": kwargs.pop("extra_link_args", []), **kwargs, } diff --git a/csrc/kernels/gqa_prefill_attn.cu b/csrc/kernels/gqa_prefill_attn.cu index bda8c9d..9ad7b0b 100644 --- a/csrc/kernels/gqa_prefill_attn.cu +++ b/csrc/kernels/gqa_prefill_attn.cu @@ -1,6 +1,29 @@ #include "gqa_prefill_attn.cuh" #include +#ifndef ASTRAI_NO_MMA +#include "gqa_prefill_attn_mma.cuh" +#endif + +template +static void dispatch_prefill(GQAParams& p) { +#ifndef ASTRAI_NO_MMA + constexpr int WARPS = 4, BC = 32, BR = 16; + dim3 grid((p.q_len + BR * WARPS - 1) / (BR * WARPS), p.q_head, p.batch); + dim3 block(WARPS * 32, 1, 1); + int smem = (2 * BC * HEAD_DIM + WARPS * BR * HEAD_DIM) * (int)sizeof(bf16); + cudaFuncSetAttribute(gqa_prefill_attn_mma_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + gqa_prefill_attn_mma_kernel<<>>(p); +#else + constexpr int G = 8, ROWS = 32, P_BC = 32; + dim3 grid((p.q_len + ROWS - 1) / ROWS, p.q_head, p.batch); + dim3 block(G, ROWS, 1); + size_t smem = 2 * P_BC * HEAD_DIM * sizeof(bf16); + gqa_prefill_attn_kernel_t<<>>(p); +#endif +} + torch::Tensor gqa_prefill_attn( torch::Tensor q, torch::Tensor k, @@ -22,7 +45,7 @@ torch::Tensor gqa_prefill_attn( p.q_len = q.size(2); p.kv_len = k.size(2); p.head_dim = q.size(3); - TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32"); + TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16"); p.use_mask = mask.has_value(); p.is_causal = (int)is_causal; p.causal_offset = (int)causal_offset; @@ -43,20 +66,15 @@ torch::Tensor gqa_prefill_attn( auto O = torch::empty_like(q); p.o = (bf16*)O.data_ptr(); - constexpr int G = 8, ROWS = 32, P_BC = 32; - dim3 grid((p.q_len + ROWS - 1) / ROWS, p.q_head, p.batch); - dim3 block(G, ROWS, 1); - size_t smem = 2 * P_BC * p.head_dim * sizeof(bf16); - switch (p.head_dim) { case 64: - gqa_prefill_attn_kernel_t<64, G, ROWS, P_BC><<>>(p); + dispatch_prefill<64>(p); break; case 128: - gqa_prefill_attn_kernel_t<128, G, ROWS, P_BC><<>>(p); + dispatch_prefill<128>(p); break; case 256: - gqa_prefill_attn_kernel_t<256, G, ROWS, P_BC><<>>(p); + dispatch_prefill<256>(p); break; default: TORCH_CHECK(false, "prefill: unsupported head_dim ", p.head_dim, @@ -74,5 +92,5 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("is_causal") = false, py::arg("causal_offset") = 0, py::arg("scale") = py::none(), - "GQA prefill (tiled, K+V smem)"); + "GQA prefill (tensor-core mma on sm_80+, scalar fallback)"); } diff --git a/csrc/kernels/gqa_prefill_attn.cuh b/csrc/kernels/gqa_prefill_attn.cuh index 5aae0c9..d0ca014 100644 --- a/csrc/kernels/gqa_prefill_attn.cuh +++ b/csrc/kernels/gqa_prefill_attn.cuh @@ -17,6 +17,19 @@ __device__ __forceinline__ float group_reduce_sum(float v, unsigned mask) { return v; } +// load 8 contiguous bf16 from (16-byte aligned) smem as one float4, unpack to +// 8 floats — cuts shared-load instructions 8x vs scalar bf16 loads. +__device__ __forceinline__ void ld8(const bf16* p, float* o) { + float4 raw = *reinterpret_cast(p); + const __nv_bfloat162* h = reinterpret_cast(&raw); +#pragma unroll + for (int j = 0; j < 4; j++) { + float2 f = __bfloat1622float2(h[j]); + o[2 * j] = f.x; + o[2 * j + 1] = f.y; + } +} + template __global__ void gqa_prefill_attn_kernel_t(GQAParams p) { constexpr int DPT = HEAD_DIM / G; @@ -83,8 +96,13 @@ __global__ void gqa_prefill_attn_kernel_t(GQAParams p) { const bf16* kr = sK + s * HEAD_DIM + gpos * DPT; float part = 0.0f; #pragma unroll - for (int i = 0; i < DPT; i++) - part = fmaf(qreg[i], __bfloat162float(kr[i]), part); + for (int i = 0; i < DPT; i += 8) { + float k8[8]; + ld8(kr + i, k8); +#pragma unroll + for (int j = 0; j < 8; j++) + part = fmaf(qreg[i + j], k8[j], part); + } float dot = group_reduce_sum(part, gmask); if (p.use_mask && p.mask && !p.mask[batch * p.kv_len + kv0 + s]) @@ -97,8 +115,13 @@ __global__ void gqa_prefill_attn_kernel_t(GQAParams p) { const bf16* vr = sV + s * HEAD_DIM + gpos * DPT; #pragma unroll - for (int i = 0; i < DPT; i++) - acc[i] = fmaf(__bfloat162float(vr[i]), be, acc[i] * al); + for (int i = 0; i < DPT; i += 8) { + float v8[8]; + ld8(vr + i, v8); +#pragma unroll + for (int j = 0; j < 8; j++) + acc[i + j] = fmaf(v8[j], be, acc[i + j] * al); + } m = nm; } __syncthreads(); diff --git a/csrc/kernels/gqa_prefill_attn_mma.cuh b/csrc/kernels/gqa_prefill_attn_mma.cuh new file mode 100644 index 0000000..25ef60c --- /dev/null +++ b/csrc/kernels/gqa_prefill_attn_mma.cuh @@ -0,0 +1,223 @@ +#pragma once +#include "gqa_common.cuh" + +// Tensor-core prefill, register-resident flash attention (raw mma.sync PTX). +// One warp owns BR=16 query rows. S = Q@K^T and O = P@V run on bf16 tensor +// cores via mma.sync.m16n8k16 (f32 accumulate). Q stays resident in registers; +// S, O, and the online-softmax stats (m, l) live in registers too — nothing is +// staged through shared memory except the cooperatively-loaded K/V tiles. The +// mma fragment layout is used directly: the S accumulator (f32) maps element- +// for-element onto the P matrix_a (bf16) operand, so softmax needs no shuffle +// repack; row reductions fold across the 4-lane thread group. Templated on +// with BC a multiple of 16. + +// mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 +// (only compiled when ASTRAI_HAS_MMA is set, i.e. built for sm_80+) +__device__ __forceinline__ void mma16816(float* d, const unsigned* a, + const unsigned* b, const float* c) { + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};" + : "=f"(d[0]), "=f"(d[1]), "=f"(d[2]), "=f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]), + "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3])); +} + +// read two adjacent bf16 from smem as one packed .b32 (elem0 low, elem1 high) +__device__ __forceinline__ unsigned ld2(const bf16* p) { + return *reinterpret_cast(p); +} +__device__ __forceinline__ unsigned pk2(float a, float b) { + __nv_bfloat162 v = __floats2bfloat162_rn(a, b); + return *reinterpret_cast(&v); +} +// pack two (non-contiguous) bf16 into one .b32 +__device__ __forceinline__ unsigned pkb(bf16 a, bf16 b) { + __nv_bfloat162 v; + v.x = a; + v.y = b; + return *reinterpret_cast(&v); +} + +template +__global__ void gqa_prefill_attn_mma_kernel(GQAParams p) { + constexpr int BR = 16; + constexpr int KD = HEAD_DIM / 16; // Q/K k-tiles + 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) + + const int warp = threadIdx.x / 32; + const int lane = threadIdx.x % 32; + const int gid = lane >> 2; // 0..7 → rows gid, gid+8 + const int tid4 = lane & 3; // 0..3 + const int nthreads = WARPS * 32; + + const int q_head = blockIdx.y; + const int batch = blockIdx.z; + const int kv_head = q_head / (p.q_head / p.kv_head); + const int qrow0 = (blockIdx.x * WARPS + warp) * BR; + + extern __shared__ __align__(16) bf16 smem[]; + bf16* sK = smem; // [BC][HEAD_DIM] + bf16* sV = sK + BC * HEAD_DIM; // [BC][HEAD_DIM] + bf16* sQ = sV + BC * HEAD_DIM + warp * (BR * HEAD_DIM); // per-warp [BR][HEAD_DIM] + + // stage Q into smem (zero-padded past q_len) + const int q_base = ((batch * p.q_head + q_head) * p.q_len) * HEAD_DIM; + for (int i = lane; i < BR * HEAD_DIM; i += 32) { + int r = i / HEAD_DIM, d = i % HEAD_DIM; + int qr = qrow0 + r; + sQ[i] = (qr < p.q_len) ? p.q[q_base + qr * HEAD_DIM + d] : __float2bfloat16(0.0f); + } + __syncwarp(); + + // Q resident A-fragments: Qa[kt][0..3] + unsigned Qa[KD][4]; +#pragma unroll + for (int kt = 0; kt < KD; kt++) { + int c0 = kt * 16 + 2 * tid4; + Qa[kt][0] = ld2(&sQ[gid * HEAD_DIM + c0]); + Qa[kt][1] = ld2(&sQ[(gid + 8) * HEAD_DIM + c0]); + Qa[kt][2] = ld2(&sQ[gid * HEAD_DIM + c0 + 8]); + Qa[kt][3] = ld2(&sQ[(gid + 8) * HEAD_DIM + c0 + 8]); + } + + 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 tiles = (p.kv_len + BC - 1) / BC; + const int qr0 = qrow0 + gid; // row for c0/c1 + const int qr1 = qrow0 + gid + 8; // row for c2/c3 + + for (int ti = 0; ti < tiles; ti++) { + int kv0 = ti * BC; + + for (int i = threadIdx.x; i < BC * HEAD_DIM; i += nthreads) { + int r = i / HEAD_DIM, d = i % HEAD_DIM; + int kc = kv0 + r; + bf16 z = __float2bfloat16(0.0f); + sK[i] = (kc < p.kv_len) ? p.k[kv_base + kc * HEAD_DIM + d] : z; + sV[i] = (kc < p.kv_len) ? p.v[kv_base + kc * HEAD_DIM + d] : z; + } + __syncthreads(); + + // S = Q @ K^T → Sacc[n8][0..3] (n8: 8 kv cols each) + 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 kv = kv0 + n8 * 8 + gid; +#pragma unroll + for (int kt = 0; kt < KD; kt++) { + unsigned b[2]; + int kr = kt * 16 + 2 * tid4; + b[0] = ld2(&sK[(n8 * 8 + gid) * HEAD_DIM + kr]); + b[1] = ld2(&sK[(n8 * 8 + gid) * HEAD_DIM + kr + 8]); + mma16816(Sacc[n8], Qa[kt], b, Sacc[n8]); + } + (void)kv; + } + + // ---- online softmax (in registers) ---- + // scale + mask, then per-row (gid, gid+8) max over held cols + float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX; +#pragma unroll + for (int n8 = 0; n8 < NC8; n8++) { + int cc = kv0 + n8 * 8 + 2 * tid4; // col for c0/c2 + bool bc0 = (cc >= p.kv_len) || + (p.use_mask && p.mask && !p.mask[batch * p.kv_len + cc]); + bool bc1 = (cc + 1 >= p.kv_len) || + (p.use_mask && p.mask && !p.mask[batch * p.kv_len + cc + 1]); + bool cz = p.is_causal; + int off = p.causal_offset; + bool bad0 = bc0 || (cz && cc > qr0 + off); + bool bad1 = bc1 || (cz && (cc + 1) > qr0 + off); + bool bad2 = bc0 || (cz && cc > qr1 + off); + bool bad3 = bc1 || (cz && (cc + 1) > qr1 + off); + float s0 = bad0 ? -FLT_MAX : Sacc[n8][0] * p.scale; + float s1 = bad1 ? -FLT_MAX : Sacc[n8][1] * p.scale; + float s2 = bad2 ? -FLT_MAX : Sacc[n8][2] * p.scale; + float s3 = bad3 ? -FLT_MAX : Sacc[n8][3] * p.scale; + 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)); + } + // reduce max across the 4-lane group (tid4) + 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; + + // rescale O accumulator by per-row correction +#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; + } + + // O += P @ V +#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]); +#pragma unroll + for (int dn8 = 0; dn8 < DN8; dn8++) { + unsigned b[2]; + int kr = kt2 * 16 + 2 * tid4; + int d = dn8 * 8 + gid; + b[0] = pkb(sV[kr * HEAD_DIM + d], sV[(kr + 1) * HEAD_DIM + d]); + b[1] = pkb(sV[(kr + 8) * HEAD_DIM + d], sV[(kr + 9) * HEAD_DIM + d]); + mma16816(Oacc[dn8], Pa, b, Oacc[dn8]); + } + } + __syncthreads(); // sK/sV reused next tile + } + + // ---- write output ---- + float rl0 = (l0 > 1e-20f) ? (1.0f / l0) : 0.0f; + float rl1 = (l1 > 1e-20f) ? (1.0f / l1) : 0.0f; + const int o_base = ((batch * p.q_head + q_head) * p.q_len) * HEAD_DIM; +#pragma unroll + for (int dn8 = 0; dn8 < DN8; dn8++) { + int d = dn8 * 8 + 2 * tid4; + if (qr0 < p.q_len) { + p.o[o_base + qr0 * HEAD_DIM + d] = __float2bfloat16(Oacc[dn8][0] * rl0); + p.o[o_base + qr0 * HEAD_DIM + d + 1] = __float2bfloat16(Oacc[dn8][1] * rl0); + } + if (qr1 < p.q_len) { + p.o[o_base + qr1 * HEAD_DIM + d] = __float2bfloat16(Oacc[dn8][2] * rl1); + p.o[o_base + qr1 * HEAD_DIM + d + 1] = __float2bfloat16(Oacc[dn8][3] * rl1); + } + } +} diff --git a/csrc/tests/gqa_prefill_test.cu b/csrc/tests/gqa_prefill_test.cu index 9b94e86..7f3e8ee 100644 --- a/csrc/tests/gqa_prefill_test.cu +++ b/csrc/tests/gqa_prefill_test.cu @@ -52,9 +52,10 @@ static float randf() { return (float)rand() / (float)RAND_MAX - 0.5f; } int main() { const int configs[][7] = { - {1,2,1,64,128,32,0}, // tiny: B,Hq,Hk,q,kv,D,causal + {1,2,1,64,128,64,0}, // tiny: B,Hq,Hk,q,kv,D,causal {1,32,4,512,512,128,0}, // standard {1,32,4,128,256,128,0}, // medium + {1,4,2,256,256,128,1}, // causal }; int n_configs = sizeof(configs) / sizeof(configs[0]); @@ -96,7 +97,6 @@ int main() { double t0=now_ms(); switch (D) { - case 32: gqa_prefill_attn_kernel_t<32, G,ROWS,P_BC><<>>(p); break; case 64: gqa_prefill_attn_kernel_t<64, G,ROWS,P_BC><<>>(p); break; case 128: gqa_prefill_attn_kernel_t<128,G,ROWS,P_BC><<>>(p); break; default: printf("unsupported D=%d\n",D); return 1;