[leonardHONG]

Motivation

MiMo DFlash can fall back to Triton during automatic backend selection on SM100-family GPUs. Use FA4 consistently for this model family. Backend A/B tests on 8×B300 show 1.85× DFlash decode throughput with FA4. Ordinary decoding also improves by 9.0%, with 34.8% lower TTFT.

Modifications

• Default MiMo attention to FA4 on SM100/SM103 through the existing model override.
• Preserve explicit attention, prefill, and decode backend choices.
• Extend the existing MiMo configuration test. The runtime change is five added lines.

Accuracy Tests

• All three MiMo configuration tests pass.
• 12/12 ordinary-decoding requests produce identical output token IDs with FA4 and Triton.
• 16 BF16 attention cases pass against an FP32 reference, covering the TP8 per-rank shapes for hidden sizes 4096 and 6144, QK=192/V=128, full attention and SWA with sinks, decode, 8-token verify, ragged page boundaries, and CUDA Graph replay.
• Ruff, formatting, and import-order checks pass.

Speed Tests and Profiling

XiaomiMiMo/MiMo-V2.6-Pro-RL@73875d00b30a, 8×B300, TP8, --moe-runner-backend flashinfer_mxfp4, BF16 KV, concurrency 1, 8192 input / 1024 output tokens. Three warmup requests followed by three rounds of four measured requests per backend; CUDA Graphs enabled, prefix caching and communication fusion disabled.

Ordinary decoding, no DFlash, page size 128 and 32,768 KV slots for both backends (benchmark revision 7faae89732):

Metric	Triton	FA4	Improvement
Decode throughput	159.47 tok/s	173.86 tok/s	+9.0%
Mean TTFT	222.52 ms	145.17 ms	−34.8%
Mean request latency	6637.36 ms	6029.08 ms	−9.2%

DFlash8 enabled, target and draft using the same attention backend:

Metric	Triton / page1	FA4 / page128	Improvement
Decode throughput	504.56 tok/s	931.40 tok/s	+84.6% (1.85×)
Mean TTFT	234.81 ms	150.93 ms	−35.7%
Mean request latency	2262.30 ms	1249.28 ms	−44.8%

Both configurations use 32,768 KV slots, and all 12 measured requests produce identical output token IDs.

Checklist

• Code formatting and style checks pass.
• Existing unit tests updated and validated.
• Correctness and speed results provided.

---

CI States

Latest PR Test (Base): ❌ Run #36716152960
Latest PR Test (Extra): 🚫 Run #36728622666
Latest PR Test (AMD ROCm 10): ❌ Run #36716152990