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 14, 05:30 UTC
refactor: Change scaling factors calculation for flashinfer FusedMoE (#22812)
Signed-off-by: Amir Klein <203507526+amirkl94@users.noreply.github.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
amirkl94 committed
Aug 15, 2025 at 09:19 UTC
b4cef5e6c7bd9ec3dbb951fd913ee36dbadf598d
4 files changed
+60
-25
vllm/model_executor/layers/fused_moe/fused_moe.py
+12
-17
@@ -1189,10 +1189,10 @@ def flashinfer_fused_moe_per_tensor_scale_fp8(
1189
hidden_states: torch.Tensor,
1190
input_scale: torch.Tensor,
1191
gemm1_weights: torch.Tensor,
1192
- gemm1_weights_scale: torch.Tensor,
1193
- activation_scale: torch.Tensor,
1192
gemm2_weights: torch.Tensor,
1195
- gemm2_weights_scale: torch.Tensor,
1193
+ output1_scales_scalar: torch.Tensor,
1194
+ output1_scales_gate_scalar: torch.Tensor,
1195
+ output2_scales_scalar: torch.Tensor,
1196
num_experts: int,
1197
top_k: int,
1198
num_expert_group: Optional[int],
@@ -1206,17 +1206,12 @@ def flashinfer_fused_moe_per_tensor_scale_fp8(
1206
num_expert_group = num_expert_group if num_expert_group is not None else 0
1207
topk_group = topk_group if topk_group is not None else 0
1208
1209
- quant_hidden_states, input_scale = moe_kernel_quantize_input(
1209
+ quant_hidden_states, _ = moe_kernel_quantize_input(
1210
hidden_states,
1211
input_scale,
1212
quant_dtype=torch.float8_e4m3fn,
1213
per_act_token_quant=False)
1214
1215
- output1_scales_scalar = gemm1_weights_scale * input_scale * (
1216
- 1.0 / activation_scale)
1217
- output1_scales_gate_scalar = gemm1_weights_scale * input_scale
1218
- output2_scales_scalar = activation_scale * gemm2_weights_scale
1219
-
1215
from vllm.utils.flashinfer import (
1216
flashinfer_trtllm_fp8_per_tensor_scale_moe)
1217
return flashinfer_trtllm_fp8_per_tensor_scale_moe(
@@ -1244,24 +1239,24 @@ def flashinfer_fused_moe_per_tensor_scale_fp8(
1239
1240
def flashinfer_fused_moe_per_tensor_scale_fp8_fake(
1241
routing_logits: torch.Tensor,
1247
- routing_bias: torch.Tensor,
1242
+ routing_bias: Optional[torch.Tensor],
1243
hidden_states: torch.Tensor,
1244
+ input_scale: torch.Tensor,
1245
gemm1_weights: torch.Tensor,
1246
+ gemm2_weights: torch.Tensor,
1247
output1_scales_scalar: torch.Tensor,
1248
output1_scales_gate_scalar: torch.Tensor,
1252
- gemm2_weights: torch.Tensor,
1249
output2_scales_scalar: torch.Tensor,
1250
num_experts: int,
1251
top_k: int,
1256
- num_expert_group: int,
1257
- topk_group: int,
1252
+ num_expert_group: Optional[int],
1253
+ topk_group: Optional[int],
1254
intermediate_size: int,
1255
local_expert_offset: int,
1256
local_num_experts: int,
1261
- routed_scaling_factor: float = 1.0,
1262
- use_routing_scales_on_input: bool = False,
1263
- tile_tokens_dim: int = 8,
1264
- routing_method_type: int = 0) -> torch.Tensor:
1257
+ use_routing_scales_on_input: bool,
1258
+ routing_method_type: int,
1259
+ routed_scaling_factor: float = 1.0) -> torch.Tensor:
1260
pass
1261
1262
vllm/model_executor/layers/quantization/fp8.py
+3
-2
@@ -24,8 +24,8 @@ from vllm.model_executor.layers.quantization.base_config import (
24
QuantizationConfig, QuantizeMethodBase)
25
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
26
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
27
- apply_flashinfer_per_tensor_scale_fp8, rotate_flashinfer_fp8_moe_weights,
28
- swap_w13_to_w31)
27
+ apply_flashinfer_per_tensor_scale_fp8, register_moe_scaling_factors,
28
+ rotate_flashinfer_fp8_moe_weights, swap_w13_to_w31)
29
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
30
get_col_major_tma_aligned_tensor, requant_weight_ue8m0_inplace)
31
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
@@ -694,6 +694,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
694
w2_weight = layer.w2_weight.data
695
w2_weight_scale_inv = layer.w2_weight_scale_inv.data
696
if not self.block_quant:
697
+ register_moe_scaling_factors(layer)
698
rotate_flashinfer_fp8_moe_weights(w13_weight, w2_weight)
699
else:
700
w13_weight = layer.w13_weight.data
vllm/model_executor/layers/quantization/modelopt.py
+3
-2
@@ -25,8 +25,8 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
25
build_flashinfer_fp4_cutlass_moe_kernel,
26
flashinfer_fp4_cutlass_moe_forward, reorder_w1w3_to_w3w1)
27
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
28
- apply_flashinfer_per_tensor_scale_fp8, rotate_flashinfer_fp8_moe_weights,
29
- swap_w13_to_w31)
28
+ apply_flashinfer_per_tensor_scale_fp8, register_moe_scaling_factors,
29
+ rotate_flashinfer_fp8_moe_weights, swap_w13_to_w31)
30
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
31
apply_fp4_marlin_linear, is_fp4_marlin_supported,
32
prepare_fp4_layer_for_marlin, prepare_moe_fp4_layer_for_marlin)
@@ -430,6 +430,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
430
layer.w13_weight.data = swap_w13_to_w31(layer.w13_weight.data)
431
rotate_flashinfer_fp8_moe_weights(layer.w13_weight,
432
layer.w2_weight)
433
+ register_moe_scaling_factors(layer)
434
435
def apply(
436
self,
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+42
-4
@@ -82,6 +82,12 @@ def apply_flashinfer_per_tensor_scale_fp8(
82
apply_router_weight_on_input: bool,
83
) -> torch.Tensor:
84
from flashinfer.fused_moe import RoutingMethodType
85
+ assert layer.output1_scales_scalar is not None, (
86
+ "Expected output1_scales_scalar to be initialized")
87
+ assert layer.output1_scales_scalar is not None, (
88
+ "Expected output1_scales_gate_scalar to be initialized")
89
+ assert layer.output1_scales_scalar is not None, (
90
+ "Expected output2_scales_scalar to be initialized")
91
92
from vllm.model_executor.models.llama4 import Llama4MoE
93
assert layer.custom_routing_function == Llama4MoE.custom_routing_function, \
@@ -92,10 +98,10 @@ def apply_flashinfer_per_tensor_scale_fp8(
98
hidden_states=hidden_states,
99
input_scale=layer.w13_input_scale,
100
gemm1_weights=layer.w13_weight,
95
- gemm1_weights_scale=layer.w13_weight_scale,
101
gemm2_weights=layer.w2_weight,
97
- gemm2_weights_scale=layer.w2_weight_scale,
98
- activation_scale=layer.w2_input_scale,
102
+ output1_scales_scalar=layer.output1_scales_scalar,
103
+ output1_scales_gate_scalar=layer.output1_scales_gate_scalar,
104
+ output2_scales_scalar=layer.output2_scales_scalar,
105
num_experts=global_num_experts,
106
top_k=top_k,
107
num_expert_group=num_expert_group,
@@ -105,4 +111,36 @@ def apply_flashinfer_per_tensor_scale_fp8(
111
local_num_experts=layer.local_num_experts,
112
use_routing_scales_on_input=apply_router_weight_on_input,
113
routing_method_type=RoutingMethodType.Llama4,
108
- )
\ No newline at end of file
114
+ )
115
+
116
+
117
+def get_moe_scaling_factors(
118
+ input_scale: torch.Tensor,
119
+ gemm1_weights_scale: torch.Tensor,
120
+ activation_scale: torch.Tensor,
121
+ gemm2_weights_scale: torch.Tensor,
122
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
123
+ output1_scales_scalar = gemm1_weights_scale * input_scale * (
124
+ 1.0 / activation_scale)
125
+ output1_scales_gate_scalar = gemm1_weights_scale * input_scale
126
+ output2_scales_scalar = activation_scale * gemm2_weights_scale
127
+
128
+ return output1_scales_scalar, output1_scales_gate_scalar, \
129
+ output2_scales_scalar
130
+
131
+
132
+def register_moe_scaling_factors(layer: torch.nn.Module) -> None:
133
+ output1_scales, output1_gate_scales, output2_scales = \
134
+ get_moe_scaling_factors(
135
+ layer.w13_input_scale, layer.w13_weight_scale,
136
+ layer.w2_input_scale, layer.w2_weight_scale
137
+ )
138
+ layer.register_parameter(
139
+ 'output1_scales_scalar',
140
+ torch.nn.Parameter(output1_scales, requires_grad=False))
141
+ layer.register_parameter(
142
+ 'output1_scales_gate_scalar',
143
+ torch.nn.Parameter(output1_gate_scales, requires_grad=False))
144
+ layer.register_parameter(
145
+ 'output2_scales_scalar',
146
+ torch.nn.Parameter(output2_scales, requires_grad=False))