perf: double-buffer KV pipeline and Q direct-to-register in decode
- Double-buffered KV (STAGES=2) for D<=128: next tile cp.async overlaps current tile MMA compute, hiding global load latency - Q loaded directly from global into mma A-operand registers, removing sQ staging and prologue syncwarp - Predicated cp.async unifies full and partial tile paths, eliminating scalar fallback branch - STAGES=1 fallback for D=256 (double-buffer would exceed smem budget) - Applied to both contiguous and paged decode MMA kernels - ~1.27x average speedup on L20 (sm_89), zero precision loss
This commit is contained in:
@@ -42,9 +42,10 @@ static void launch_paged_decode(PagedAttentionParams<bf16, float>& p) {
|
||||
int G_check = p.q_head / p.kv_head;
|
||||
bool use_mma = !p.use_mask && G_check >= 1 && G_check <= 16 && p.page_size >= 32;
|
||||
if (use_mma) {
|
||||
constexpr int STAGES = (HEAD_DIM <= 128) ? 2 : 1;
|
||||
int tiles_total = (p.kv_len + 32 - 1) / 32;
|
||||
p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total);
|
||||
paged_attn_decode_split_kv_mma_kernel<HEAD_DIM, 32>
|
||||
paged_attn_decode_split_kv_mma_kernel<HEAD_DIM, 32, STAGES>
|
||||
<<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
|
||||
} else
|
||||
#endif
|
||||
|
||||
Reference in New Issue
Block a user