Skip to content

MiMoV2: do not force flashinfer_trtllm when an EP a2a backend is active - #39079

Open
whn09 wants to merge 1 commit into
sgl-project:mainfrom
whn09:fix/mimo-v2-trtllm-override-respects-a2a
Open

MiMoV2: do not force flashinfer_trtllm when an EP a2a backend is active#39079
whn09 wants to merge 1 commit into
sgl-project:mainfrom
whn09:fix/mimo-v2-trtllm-override-respects-a2a

Conversation

@whn09

@whn09 whn09 commented Sep 11, 2026

Copy link
Copy Markdown
Contributor

Motivation

The MiMoV2 config-time override forces moe_runner_backend=flashinfer_trtllm on sm100 FP8
without checking whether an EP all-to-all backend is active:

if (
    get_platform().is_sm100
    and cfg.moe_runner_backend == "auto"
    and get_quantization_config(hf_config) == "fp8"
):
    overrides["moe_runner_backend"] = "flashinfer_trtllm"

FusedOpPool has no flashinfer_trtllm fused func registered for deepep or deepep_v2,
and flashinfer_trtllm has no runner_core (fused path only), so MoeRunner.__init__ hits
runner_core is None and fused_func is None:

NotImplementedError: Runner backend MoeRunnerBackend.FLASHINFER_TRTLLM requires a fused func
for a2a backend deepep_v2, but none is registered.

The user asked for auto. auto resolves to deep_gemm, which works. The override turns a
working default into a hard failure at model init, and because the substitution is silent the
traceback names a backend the command line never mentioned.

Sibling overrides already carry exactly this guard — glm4_moe.py:39, qwen3_moe.py:43,
deepseek_v4.py:62, nemotron_h.py:86 all test cfg.moe_a2a_backend == "none" before
forcing a runner. mimo_v2.py is the one that does not, so this brings it in line rather
than introducing a new convention.

Modifications

One condition, and cfg.moe_a2a_backend == "none", plus a comment explaining why (the next
person to read this will otherwise wonder whether the guard is load-bearing). No behaviour
change without an a2a backend: the sm100 FP8 fast path is untouched for the plain-TP case
the override was written for.

Accuracy Tests

No numerical change. With --moe-a2a-backend none the resolved runner is flashinfer_trtllm
before and after — verified via /get_server_info on 8× B200 with an FP8 MiMoV2 checkpoint.
With --moe-a2a-backend deepep_v2 --moe-runner-backend auto the resolved runner goes from
"init raises" to deep_gemm, whose numerics against the deepep (v1) path are 10/12 prompts
bit-identical on a fixed logprob set (details in the companion PR).

Speed Tests and Profiling

Unchanged where the override still fires. Where it no longer fires, the comparison is
NotImplementedError versus a running server.

Worth stating plainly for anyone weighing it: flashinfer_trtllm really is the faster runner
where it is usable — on decode at global bs=32 it is 24.58 ms/token versus 34.90 for
deep_gemm, ~30%. That is an argument for registering a deepep/deepep_v2 fused func, not
for forcing a runner that cannot load. (The one existing deepep+flashinfer_cutedsl fused
func takes CuteDslFp4MoeQuantInfo, i.e. NVFP4 only, so it does not help an FP8 checkpoint.)

Checklist

Review and Merge Process

  1. Ping Merge Oncalls to start the process. See the PR Merge Process.
  2. Get approvals from CODEOWNERS and other reviewers.
  3. Trigger CI tests with comments or contact authorized users to do so.
  4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

I cannot add the run-ci label myself, so this needs someone with write access to start CI.


Companions, independent of each other: #39080 (enable deepep_v2 for MiMoV2 — the PR that
makes this guard reachable), #39078 (ElasticBuffer vs CUDA graph capture), #37211
(ep_scatter_from_psum kernel args).


CI States

Latest PR Test (Base): ❌ Run #34590432083
Latest PR Test (Extra): ❌ Run #34590431824
Latest PR Test (AMD ROCm 10): ❌ Run #34590431968

The sm100 override rewrites moe_runner_backend auto -> flashinfer_trtllm without
looking at moe_a2a_backend. FusedOpPool has no ("deepep", "flashinfer_trtllm") or
("deepep_v2", "flashinfer_trtllm") entry and flashinfer_trtllm's runner_core is
None, so every DeepEP arm on Blackwell dies at model init with

    NotImplementedError: Runner backend MoeRunnerBackend.FLASHINFER_TRTLLM requires
    a fused func for a2a backend deepep, but none is registered

Left alone, "auto" would have resolved to deep_gemm at runtime, which is the correct
and working choice for deepep; the config-time override just gets there first. The
glm4_moe, qwen3_moe, deepseek_v4 and nemotron_h overrides all gate the same sm100
trtllm rewrite on moe_a2a_backend == "none"; this adds the missing clause. Invisible
on Hopper, where the override does not fire.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant