FlashMLA sm_120 kernel build with 2-3x performance increase from SPDA

Hey guys I have been working on a open source LLM development repo. Anyone can train any amount of parameter against any arch (GQA, MLA, Dense...) When I tried to implement MLA I read the technical research paper on deepseek implementation and found they used FlashMLA for fast training and inference. The only thing is that it was compiled for sm_100 and sm_90 only and I couldn't find anyone that had tried to build it for consumer grade blackwell sm_120. https://github.com/IISuperluminaLII/FlashMLA_Windows_Linux_sm120 *I am not creative with names Inference FlashMLA vs PyTorch SDPA benchmarks Inference / Serving Workload FlashMLA SDPA Speedup Sparse FP8 decode — b=128, s_q=2, topk=2048 0.809 ms 2.118 ms (gather + math) 2.62× Sparse serving — b=4, s_q=1 (CFG=4, warm) 0.050 ms 0.257 ms (CFG=1 as proxy) ~5× Sparse prefill forward — s_q=512, s_kv=8192 1.240 ms 3.232 ms 2.61× Dense decode — H=22, s_q=1, 4K cache (CFG=4) 0.440 ms / 1394 GB/s No equivalent PyTorch path — Model-level BF16-cache decode step ~Parity ~Parity ~1.0× FP8 KV cache, with FlashMLA engaged 8.0% lower latency, 1.84× lower cache memory 1.8% higher latency — Training For my use case — and probably a lot more people soon. Forward + backward using the model's actual attention shape: 192/128, H=22. Workload FlashMLA SDPA Speedup Dense S=4096 3.630 ms 8.696 ms 2.40× Dense S=8192 9.911 ms 30.105 ms 3.04× Dense S=1024 (warm-clock run) 0.306 ms 1.007 ms 3.29× Sparse prefill — s_q=512, topk=2048 6.651 ms 20.054 ms 3.01× At the full-model level, BF16-cache decode is basically at parity, so I wouldn't interpret the kernel-level numbers as an automatic end-to-end 3× model speedup. But for attention-heavy workloads — particularly long-context training and sparse prefill — the difference is substantial.