Fix EPLB + NVFP4: exclude broadcast scales and fix stale quant config

Cherry-pick of PR #34646 (elvircrn/fix-eplb-nvfp4-contiguous).

Pre-compute g1/g2 alphas as registered parameters so EPLB rearranges
them alongside expert weights. Without this, the quant config caches
g1_alphas = a_scale * w_scale_2 once at init, and EPLB's in-place
rearrangement of w_scale_2 leaves the cached product stale.

Also excludes broadcast activation scales (w13_input_scale,
w2_input_scale) from EPLB to prevent contiguity assertion crash
(these are expanded stride-0 tensors from .max().expand()).

Co-Authored-By: Elvir Crncevic <elvircrn@users.noreply.github.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

vllm/model_executor/layers/fused_moe/layer.py

@@ -1392,19 +1392,23 @@ class FusedMoE(CustomOp):
         weights = list(self.named_parameters())
         weights = [(name, _maybe_make_contiguous(name, p)) for name, p in weights]
 
+        # `w13_input_scale` and `w2_input_scale` are global per-tensor
+        # activation scales shared across all experts (e.g. NVFP4).
+        # They are broadcast views (stride 0) from .expand() and are
+        # not actual expert weights, so exclude them from EPLB.
+        NON_EXPERT_WEIGHTS = {
+            "e_score_correction_bias",
+            "w13_input_scale",
+            "w2_input_scale",
+        }
+
         assert all(
             weight.is_contiguous()
             for name, weight in weights
             if not (name.startswith("_shared_experts.") or name.startswith("_gate."))
+            and name not in NON_EXPERT_WEIGHTS
         )
 
