@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / af8fd73051

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 10, 16:30 UTC

[MoE Refactor][14/N] Clean Up FI Quant Config Smuggling (#31593)

Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>

Robert Shaw committed Jan 6, 2026 at 10:47 UTC af8fd730512341773cb43c6f1763f14e120a781d

7 files changed +173 -84

tests/kernels/moe/test_flashinfer.py

+12 -7

@@ -15,7 +15,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts

15 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

16 apply_flashinfer_per_tensor_scale_fp8,

17 flashinfer_cutlass_moe_fp8,

18 - register_moe_scaling_factors,

18 + register_scales_for_trtllm_fp8_per_tensor_moe,

19 rotate_flashinfer_fp8_moe_weights,

20 swap_w13_to_w31,

21 )

@@ -85,7 +85,7 @@ class TestData:

86 @staticmethod

87 def make_moe_tensors_8bit(

88 - m: int, k: int, n: int, e: int, reorder: bool, activation: str = "silu"

88 + m: int, k: int, n: int, e: int, is_trtllm: bool, activation: str = "silu"

89 ) -> "TestData":

90 is_gated = activation != "relu2_no_mul"

@@ -123,12 +123,17 @@ class TestData:

123 all2all_backend="naive",

124 )

125

126 - register_moe_scaling_factors(layer)

127 -

126 # flashinfer expects swapped rows for w13

127 layer.w13_weight.data = swap_w13_to_w31(layer.w13_weight.data)

130 - if reorder:

128 + if is_trtllm:

129 rotate_flashinfer_fp8_moe_weights(layer.w13_weight, layer.w2_weight)

130 + register_scales_for_trtllm_fp8_per_tensor_moe(

131 + layer,

132 + layer.w13_weight_scale,

133 + layer.w13_input_scale,

134 + layer.w2_weight_scale,

135 + layer.w2_input_scale,

136 + )

137 layer.custom_routing_function = Llama4MoE.custom_routing_function

138 layer.intermediate_size_per_partition = n

139 layer.ep_rank = 0

@@ -162,7 +167,7 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(

167 set_random_seed(7)

168 monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")

169 with set_current_vllm_config(vllm_config):

165 - td = TestData.make_moe_tensors_8bit(m, k, n, e, reorder=True)

170 + td = TestData.make_moe_tensors_8bit(m, k, n, e, is_trtllm=True)

171

172 score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16)

173 topk_weights, topk_ids = Llama4MoE.custom_routing_function(

@@ -227,7 +232,7 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(

232 monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")

233 with set_current_vllm_config(vllm_config):

234 td = TestData.make_moe_tensors_8bit(

230 - m, k, n, e, reorder=False, activation=activation

235 + m, k, n, e, is_trtllm=False, activation=activation

236 )

237

238 score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16)

vllm/model_executor/layers/fused_moe/config.py

+7 -4

@@ -452,11 +452,14 @@ class FusedMoEQuantConfig:

452 - a1_scale: Optional scale to be used for a1.

453 - a2_scale: Optional scale to be used for a2.

454 - g1_alphas: Optional global quantization scales for w1 (for nvfp4).

455 - per-channel scales for w1 (for W4A8 FP8).

455 + Optional per-channel scales for w1 (for W4A8 FP8).

456 + Optional dq scale i.e. w_scale * a_scale (for W8A8 fp8).

457 - g2_alphas: Optional global quantization scales for w2 (for nvfp4).

457 - per-channel scales for w2 (for W4A8 FP8).

458 - - a1_gscale: Optional global quantization scales for a1 (for nvfp4).

459 - - a2_gscale: Optional global quantization scales for a2 (for nvfp4).

458 + Optional per-channel scales for w2 (for W4A8 FP8).

459 + Optional dq scale i.e. w_scale * a_scale (for W8A8 fp8).

460 + - a1_gscale: Optional global quantization scales for a1 (1.0 /a2_scale).

461 + - a2_gscale: Optional global quantization scales for a2 (1.0 /a2_scale).

462 +

463 - w1_bias: Optional biases for w1 (GPT OSS Triton).

464 - w2_bias: Optional biases for w1 (GPT OSS Triton).

465 - w1_zp: Optional w1 zero points for int4/int8 quantization.

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py

+4 -4

@@ -165,10 +165,10 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):

165 ):

