fix(modelopt): dispatch NVFP4 MoE on the cached backend, not the live global

[divyvasal]

Motivation

Fixes #38795.

ModelOptNvFp4FusedMoEMethod.apply resolves the MoE backend twice from two different sources, and under speculative decoding they disagree.

create_weights caches it:

self._moe_runner_backend = moe_runner_backend

apply reads that cache for the error message and for the marlin/megamoe checks:

moe_runner_backend = getattr(self, "_moe_runner_backend", get_moe_runner_backend())

but the flashinfer dispatch guards re-read the live global instead:

@property
def enable_flashinfer_cutlass_moe(self) -> bool:
    return get_moe_runner_backend().is_flashinfer_cutlass()

With --speculative-moe-runner-backend, the global is the draft model's backend while the target layer's cache is still the user's choice. Every flashinfer branch is skipped and the method falls through to:

NotImplementedError: Unsupported moe_runner_backend for NVFP4 MoE:
MoeRunnerBackend.FLASHINFER_CUTLASS. Use --moe-runner-backend flashinfer_cutlass instead.

The message names the backend already in use, because it prints the cache while the guard read the global.

Repro

nvidia/GLM-5.3-Flash-NVFP4, 8×B300 (sm_103), TP8/EP8:

--quantization modelopt_fp4 --moe-runner-backend flashinfer_cutlass
--speculative-algorithm EAGLE --speculative-num-steps 5 --speculative-eagle-topk 1
--speculative-num-draft-tokens 6 --speculative-adaptive
--speculative-moe-runner-backend triton

Dies at build_adaptive_runtime_state -> decode_cuda_graph_runner -> run_moe_core -> modelopt_quant.apply. Dropping --speculative-adaptive (fixed-step EAGLE, everything else identical) boots healthy and serves, which is why this only surfaces on the adaptive path.

Modifications

Point the three flashinfer guards at the local moe_runner_backend the function already computes, matching the is_marlin() and is_flashinfer_megamoe() checks a few lines above:

-        if self.enable_flashinfer_trtllm_moe and hasattr(layer, "g1_scale_c"):
+        if moe_runner_backend.is_flashinfer_trtllm() and hasattr(layer, "g1_scale_c"):
-        if self.enable_flashinfer_cutedsl_moe:
+        if moe_runner_backend.is_flashinfer_cutedsl():
-        if self.enable_flashinfer_cutlass_moe:
+        if moe_runner_backend.is_flashinfer_cutlass():

The branch taken and the backend reported are now the same value by construction. A comment at the local records the invariant so the global does not creep back in.

The properties are left in place; they have other callers outside this method.

Checklist

