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

- 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

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).

csrc/kernels/attention/dispatchers.cuh

@@ -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;
+        // GQA head packing: HB = min(G, WARPS) q-heads of one kv-head group
-        dim3 grid(QSchedule::host_q_blocks(p, ROWS), p.q_head,
+        // 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

@@ -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

@@ -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;
+    const int G = p.q_head / p.kv_head;
-    int batch, q_tile;
+    const int HB = min(G, Traits::WARPS);   // q heads packed per block
-    QSchedule::map_block(p, batch, q_tile);
+    const int WPH = Traits::WARPS / HB;     // 16-row chunks per head
-    const int kv_head = q_head / (p.q_head / p.kv_head);
+    const int BPG = (G + HB - 1) / HB;      // blocks per GQA group
-    const int qrow0 = (q_tile * Traits::WARPS + warp) * Traits::BR;
+    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 =
+    const int block_max_kv = row_base + WPH * Traits::BR - 1 + causal_off;
-        q_tile * Traits::WARPS * Traits::BR + Traits::WARPS * 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*>(

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.x = num_q_tiles * HB   # HB = min(G, WARPS): q heads packed per block
-grid.y = q_heads
+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];
+host_tile = blockIdx.x / HB;
-q_tile = q_tile_to_index[blockIdx.x];
+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