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
This commit is contained in:
2026-08-31 13:39:47 +08:00
parent 7dd184a4e5
commit a1a1a6bf0f
5 changed files with 91 additions and 20 deletions
+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).