refactor: extract load_q_mma_frags template, unify comment style
- Add load_q_mma_frags<KD>() shared template in attn_mma_utils.cuh - Replace ~15 duplicated Q-load lines in 3 MMA kernels - Unify section header comment style to // ---- Section ---- - Remove duplicate separator line in attn_mma_utils.cuh
This commit is contained in:
parent
bb175fda91
commit
1f0be382ad
|
|
@ -8,39 +8,21 @@ using bf16 = __nv_bfloat16;
|
|||
|
||||
// Split-K (FlashDecoding) tensor-core decode via GQA head-packing.
|
||||
//
|
||||
// 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
|
||||
// one kv_head. We pack those G heads into the M=16 rows of mma.sync.m16n8k16,
|
||||
// turning G independent GEMVs into a single GEMM that reuses each loaded K/V tile
|
||||
// across all G heads (K/V load is the decode bottleneck, so the reuse is the win,
|
||||
// not the flops). 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,
|
||||
// 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 one kv_head. We pack those G heads into the M=16 rows of
|
||||
// mma.sync.m16n8k16, turning G independent GEMVs into a single GEMM that
|
||||
// reuses each loaded K/V tile across all G heads (K/V load is the decode
|
||||
// bottleneck, so the reuse is the win, not the flops). 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)
|
||||
//
|
||||
// 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
|
||||
// - Q loaded directly from global into mma A-operand registers (no sQ staging,
|
||||
// no prologue syncwarp) — frees shared memory for double-buffering
|
||||
// - Double-buffered KV (STAGES=2): next tile's cp.async overlaps current
|
||||
// tile's MMA compute — hides global load latency / boosts bandwidth
|
||||
// utilization for small-batch (low-occupancy) decode
|
||||
// - Predicated cp.async (cp_async_16_pred) for full AND partial tiles on one
|
||||
// uniform path — eliminates the scalar fallback branch
|
||||
//
|
||||
// Smem footprint (BC=32): STAGES=2 → 2*(sK+sV) = 2*2*32*HEAD_DIM*2 bytes.
|
||||
// D=128: 16 KB (fits 48 KB static cap). D=256: 32 KB (also fits).
|
||||
// STAGES=1 fallback (4/8 KB) for smem-constrained configs.
|
||||
|
||||
template <int HEAD_DIM, int BC, int STAGES = 2>
|
||||
__global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
|
||||
constexpr int BR = 16;
|
||||
constexpr int KD = HEAD_DIM / 16;
|
||||
constexpr int NC8 = BC / 8;
|
||||
constexpr int KT2 = BC / 16;
|
||||
|
|
@ -66,26 +48,13 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
|
|||
__shared__ __align__(16) bf16 sV[STAGES * BC * LD];
|
||||
|
||||
// ---- Load Q directly from global into mma A-operand registers ----
|
||||
// Same layout as prefill: frag[0]/[2] = row gid, frag[1]/[3] = row gid+8
|
||||
// cols kt*16 + tid4*2 + {0,1} / +{8,9}. pau[0]=cols c,c+1; pau[4]=c+8,c+9.
|
||||
// Stride-based: Q is [batch, q_head, q_len=1, head_dim]
|
||||
const int q_base = batch * p.q_stride_b + q_head0 * p.q_stride_h;
|
||||
const int qra = gid;
|
||||
const int qrb = gid + 8;
|
||||
const bool va = qra < G, vb = qrb < G;
|
||||
unsigned Qa[KD][4];
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < KD; kt++) {
|
||||
int c = kt * 16 + tid4 * 2;
|
||||
const unsigned* pau = reinterpret_cast<const unsigned*>(
|
||||
&p.q[q_base + qra * p.q_stride_h + c * p.q_stride_d]);
|
||||
const unsigned* pbu = reinterpret_cast<const unsigned*>(
|
||||
&p.q[q_base + qrb * p.q_stride_h + c * p.q_stride_d]);
|
||||
Qa[kt][0] = va ? pau[0] : 0u;
|
||||
Qa[kt][1] = vb ? pbu[0] : 0u;
|
||||
Qa[kt][2] = va ? pau[4] : 0u;
|
||||
Qa[kt][3] = vb ? pbu[4] : 0u;
|
||||
}
|
||||
load_q_mma_frags<KD>(p.q + q_base, p.q_stride_h, p.q_stride_d,
|
||||
qra, qrb, va, vb, tid4, Qa);
|
||||
|
||||
float Oacc[DN8][4];
|
||||
#pragma unroll
|
||||
|
|
|
|||
|
|
@ -108,6 +108,36 @@ __device__ __forceinline__ void cp_async_wait_group() {
|
|||
asm volatile("cp.async.wait_group %0;" :: "n"(N));
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Q-load: load query rows directly from global memory into mma A-operand
|
||||
// register layout. One call replaces ~15 duplicated lines in each MMA kernel.
|
||||
// stride_row is p.q_stride_h for decode (q_len=1, G heads) or
|
||||
// p.q_stride_l for prefill (multi-q rows).
|
||||
// ---------------------------------------------------------------------------
|
||||
template <int KD>
|
||||
__device__ inline void load_q_mma_frags(
|
||||
const bf16* __restrict__ q,
|
||||
int stride_row,
|
||||
int stride_d,
|
||||
int qra, int qrb,
|
||||
bool va, bool vb,
|
||||
int tid4,
|
||||
unsigned Qa[KD][4])
|
||||
{
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < KD; kt++) {
|
||||
int c = kt * 16 + tid4 * 2;
|
||||
const unsigned* pau = reinterpret_cast<const unsigned*>(
|
||||
&q[qra * stride_row + c * stride_d]);
|
||||
const unsigned* pbu = reinterpret_cast<const unsigned*>(
|
||||
&q[qrb * stride_row + c * stride_d]);
|
||||
Qa[kt][0] = va ? pau[0] : 0u;
|
||||
Qa[kt][1] = vb ? pbu[0] : 0u;
|
||||
Qa[kt][2] = va ? pau[4] : 0u;
|
||||
Qa[kt][3] = vb ? pbu[4] : 0u;
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Shared MMA compute functions — used by both decode and prefill MMA kernels.
|
||||
// Extracted because S=Q@K^T, online softmax, and P@V are structurally identical
|
||||
|
|
|
|||
|
|
@ -11,14 +11,9 @@ using bf16 = __nv_bfloat16;
|
|||
// directly from the page pool through a page table, eliminating the gather
|
||||
// copy. Each tile (BC=32) fits within a single page (page_size >= 32), so
|
||||
// the page-table lookup happens once per tile for cp.async.
|
||||
//
|
||||
// Optimizations mirror attn_decode_split_kv_mma_kernel:
|
||||
// - Q loaded directly from global into mma A-operand registers (no sQ)
|
||||
// - Double-buffered KV (STAGES=2) for D<=128, single-buffer for D=256
|
||||
// - Predicated cp.async for unified full/partial tile path
|
||||
|
||||
template <int HEAD_DIM, int BC, int STAGES = (HEAD_DIM <= 128) ? 2 : 1>
|
||||
__global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16> p) {
|
||||
constexpr int BR = 16;
|
||||
constexpr int KD = HEAD_DIM / 16;
|
||||
constexpr int NC8 = BC / 8;
|
||||
constexpr int KT2 = BC / 16;
|
||||
|
|
@ -32,34 +27,23 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
|
|||
const int gid = lane >> 2;
|
||||
const int tid4 = lane & 3;
|
||||
|
||||
const int kv_head_idx = blockIdx.x;
|
||||
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_idx * G;
|
||||
const int q_head0 = kv_head * G;
|
||||
|
||||
__shared__ __align__(16) bf16 sK[STAGES * BC * LD];
|
||||
__shared__ __align__(16) bf16 sV[STAGES * BC * LD];
|
||||
|
||||
// ---- Load Q directly from global into mma A-operand registers ----
|
||||
// Q stride-based: [batch, q_head, q_len=1, head_dim]
|
||||
const int q_base = batch * p.q_stride_b + q_head0 * p.q_stride_h;
|
||||
const int qra = gid;
|
||||
const int qrb = gid + 8;
|
||||
const bool va = qra < G, vb = qrb < G;
|
||||
unsigned Qa[KD][4];
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < KD; kt++) {
|
||||
int c = kt * 16 + tid4 * 2;
|
||||
const unsigned* pau = reinterpret_cast<const unsigned*>(
|
||||
&p.q[q_base + qra * p.q_stride_h + c * p.q_stride_d]);
|
||||
const unsigned* pbu = reinterpret_cast<const unsigned*>(
|
||||
&p.q[q_base + qrb * p.q_stride_h + c * p.q_stride_d]);
|
||||
Qa[kt][0] = va ? pau[0] : 0u;
|
||||
Qa[kt][1] = vb ? pbu[0] : 0u;
|
||||
Qa[kt][2] = va ? pau[4] : 0u;
|
||||
Qa[kt][3] = vb ? pbu[4] : 0u;
|
||||
}
|
||||
load_q_mma_frags<KD>(p.q + q_base, p.q_stride_h, p.q_stride_d,
|
||||
qra, qrb, va, vb, tid4, Qa);
|
||||
|
||||
float Oacc[DN8][4];
|
||||
#pragma unroll
|
||||
|
|
@ -77,7 +61,7 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
|
|||
// Paged strides (constant for the block)
|
||||
const int64_t page_stride = (int64_t)p.page_size * p.kv_head * HEAD_DIM;
|
||||
const int64_t pos_stride = (int64_t)p.kv_head * HEAD_DIM;
|
||||
const int64_t head_off = (int64_t)kv_head_idx * HEAD_DIM;
|
||||
const int64_t head_off = (int64_t)kv_head * HEAD_DIM;
|
||||
|
||||
// ---- Load tile lambda: predicated cp.async, paged addressing ----
|
||||
auto load_tile = [&](int ti, int buf) {
|
||||
|
|
|
|||
|
|
@ -57,7 +57,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
|
|||
const int kv_head = q_head / (p.q_head / p.kv_head);
|
||||
const int qrow0 = (blockIdx.x * WARPS + warp) * BR;
|
||||
|
||||
// Static shared memory — sized by template parameters at compile time.
|
||||
// ---- Static shared memory: double-buffered K/V ----
|
||||
// K/V are double-buffered (STAGES=2): the next tile's cp.async load runs
|
||||
// while the current tile's tensor-core math executes, hiding global-load
|
||||
// latency (FA2-style software pipeline). No dynamic smem / carveout opt-in.
|
||||
|
|
@ -65,31 +65,16 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
|
|||
__shared__ __align__(16) bf16 sK[STAGES * BC * LD];
|
||||
__shared__ __align__(16) bf16 sV[STAGES * BC * LD];
|
||||
|
||||
// Load the Q fragments straight from global into the mma A-operand layout
|
||||
// (m16n8k16, row-major): no sQ staging area and no serialized per-warp
|
||||
// prologue barriers. Each lane reads exactly the 8 Q elements ldmatrix
|
||||
// would have produced, pre-scaled by the attention scale. Kept resident in
|
||||
// registers across the tile loop.
|
||||
// frag[0]/[2]: row = qrow0 + gid ; frag[1]/[3]: row = qrow0 + gid + 8
|
||||
// frag[0]/[1]: cols kt*16 + tid4*2 + {0,1} ; frag[2]/[3]: + 8
|
||||
// Q stride-based: [batch, q_head, q_len, head_dim]
|
||||
// Load Q fragments straight from global into mma A-operand layout.
|
||||
// stride_row = p.q_stride_l for prefill (multi-q rows across q_len).
|
||||
// See attn_mma_utils.cuh for the shared template.
|
||||
const int q_base = batch * p.q_stride_b + q_head * p.q_stride_h;
|
||||
const int qra = qrow0 + gid;
|
||||
const int qrb = qrow0 + gid + 8;
|
||||
const bool va = qra < p.q_len, vb = qrb < p.q_len;
|
||||
unsigned Qa[KD][4];
|
||||
#pragma unroll
|
||||
for (int kt = 0; kt < KD; kt++) {
|
||||
int c = kt * 16 + tid4 * 2;
|
||||
const unsigned* pau = reinterpret_cast<const unsigned*>(
|
||||
&p.q[q_base + qra * p.q_stride_l + c * p.q_stride_d]);
|
||||
const unsigned* pbu = reinterpret_cast<const unsigned*>(
|
||||
&p.q[q_base + qrb * p.q_stride_l + c * p.q_stride_d]);
|
||||
Qa[kt][0] = va ? pau[0] : 0u;
|
||||
Qa[kt][1] = vb ? pbu[0] : 0u;
|
||||
Qa[kt][2] = va ? pau[4] : 0u;
|
||||
Qa[kt][3] = vb ? pbu[4] : 0u;
|
||||
}
|
||||
load_q_mma_frags<KD>(p.q + q_base, p.q_stride_l, p.q_stride_d,
|
||||
qra, qrb, va, vb, tid4, Qa);
|
||||
|
||||
float Oacc[DN8][4];
|
||||
#pragma unroll
|
||||
|
|
@ -122,6 +107,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
|
|||
constexpr int VEC = 8; // bf16 per cp.async unit (16 bytes)
|
||||
constexpr int TOTAL = BC * HEAD_DIM;
|
||||
|
||||
// ---- Load tile lambda: predicated cp.async ----
|
||||
// Issue cp.async loads for tile `ti` into shared buffer `buf`. Predicated
|
||||
// loads zero-fill rows past kv_len, so partial tiles need no scalar path.
|
||||
auto load_tile = [&](int ti, int buf) {
|
||||
|
|
@ -141,7 +127,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
|
|||
cp_async_commit();
|
||||
};
|
||||
|
||||
// Prologue: kick off the first tile's load.
|
||||
// ---- Prologue: issue first tile load ----
|
||||
load_tile(0, 0);
|
||||
|
||||
for (int ti = 0; ti <= t_end; ti++) {
|
||||
|
|
|
|||
Loading…
Reference in New Issue