From c1637cfb68518a71bc003189e1d4461d7865d53d Mon Sep 17 00:00:00 2001 From: leavelet Date: Sun, 21 Jun 2026 17:48:06 +0000 Subject: [PATCH] =?UTF-8?q?B300=20NextN=20fp4:=20review=20cleanups=20?= =?UTF-8?q?=E2=80=94=20keep=20else-slice,=20non-routed=20draft=20trtllm,?= =?UTF-8?q?=20weight=5Fscale=5F2=20probe?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adversarial review of b3a4f4174f found 3 items (Bug B 2b8593ff24 accepted as-is): 1) Re-add the modelopt_quant else-branch input-scale EP-slice (reverting it re-opened a general latent EP>1 NVFP4 crash on the plain/triton path: --moe-runner-backend triton --quantization modelopt_fp4 --ep>1). No-op at ep==1; orthogonal defense-in-depth (draft now takes the trtllm scalar branch). 2) Draft uses the NON-routed flashinfer_trtllm (FlashInferFP4MoE only does the non-routed kernel; the _routed impl-class isn't wired). Reject explicit --speculative-moe-runner-backend flashinfer_trtllm_routed. 3) Tighten the checkpoint probe to key on weight_scale_2 (fp4-specific) instead of the generic weight_scale substring (which also matches fp8 weight_scale_inv). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../srt/layers/quantization/modelopt_quant.py | 13 +++++++++++++ python/sglang/srt/model_loader/weight_utils.py | 12 +++++++----- python/sglang/srt/server_args.py | 14 ++++++++++++-- 3 files changed, 32 insertions(+), 7 deletions(-) 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"