B300 NextN fp4: review cleanups — keep else-slice, non-routed draft trtllm, weight_scale_2 probe

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) <noreply@anthropic.com>

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

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
     )
 

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"