diff --git a/csrc/kernels/attn_decode_split_kv_mma.cuh b/csrc/kernels/attn_decode_split_kv_mma.cuh index 07ba795..30be5bb 100644 --- a/csrc/kernels/attn_decode_split_kv_mma.cuh +++ b/csrc/kernels/attn_decode_split_kv_mma.cuh @@ -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 __global__ void attn_decode_split_kv_mma_kernel(AttentionParams 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 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( - &p.q[q_base + qra * p.q_stride_h + c * p.q_stride_d]); - const unsigned* pbu = reinterpret_cast( - &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(p.q + q_base, p.q_stride_h, p.q_stride_d, + qra, qrb, va, vb, tid4, Qa); float Oacc[DN8][4]; #pragma unroll diff --git a/csrc/kernels/attn_mma_utils.cuh b/csrc/kernels/attn_mma_utils.cuh index 47f95e2..4baefa6 100644 --- a/csrc/kernels/attn_mma_utils.cuh +++ b/csrc/kernels/attn_mma_utils.cuh @@ -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 +__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( + &q[qra * stride_row + c * stride_d]); + const unsigned* pbu = reinterpret_cast( + &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 diff --git a/csrc/kernels/attn_paged_decode_split_kv_mma.cuh b/csrc/kernels/attn_paged_decode_split_kv_mma.cuh index 64f0d65..dca6b7e 100644 --- a/csrc/kernels/attn_paged_decode_split_kv_mma.cuh +++ b/csrc/kernels/attn_paged_decode_split_kv_mma.cuh @@ -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 __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams 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 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( - &p.q[q_base + qra * p.q_stride_h + c * p.q_stride_d]); - const unsigned* pbu = reinterpret_cast( - &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(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 // 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) { diff --git a/csrc/kernels/attn_prefill_split_q_mma.cuh b/csrc/kernels/attn_prefill_split_q_mma.cuh index edc62db..6374162 100644 --- a/csrc/kernels/attn_prefill_split_q_mma.cuh +++ b/csrc/kernels/attn_prefill_split_q_mma.cuh @@ -57,7 +57,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams 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 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( - &p.q[q_base + qra * p.q_stride_l + c * p.q_stride_d]); - const unsigned* pbu = reinterpret_cast( - &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(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 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 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++) {