EPLB: add post-rearrangement weight invariant checks

After each expert rearrangement, verify that:
- g1_alphas[i] == a13_scale * w13_weight_scale_2[i]
- g2_alphas[i] == a2_scale * w2_weight_scale_2[i]

Also logs per-weight checksums (rank 0, every 20th layer) to
track cumulative corruption across rearrangements.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

vllm/distributed/eplb/eplb_state.py

@@ -56,6 +56,105 @@ from .rebalance_execute import (
 logger = init_logger(__name__)
 
 
+def _verify_expert_weights_after_rearrange(
+    model_state: "EplbModelState",
+    ep_rank: int,
+) -> None:
+    """
+    Post-rearrangement diagnostic: verify expert weight consistency.
+
+    Checks:
+    1. g1_alphas[i] == a13_scale_val * w13_weight_scale_2[i] for all experts
+    2. g2_alphas[i] == a2_scale_val * w2_weight_scale_2[i] for all experts
+    3. Per-weight checksums for tracking corruption across rearrangements
+    """
+    torch.cuda.synchronize()
+    model = model_state.model
+    g1_broken = 0
+    g2_broken = 0
+    g1_max_diff_all = 0.0
+    g2_max_diff_all = 0.0
+    for layer_idx, layer in enumerate(model.moe_layers):
+        g1 = getattr(layer, "g1_alphas", None)
+        g2 = getattr(layer, "g2_alphas", None)
+        s2_13 = getattr(layer, "w13_weight_scale_2", None)
+        s2_2 = getattr(layer, "w2_weight_scale_2", None)
+        a13 = getattr(layer, "w13_input_scale", None)
+        a2 = getattr(layer, "w2_input_scale", None)
+
+        # Invariant checks
+        if g1 is not None and s2_13 is not None and a13 is not None:
+            a13_val = a13.float().max().item()
+            expected_g1 = a13_val * s2_13.float()
+            diff = (g1.float() - expected_g1).abs()
+            max_diff = diff.max().item()
+            g1_max_diff_all = max(g1_max_diff_all, max_diff)
+            if max_diff > 1e-6:
+                g1_broken += 1
+                bad = (diff > 1e-6).nonzero(as_tuple=True)[0]
+                logger.error(
+                    "EPLB INVARIANT BROKEN rank %d layer %d: "
+                    "g1_alphas != a13_scale * w13_scale_2, "
+                    "max_diff=%.6e, broken_slots=%s "
+                    "(g1=%s, expected=%s)",
+                    ep_rank, layer_idx, max_diff,
+                    bad[:8].tolist(),
+                    g1.float()[bad[:4]].tolist(),
+                    expected_g1[bad[:4]].tolist(),
+                )
+
+        if g2 is not None and s2_2 is not None and a2 is not None:
+            a2_val = a2.float().max().item()
+            expected_g2 = a2_val * s2_2.float()
+            diff = (g2.float() - expected_g2).abs()
+            max_diff = diff.max().item()
+            g2_max_diff_all = max(g2_max_diff_all, max_diff)
+            if max_diff > 1e-6:
+                g2_broken += 1
+                bad = (diff > 1e-6).nonzero(as_tuple=True)[0]
+                logger.error(
+                    "EPLB INVARIANT BROKEN rank %d layer %d: "
+                    "g2_alphas != a2_scale * w2_scale_2, "
+                    "max_diff=%.6e, broken_slots=%s",
+                    ep_rank, layer_idx, max_diff,
+                    bad[:8].tolist(),
+                )
+
+        # Per-weight checksums (rank 0 only, sample layers)
+        if ep_rank == 0 and layer_idx % 20 == 0:
+            checksums = []
+            for name, param in layer.named_parameters():
+                if name in {"w13_input_scale", "w2_input_scale",
+                            "e_score_correction_bias"}:
+                    continue
+                if (name.startswith("_shared_experts.")
+                        or name.startswith("_gate.")):
+                    continue
+                cs = param.float().abs().sum().item()
+                checksums.append(f"{name}={cs:.4f}")
+            logger.info(
+                "EPLB checksums rank %d layer %d: %s",
+                ep_rank, layer_idx, ", ".join(checksums),
+            )
+
+    num_layers = model.num_moe_layers
+    if g1_broken > 0 or g2_broken > 0:
+        logger.error(
+            "EPLB VERIFY rank %d: %d/%d layers g1 broken, "
+            "%d/%d layers g2 broken (g1_max=%.2e, g2_max=%.2e)",
+            ep_rank, g1_broken, num_layers,
+            g2_broken, num_layers,
+            g1_max_diff_all, g2_max_diff_all,
+        )
+    else:
+        logger.info(
+            "EPLB VERIFY rank %d: all %d layers OK "
+            "(g1_max=%.2e, g2_max=%.2e)",
+            ep_rank, num_layers,
+            g1_max_diff_all, g2_max_diff_all,
+        )
+
+
 @dataclass
 class EplbStats:
     """
@@ -571,18 +670,22 @@ class EplbState:
                     .float()
                 )
 
-                # Compute balancedness ratio:
-                # for each layer:
-                #   (mean load across ranks) / (max load across ranks)
-                avg_tokens_tensor = num_tokens_per_rank.mean(dim=0).sum(dim=0)
-                max_tokens_tensor = num_tokens_per_rank.max(dim=0).values.sum(dim=0)
+                # Compute per-layer balancedness ratio:
+                #   for each layer: (mean across ranks) / (max across ranks)
+                #   then average across layers.
+                # dim=-1 is the rank dimension.
+                avg_per_layer = num_tokens_per_rank.mean(dim=-1)
+                max_per_layer = num_tokens_per_rank.max(dim=-1).values
+                per_layer_balance = torch.where(
+                    max_per_layer > 0,
+                    avg_per_layer / max_per_layer,
+                    torch.ones_like(max_per_layer),
+                )
+                balancedness = float(per_layer_balance.mean().item())
 
-                # Just to make type checker happy
-                tokens_tensors: list[float] = torch.stack(
-                    [avg_tokens_tensor, max_tokens_tensor]
-                ).tolist()
-                avg_tokens, max_tokens = tokens_tensors
-                balancedness = avg_tokens / max_tokens if max_tokens > 0 else 0.0
+                # Summary stats for logging
+                avg_tokens = float(avg_per_layer.sum().item())
+                max_tokens = float(max_per_layer.sum().item())
 
                 if ep_group.rank() == 0:
                     logger.info(
@@ -762,6 +865,9 @@ class EplbState:
                 )
 
                 if not is_profile:
+                    _verify_expert_weights_after_rearrange(
+                        eplb_model_state, ep_rank
+                    )
                     if (
                         eplb_model_state.physical_to_logical_map.shape[1]
                         != new_physical_to_logical_map.shape[1]

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,29 @@ 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.
+        #
+        # Use direct Parameter assignment (not replace_parameter) because
+        # g1_alphas/g2_alphas are not pre-registered in create_weights.
+        if self.nvfp4_backend not in (
+            NvFp4MoeBackend.FLASHINFER_TRTLLM,
+            NvFp4MoeBackend.MARLIN,
+        ):
+            layer.g1_alphas = torch.nn.Parameter(
+                (a13_scale * w13_scale_2).contiguous(), requires_grad=False
+            )
+            layer.g2_alphas = torch.nn.Parameter(
+                (a2_scale * w2_scale_2).contiguous(), 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(
@@ -1385,7 +1408,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 +1416,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):