diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 6200477092a..94d3a66e7ea 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/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 diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index dd1a24d863d..75e902c3b43 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/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, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py index f3ed9a62865..c306901d43c 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py +++ b/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, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index f167e213447..d0c96b688f3 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/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: diff --git a/vllm/tool_parsers/abstract_tool_parser.py b/vllm/tool_parsers/abstract_tool_parser.py index 75cffd3297f..81ee4ea671e 100644 --- a/vllm/tool_parsers/abstract_tool_parser.py +++ b/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):