f8848b2f2d · tduncan/Nvidia-Optimization-Stack-VLLM

tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py

+47

@@ -66,6 +66,24 @@ def reference(x: torch.Tensor, use_ue8m0: bool) -> tuple[torch.Tensor, torch.Ten

66 return reference_quant(ref_act_out, use_ue8m0)

67

68

69 +def reference_with_clamp(

70 + x: torch.Tensor, use_ue8m0: bool, clamp_limit: float

71 +) -> tuple[torch.Tensor, torch.Tensor]:

72 + """Pre-clamp inputs (gate from above, up symmetric) at the input dtype to

73 + match the C++ compute() template, then run the standard silu_and_mul +

74 + quant reference."""

75 + N_2 = x.size(1) // 2

76 + dtype = x.dtype

77 + gate = x[..., :N_2].to(torch.float32).clamp(max=clamp_limit).to(dtype)

78 + up = (

79 + x[..., N_2:]

80 + .to(torch.float32)

81 + .clamp(min=-clamp_limit, max=clamp_limit)

82 + .to(dtype)

83 + )

84 + return reference(torch.cat([gate, up], dim=-1), use_ue8m0)

85 +

86 +

87 @pytest.mark.parametrize("T", [128, 256, 512])

88 @pytest.mark.parametrize("N", [128 * 2, 256 * 2, 768 * 2, 2048 * 2, 7168 * 2])

89 @pytest.mark.skipif(

@@ -89,3 +107,32 @@ def test_silu_mul_fp8_quant_deep_gemm(T: int, N: int):

107

108 torch.testing.assert_close(output.to(torch.float32), ref_output.to(torch.float32))

109 torch.testing.assert_close(output_scales, ref_output_scales)

110 +

111 +

112 +@pytest.mark.parametrize("T", [128, 256, 512])

113 +@pytest.mark.parametrize("N", [128 * 2, 256 * 2, 768 * 2, 2048 * 2, 7168 * 2])

114 +@pytest.mark.parametrize("clamp_limit", [7.0, 10.0])

115 +@pytest.mark.skipif(

116 + current_platform.is_rocm(),

117 + reason="ROCm does not support DeepGemm.",

118 +)

119 +def test_silu_mul_fp8_quant_deep_gemm_clamp(T: int, N: int, clamp_limit: float):

120 + set_random_seed(42)

121 +

122 + # Use a wide distribution so values routinely exceed both clamp limits and

123 + # the clamp branch is actually exercised (uniform [0, 1) inputs would never

124 + # trigger it).

125 + input = torch.randn((T, N), dtype=torch.bfloat16, device="cuda") * 8.0

126 +

127 + use_ue8m0 = is_deep_gemm_e8m0_used()

128 +

129 + # Test

130 + output, output_scales = silu_mul_per_token_group_quant_fp8_colmajor(

131 + input, use_ue8m0=use_ue8m0, clamp_limit=clamp_limit

132 + )

133 +

134 + # Reference

135 + ref_output, ref_output_scales = reference_with_clamp(input, use_ue8m0, clamp_limit)

136 +

137 + torch.testing.assert_close(output.to(torch.float32), ref_output.to(torch.float32))

138 + torch.testing.assert_close(output_scales, ref_output_scales)

vllm/model_executor/layers/fused_moe/config.py

+4

@@ -598,6 +598,7 @@ def fp8_w8a8_moe_quant_config(

598 a2_gscale: torch.Tensor | None = None,

599 g1_alphas: torch.Tensor | None = None,

600 g2_alphas: torch.Tensor | None = None,

601 + gemm1_clamp_limit: float | None = None,

602 ) -> FusedMoEQuantConfig:

603 """

604 Construct a quant config for fp8 activations and fp8 weights.

@@ -617,6 +618,7 @@ def fp8_w8a8_moe_quant_config(

618 per_act_token_quant=per_act_token_quant,

619 per_out_ch_quant=per_out_ch_quant,

620 block_shape=block_shape,

621 + gemm1_clamp_limit=gemm1_clamp_limit,

622 )

623

624

@@ -743,6 +745,7 @@ def mxfp4_w4a8_moe_quant_config(

745 w1_bias: torch.Tensor | None = None,

746 w2_bias: torch.Tensor | None = None,

747 block_shape: list[int] | None = None,

748 + gemm1_clamp_limit: float | None = None,

749 ) -> FusedMoEQuantConfig:

750 """

751 Construct a quant config for fp8 activations and mxfp4 weights.

@@ -752,6 +755,7 @@ def mxfp4_w4a8_moe_quant_config(

755 _a2=FusedMoEQuantDesc("fp8", None, a2_scale, None, None, None),

756 _w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias),

757 _w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias),

758 + gemm1_clamp_limit=gemm1_clamp_limit,

759 )

