ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf: pack gqa q-heads per prefill block to reuse kv tiles

Browse Source 
- 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
This commit is contained in:
ViperEkura
2026-08-31 13:39:47 +08:00
parent 7dd184a4e5
commit a1a1a6bf0f
5 changed files with 91 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/attention/common.h
+5
View File
@@ -16,6 +16,11 @@ enum TensorLayout : int {
// Split-KV workspace cap: max decode splits per (batch, q_head).
constexpr int MAX_SPLITS = 32;
// Paged-prefill host Q-tile granularity in q rows: one q_tile_to_index unit
// covers this many query rows of one request. Must match Q_TILE_ROWS in
// astrai/inference/workspace.py, which builds the device-side tile maps.
constexpr int HOST_Q_TILE_ROWS = 64;
// Unified attention params covering BOTH addressing modes:
// - Contiguous K/V: dense [batch, kv_head, kv_len, head_dim] tensors (k/v).
csrc/kernels/attention/dispatchers.cuh
+11 -2
View File
@@ -88,8 +88,17 @@ struct PrefillLauncherMMA {
static void launch(AttentionParams<bf16>& p, cudaStream_t stream) {
using Config = PrefillConfigMap<HEAD_DIM, IsCausal>;
using Traits = KernelTraits<HEAD_DIM, Config::BC, Config::WARPS, Config::STAGES>;
constexpr int ROWS = Traits::BR * Config::WARPS;
dim3 grid(QSchedule::host_q_blocks(p, ROWS), p.q_head,
// GQA head packing: HB = min(G, WARPS) q-heads of one kv-head group
// share each block's K/V stream (~HB× less global K/V traffic).
// Each head gets WPH = WARPS/HB 16-row chunks per block, so per-head
// rows drop from 64 to BR*WPH while total mma work per K/V byte is
// unchanged. G=1 (MHA) reproduces the historical grid exactly.
const int G = p.q_head / p.kv_head;
const int HB = std::min(G, Config::WARPS);
const int WPH = Config::WARPS / HB;
constexpr int BR = Traits::BR;
dim3 grid(QSchedule::packed_grid_x(p, BR * WPH),
p.kv_head * ((G + HB - 1) / HB),
QSchedule::host_grid_batch(p));
dim3 block(Traits::NUM_THREADS);
attn_prefill_split_q_mma_kernel<Traits, QSchedule, KV, IsCausal, HasMask>
csrc/kernels/attention/layout_policies.cuh
+31
View File
@@ -56,6 +56,20 @@ struct DenseQSchedule {
q_tile = blockIdx.x;
}
// GQA-packed prefill mapping: HB q-heads of one kv-head group share a
// block's K/V stream, each head owning `rows` = BR*WPH consecutive q rows
// per block. Dense tensors tile q_len directly, one block per range.
HOST_FORCEINLINE int packed_grid_x(
const AttentionParams<bf16>& p, int rows) {
return (p.q_len + rows - 1) / rows;
}
DEVICE_FORCEINLINE void map_packed_block(
const AttentionParams<bf16>&, int rows, int& batch, int& row_base) {
batch = blockIdx.z;
row_base = blockIdx.x * rows;
}
DEVICE_FORCEINLINE int q_len(
const AttentionParams<bf16>& p, int) {
return p.q_len;
@@ -84,6 +98,23 @@ struct PackedQSchedule {
q_tile = p.q_tile_to_index[blockIdx.x];
}
// GQA-packed prefill mapping: the host tile maps are built in
// HOST_Q_TILE_ROWS granularity, so each host tile splits into
// HOST_Q_TILE_ROWS / rows packed blocks along blockIdx.x.
HOST_FORCEINLINE int packed_grid_x(
const AttentionParams<bf16>& p, int rows) {
return p.num_q_tiles * (HOST_Q_TILE_ROWS / rows);
}
DEVICE_FORCEINLINE void map_packed_block(
const AttentionParams<bf16>& p, int rows, int& batch, int& row_base) {
const int hb = HOST_Q_TILE_ROWS / rows;
const int host_tile = blockIdx.x / hb;
batch = p.q_tile_to_batch[host_tile];
row_base = p.q_tile_to_index[host_tile] * HOST_Q_TILE_ROWS
+ (blockIdx.x - host_tile * hb) * rows;
}
DEVICE_FORCEINLINE int q_len(
const AttentionParams<bf16>& p, int batch) {
return p.qo_indptr[batch + 1] - p.qo_indptr[batch];
csrc/kernels/attention/prefill_split_q_mma.cuh
+30 -11
View File
@@ -13,6 +13,13 @@ namespace attention {
// 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
@@ -26,11 +33,23 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
const int gid = lane >> 2; // 0..7
const int tid4 = lane & 3; // 0..3
const int q_head = blockIdx.y;
int batch, q_tile;
QSchedule::map_block(p, batch, q_tile);
const int kv_head = q_head / (p.q_head / p.kv_head);
const int qrow0 = (q_tile * Traits::WARPS + warp) * Traits::BR;
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);
@@ -62,11 +81,11 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
const int qr0 = qrow0 + gid;
const int qr1 = qrow0 + gid + 8;
// Causal tile-skip bounds (dead code when IsCausal == false)
// 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 =
q_tile * Traits::WARPS * Traits::BR + Traits::WARPS * 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) {
@@ -144,13 +163,13 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
#pragma unroll
for (int dn8 = 0; dn8 < Traits::DN8; dn8++) {
int d = dn8 * 8 + 2 * tid4;
if (qr0 < q_len) {
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 (qr1 < q_len) {
if (active && qr1 < q_len) {
__nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][2] * rl1,
Oacc[dn8][3] * rl1);
*reinterpret_cast<__nv_bfloat162*>(