diff --git a/csrc/kernels/attention/common.h b/csrc/kernels/attention/common.h index 88ae20b..1222aae 100644 --- a/csrc/kernels/attention/common.h +++ b/csrc/kernels/attention/common.h @@ -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). diff --git a/csrc/kernels/attention/dispatchers.cuh b/csrc/kernels/attention/dispatchers.cuh index f6b78a7..aad4052 100644 --- a/csrc/kernels/attention/dispatchers.cuh +++ b/csrc/kernels/attention/dispatchers.cuh @@ -88,8 +88,17 @@ struct PrefillLauncherMMA { static void launch(AttentionParams& p, cudaStream_t stream) { using Config = PrefillConfigMap; using Traits = KernelTraits; - 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 diff --git a/csrc/kernels/attention/layout_policies.cuh b/csrc/kernels/attention/layout_policies.cuh index 0354c1d..257dc32 100644 --- a/csrc/kernels/attention/layout_policies.cuh +++ b/csrc/kernels/attention/layout_policies.cuh @@ -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& p, int rows) { + return (p.q_len + rows - 1) / rows; + } + + DEVICE_FORCEINLINE void map_packed_block( + const AttentionParams&, int rows, int& batch, int& row_base) { + batch = blockIdx.z; + row_base = blockIdx.x * rows; + } + DEVICE_FORCEINLINE int q_len( const AttentionParams& 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& p, int rows) { + return p.num_q_tiles * (HOST_Q_TILE_ROWS / rows); + } + + DEVICE_FORCEINLINE void map_packed_block( + const AttentionParams& 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& p, int batch) { return p.qo_indptr[batch + 1] - p.qo_indptr[batch]; diff --git a/csrc/kernels/attention/prefill_split_q_mma.cuh b/csrc/kernels/attention/prefill_split_q_mma.cuh index 4692529..99c0f5a 100644 --- a/csrc/kernels/attention/prefill_split_q_mma.cuh +++ b/csrc/kernels/attention/prefill_split_q_mma.cuh @@ -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 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 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 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*>( diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index 845a1da..2b878d4 100644 --- a/docs/developer/cuda_kernels.md +++ b/docs/developer/cuda_kernels.md @@ -374,19 +374,26 @@ q_tile_to_batch = [0, 0, 1, 2, 2, 2] q_tile_to_index = [0, 1, 0, 0, 1, 2] ``` -Paged prefill launches: +Paged prefill launches (MMA path, GQA head packing): ```text -grid.x = num_q_tiles # 6, exactly the valid ragged work items -grid.y = q_heads +grid.x = num_q_tiles * HB # HB = min(G, WARPS): q heads packed per block +grid.y = kv_heads * ceil(G / HB) grid.z = 1 ``` -Each block resolves its request and request-local tile in O(1): +The tensor-core prefill kernel packs `HB = min(G, WARPS)` query heads of one +kv-head group into a block, so K/V tiles stream once per block instead of once +per q head (~HB× less global K/V traffic). Warp `w` handles head slot `w / WPH` +and 16-row chunk `w % WPH`, where `WPH = WARPS / HB`; `G = q_heads / kv_heads` +and `G = 1` (MHA) degenerates to the historical one-head-per-block layout. +Each host Q tile (64 rows, `Q_TILE_ROWS`) splits into `HB` packed blocks along +`grid.x`. Each block resolves its request and request-local row range in O(1): ```cpp -batch = q_tile_to_batch[blockIdx.x]; -q_tile = q_tile_to_index[blockIdx.x]; +host_tile = blockIdx.x / HB; +batch = q_tile_to_batch[host_tile]; +row_base = q_tile_to_index[host_tile] * 64 + (blockIdx.x % HB) * (64 / HB); ``` The kernel then uses `qo_indptr[batch]` for the packed Q base and adjacent @@ -448,7 +455,7 @@ csrc/ │ │ ├── decode_split_kv.cuh # decode kernel, scalar (split-KV) │ │ ├── decode_split_kv_mma.cuh # decode kernel, MMA + split-K │ │ ├── prefill_split_q.cuh # prefill kernel, scalar (split-Q) -│ │ ├── prefill_split_q_mma.cuh # prefill kernel, MMA (split-Q, packed/ragged Q schedule) +│ │ ├── prefill_split_q_mma.cuh # prefill kernel, MMA (split-Q, GQA head packing, packed/ragged Q schedule) │ │ ├── decode.cu # → module attn_decode │ │ ├── prefill.cu # → module attn_prefill │ │ ├── paged_decode.cu # → module attn_paged_decode