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

perf: enlarge prefill KV tile to BC=32 for D<=128

Browse Source 
- Kernel is latency-bound (25% occupancy), not compute/bandwidth-bound
- BC=16 wasted a cp.async wait + barrier + loop overhead per tiny tile
- Double KV tile to BC=32 for D<=128; D=256 stays 16 (64KB > 48KB smem cap)
- Retune MIN_BLOCKS per head_dim (32->6, 64->4, 128->3, 256->2)
- Result: ~6-8% faster on L20, 0.93-1.20x vs torch SDPA, correctness unchanged
This commit is contained in:
ViperEkura
2026-07-10 17:43:23 +08:00
parent 9027fdc546
commit 88f8dca2c2
3 changed files with 18 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/tests/gqa_prefill_test.cu
+4 -2
View File
@@ -25,8 +25,10 @@ static double now_ms() {
template <int HEAD_DIM>
static void launch_prefill(GQAParams& p) {
#ifndef ASTRAI_NO_MMA
constexpr int WARPS = 4, BC = 16, BR = 16;
constexpr int MIN_BLOCKS = (HEAD_DIM <= 64) ? 6 : (HEAD_DIM <= 128) ? 3 : 2;
constexpr int WARPS = 4, BR = 16;
constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
constexpr int MIN_BLOCKS = (HEAD_DIM <= 32) ? 6 : (HEAD_DIM <= 64) ? 4
: (HEAD_DIM <= 128) ? 3 : 2;
dim3 grid((p.q_len + BR * WARPS - 1) / (BR * WARPS), p.q_head, p.batch);
dim3 block(WARPS * 32, 1, 1);
gqa_prefill_attn_mma_kernel<HEAD_DIM, WARPS, BC, MIN_BLOCKS><<<grid, block>>>(p);