diff --git a/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py b/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py index cca02928b49..cb01db44f8c 100644 --- a/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py +++ b/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py @@ -66,6 +66,24 @@ def reference(x: torch.Tensor, use_ue8m0: bool) -> tuple[torch.Tensor, torch.Ten return reference_quant(ref_act_out, use_ue8m0) +def reference_with_clamp( + x: torch.Tensor, use_ue8m0: bool, clamp_limit: float +) -> tuple[torch.Tensor, torch.Tensor]: + """Pre-clamp inputs (gate from above, up symmetric) at the input dtype to + match the C++ compute() template, then run the standard silu_and_mul + + quant reference.""" + N_2 = x.size(1) // 2 + dtype = x.dtype + gate = x[..., :N_2].to(torch.float32).clamp(max=clamp_limit).to(dtype) + up = ( + x[..., N_2:] + .to(torch.float32) + .clamp(min=-clamp_limit, max=clamp_limit) + .to(dtype) + ) + return reference(torch.cat([gate, up], dim=-1), use_ue8m0) + + @pytest.mark.parametrize("T", [128, 256, 512]) @pytest.mark.parametrize("N", [128 * 2, 256 * 2, 768 * 2, 2048 * 2, 7168 * 2]) @pytest.mark.skipif( @@ -89,3 +107,32 @@ def test_silu_mul_fp8_quant_deep_gemm(T: int, N: int): torch.testing.assert_close(output.to(torch.float32), ref_output.to(torch.float32)) torch.testing.assert_close(output_scales, ref_output_scales) + + +@pytest.mark.parametrize("T", [128, 256, 512]) +@pytest.mark.parametrize("N", [128 * 2, 256 * 2, 768 * 2, 2048 * 2, 7168 * 2]) +@pytest.mark.parametrize("clamp_limit", [7.0, 10.0]) +@pytest.mark.skipif( + current_platform.is_rocm(), + reason="ROCm does not support DeepGemm.", +) +def test_silu_mul_fp8_quant_deep_gemm_clamp(T: int, N: int, clamp_limit: float): + set_random_seed(42) + + # Use a wide distribution so values routinely exceed both clamp limits and + # the clamp branch is actually exercised (uniform [0, 1) inputs would never + # trigger it). + input = torch.randn((T, N), dtype=torch.bfloat16, device="cuda") * 8.0 + + use_ue8m0 = is_deep_gemm_e8m0_used() + + # Test + output, output_scales = silu_mul_per_token_group_quant_fp8_colmajor( + input, use_ue8m0=use_ue8m0, clamp_limit=clamp_limit + ) + + # Reference + ref_output, ref_output_scales = reference_with_clamp(input, use_ue8m0, clamp_limit) + + torch.testing.assert_close(output.to(torch.float32), ref_output.to(torch.float32)) + torch.testing.assert_close(output_scales, ref_output_scales) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index d5435261327..1b21f75ea12 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -598,6 +598,7 @@ def fp8_w8a8_moe_quant_config( a2_gscale: torch.Tensor | None = None, g1_alphas: torch.Tensor | None = None, g2_alphas: torch.Tensor | None = None, + gemm1_clamp_limit: float | None = None, ) -> FusedMoEQuantConfig: """ Construct a quant config for fp8 activations and fp8 weights. @@ -617,6 +618,7 @@ def fp8_w8a8_moe_quant_config( per_act_token_quant=per_act_token_quant, per_out_ch_quant=per_out_ch_quant, block_shape=block_shape, + gemm1_clamp_limit=gemm1_clamp_limit, ) @@ -743,6 +745,7 @@ def mxfp4_w4a8_moe_quant_config( w1_bias: torch.Tensor | None = None, w2_bias: torch.Tensor | None = None, block_shape: list[int] | None = None, + gemm1_clamp_limit: float | None = None, ) -> FusedMoEQuantConfig: """ Construct a quant config for fp8 activations and mxfp4 weights. @@ -752,6 +755,7 @@ def mxfp4_w4a8_moe_quant_config( _a2=FusedMoEQuantDesc("fp8", None, a2_scale, None, None, None), _w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias), _w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias), + gemm1_clamp_limit=gemm1_clamp_limit, ) diff --git a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py index b4394b5fd38..e3e15e31618 100644 --- a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py @@ -128,6 +128,8 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular): assert not quant_config.per_act_token_quant assert not quant_config.per_out_ch_quant + self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit + @staticmethod def activation_format() -> mk.FusedMoEActivationFormat: return mk.FusedMoEActivationFormat.Standard @@ -209,6 +211,7 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular): input=input, output_q=output, group_size=block_k, + clamp_limit=self.gemm1_clamp_limit, ) act_out = torch.empty( (M_sum, activation_out_dim), dtype=input.dtype, device=input.device @@ -228,6 +231,7 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular): input=input, output=output, use_ue8m0=use_ue8m0, + clamp_limit=self.gemm1_clamp_limit, ) # 3. fallback path for non-SiLU activations in non‑UE8M0 cases. @@ -437,6 +441,7 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular): input=input, output=output, use_ue8m0=use_ue8m0, + clamp_limit=self.gemm1_clamp_limit, ) act_out = torch.empty( diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index ed144a992c8..a64c99bdb9f 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -511,6 +511,7 @@ def make_fp8_moe_quant_config( block_shape: list[int] | None = None, per_act_token_quant: bool = False, per_out_ch_quant: bool = False, + swiglu_limit: float | None = None, ) -> FusedMoEQuantConfig: """ Create FusedMoEQuantConfig for the specified FP8 Backend. @@ -554,6 +555,7 @@ def make_fp8_moe_quant_config( a2_gscale=(1.0 / a2_scale), g1_alphas=(w1_scale * a1_scale).squeeze(), g2_alphas=(w2_scale * a2_scale).squeeze(), + gemm1_clamp_limit=swiglu_limit, ) # MXFP8 uses "mxfp8" quant_dtype so the prepare step dispatches to # _mxfp8_e4m3_quantize rather than standard FP8 block quantization. @@ -568,6 +570,7 @@ def make_fp8_moe_quant_config( a2_scale=a2_scale, block_shape=block_shape, is_scale_swizzled=False, + gemm1_clamp_limit=swiglu_limit, ) # All other backends use normal config. @@ -579,6 +582,7 @@ def make_fp8_moe_quant_config( block_shape=block_shape, per_act_token_quant=per_act_token_quant, per_out_ch_quant=per_out_ch_quant, + gemm1_clamp_limit=swiglu_limit, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 6e233f4c392..80408998397 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -1488,6 +1488,7 @@ def make_mxfp4_moe_quant_config( w1_bias=w1_bias, w2_bias=w2_bias, block_shape=None, + gemm1_clamp_limit=swiglu_limit, ) elif mxfp4_backend in ( Mxfp4MoeBackend.MARLIN, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py index 1e59f357f52..da5d85e4abc 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py @@ -360,6 +360,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod): per_act_token_quant=is_per_token, per_out_ch_quant=is_per_token, block_shape=self.weight_block_size, + swiglu_limit=getattr(layer, "swiglu_limit", None), ) def apply_monolithic( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py index 97e1e92c2c9..dc851cc1313 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_mxfp8.py @@ -152,6 +152,7 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod): a1_scale=layer.w13_input_scale, a2_scale=layer.w2_input_scale, block_shape=self.weight_block_size, + swiglu_limit=getattr(layer, "swiglu_limit", None), ) def maybe_make_prepare_finalize( diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index fe8522a4f4b..07261648303 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -856,6 +856,7 @@ class Fp8MoEMethod(FusedMoEMethodBase): a1_scale=a1_scale, a2_scale=a2_scale, block_shape=self.weight_block_size, + swiglu_limit=getattr(layer, "swiglu_limit", None), ) # Inject biases into the quant config if the model has them diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 092346cb344..4a3b7619b94 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -944,6 +944,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): w2_scale=w2_scale, a1_scale=a1_scale, a2_scale=a2_scale, + swiglu_limit=getattr(layer, "swiglu_limit", None), ) def apply_monolithic( diff --git a/vllm/model_executor/layers/quantization/online/fp8.py b/vllm/model_executor/layers/quantization/online/fp8.py index 7f8bd11cb7a..add28f1a8c6 100644 --- a/vllm/model_executor/layers/quantization/online/fp8.py +++ b/vllm/model_executor/layers/quantization/online/fp8.py @@ -376,6 +376,7 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase): a1_scale=a1_scale, a2_scale=a2_scale, block_shape=self.weight_block_size, + swiglu_limit=getattr(layer, "swiglu_limit", None), ) self._maybe_inject_biases(quant_config, layer) diff --git a/vllm/model_executor/layers/quantization/online/mxfp8.py b/vllm/model_executor/layers/quantization/online/mxfp8.py index f2ff5b043b9..28d0706b11e 100644 --- a/vllm/model_executor/layers/quantization/online/mxfp8.py +++ b/vllm/model_executor/layers/quantization/online/mxfp8.py @@ -221,6 +221,7 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase): a1_scale=a1_scale, a2_scale=a2_scale, block_shape=self.weight_block_size, + swiglu_limit=getattr(layer, "swiglu_limit", None), ) self._maybe_inject_biases(quant_config, layer) diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 2dd38c07fe6..ad2b842e782 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -449,6 +449,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod): w2_bias=layer.w2_bias, per_act_token_quant=self.input_qscheme == "per_channel", per_out_ch_quant=self.weight_qscheme == "per_channel", + gemm1_clamp_limit=getattr(layer, "swiglu_limit", None), ) def apply( @@ -941,6 +942,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod): w1_scale=layer.w13_weight_scale_2, w2_scale=layer.w2_weight_scale_2, per_out_ch_quant=True, + gemm1_clamp_limit=getattr(layer, "swiglu_limit", None), ) def apply( diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index a5029a3ae68..09b31add3b4 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -302,9 +302,11 @@ def _silu_mul_per_token_group_quant_fp8_colmajor( y_s_col_stride: tl.int64, # Information for float8 eps, + clamp_limit, fp8_min: tl.constexpr, fp8_max: tl.constexpr, use_ue8m0: tl.constexpr, + HAS_CLAMP: tl.constexpr, # Meta-parameters GROUP_SIZE: tl.constexpr, BLOCK_M: tl.constexpr, @@ -336,7 +338,16 @@ def _silu_mul_per_token_group_quant_fp8_colmajor( act_in = tl.load(act_in_ptrs) mul_in = tl.load(act_in_ptrs + N_2) - # silu & mul + # silu & mul — match C++ silu_and_mul: clamp in fp32 then store back to the + # input dtype, run silu in fp32 then narrow, and do the mul at input + # precision so HAS_CLAMP True/False share the same multiplication path. + if HAS_CLAMP: + act_in = tl.minimum(act_in.to(tl.float32), clamp_limit).to( + y_ptr.dtype.element_ty + ) + mul_in = tl.clamp(mul_in.to(tl.float32), -clamp_limit, clamp_limit).to( + y_ptr.dtype.element_ty + ) act_in = act_in.to(tl.float32) one_f32 = tl.cast(1, tl.float32) silu_out = (act_in / (one_f32 + tl.exp(-act_in))).to(y_ptr.dtype.element_ty) @@ -367,6 +378,7 @@ def silu_mul_per_token_group_quant_fp8_colmajor( output: torch.Tensor | None = None, # [M, N // 2] use_ue8m0: bool | None = None, eps: float = 1e-10, + clamp_limit: float | None = None, ): """ silu+mul + block-fp8 quant with group size 128. @@ -409,6 +421,7 @@ def silu_mul_per_token_group_quant_fp8_colmajor( assert N_2 % BLOCK_N == 0 grid = (M // BLOCK_M, N_2 // BLOCK_N) + has_clamp = clamp_limit is not None _silu_mul_per_token_group_quant_fp8_colmajor[grid]( input, output, @@ -417,9 +430,11 @@ def silu_mul_per_token_group_quant_fp8_colmajor( N, output_scales.stride(-1), eps, + clamp_limit if has_clamp else 0.0, fp8_min, fp8_max, use_ue8m0, + has_clamp, GROUP_SIZE, BLOCK_M, BLOCK_N,