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:
@@ -27,14 +27,14 @@ static bool decode_use_mma(const AttentionParams<bf16>& p) {
|
||||
return !p.use_mask && G > 1 && G <= 16;
|
||||
}
|
||||
|
||||
template <int HEAD_DIM, int BC>
|
||||
template <int HEAD_DIM, int BC, int STAGES = (HEAD_DIM <= 128) ? 2 : 1>
|
||||
static void launch_mma_decode(AttentionParams<bf16>& p, DecodeScratch& sc) {
|
||||
int tiles_total = (p.kv_len + BC - 1) / BC;
|
||||
p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total);
|
||||
p.o_part = sc.o_part;
|
||||
p.ml_part = sc.ml_part;
|
||||
|
||||
attn_decode_split_kv_mma_kernel<HEAD_DIM, BC>
|
||||
attn_decode_split_kv_mma_kernel<HEAD_DIM, BC, STAGES>
|
||||
<<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
|
||||
attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user