760

761

vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py

+5

@@ -128,6 +128,8 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):

128 assert not quant_config.per_act_token_quant

129 assert not quant_config.per_out_ch_quant

130

131 + self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit

132 +

133 @staticmethod

134 def activation_format() -> mk.FusedMoEActivationFormat:

135 return mk.FusedMoEActivationFormat.Standard

@@ -209,6 +211,7 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):

211 input=input,

212 output_q=output,

213 group_size=block_k,

214 + clamp_limit=self.gemm1_clamp_limit,

215 )

216 act_out = torch.empty(

217 (M_sum, activation_out_dim), dtype=input.dtype, device=input.device

@@ -228,6 +231,7 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):

231 input=input,

232 output=output,

233 use_ue8m0=use_ue8m0,

234 + clamp_limit=self.gemm1_clamp_limit,

235 )

236

237 # 3. fallback path for non-SiLU activations in non‑UE8M0 cases.

@@ -437,6 +441,7 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular):

441 input=input,

442 output=output,

443 use_ue8m0=use_ue8m0,

444 + clamp_limit=self.gemm1_clamp_limit,

445 )

446

447 act_out = torch.empty(

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

+4

@@ -511,6 +511,7 @@ def make_fp8_moe_quant_config(

511 block_shape: list[int] | None = None,

512 per_act_token_quant: bool = False,

513 per_out_ch_quant: bool = False,

514 + swiglu_limit: float | None = None,

515 ) -> FusedMoEQuantConfig:

516 """

517 Create FusedMoEQuantConfig for the specified FP8 Backend.

@@ -554,6 +555,7 @@ def make_fp8_moe_quant_config(

555 a2_gscale=(1.0 / a2_scale),

556 g1_alphas=(w1_scale * a1_scale).squeeze(),

557 g2_alphas=(w2_scale * a2_scale).squeeze(),

558 + gemm1_clamp_limit=swiglu_limit,

559 )

560 # MXFP8 uses "mxfp8" quant_dtype so the prepare step dispatches to

561 # _mxfp8_e4m3_quantize rather than standard FP8 block quantization.

@@ -568,6 +570,7 @@ def make_fp8_moe_quant_config(

570 a2_scale=a2_scale,

571 block_shape=block_shape,

572 is_scale_swizzled=False,

573 + gemm1_clamp_limit=swiglu_limit,

574 )

575

576 # All other backends use normal config.

@@ -579,6 +582,7 @@ def make_fp8_moe_quant_config(

582 block_shape=block_shape,

583 per_act_token_quant=per_act_token_quant,

584 per_out_ch_quant=per_out_ch_quant,

585 + gemm1_clamp_limit=swiglu_limit,

586 )

587

588

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

+1

@@ -1488,6 +1488,7 @@ def make_mxfp4_moe_quant_config(

1488 w1_bias=w1_bias,

1489 w2_bias=w2_bias,

1490 block_shape=None,

1491 + gemm1_clamp_limit=swiglu_limit,

1492 )

1493 elif mxfp4_backend in (

1494 Mxfp4MoeBackend.MARLIN,

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

+1

@@ -360,6 +360,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

360 per_act_token_quant=is_per_token,

361 per_out_ch_quant=is_per_token,

362 block_shape=self.weight_block_size,

363 + swiglu_limit=getattr(layer, "swiglu_limit", None),

364 )

365

366 def apply_monolithic(

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

+1

@@ -152,6 +152,7 @@ class CompressedTensorsW8A8Mxfp8MoEMethod(CompressedTensorsMoEMethod):

152 a1_scale=layer.w13_input_scale,

153 a2_scale=layer.w2_input_scale,

154 block_shape=self.weight_block_size,

155 + swiglu_limit=getattr(layer, "swiglu_limit", None),

156 )

157

158 def maybe_make_prepare_finalize(

vllm/model_executor/layers/quantization/fp8.py

+1

@@ -856,6 +856,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

856 a1_scale=a1_scale,

857 a2_scale=a2_scale,

858 block_shape=self.weight_block_size,

859 + swiglu_limit=getattr(layer, "swiglu_limit", None),

860 )

861

862 # Inject biases into the quant config if the model has them

vllm/model_executor/layers/quantization/modelopt.py

+1

@@ -944,6 +944,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

944 w2_scale=w2_scale,

945 a1_scale=a1_scale,

946 a2_scale=a2_scale,

947 + swiglu_limit=getattr(layer, "swiglu_limit", None),

948 )

949

950 def apply_monolithic(

vllm/model_executor/layers/quantization/online/fp8.py

+1

@@ -376,6 +376,7 @@ class _Fp8OnlineMoEBase(OnlineMoEMethodBase):

376 a1_scale=a1_scale,

377 a2_scale=a2_scale,

378 block_shape=self.weight_block_size,

379 + swiglu_limit=getattr(layer, "swiglu_limit", None),

380 )

381

382 self._maybe_inject_biases(quant_config, layer)

vllm/model_executor/layers/quantization/online/mxfp8.py

+1

@@ -221,6 +221,7 @@ class Mxfp8OnlineMoEMethod(OnlineMoEMethodBase):

221 a1_scale=a1_scale,

222 a2_scale=a2_scale,

223 block_shape=self.weight_block_size,

224 + swiglu_limit=getattr(layer, "swiglu_limit", None),

225 )

226

227 self._maybe_inject_biases(quant_config, layer)

vllm/model_executor/layers/quantization/quark/quark_moe.py

+2

@@ -449,6 +449,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

449 w2_bias=layer.w2_bias,

450 per_act_token_quant=self.input_qscheme == "per_channel",

451 per_out_ch_quant=self.weight_qscheme == "per_channel",

452 + gemm1_clamp_limit=getattr(layer, "swiglu_limit", None),

453 )

454

455 def apply(

@@ -941,6 +942,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):

942 w1_scale=layer.w13_weight_scale_2,

943 w2_scale=layer.w2_weight_scale_2,

944 per_out_ch_quant=True,

945 + gemm1_clamp_limit=getattr(layer, "swiglu_limit", None),

946 )

947

948 def apply(

vllm/model_executor/layers/quantization/utils/fp8_utils.py

+16 -1

@@ -302,9 +302,11 @@ def _silu_mul_per_token_group_quant_fp8_colmajor(

302 y_s_col_stride: tl.int64,

303 # Information for float8

304 eps,

305 + clamp_limit,

306 fp8_min: tl.constexpr,

307 fp8_max: tl.constexpr,

308 use_ue8m0: tl.constexpr,

309 + HAS_CLAMP: tl.constexpr,

310 # Meta-parameters

311 GROUP_SIZE: tl.constexpr,

312 BLOCK_M: tl.constexpr,

@@ -336,7 +338,16 @@ def _silu_mul_per_token_group_quant_fp8_colmajor(

338 act_in = tl.load(act_in_ptrs)

339 mul_in = tl.load(act_in_ptrs + N_2)

340

339 - # silu & mul

341 + # silu & mul — match C++ silu_and_mul: clamp in fp32 then store back to the

342 + # input dtype, run silu in fp32 then narrow, and do the mul at input

343 + # precision so HAS_CLAMP True/False share the same multiplication path.

344 + if HAS_CLAMP:

345 + act_in = tl.minimum(act_in.to(tl.float32), clamp_limit).to(

346 + y_ptr.dtype.element_ty

347 + )

348 + mul_in = tl.clamp(mul_in.to(tl.float32), -clamp_limit, clamp_limit).to(

349 + y_ptr.dtype.element_ty

350 + )

351 act_in = act_in.to(tl.float32)

352 one_f32 = tl.cast(1, tl.float32)

353 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(

378 output: torch.Tensor | None = None, # [M, N // 2]

379 use_ue8m0: bool | None = None,

380 eps: float = 1e-10,

381 + clamp_limit: float | None = None,

382 ):

383 """

384 silu+mul + block-fp8 quant with group size 128.

@@ -409,6 +421,7 @@ def silu_mul_per_token_group_quant_fp8_colmajor(

421 assert N_2 % BLOCK_N == 0

422 grid = (M // BLOCK_M, N_2 // BLOCK_N)

423

424 + has_clamp = clamp_limit is not None

425 _silu_mul_per_token_group_quant_fp8_colmajor[grid](

426 input,

427 output,

@@ -417,9 +430,11 @@ def silu_mul_per_token_group_quant_fp8_colmajor(

430 N,

431 output_scales.stride(-1),

432 eps,

433 + clamp_limit if has_clamp else 0.0,

434 fp8_min,

435 fp8_max,

436 use_ue8m0,

437 + has_clamp,

438 GROUP_SIZE,

439 BLOCK_M,

440 BLOCK_N,

[Bugfix] Add swiglu limits to deepgemm fp8 methods (#41986)