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/kernels/attn_paged_decode.cu

@@ -33,7 +33,7 @@ static void launch_paged_scalar_decode(PagedAttentionParams<bf16>& p) {
 }
 
 #ifndef ASTRAI_NO_MMA
-template <int HEAD_DIM, int BC>
+template <int HEAD_DIM, int BC, int STAGES = (HEAD_DIM <= 128) ? 2 : 1>
 static void launch_paged_mma_decode(PagedAttentionParams<bf16>& p) {
     int tiles_total = (p.kv_len + BC - 1) / BC;
     p.num_splits = paged_decode_num_splits(p.batch * p.kv_head, tiles_total);
@@ -44,8 +44,7 @@ static void launch_paged_mma_decode(PagedAttentionParams<bf16>& p) {
     p.o_part = o_part.data_ptr<float>();
     p.ml_part = ml_part.data_ptr<float>();
 
-    paged_attn_decode_split_kv_mma_kernel<HEAD_DIM, BC>
-        <<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
+    paged_attn_decode_split_kv_mma_kernel<HEAD_DIM, BC, STAGES><<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
     paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
 }
 #endif