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

csrc/tests/attn_decode_test.cu

@@ -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);
 }

csrc/tests/attn_paged_decode_test.cu

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