diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index efa58209b..599d199f5 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -2027,6 +2027,19 @@ class ModelOptNvFp4FusedMoEMethod(FusedMoEMethodBase): else: w13_input_scale = layer.w13_input_scale.max(dim=-1).values.to(torch.float32) w2_input_scale = layer.w2_input_scale + # w13_input_scale/w2_input_scale are loaded GLOBAL (num_experts) because + # create_weights tags them _sglang_require_global_experts, while + # w13_weight_scale_2 is EP-local (num_local_experts). Under EP>1 this + # non-flashinfer path must slice the global input scales down to this + # rank's local experts so the gemm1/gemm2 alpha products line up (the + # flashinfer branches above .max()-collapse or slice instead). No-op for + # moe_ep_size == 1. (The fp4 EAGLE draft now takes the trtllm scalar + # branch above; this keeps the plain/triton EP>1 NVFP4 path correct too.) + if layer.moe_ep_size > 1 and w13_input_scale.shape[0] == layer.num_experts: + start = layer.moe_ep_rank * layer.num_local_experts + end = start + layer.num_local_experts + w13_input_scale = w13_input_scale[start:end] + w2_input_scale = w2_input_scale[start:end] if self.quant_config.use_per_token_activation: # FlashInfer computes activation scales dynamically per token, so diff --git a/python/sglang/srt/model_loader/weight_utils.py b/python/sglang/srt/model_loader/weight_utils.py index c54a77308..d1681aaa8 100644 --- a/python/sglang/srt/model_loader/weight_utils.py +++ b/python/sglang/srt/model_loader/weight_utils.py @@ -603,10 +603,12 @@ def nextn_moe_is_fp4_quantized_in_checkpoint( """True/False if the checkpoint's NextN/MTP MoE experts ship fp4 scales, else None. The MTP/NextN decoder lives at HF prefix ``model.layers.{num_hidden_layers}``. - fp4-quantized MoE experts always carry ``weight_scale`` / ``weight_scale_2`` - tensors; an unquantized (bf16) MTP MoE (DeepSeek-R1-FP4 / V3-0324-FP4) does not. - Used to decide both whether to keep modelopt_fp4 for the NextN module and - whether the EAGLE draft should run the fp4-capable trtllm MoE backend. + NVFP4-quantized MoE experts carry a per-tensor ``weight_scale_2`` tensor; an + unquantized (bf16) MTP MoE (DeepSeek-R1-FP4 / V3-0324-FP4) does not. We key on + ``weight_scale_2`` specifically (not the generic ``weight_scale`` substring, + which would also match the fp8 block-quant ``weight_scale_inv``). Used to decide + both whether to keep modelopt_fp4 for the NextN module and whether the EAGLE + draft should run the fp4-capable trtllm MoE backend. """ if not model_dir or num_hidden_layers is None: return None @@ -615,7 +617,7 @@ def nextn_moe_is_fp4_quantized_in_checkpoint( return None expert_prefix = f"model.layers.{num_hidden_layers}.mlp.experts." return any( - name.startswith(expert_prefix) and "weight_scale" in name + name.startswith(expert_prefix) and "weight_scale_2" in name for name in weight_map ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 9f1e54fc8..f4626cb57 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -3318,8 +3318,10 @@ class ServerArgs: if self.speculative_moe_runner_backend is None: if _nextn_is_fp4: - # fp4 NextN -> draft uses the same fp4 trtllm backend as main. - self.speculative_moe_runner_backend = self.moe_runner_backend + # fp4 NextN -> the draft runs the same fp4 trtllm MoE path as main. + # Always the NON-routed kernel: the draft's FlashInferFP4MoE only + # implements trtllm_fp4_block_scale_moe, not the _routed variant. + self.speculative_moe_runner_backend = "flashinfer_trtllm" elif _main_is_trtllm: self.speculative_moe_runner_backend = "auto" else: @@ -3332,6 +3334,14 @@ class ServerArgs: "an fp4 NextN/MTP layer (the bf16 trtllm MoE requires RenormalizeNaive " "routing); use triton or auto for a bf16 MTP." ) + elif MoeRunnerBackend( + self.speculative_moe_runner_backend + ).is_flashinfer_trtllm_routed(): + raise ValueError( + "flashinfer_trtllm_routed is not supported as a speculative MoE runner " + "backend (the draft NextN MoE uses the non-routed fp4 kernel); use " + "flashinfer_trtllm, triton, or auto instead." + ) if self.speculative_algorithm == "NEXTN": self.speculative_algorithm = "EAGLE"