[Fix] Route asymmetric-KV models to fa4 on SM100 and pin MiMoV2 FP8 MoE to flashinfer_trtllm (#32818)

This commit is contained in:
Liangsheng Yin
2026-07-29 16:37:43 -07:00
committed by GitHub
parent 3c1717d9b6
commit e5c46ff07d
6 changed files with 88 additions and 21 deletions
@@ -79,7 +79,9 @@ import { MiMoV25Deployment } from '/src/snippets/autoregressive/mimo-v25-deploym
**MiMo-V2.5 (310B):**
- The checkpoint has a TP=4-interleaved fused `qkv_proj`; attention-TP per DP group **must** be 4. Use `--dp = TP / 4`; for TP > 4 this also requires DP-attention. Total GPUs must be a multiple of 4. A bare `--tp 8` without `--dp 2` will fail to load with `MiMoV2 fused qkv_proj checkpoint is TP=4-interleaved; got attention tp_size=8`.
- Single-node deployments: H100/H200 8× GPUs (`--tp 8 --dp 2`), B200 4× GPUs (`--tp 4`, dp=1, no DP-attn flag needed), GB300 4× GPUs (`--tp 4`, single NVL4 node). FP8 quantization.
- On Blackwell, pass `--attention-backend fa4`: MiMoV2's asymmetric KV (`head_dim` 192 / `v_head_dim` 128) fails on the SM100 default `trtllm_mha`, which requires equal K/V widths.
- On Blackwell, pass `--mm-attention-backend fa4` for the V2.5 vision encoder. The checkpoint config requests FlashAttention-3 internally, but SGLang rejects FA3 on Blackwell and expects FA4 for multimodal attention.
- On Blackwell, pass `--moe-runner-backend flashinfer_trtllm`; the default `auto` falls through to the triton fused-MoE runner, ~12% slower at bs=1 decode.
- `--enable-dp-lm-head` and `--mm-enable-dp-encoder` are required whenever `--enable-dp-attention` is on, to keep LM head and encoder sharding consistent.
- EAGLE MTP uses the checkpoint's MTP weights. Enable with `--speculative-algorithm EAGLE` and `--enable-multi-layer-eagle` (both Hopper and Blackwell).
- **Multimodal**: Supports image, video, and audio understanding; see Section 4.3 for invocation examples.
@@ -321,9 +321,9 @@ export const MiMoV25Deployment = () => {
if (multinode) flags.push(...multiNodeFlags(nnodes));
// MoE backend: Blackwell uses flashinfer_trtllm (hardware-driven); Hopper
// MoE backend: Blackwell uses flashinfer_trtllm (both variants); Hopper
// optionally uses DeepEP (toggle).
if (isPro && blackwell) {
if (blackwell) {
flags.push(" --moe-runner-backend flashinfer_trtllm");
} else if (useDeepep) {
flags.push(" --moe-a2a-backend deepep");
@@ -349,7 +349,11 @@ export const MiMoV25Deployment = () => {
flags.push(` --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 64}'`);
}
} else {
if (blackwell) flags.push(" --mm-attention-backend fa4");
if (blackwell) {
// fa4 is required, not tuning: trtllm_mha rejects MiMoV2's 192/128 KV.
flags.push(" --attention-backend fa4");
flags.push(" --mm-attention-backend fa4");
}
flags.push(" --mem-fraction-static 0.65");
flags.push(" --chunked-prefill-size 16384");
}