166 # FP8 per-tensor path: use global alphas/scales; do not pass input_sf

167 quant_scales = [

168 - self.g1_alphas,

169 - self.a2_gscale,

170 - self.g2_alphas,

171 - self.a1_gscale,

168 + self.g1_alphas, # w13_weight_scale * w13_input_scale

169 + self.a2_gscale, # 1.0 / w2_input_scale

170 + self.g2_alphas, # w2_weight_scale * w2_input_scale

171 + self.a1_scale,

172 ]

173

174 a1q_scale = None # not passing input_sf in fp8

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py

+4 -3

@@ -184,13 +184,14 @@ class FlashInferAllGatherMoEPrepareAndFinalize(FlashInferCutlassMoEPrepareAndFin

184 self._apply_router_weight_on_input(

185 a1, topk_weights, topk_ids, apply_router_weight_on_input

186 )

187 - if not self.use_dp and quant_config.quant_dtype == "nvfp4":

187 + is_nvfp4 = quant_config.quant_dtype == "nvfp4"

188 + if not self.use_dp and is_nvfp4:

189 return a1, None, None, topk_ids, topk_weights

190

191 if not self.use_deepseek_fp8_block_scale:

192 a1q, a1q_scale = moe_kernel_quantize_input(

193 a1,

193 - quant_config.a1_gscale,

194 + quant_config.a1_gscale if is_nvfp4 else quant_config.a1_scale,

195 quant_config.quant_dtype,

196 quant_config.per_act_token_quant,

197 quant_config.block_shape,

@@ -222,7 +223,7 @@ class FlashInferAllGatherMoEPrepareAndFinalize(FlashInferCutlassMoEPrepareAndFin

223 topk_weights, topk_ids, a1q = gathered

224 a1q_scale = None

225

225 - if quant_config.quant_dtype == "nvfp4" and a1q_scale is not None:

226 + if is_nvfp4 and a1q_scale is not None:

227 a1q_scale = nvfp4_block_scale_interleave(a1q_scale)

228

229 return a1q, a1q_scale, None, topk_ids, topk_weights

vllm/model_executor/layers/quantization/fp8.py

+74 -10

@@ -50,7 +50,8 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

50 apply_flashinfer_per_tensor_scale_fp8,

51 build_flashinfer_fp8_cutlass_moe_prepare_finalize,

52 get_flashinfer_moe_backend,

53 - register_moe_scaling_factors,

53 + make_fp8_moe_alpha_scales_for_fi,

54 + register_scales_for_trtllm_fp8_per_tensor_moe,

55 rotate_flashinfer_fp8_moe_weights,

56 select_cutlass_fp8_gemm_impl,

57 swap_w13_to_w31,

@@ -774,6 +775,14 @@ class Fp8MoEMethod(FusedMoEMethodBase):

775 "FlashInfer CUTLASS FP8 MoE backend only supports SiLU "

776 "activation function, but got {layer.activation}."

777 )

778 + dynamic_per_token = (

779 + not self.block_quant and self.quant_config.activation_scheme != "static"

780 + )

781 + if self.flashinfer_moe_backend is not None and dynamic_per_token:

782 + raise NotImplementedError(

783 + "FlashInfer FP8 MoE backend does not support dynamic per token "

784 + "activation quantization."

785 + )

786

787 def create_weights(

788 self,

@@ -905,6 +914,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):

914 w2_weight: torch.Tensor,

915 w13_weight_scale: torch.Tensor,

916 w2_weight_scale: torch.Tensor,

917 + w13_input_scale: torch.Tensor | None,

918 + w2_input_scale: torch.Tensor | None,

919 ) -> None:

920 if self.fp8_backend == Fp8MoeBackend.DEEPGEMM:

921 assert self.block_quant

@@ -949,11 +960,16 @@ class Fp8MoEMethod(FusedMoEMethodBase):

960 if self.block_quant:

961 w13_weight_scale = swap_w13_to_w31(w13_weight_scale)

962 else:

952 - # TODO(rob): this function is a hack that renames the scaling

953 - # factors in the Module. This is a hack we should clean up.

954 - register_moe_scaling_factors(layer)

963 if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

964 rotate_flashinfer_fp8_moe_weights(w13_weight, w2_weight)

965 + register_scales_for_trtllm_fp8_per_tensor_moe(

966 + layer=layer,

967 + w13_weight_scale=w13_weight,

968 + w13_input_scale=w13_input_scale,

969 + w2_weight_scale=w2_weight,

970 + w2_input_scale=w2_input_scale,

971 + )

972 +

973 elif self.fp8_backend == Fp8MoeBackend.AITER:

974 w13_weight, w2_weight = rocm_aiter_ops.shuffle_weights(

975 w13_weight, w2_weight

@@ -990,6 +1006,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1006 AiterExperts,

1007 )

1008

1009 + # Flashinfer TRTLLM does not use the modular kernel abstraction.

1010 + if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

1011 + return

1012 +

1013 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

1014 assert self.moe_quant_config is not None

1015 self.use_inplace = True

@@ -1087,7 +1107,13 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1107

1108 # Shuffle weights into the runtime format.

1109 self._convert_weights_to_kernel_format(

1090 - layer, w13_weight, w2_weight, w13_weight_scale, w2_weight_scale

1110 + layer=layer,

1111 + w13_weight=w13_weight,

1112 + w2_weight=w2_weight,

1113 + w13_weight_scale=w13_weight_scale,

1114 + w2_weight_scale=w2_weight_scale,

1115 + w13_input_scale=w13_input_scale,

1116 + w2_input_scale=w2_input_scale,

1117 )

1118

1119 # Setup modular kernel for TP case.

@@ -1182,6 +1208,11 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1208 def get_fused_moe_quant_config(

1209 self, layer: torch.nn.Module

1210 ) -> FusedMoEQuantConfig | None:

1211 + # TRTLLM does not use Modular Kernel.

1212 + if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:

1213 + return None

1214 +

1215 + # MARLIN uses mixed precision W8A16 config.

1216 if self.fp8_backend == Fp8MoeBackend.MARLIN:

1217 return fp8_w8a16_moe_quant_config(

1218 w1_scale=getattr(layer, f"w13_{self.weight_scale_name}"),

@@ -1189,11 +1220,38 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1220 block_shape=self.weight_block_size,

1221 )

1222

1223 + w1_scale = getattr(layer, f"w13_{self.weight_scale_name}")

1224 + w2_scale = getattr(layer, f"w2_{self.weight_scale_name}")

1225 + a1_scale = layer.w13_input_scale

1226 + a2_scale = layer.w2_input_scale

1227 +

1228 + # Flashinfer CUTLASS per-tensor uses single dq scale

1229 + # (alpha = w_scale * a_scale) and inverse a2 scale.

1230 + if (

1231 + self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS

1232 + and not self.block_quant

1233 + ):

1234 + g1_alphas, g2_alphas = make_fp8_moe_alpha_scales_for_fi(

1235 + w1_scale,

1236 + a1_scale,

1237 + w2_scale,

1238 + a2_scale,

1239 + )

1240 + return fp8_w8a8_moe_quant_config(

1241 + w1_scale=w1_scale,

1242 + w2_scale=w2_scale,

1243 + a1_scale=a1_scale,

1244 + a2_scale=(1.0 / a2_scale),

1245 + g1_alphas=g1_alphas,

1246 + g2_alphas=g2_alphas,

1247 + )

1248 +

1249 + # All other backends use normal config.

1250 return fp8_w8a8_moe_quant_config(

1193 - w1_scale=getattr(layer, f"w13_{self.weight_scale_name}"),

1194 - w2_scale=getattr(layer, f"w2_{self.weight_scale_name}"),

1195 - a1_scale=layer.w13_input_scale,

1196 - a2_scale=layer.w2_input_scale,

1251 + w1_scale=w1_scale,

1252 + w2_scale=w2_scale,

1253 + a1_scale=a1_scale,

1254 + a2_scale=a2_scale,

1255 block_shape=self.weight_block_size,

1256 )

1257

@@ -1414,7 +1472,13 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):