-        # Filter out the non-expert weights.
-        # `e_score_correction_bias` is a bias for each logical expert,
-        # with shape (num_logical_experts,), not an expert weight.
-        NON_EXPERT_WEIGHTS = {
-            "e_score_correction_bias",
-        }
-
         return [
             weight.view(self.local_num_experts, -1)
             for name, weight in weights

vllm/model_executor/layers/fused_moe/oracle/nvfp4.py

@@ -365,6 +365,8 @@ def make_nvfp4_moe_quant_config(
     w2_scale_2: torch.Tensor,
     a13_scale: torch.Tensor,
     a2_scale: torch.Tensor,
+    g1_alphas: torch.Tensor | None = None,
+    g2_alphas: torch.Tensor | None = None,
 ) -> FusedMoEQuantConfig:
     if backend == NvFp4MoeBackend.MARLIN:
         return nvfp4_w4a16_moe_quant_config(
@@ -374,8 +376,10 @@ def make_nvfp4_moe_quant_config(
             w2_scale=w2_scale,
         )
 
-    g1_alphas = a13_scale * w13_scale_2
-    g2_alphas = a2_scale * w2_scale_2
+    if g1_alphas is None:
+        g1_alphas = a13_scale * w13_scale_2
+    if g2_alphas is None:
+        g2_alphas = a2_scale * w2_scale_2
     return nvfp4_moe_quant_config(
         g1_alphas=g1_alphas,
         g2_alphas=g2_alphas,

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

@@ -554,7 +554,23 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
         layer.w13_input_scale = a13_scale
         layer.w2_input_scale = a2_scale
 
-        # Setup modular kernel.
+        # Pre-compute g1/g2 alphas as registered parameters so EPLB
+        # rearranges them alongside expert weights (see modelopt.py).
+        if self.nvfp4_backend not in (
+            NvFp4MoeBackend.FLASHINFER_TRTLLM,
+            NvFp4MoeBackend.MARLIN,
+        ):
+            layer.g1_alphas = torch.nn.Parameter(
+                a13_scale * w13_scale_2, requires_grad=False
+            )
+            layer.g2_alphas = torch.nn.Parameter(
+                a2_scale * w2_scale_2, requires_grad=False
+            )
+
+        # Setup modular kernel for TP case and naive DP/EP case.
+        # In non-naive DP/EP case, we will create a ModularKernelMethod.
+        # TODO(rob): unify these so FP8MoEMethod owns the ModularKernel
+        # in both cases.
         self.moe_quant_config = self.get_fused_moe_quant_config(layer)
         assert self.experts_cls is not None
         self.moe_kernel = make_nvfp4_moe_kernel(
@@ -575,7 +591,7 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
         )
 
     def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
-        return make_nvfp4_moe_quant_config(
+        result = make_nvfp4_moe_quant_config(
             backend=self.nvfp4_backend,
             w13_scale=layer.w13_weight_scale,
             w2_scale=layer.w2_weight_scale,
@@ -583,7 +599,11 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
             w2_scale_2=layer.w2_weight_scale_2,
             a13_scale=layer.w13_input_scale,
             a2_scale=layer.w2_input_scale,
+            g1_alphas=getattr(layer, "g1_alphas", None),
+            g2_alphas=getattr(layer, "g2_alphas", None),
         )
+        assert result is not None
+        return result
 
     def apply_monolithic(
         self,

vllm/model_executor/layers/quantization/modelopt.py

@@ -29,6 +29,7 @@ from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
     select_fp8_moe_backend,
 )
 from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
+    NvFp4MoeBackend,
     convert_to_nvfp4_moe_kernel_format,
     is_global_sf_supported_for_nvfp4_backend,
     make_nvfp4_moe_kernel,
@@ -1373,7 +1374,22 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
         replace_parameter(layer, "w2_weight_scale_2", w2_scale_2)
         replace_parameter(layer, "w2_input_scale", a2_scale)
 
-        # Setup modular kernel.
+        # Pre-compute g1/g2 alphas as registered parameters so EPLB
+        # rearranges them alongside expert weights. Without this, the
+        # quant config caches g1_alphas = a_scale * w_scale_2 once at
+        # init, and EPLB's in-place rearrangement of w_scale_2 leaves
+        # the cached product stale, corrupting dequantization.
+        if self.nvfp4_backend not in (
+            NvFp4MoeBackend.FLASHINFER_TRTLLM,
+            NvFp4MoeBackend.MARLIN,
+        ):
+            replace_parameter(layer, "g1_alphas", a13_scale * w13_scale_2)
+            replace_parameter(layer, "g2_alphas", a2_scale * w2_scale_2)
+
+        # Setup modular kernel for TP case and naive DP/EP case.
+        # In non-naive DP/EP case, we will create a ModularKernelMethod.
+        # TODO(rob): unify these so FP8MoEMethod owns the ModularKernel
+        # in both cases.
         self.moe_quant_config = self.get_fused_moe_quant_config(layer)
         assert self.experts_cls is not None
         self.moe_kernel = make_nvfp4_moe_kernel(
@@ -1385,7 +1401,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
         )
 
     def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
-        return make_nvfp4_moe_quant_config(
+        result = make_nvfp4_moe_quant_config(
             backend=self.nvfp4_backend,
             w13_scale=layer.w13_weight_scale,
             w2_scale=layer.w2_weight_scale,
@@ -1393,7 +1409,11 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
             w2_scale_2=layer.w2_weight_scale_2,
             a13_scale=layer.w13_input_scale,
             a2_scale=layer.w2_input_scale,
+            g1_alphas=getattr(layer, "g1_alphas", None),
+            g2_alphas=getattr(layer, "g2_alphas", None),
         )
+        assert result is not None
+        return result
 
     @property
     def supports_eplb(self) -> bool:

vllm/tool_parsers/abstract_tool_parser.py

@@ -68,7 +68,7 @@ class ToolParser:
                 # tool_choice: "Forced Function" or "required" will override
                 # structured output json settings to make tool calling work correctly
                 request.structured_outputs = StructuredOutputsParams(
-                    json=json_schema_from_tool
+                    json=json_schema_from_tool  # type: ignore[call-arg]
                 )
                 request.response_format = None
             if isinstance(request, ResponsesRequest):