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

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com>
This commit is contained in:
Divy
2026-09-20 19:28:54 -04:00
committed by GitHub
co-authored by Claude Opus 5 Po-Han Huang
parent 2fa6b94e34
commit 42875bcd2a
2 changed files with 194 additions and 4 deletions
@@ -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,
)