1472

1473 # Shuffle weights into the runtime format.

1474 self._convert_weights_to_kernel_format(

1417 - layer, w13_weight, w2_weight, layer.w13_weight_scale, layer.w2_weight_scale

1475 + layer=layer,

1476 + w13_weight=w13_weight,

1477 + w2_weight=w2_weight,

1478 + w13_weight_scale=layer.w13_weight_scale,

1479 + w2_weight_scale=layer.w2_weight_scale,

1480 + w13_input_scale=None,

1481 + w2_input_scale=None,

1482 )

1483

1484 # Setup modular kernel for TP case.

vllm/model_executor/layers/quantization/modelopt.py

+38 -13

@@ -50,7 +50,8 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

50 flashinfer_cutlass_moe_fp8,

51 get_flashinfer_moe_backend,

52 is_flashinfer_supporting_global_sf,

53 - register_moe_scaling_factors,

53 + make_fp8_moe_alpha_scales_for_fi,

54 + register_scales_for_trtllm_fp8_per_tensor_moe,

55 rotate_flashinfer_fp8_moe_weights,

56 select_cutlass_fp8_gemm_impl,

57 swap_w13_to_w31,

@@ -947,9 +948,18 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

948 if self.flashinfer_moe_backend is not None:

949 if self.moe.is_act_and_mul:

