diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index 24cc6adec..5a15eeecf 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -3045,6 +3045,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): # tuple). Defer per-attribute access to the branches that actually # consume them. activation = self.moe_runner_config.activation + # Use the cached backend: the global differs under speculative decoding. moe_runner_backend = getattr( self, "_moe_runner_backend", get_moe_runner_backend() ) @@ -3077,8 +3078,11 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): quant_info = self.get_marlin_quant_info(layer) return self.runner.run(dispatch_output, quant_info) - # FlashInfer TRTLLM FP4 path - if self.enable_flashinfer_trtllm_moe 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, ) @@ -3114,7 +3118,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): 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, @@ -3170,7 +3174,7 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): ) 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, ) 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 000000000..99c203176 --- /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()