From 3d9b915c17beb2356ac359f4bfd1e22d4834e416 Mon Sep 17 00:00:00 2001 From: Divy Date: Thu, 10 Sep 2026 11:34:11 -0700 Subject: [PATCH 1/3] fix(modelopt): dispatch NVFP4 MoE on the cached backend, not the live global `ModelOptNvFp4FusedMoEMethod.apply` resolves the MoE backend twice from two different sources. `create_weights` caches it as `self._moe_runner_backend`, and `apply` reads that cache for its error message and for the marlin/megamoe checks, but the flashinfer dispatch guards use properties that re-read the *global* `get_moe_runner_backend()`. Under speculative decoding those disagree. `--speculative-moe-runner-backend` makes the global the draft model's backend while the target layer's cache is still the one the user asked for, so 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. which names the backend already in use, because the message prints the cache while the guard read the global. Reported in #38795. Reproduces on nvidia/GLM-5.3-Flash-NVFP4, 8xB300, TP8, with `--moe-runner-backend flashinfer_cutlass --speculative-adaptive`, at adaptive CUDA-graph capture (`build_adaptive_runtime_state`). Fixed-step EAGLE with the same backend is unaffected, which is why this only shows up with adaptive. Point the three flashinfer guards at the local `moe_runner_backend` the function already computes, matching the marlin and megamoe checks a few lines above. The branch taken and the backend reported are now the same value by construction. Co-Authored-By: Claude Opus 5 (1M context) --- .../sglang/srt/layers/quantization/modelopt_quant.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index bcad61e6fc83..f5e9bba1459e 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -2952,6 +2952,11 @@ def apply( # tuple). Defer per-attribute access to the branches that actually # consume them. activation = self.moe_runner_config.activation + # Resolved once in create_weights and cached here. Dispatch MUST use this, not + # the live global: under speculative decoding the global is the draft model's + # backend (--speculative-moe-runner-backend), so a global read disagrees with + # this layer's own backend, falls through to the NotImplementedError below, and + # then reports this cached value -- naming the backend already in use. moe_runner_backend = getattr( self, "_moe_runner_backend", get_moe_runner_backend() ) @@ -2985,7 +2990,7 @@ def apply( return self.runner.run(dispatch_output, quant_info) # FlashInfer TRTLLM FP4 path - 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"): from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( FlashInferTrtllmFp4MoeQuantInfo, ) @@ -3021,7 +3026,7 @@ def apply( return self.runner.run(dispatch_output, quant_info) - if self.enable_flashinfer_cutedsl_moe: + if moe_runner_backend.is_flashinfer_cutedsl(): from sglang.srt.layers.moe.moe_runner.flashinfer_cutedsl import ( CuteDslFp4MoeQuantInfo, ensure_cutedsl_wrapper, @@ -3077,7 +3082,7 @@ def apply( ) return self.runner.run(dispatch_output, quant_info) - if self.enable_flashinfer_cutlass_moe: + if moe_runner_backend.is_flashinfer_cutlass(): from sglang.srt.layers.moe.moe_runner.flashinfer_cutlass import ( FlashInferCutlassMoeQuantInfo, ) From 3338742ff080e4bd32c3af1c99e34936d301be24 Mon Sep 17 00:00:00 2001 From: Divy Date: Fri, 11 Sep 2026 09:44:16 -0700 Subject: [PATCH 2/3] fix(modelopt): shorten the cached-backend comment Review nit: trim the explanation to one line. Co-Authored-By: Claude Opus 5 (1M context) --- python/sglang/srt/layers/quantization/modelopt_quant.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index f5e9bba1459e..05e67aac7fd7 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -2952,11 +2952,7 @@ def apply( # tuple). Defer per-attribute access to the branches that actually # consume them. activation = self.moe_runner_config.activation - # Resolved once in create_weights and cached here. Dispatch MUST use this, not - # the live global: under speculative decoding the global is the draft model's - # backend (--speculative-moe-runner-backend), so a global read disagrees with - # this layer's own backend, falls through to the NotImplementedError below, and - # then reports this cached value -- naming the backend already in use. + # Use the cached backend: the global differs under speculative decoding. moe_runner_backend = getattr( self, "_moe_runner_backend", get_moe_runner_backend() ) From 45ae650d6e47da4be4b09fde67421da90a8df6c4 Mon Sep 17 00:00:00 2001 From: Divy Date: Sun, 13 Sep 2026 20:10:48 -0700 Subject: [PATCH 3/3] fix(modelopt): keep routed TRT-LLM on the TRT-LLM path; drop the hasattr guard The cached-backend guard replaced enable_flashinfer_trtllm_moe, which covered both flashinfer_trtllm and flashinfer_trtllm_routed, with is_flashinfer_trtllm() alone. A layer whose cached backend is the routed one then fell through every branch and apply() raised NotImplementedError. The guard now accepts both, as __init__ and the weight prep already do. The hasattr(layer, "g1_scale_c") check goes too. The TRT-LLM weight prep (align_fp4_moe_weights_for_flashinfer_trtllm) always registers g1_scale_c, so the branch is selected on the backend predicates only and reads layer.g1_scale_c directly; a layer that skipped the prep now fails on the missing field instead of reporting the backend as unsupported. Adds a CPU unit test that pins the dispatch: routed and regular TRT-LLM select FlashInferTrtllmFp4MoeQuantInfo, the cached backend wins over the live global, and a missing prep names g1_scale_c. Co-Authored-By: Claude Fable 5.1 --- .../srt/layers/quantization/modelopt_quant.py | 7 +- .../test_modelopt_nvfp4_moe_dispatch.py | 186 ++++++++++++++++++ 2 files changed, 191 insertions(+), 2 deletions(-) create mode 100644 test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_dispatch.py diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 05e67aac7fd7..120a160d4564 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -2985,8 +2985,11 @@ def apply( quant_info = self.get_marlin_quant_info(layer) return self.runner.run(dispatch_output, quant_info) - # FlashInfer TRTLLM FP4 path - if moe_runner_backend.is_flashinfer_trtllm() and hasattr(layer, "g1_scale_c"): + # FlashInfer TRTLLM FP4 path (routed shares the weight prep and the runner) + if ( + moe_runner_backend.is_flashinfer_trtllm() + or moe_runner_backend.is_flashinfer_trtllm_routed() + ): from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( FlashInferTrtllmFp4MoeQuantInfo, ) diff --git a/test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_dispatch.py b/test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_dispatch.py new file mode 100644 index 000000000000..99c203176d18 --- /dev/null +++ b/test/registered/unit/layers/quantization/test_modelopt_nvfp4_moe_dispatch.py @@ -0,0 +1,186 @@ +"""CPU unit tests for NVFP4 fused-MoE backend dispatch in ModelOptNvFp4FusedMoEMethod.apply. + +apply() picks the kernel path from the backend the method cached in +create_moe_runner, not from the process-wide MoE runner backend, which +speculative decoding changes after the weights were prepared. These tests pin +that for the FlashInfer TRT-LLM path, which serves the regular and the routed +TRT-LLM backend from one weight prep. + +The platform check and the runner are stubbed and the layer is a bag of small +tensors, so the tests stay on CPU; the kernels are covered on-device. +""" + +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=6, suite="base-a-test-cpu") + +import unittest +from contextlib import contextmanager +from types import SimpleNamespace +from unittest import mock + +import torch + +# Import modelopt_quant before flashinfer_trtllm (see +# test_modelopt_nvfp4_moe_scales.py for the circular-import reason). +# isort: off +from sglang.srt.layers.quantization import modelopt_quant +from sglang.srt.layers.quantization.modelopt_quant import ModelOptNvFp4FusedMoEMethod +from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( + FlashInferTrtllmFp4MoeQuantInfo, +) + +# isort: on +from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig +from sglang.srt.layers.moe.utils import MoeRunnerBackend +from sglang.test.test_utils import CustomTestCase + +NUM_EXPERTS = 8 +HIDDEN = 64 +INTERMEDIATE = 32 + + +class _Runner: + """Records the quant_info apply() hands to the runner.""" + + def __init__(self): + self.calls = [] + + def run(self, dispatch_output, quant_info): + self.calls.append((dispatch_output, quant_info)) + return "combine-input" + + +def _trtllm_prepared_layer() -> SimpleNamespace: + """A layer as align_fp4_moe_weights_for_flashinfer_trtllm leaves it: packed FP4 + weights, FP8 block scales and the TRT-LLM output scalars, incl. g1_scale_c.""" + + def p(t: torch.Tensor) -> torch.nn.Parameter: + return torch.nn.Parameter(t, requires_grad=False) + + return SimpleNamespace( + w13_weight=p( + torch.zeros(NUM_EXPERTS, 2 * INTERMEDIATE, HIDDEN // 2, dtype=torch.uint8) + ), + w2_weight=p( + torch.zeros(NUM_EXPERTS, HIDDEN, INTERMEDIATE // 2, dtype=torch.uint8) + ), + w13_weight_scale=p( + torch.zeros( + NUM_EXPERTS, 2 * INTERMEDIATE, HIDDEN // 16, dtype=torch.float8_e4m3fn + ) + ), + w2_weight_scale=p( + torch.zeros( + NUM_EXPERTS, HIDDEN, INTERMEDIATE // 16, dtype=torch.float8_e4m3fn + ) + ), + g1_scale_c=p(torch.full((NUM_EXPERTS,), 0.25, dtype=torch.float32)), + g1_alphas=p(torch.full((NUM_EXPERTS,), 0.5, dtype=torch.float32)), + g2_alphas=p(torch.full((NUM_EXPERTS,), 0.75, dtype=torch.float32)), + w13_input_scale_quant=torch.tensor(2.0, dtype=torch.float32), + num_experts=NUM_EXPERTS, + num_local_experts=NUM_EXPERTS, + moe_ep_rank=0, + intermediate_size_per_partition=INTERMEDIATE, + ) + + +@contextmanager +def _live_backend(backend: MoeRunnerBackend): + """The process-wide MoE runner backend, as the method sees it.""" + with mock.patch.object( + modelopt_quant, "get_moe_runner_backend", return_value=backend + ): + yield + + +def _method_set_up_for(backend: MoeRunnerBackend) -> ModelOptNvFp4FusedMoEMethod: + """A method constructed and given its runner while ``backend`` was the live one.""" + with ( + _live_backend(backend), + mock.patch.object( + modelopt_quant, + "get_platform", + return_value=SimpleNamespace(is_blackwell=True), + ), + mock.patch.object(modelopt_quant, "is_cuda", return_value=False), + ): + method = ModelOptNvFp4FusedMoEMethod( + SimpleNamespace(use_per_token_activation=False) + ) + method.create_moe_runner( + SimpleNamespace(), + MoeRunnerConfig( + num_experts=NUM_EXPERTS, + num_local_experts=NUM_EXPERTS, + hidden_size=HIDDEN, + intermediate_size_per_partition=INTERMEDIATE, + activation="silu", + is_gated=True, + ), + ) + method.runner = _Runner() + return method + + +class TestNvFp4MoeDispatch(CustomTestCase): + def test_routed_trtllm_takes_the_trtllm_path(self): + for backend in ( + MoeRunnerBackend.FLASHINFER_TRTLLM_ROUTED, + MoeRunnerBackend.FLASHINFER_TRTLLM, + ): + with self.subTest(backend=backend.value): + method = _method_set_up_for(backend) + layer = _trtllm_prepared_layer() + dispatch_output = object() + + with _live_backend(backend): + out = method.apply(layer, dispatch_output) + + self.assertEqual(out, "combine-input") + ((seen_dispatch, quant_info),) = method.runner.calls + self.assertIs(seen_dispatch, dispatch_output) + self.assertIsInstance(quant_info, FlashInferTrtllmFp4MoeQuantInfo) + self.assertEqual( + quant_info.g1_scale_c.data_ptr(), layer.g1_scale_c.data_ptr() + ) + self.assertEqual( + quant_info.g1_alphas.data_ptr(), layer.g1_alphas.data_ptr() + ) + self.assertEqual(quant_info.local_num_experts, NUM_EXPERTS) + self.assertEqual( + quant_info.intermediate_size_per_partition, INTERMEDIATE + ) + + def test_dispatch_follows_the_backend_the_runner_was_created_for(self): + method = _method_set_up_for(MoeRunnerBackend.FLASHINFER_TRTLLM_ROUTED) + for live in ( + MoeRunnerBackend.AUTO, + MoeRunnerBackend.FLASHINFER_CUTLASS, + MoeRunnerBackend.FLASHINFER_CUTEDSL, + ): + with self.subTest(live=live.value): + method.runner = _Runner() + + with _live_backend(live): + method.apply(_trtllm_prepared_layer(), object()) + + ((_, quant_info),) = method.runner.calls + self.assertIsInstance(quant_info, FlashInferTrtllmFp4MoeQuantInfo) + + def test_missing_trtllm_prep_names_the_missing_field(self): + method = _method_set_up_for(MoeRunnerBackend.FLASHINFER_TRTLLM_ROUTED) + layer = _trtllm_prepared_layer() + del layer.g1_scale_c + + with _live_backend(MoeRunnerBackend.FLASHINFER_TRTLLM_ROUTED): + with self.assertRaises(AttributeError) as ctx: + method.apply(layer, object()) + + self.assertIn("g1_scale_c", str(ctx.exception)) + self.assertEqual(method.runner.calls, []) + + +if __name__ == "__main__": + unittest.main()