Decode has only batch*kv_head independent tasks, so the grid was tiny (e.g. 16 blocks) leaving most SMs idle (ncu: 0.04 waves/SM, 11% DRAM).
- Partition KV across gridDim.z blocks emitting unnormalised (O, m, l) partials, reduced by a new combine kernel
- Choose split count to fill the device (~2 blocks/SM), capped by tile count and 32; fall back to single-pass direct-write when batch*kv_head already saturates the SMs
- Refactor decode dispatch into named helpers, de-duplicate scalar fallback
Result: now DRAM-bound at 63% (99->543 GB/s), 2.1-2.5x over torch SDPA in the low-parallelism regime, on par at high parallelism
ViperEkura
7ba43a7c6f