• Behaviour is unchanged when the global and the cache agree, which is every non-speculative path and every fixed-step speculative path.
• No new lint or formatting findings on the changed lines (the file's existing black/ruff findings are pre-existing on main).
• I could not add a unit test that exercises this without a Blackwell GPU: the divergence only appears during adaptive CUDA-graph capture. Happy to add one if a maintainer can point me at the right harness.

---

CI States

Latest PR Test (Base): ✅ Run #35291501780
Latest PR Test (Extra): ✅ Run #35360849765
Latest PR Test (AMD ROCm 10): ❌ Run #35291501721

[divyvasal]

Could a maintainer add the run-ci label? All 16 red checks here are the pr-gate job exiting on Missing required label 'run-ci', and the *-finish jobs cascading from it — no test has actually run against the change yet.

The diff is 3 lines in modelopt_quant.py, pointing the flashinfer dispatch guards at the local moe_runner_backend the function already computes, so it matches the is_marlin() / is_flashinfer_megamoe() checks just above. Behaviour is identical wherever the cached backend and the live global agree, which is every non-speculative path and every fixed-step speculative path.

Happy to rebase if that helps.

[nvpohanh]

Could you make the comment less verbose? Thanks

[divyvasal]

Shortened to one line in 3338742, thanks.

[nvpohanh]

[by Codex] Posting two inline review findings from the revalidated dry-run review.

[nvpohanh]

[by Codex] Severity: functional | Confidence: High

Issue

This guard no longer recognizes flashinfer_trtllm_routed. ModelOptNvFp4FusedMoEMethod.__init__ deliberately enables the same weight-preparation path for both regular and routed TRT-LLM, and the runner registers FP4 support for the routed backend. With a cached FLASHINFER_TRTLLM_ROUTED backend, this new condition is false, every later backend condition is also false, and apply() raises NotImplementedError.

Fix

Keep using the cached backend, but also accept moe_runner_backend.is_flashinfer_trtllm_routed(). Add a regression case that verifies routed TRT-LLM selects this quant-info branch.

[divyvasal]

Fixed in 45ae650: the guard accepts is_flashinfer_trtllm_routed() too, with a CPU regression test for the routed dispatch.

[nvpohanh]

[by Codex] Severity: style | Confidence: High

Issue

The modified guard retains hasattr, which is not allowed in this code path. It also turns a missing required TRT-LLM preparation field into a misleading unsupported-backend error.

Fix

Declare and initialize g1_scale_c as part of the TRT-LLM layer/setup structure, then select the branch only from the explicit backend predicates and access layer.g1_scale_c directly.

[divyvasal]

Dropped the hasattr in 45ae650: the branch selects on the backend predicates only and reads layer.g1_scale_c directly; a layer that skipped the TRT-LLM prep now fails on the missing field.

[divyvasal]

Thanks, both findings were right. Fixed in 45ae650.

Routed TRT-LLM: enable_flashinfer_trtllm_moe had covered both flashinfer_trtllm and flashinfer_trtllm_routed; my guard kept only the former, so a cached routed backend fell through to NotImplementedError. The guard now accepts both predicates, matching __init__ and the weight prep.

hasattr: dropped. align_fp4_moe_weights_for_flashinfer_trtllm always registers g1_scale_c, so the branch selects on the backend predicates only and reads layer.g1_scale_c directly. A layer that skipped the TRT-LLM prep now fails on the missing field rather than reporting the backend as unsupported.

Regression test: test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_dispatch.py (CPU, base-a-test-cpu): routed and regular TRT-LLM select FlashInferTrtllmFp4MoeQuantInfo with the layer's g1_scale_c; the backend cached by create_moe_runner wins over the live global (auto / cutlass / cutedsl); a missing prep raises on g1_scale_c. Passes on CPU against this branch and fails with the NotImplementedError above against the previous guard.

[nvpohanh]

/tag-and-rerun-ci

[nvpohanh]

/tag-and-rerun-ci

[divyvasal]

Thanks for the rerun. On "PR Test Base" it retried only the failed jobs (wait-for-base-b, pr-test-finish), so the two shards that were cancelled when the run aged out — base-c-test-4-gpu-b200 (4) and base-c-test-8-gpu-h20 (2) — did not run again, and pr-test-finish is red only because of them. Everything else in Base is green (95/97 GPU shards, including the other four B200 shards, B300, GB300, H100, H200 and H20 0–1). A full rerun of that workflow, or of those two jobs, should clear it. The PR Test Extra gate just wants the run-ci-extra label if you want that suite too.

[nvpohanh]

/rerun-failed-ci

[divyvasal]

Status after the rerun, in case it helps:

Base (attempt 2): all 17 base-c-* shards were cancelled at 06:35:06Z while still queued — they never picked up a runner (queued 06:31:17Z). pr-test-finish then exits on the first non-success, so Base is red purely from that. Everything that actually ran is green: base-a-test-cpu (10/10 shards), base-a-test-1-gpu-small, base-b 1-/2-/4-GPU including B200, plus NPU, XPU, Xeon, MUSA, Arm64, sgl-router and Lint. The CPU unit test this PR adds passed: test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_dispatch.py — {"passed": true, "elapsed": 6}.

Extra: the gate exits by design — Missing required label 'run-ci-extra'. Could you run /tag-and-rerun-ci extra so that suite is covered before merge? This PR changes the NVFP4 MoE dispatch path, so the extra Blackwell coverage seems worth having on the record.

MLX / ROCm: both are currently failing on main independently of this PR — MLX on TestOverlapLoopGracefulExit.test_loop_exits_after_shutdown_req (test/registered/unit/hardware_backend/mlx/test_scheduler_mixin.py), ROCm on test_basic_sanity.py::test_accuracy_floor, which 404s fetching hellaswag_val.jsonl from raw.githubusercontent.com/rowanz/hellaswag.

So the two things outstanding are a Base run where the base-c shards get runners, and the run-ci-extra label. Happy to rebase if that makes the GPU shards easier to schedule.

[nvpohanh]

/rerun-failed-ci

[divyvasal]

The two red B200 shards aren't test failures — both died in actions/checkout before any test ran:

fatal: remote error: upload-pack: not our ref b8858a5904c6536435f8ed3ecbf1e7eab0d30747

That SHA is the old refs/pull/38932/merge commit. It has since been regenerated — the merge ref is now 6c9a25f7 and b8858a59 no longer exists (the API returns "No commit found"). The cause looks like queue latency: every base-c shard was queued at 07:59:31Z, and these two didn't get runners until 14:56Z and 16:47Z, by which point main had moved on and the merge commit they were pinned to had been garbage-collected. Each retried the fetch three times and got the same error.

Everything that picked up a runner before ~12:30Z passed, including the on-device Blackwell coverage this PR cares about: 4-gpu-b200 (0) and (1), 8-gpu-b300 (0)/(1), 4-gpu-gb300, 4-gpu-h100 (0)-(2), 8-gpu-h200 (0)-(2), 8-gpu-h20 (0)-(2).

4-gpu-b200 (4) is still queued from 07:59Z and will hit the same checkout error whenever it starts. A /rerun-failed-ci should be cheap now — only those two (soon three) jobs are in a failed state, so it would re-run just them against the current merge ref rather than the whole matrix.

Also, when you get a chance: could you add run-ci-extra (/tag-and-rerun-ci extra)? That gate fails by design without the label, and since this PR changes the NVFP4 MoE dispatch path it seems worth having the extra Blackwell coverage on the record before merge.

@nvpohanh

[nvpohanh]

this small change doesn't need extra-ci. You can ignore the extra-ci failures.

[b8zhong]

/rerun-failed-ci

[b8zhong]

/rerun-test test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_dispatch.py

[github-actions]

Results for /rerun-test test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_dispatch.py:

🚀 ubuntu-latest (1 test): ✅ View workflow run

cd test/ && python3 registered/unit/layers/quantization/test_modelopt_nvfp4_moe_dispatch.py

[divyvasal]

@b8zhong @nvpohanh Could one of you merge this when CI allows? It is approved and mergeable; remaining red checks are unrelated runner/main failures, and @nvpohanh already noted extra-CI is not needed. Happy to rebase if that helps scheduling.