ViperEkura
a1a1a6bf0f
- pack HB = min(G, WARPS) q heads per block; K/V tiles stream once per block instead of once per q head
- G=1 keeps the old grid; paged path splits 64-row host Q tiles into HB blocks along grid.x (host maps unchanged)
Benchmark: NVIDIA RTX 6000D, short-q/long-kv prefill 1.4-3.4x (G=8 B=16 q=16 kv=16k 4.22 -> 1.26 ms); full prefill/MHA/paged unchanged (compute-bound); verified vs SDPA G in {1,2,3,4,8,32}, 99 tests pass
183 lines
7.9 KiB
Plaintext
183 lines
7.9 KiB
Plaintext
#pragma once
|
||
#include <cfloat>
|
||
#include <cuda_bf16.h>
|
||
#include "common.h"
|
||
#include "layout_policies.cuh"
|
||
#include "mma_utils.cuh"
|
||
|
||
namespace astrai {
|
||
namespace attention {
|
||
|
||
// Tensor-core prefill flash attention (raw mma.sync PTX), unified across
|
||
// contiguous and paged (SGLang flat-pool) K/V via the KV template parameter.
|
||
// 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).
|
||
//
|
||
// GQA head packing (FA2/FA3-style): HB = min(G, WARPS) query heads of one
|
||
// kv-head group share a block's K/V tiles, so each K/V element is read from
|
||
// global memory once per block instead of once per q head (~HB× less K/V
|
||
// traffic). WARPS = WPH × HB: warp w handles head slot w/WPH, chunk w%WPH;
|
||
// all warps of a block cover the same token range, keeping the causal sweep
|
||
// end block-uniform. G=1 (MHA) degenerates to the unpadded layout.
|
||
//
|
||
// KV = ContigKV (dense [batch, kv_head, kv_len, head_dim]) or PagedKV
|
||
// (flat pool + req_to_token, ragged batches via qo_indptr/kv_indptr).
|
||
// IsCausal and HasMask are compile-time bools — the compiler eliminates all
|
||
// dead branches in the inner compute loop (FA2-style).
|
||
//
|
||
// Traits = KernelTraits<HEAD_DIM, BC, WARPS=4, STAGES=2>.
|
||
template <typename Traits, typename QSchedule, typename KV, bool IsCausal, bool HasMask>
|
||
__global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
|
||
const int warp = threadIdx.x / 32;
|
||
const int lane = threadIdx.x % 32;
|
||
const int gid = lane >> 2; // 0..7
|
||
const int tid4 = lane & 3; // 0..3
|
||
|
||
const int G = p.q_head / p.kv_head;
|
||
const int HB = min(G, Traits::WARPS); // q heads packed per block
|
||
const int WPH = Traits::WARPS / HB; // 16-row chunks per head
|
||
const int BPG = (G + HB - 1) / HB; // blocks per GQA group
|
||
const int chunk = warp % WPH;
|
||
|
||
int batch, row_base;
|
||
QSchedule::map_packed_block(p, Traits::BR * WPH, batch, row_base);
|
||
const int kv_head = blockIdx.y / BPG;
|
||
const int slot = blockIdx.y - kv_head * BPG;
|
||
const int head_idx = slot * HB + warp / WPH;
|
||
// G % HB tail blocks have idle head slots: clamp to the last head so all
|
||
// warps do valid work (cp.async + __syncthreads stay block-uniform) and
|
||
// just skip the O store via `active`.
|
||
const bool active = head_idx < G;
|
||
const int q_head = kv_head * G + min(head_idx, G - 1);
|
||
const int qrow0 = row_base + chunk * Traits::BR;
|
||
|
||
// Per-request dims (from KV policy — paged reads kv_indptr/qo_indptr).
|
||
const int seq_len = KV::kv_len(p, batch);
|
||
const int q_len = QSchedule::q_len(p, batch);
|
||
const int causal_off = KV::causal_offset(p, batch, q_len);
|
||
const KVContext kctx = KV::template make_ctx<Traits::HEAD_DIM>(p, batch, kv_head);
|
||
|
||
// Static shared memory: double-buffered K/V (no sQ — Q goes direct
|
||
// to registers in mma A-operand layout).
|
||
__shared__ __align__(16) bf16 sK[Traits::STAGES * Traits::BC * Traits::LD];
|
||
__shared__ __align__(16) bf16 sV[Traits::STAGES * Traits::BC * Traits::LD];
|
||
|
||
// Load Q fragments straight from global into mma A-operand layout.
|
||
const int q_base = QSchedule::q_base(p, batch, q_head);
|
||
const int qra = qrow0 + gid;
|
||
const int qrb = qrow0 + gid + 8;
|
||
const bool va = qra < q_len, vb = qrb < q_len;
|
||
unsigned Qa[Traits::KD][4];
|
||
load_q_mma_frags<Traits::KD>(p.q_ptr + q_base, p.q_l_stride, p.q_d_stride,
|
||
qra, qrb, va, vb, tid4, Qa);
|
||
|
||
float Oacc[Traits::DN8][4];
|
||
#pragma unroll
|
||
for (int j = 0; j < Traits::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 tiles = (seq_len + Traits::BC - 1) / Traits::BC;
|
||
const int qr0 = qrow0 + gid;
|
||
const int qr1 = qrow0 + gid + 8;
|
||
|
||
// Causal tile-skip bounds (dead code when IsCausal == false).
|
||
// max_kv is per-warp (its own 16 rows); block_max_kv is the last row of
|
||
// the whole block's range and must be uniform for the shared sweep loop.
|
||
const int max_kv = qrow0 + Traits::BR - 1 + causal_off;
|
||
const int block_max_kv = row_base + WPH * Traits::BR - 1 + causal_off;
|
||
|
||
int t_end = tiles - 1;
|
||
if constexpr (IsCausal) {
|
||
int bt = block_max_kv / Traits::BC;
|
||
if (bt < t_end) t_end = bt;
|
||
}
|
||
|
||
// ---- Load tile lambda: predicated cp.async (addressing via KV policy) ----
|
||
auto load_tile = [&](int ti, int buf) {
|
||
int kv0 = ti * Traits::BC;
|
||
bf16* dK = sK + buf * Traits::BC * Traits::LD;
|
||
bf16* dV = sV + buf * Traits::BC * Traits::LD;
|
||
#pragma unroll
|
||
for (int i = threadIdx.x * Traits::VEC; i < Traits::TOTAL;
|
||
i += Traits::NUM_THREADS * Traits::VEC) {
|
||
int r = i / Traits::HEAD_DIM, d = i % Traits::HEAD_DIM;
|
||
int kc = kv0 + r;
|
||
bool valid = kc < seq_len;
|
||
int token = KV::resolve_token(p, kctx, kc, valid);
|
||
KVAddr a = KV::kv_addr_from_token(p, kctx, token, d);
|
||
int off = r * Traits::LD + swiz_col(d, r, Traits::SWIZ_MASK);
|
||
astrai::cp_async_16(&dK[off], a.k, a.valid);
|
||
astrai::cp_async_16(&dV[off], a.v, a.valid);
|
||
}
|
||
astrai::cp_async_commit_group();
|
||
};
|
||
|
||
// ---- Prologue: issue first tile load ----
|
||
load_tile(0, 0);
|
||
|
||
for (int ti = 0; ti <= t_end; ti++) {
|
||
int buf = ti & 1;
|
||
|
||
// Wait for current tile, then publish cross-warp + guard buffer reuse.
|
||
astrai::cp_async_wait_group<0>();
|
||
__syncthreads();
|
||
if (ti < t_end) load_tile(ti + 1, (ti + 1) & 1);
|
||
|
||
const bf16* bK = sK + buf * Traits::BC * Traits::LD;
|
||
const bf16* bV = sV + buf * Traits::BC * Traits::LD;
|
||
int kv0 = ti * Traits::BC;
|
||
|
||
// Warp-level causal skip (dead branch eliminated when IsCausal == false)
|
||
if (!IsCausal || kv0 <= max_kv) {
|
||
|
||
float Sacc[Traits::NC8][4];
|
||
mma_compute_scores<Traits>(Qa, bK, lane, Sacc);
|
||
|
||
// Post-multiply scale in float (no bf16 precision loss)
|
||
#pragma unroll
|
||
for (int n8 = 0; n8 < Traits::NC8; n8++)
|
||
Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
|
||
Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale;
|
||
|
||
int maxc0 = IsCausal ? min(seq_len, causal_off + qr0 + 1)
|
||
: seq_len;
|
||
int maxc1 = IsCausal ? min(seq_len, causal_off + qr1 + 1)
|
||
: seq_len;
|
||
mma_softmax_tile<Traits, HasMask>(kv0, maxc0, maxc1,
|
||
qr0, qr1,
|
||
p.mask_b_stride, p.mask_h_stride, p.mask_l_stride,
|
||
batch, q_head, q_head,
|
||
p.mask,
|
||
va, vb,
|
||
Sacc, Oacc, m0, m1, l0, l1, lane);
|
||
|
||
mma_pv_accumulate<Traits>(Sacc, bV, lane, Oacc);
|
||
}
|
||
}
|
||
|
||
// ---- write output: packed bf16x2 stores ----
|
||
float rl0 = (l0 > 1e-20f) ? (1.0f / l0) : 0.0f;
|
||
float rl1 = (l1 > 1e-20f) ? (1.0f / l1) : 0.0f;
|
||
const int o_base = QSchedule::q_base(p, batch, q_head);
|
||
#pragma unroll
|
||
for (int dn8 = 0; dn8 < Traits::DN8; dn8++) {
|
||
int d = dn8 * 8 + 2 * tid4;
|
||
if (active && qr0 < q_len) {
|
||
__nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][0] * rl0,
|
||
Oacc[dn8][1] * rl0);
|
||
*reinterpret_cast<__nv_bfloat162*>(
|
||
&p.o_ptr[o_base + qr0 * p.q_l_stride + d * p.q_d_stride]) = v;
|
||
}
|
||
if (active && qr1 < q_len) {
|
||
__nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][2] * rl1,
|
||
Oacc[dn8][3] * rl1);
|
||
*reinterpret_cast<__nv_bfloat162*>(
|
||
&p.o_ptr[o_base + qr1 * p.q_l_stride + d * p.q_d_stride]) = v;
|
||
}
|
||
}
|
||
}
|
||
|
||
} // namespace attention
|
||
} // namespace astrai
|