950 layer.w13_weight.data = swap_w13_to_w31(layer.w13_weight.data)

951 +

952 + # NOTE: this adds some attributes used by the trtllm kernel,

953 + # which does not conform to the modular kernels abstraction (yet).

954 if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:

955 rotate_flashinfer_fp8_moe_weights(layer.w13_weight, layer.w2_weight)

952 - register_moe_scaling_factors(layer)

956 + register_scales_for_trtllm_fp8_per_tensor_moe(

957 + layer=layer,

958 + w13_weight_scale=layer.w13_weight_scale,

959 + w13_input_scale=layer.w13_input_scale,

960 + w2_weight_scale=layer.w2_weight_scale,

961 + w2_input_scale=layer.w2_input_scale,

962 + )

963

964 def _maybe_pad_intermediate_for_flashinfer(self, layer: torch.nn.Module) -> None:

965 """Pad intermediate size so FlashInfer kernels' alignment constraints hold.

@@ -999,19 +1009,34 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

1009 self, layer: torch.nn.Module

1010 ) -> FusedMoEQuantConfig | None:

1011 if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:

1012 + # TRTLLM does not use modular kernels

1013 return None

1014

1004 - return fp8_w8a8_moe_quant_config(

1005 - w1_scale=layer.w13_weight_scale,

1006 - g1_alphas=layer.output1_scales_gate_scalar.squeeze(),

1007 - w2_scale=layer.w2_weight_scale,

1008 - g2_alphas=layer.output2_scales_scalar.squeeze(),

1009 - a1_scale=layer.w13_input_scale,

1010 - a1_gscale=layer.w13_input_scale,

1011 - a2_scale=layer.w2_input_scale,

1012 - a2_gscale=layer.w2_input_scale_inv,

1013 - per_act_token_quant=False,

1014 - )

1015 + elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:

1016 + g1_alphas, g2_alphas = make_fp8_moe_alpha_scales_for_fi(

1017 + layer.w13_weight_scale,

1018 + layer.w13_input_scale,

1019 + layer.w2_weight_scale,

1020 + layer.w2_input_scale,

1021 + )

1022 + return fp8_w8a8_moe_quant_config(

1023 + w1_scale=layer.w13_weight_scale,

1024 + w2_scale=layer.w2_weight_scale,

1025 + a1_scale=layer.w13_input_scale,

1026 + a2_scale=layer.w2_input_scale,

1027 + a1_gscale=(1.0 / layer.w13_input_scale),

1028 + a2_gscale=(1.0 / layer.w2_input_scale),

1029 + g1_alphas=g1_alphas,

1030 + g2_alphas=g2_alphas,

1031 + )

1032 + else:

1033 + assert self.flashinfer_moe_backend is None

1034 + return fp8_w8a8_moe_quant_config(

1035 + w1_scale=layer.w13_weight_scale,

1036 + w2_scale=layer.w2_weight_scale,

1037 + a1_scale=layer.w13_input_scale,

1038 + a2_scale=layer.w2_input_scale,

1039 + )

1040

1041 def apply(

1042 self,

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

+34 -43

@@ -103,6 +103,26 @@ def rotate_flashinfer_fp8_moe_weights(

103 )

104

105

106 +def register_scales_for_trtllm_fp8_per_tensor_moe(

107 + layer: torch.nn.Module,

108 + w13_weight_scale: torch.Tensor,

109 + w13_input_scale: torch.Tensor,

110 + w2_weight_scale: torch.Tensor,

111 + w2_input_scale: torch.Tensor,

112 +) -> None:

113 + """Register necessary scales for FlashInfer TRTLLM FP8 MoE kernel"""

114 + g1_alphas, g2_alphas = make_fp8_moe_alpha_scales_for_fi(

115 + w13_scale=w13_weight_scale,

116 + w13_input_scale=w13_input_scale,

117 + w2_scale=w2_weight_scale,

118 + w2_input_scale=w2_input_scale,

119 + )

120 + layer.w2_input_scale_inv = 1.0 / w2_input_scale

121 + layer.output1_scales_gate_scalar = g1_alphas

122 + layer.output1_scales_scalar = g1_alphas * layer.w2_input_scale_inv

123 + layer.output2_scales_scalar = g2_alphas

124 +

125 +

126 def apply_flashinfer_per_tensor_scale_fp8(

127 layer: torch.nn.Module,

128 hidden_states: torch.Tensor,

@@ -117,18 +137,13 @@ def apply_flashinfer_per_tensor_scale_fp8(

137 from flashinfer.fused_moe import RoutingMethodType

138

139 import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401

140 + from vllm.model_executor.models.llama4 import Llama4MoE

141

121 - assert layer.output1_scales_scalar is not None, (

122 - "Expected output1_scales_scalar to be initialized"

123 - )

124 - assert layer.output1_scales_scalar is not None, (

125 - "Expected output1_scales_gate_scalar to be initialized"

142 + assert (

143 + hasattr(layer, "output1_scales_scalar")

144 + and hasattr(layer, "output1_scales_gate_scalar")

145 + and hasattr(layer, "output2_scales_scalar")

146 )

127 - assert layer.output1_scales_scalar is not None, (

128 - "Expected output2_scales_scalar to be initialized"

129 - )

130 -

131 - from vllm.model_executor.models.llama4 import Llama4MoE

147

148 assert layer.custom_routing_function == Llama4MoE.custom_routing_function, (

149 "FusedMoE flashinfer kernels are only supported for Llama4"

@@ -155,40 +170,16 @@ def apply_flashinfer_per_tensor_scale_fp8(

170 )

171

172

158 -def get_moe_scaling_factors(

159 - input_scale: torch.Tensor,

160 - gemm1_weights_scale: torch.Tensor,

161 - activation_scale: torch.Tensor,

162 - gemm2_weights_scale: torch.Tensor,

163 -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

164 - output1_scales_scalar = gemm1_weights_scale * input_scale * (1.0 / activation_scale)

165 - output1_scales_gate_scalar = gemm1_weights_scale * input_scale

166 - output2_scales_scalar = activation_scale * gemm2_weights_scale

167 -

168 - return output1_scales_scalar, output1_scales_gate_scalar, output2_scales_scalar

173 +def make_fp8_moe_alpha_scales_for_fi(

174 + w13_scale: torch.Tensor,

175 + w13_input_scale: torch.Tensor,

176 + w2_scale: torch.Tensor,

177 + w2_input_scale: torch.Tensor,

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

179 + g1_alphas = (w13_scale * w13_input_scale).squeeze()

180 + g2_alphas = (w2_scale * w2_input_scale).squeeze()

181

170 -

171 -def register_moe_scaling_factors(layer: torch.nn.Module) -> None:

172 - output1_scales, output1_gate_scales, output2_scales = get_moe_scaling_factors(

173 - layer.w13_input_scale,

174 - layer.w13_weight_scale,

175 - layer.w2_input_scale,

176 - layer.w2_weight_scale,

177 - )

178 - layer.register_parameter(

179 - "output1_scales_scalar", torch.nn.Parameter(output1_scales, requires_grad=False)

180 - )

181 - layer.register_parameter(

182 - "output1_scales_gate_scalar",

183 - torch.nn.Parameter(output1_gate_scales, requires_grad=False),

184 - )

185 - layer.register_parameter(

186 - "output2_scales_scalar", torch.nn.Parameter(output2_scales, requires_grad=False)

187 - )

188 - layer.register_parameter(

189 - "w2_input_scale_inv",

190 - torch.nn.Parameter(1.0 / layer.w2_input_scale, requires_grad=False),

191 - )

182 + return g1_alphas, g2_alphas

183

184

185 def build_flashinfer_fp8_cutlass_moe_prepare_finalize(