mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-06-29 19:28:33 +00:00
Enable 16x16 MFMA for decode by making the softmax cross-warp reduction conditional on the warp gemm M dimension: use permlane32_swap for 32x32 MFMA (2 lanes per row), fall back to block_tile_reduce_sync for 16x16 MFMA (4 lanes per row). New tiny decode traits: 1 warp, sequence<1,1,1>, warp_gemm 16x16x32, kBlockM=16, kBlockQ=2 for GQA-8. This matches Triton's BLOCK_M=16 / BLOCK_Q=2 decode configuration exactly. Also adds 4-tier dispatch: tiny (avg_q<=2) -> small (avg_q<=8) -> medium (avg_q<=128) -> large (prefill). Benchmark results (d64 GQA-8 via aiter, 363 shapes): Before: CK faster 135 (37.2%), Triton faster 228 (62.8%) After: CK faster 247 (68.0%), Triton faster 116 (32.0%) Key shapes: 1-seq decode: 0.021ms (CK 0.75x, wins 25%) 64-seq decode: 0.025ms vs Triton 0.029ms (CK wins 14%) 512-seq decode: 0.018ms vs Triton 0.021ms (CK wins) Weighted end-to-end: CK/Triton = 0.999x (tied) Verified correct on 10 shapes: bf16+fp16, d64 GQA-8 + d128 MHA, batch 1-64, all 4 dispatch tiers. Made-with: Cursor