From de726b4e7e902cfba6e7841938b8a49a8242383f Mon Sep 17 00:00:00 2001 From: leavelet Date: Sun, 21 Jun 2026 17:02:43 +0000 Subject: [PATCH] B300 NVFP4 MoE: slice global input scales to EP-local on the non-flashinfer process_weights path The NextN/EAGLE draft runs under speculative_moe_backend_context -> MOE_RUNNER_BACKEND=AUTO (trtllm is forbidden as a speculative backend), so its NVFP4 MoE takes the plain-FusedMoE 'else' branch of process_weights_after_loading. That branch left w13_input_scale/w2_input_scale GLOBAL (num_experts=256, they are tagged _sglang_require_global_experts at load) while w13_weight_scale_2 is EP-local (64) -> _compute_gemm1_alphas did 256*64 -> RuntimeError. trtllm/cutlass collapse to a scalar and cutedsl already slices; the else/triton path was a latent EP>1 NVFP4 bug. Slice the global input scales to this rank's local experts (mirrors the cutedsl _slice_scale). No-op for moe_ep_size==1; main model unaffected (never takes this branch). opus-diagnosed. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../sglang/srt/layers/quantization/modelopt_quant.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/python/sglang/srt/layers/quantization/modelopt_quant.py b/python/sglang/srt/layers/quantization/modelopt_quant.py index efa58209b..330fa6c21 100755 --- a/python/sglang/srt/layers/quantization/modelopt_quant.py +++ b/python/sglang/srt/layers/quantization/modelopt_quant.py @@ -2027,6 +2027,17 @@ 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 + # (mirrors the cutedsl branch above). No-op for moe_ep_size == 1. + 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