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, 22:35 UTC
[NVIDIA] Add SM100 Flashinfer MoE per tensor scale fp8 backend (#21458)
Signed-off-by: Amir Klein <203507526+amirkl94@users.noreply.github.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: mgoin <mgoin64@gmail.com>
amirkl94 committed
Jul 31, 2025 at 16:00 UTC
207b750e194829c4bcd4df0450f5f93d71755dae
5 files changed
+269
-49
vllm/model_executor/layers/fused_moe/fused_moe.py
+95
-18
@@ -30,6 +30,8 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
30
TopKWeightAndReduceNoOP)
31
from vllm.model_executor.layers.fused_moe.utils import (
32
_resize_cache, moe_kernel_quantize_input, per_token_group_quant_fp8)
33
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
34
+ calculate_tile_tokens_dim)
35
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
36
dequant_mxfp4)
37
from vllm.platforms import current_platform
@@ -1065,22 +1067,6 @@ direct_register_custom_op(
1067
)
1068
1069
1068
-def next_positive_power_of_2(x: int) -> int:
1069
- if x < 1:
1070
- return 1
1071
- return 1 << (x - 1).bit_length()
1072
-
1073
-
1074
-def _get_tile_tokens_dim(num_tokens, top_k, num_experts):
1075
- # Guess tokens per expert assuming perfect expert distribution first.
1076
- num_tokens_per_expert = (num_tokens * top_k) // num_experts
1077
- # And pad the number to the next power of 2.
1078
- tile_tokens_dim = next_positive_power_of_2(num_tokens_per_expert)
1079
- # Cap to 8-64 tokens per CTA tile as it's the range supported by the kernel.
1080
- tile_tokens_dim = min(max(tile_tokens_dim, 8), 64)
1081
- return tile_tokens_dim
1082
-
1083
-
1070
def flashinfer_fused_moe_blockscale_fp8(
1071
routing_logits: torch.Tensor,
1072
routing_bias: torch.Tensor,
@@ -1128,8 +1114,8 @@ def flashinfer_fused_moe_blockscale_fp8(
1114
local_expert_offset=expert_offset,
1115
local_num_experts=local_num_experts,
1116
routed_scaling_factor=routed_scaling,
1131
- tile_tokens_dim=_get_tile_tokens_dim(x.shape[0], top_k,
1132
- global_num_experts),
1117
+ tile_tokens_dim=calculate_tile_tokens_dim(x.shape[0], top_k,
1118
+ global_num_experts),
1119
routing_method_type=2, # DeepSeek-styled routing method
1120
use_shuffled_weight=False,
1121
)
@@ -1164,6 +1150,97 @@ direct_register_custom_op(
1150
)
1151
1152
1153
+def flashinfer_fused_moe_per_tensor_scale_fp8(
1154
+ routing_logits: torch.Tensor,
1155
+ routing_bias: Optional[torch.Tensor],
1156
+ hidden_states: torch.Tensor,
1157
+ input_scale: torch.Tensor,
1158
+ gemm1_weights: torch.Tensor,
1159
+ gemm1_weights_scale: torch.Tensor,
1160
+ activation_scale: torch.Tensor,
1161
+ gemm2_weights: torch.Tensor,
1162
+ gemm2_weights_scale: torch.Tensor,
1163
+ num_experts: int,
1164
+ top_k: int,
1165
+ num_expert_group: Optional[int],
1166
+ topk_group: Optional[int],
1167
+ intermediate_size: int,
1168
+ local_expert_offset: int,
1169
+ local_num_experts: int,
1170
+ use_routing_scales_on_input: bool,
1171
+ routing_method_type: int,
1172
+ routed_scaling_factor: float = 1.0) -> torch.Tensor:
1173
+ num_expert_group = num_expert_group if num_expert_group is not None else 0
1174
+ topk_group = topk_group if topk_group is not None else 0
1175
+
1176
+ quant_hidden_states, input_scale = moe_kernel_quantize_input(
1177
+ hidden_states,
1178
+ input_scale,
1179
+ quant_dtype=torch.float8_e4m3fn,
1180
+ per_act_token_quant=False)
1181
+
1182
+ output1_scales_scalar = gemm1_weights_scale * input_scale * (
1183
+ 1.0 / activation_scale)
1184
+ output1_scales_gate_scalar = gemm1_weights_scale * input_scale
1185
+ output2_scales_scalar = activation_scale * gemm2_weights_scale
1186
+
1187
+ from vllm.utils.flashinfer import (
1188
+ flashinfer_trtllm_fp8_per_tensor_scale_moe)
1189
+ return flashinfer_trtllm_fp8_per_tensor_scale_moe(
1190
+ routing_logits=routing_logits,
1191
+ routing_bias=routing_bias,
1192
+ hidden_states=quant_hidden_states,
1193
+ gemm1_weights=gemm1_weights,
1194
+ output1_scales_scalar=output1_scales_scalar,
1195
+ output1_scales_gate_scalar=output1_scales_gate_scalar,
1196
+ gemm2_weights=gemm2_weights,
1197
+ output2_scales_scalar=output2_scales_scalar,
1198
+ num_experts=num_experts,
1199
+ top_k=top_k,
1200
+ n_group=num_expert_group,
1201
+ topk_group=topk_group,
1202
+ intermediate_size=intermediate_size,
1203
+ local_expert_offset=local_expert_offset,
1204
+ local_num_experts=local_num_experts,
1205
+ routed_scaling_factor=routed_scaling_factor,
1206
+ use_routing_scales_on_input=use_routing_scales_on_input,
1207
+ tile_tokens_dim=calculate_tile_tokens_dim(hidden_states.shape[0],
1208
+ top_k, num_experts),
1209
+ routing_method_type=routing_method_type)
1210
+
1211
+
1212
+def flashinfer_fused_moe_per_tensor_scale_fp8_fake(
1213
+ routing_logits: torch.Tensor,
1214
+ routing_bias: torch.Tensor,
1215
+ hidden_states: torch.Tensor,
1216
+ gemm1_weights: torch.Tensor,
1217
+ output1_scales_scalar: torch.Tensor,
1218
+ output1_scales_gate_scalar: torch.Tensor,
1219
+ gemm2_weights: torch.Tensor,
1220
+ output2_scales_scalar: torch.Tensor,
1221
+ num_experts: int,
1222
+ top_k: int,
1223
+ num_expert_group: int,
1224
+ topk_group: int,
1225
+ intermediate_size: int,
1226
+ local_expert_offset: int,
1227
+ local_num_experts: int,
1228
+ routed_scaling_factor: float = 1.0,
1229
+ use_routing_scales_on_input: bool = False,
1230
+ tile_tokens_dim: int = 8,
1231
+ routing_method_type: int = 0) -> torch.Tensor:
1232
+ pass
1233
+
1234
+
1235
+direct_register_custom_op(
1236
+ op_name="flashinfer_fused_moe_per_tensor_scale_fp8",
1237
+ op_func=flashinfer_fused_moe_per_tensor_scale_fp8,
1238
+ mutates_args=["hidden_states"],
1239
+ fake_impl=flashinfer_fused_moe_per_tensor_scale_fp8_fake,
1240
+ tags=(torch.Tag.needs_fixed_stride_order, ),
1241
+)
1242
+
1243
+
1244
def outplace_fused_experts(
1245
hidden_states: torch.Tensor,
1246
w1: torch.Tensor,
vllm/model_executor/layers/quantization/fp8.py
+44
-31
@@ -23,6 +23,9 @@ from vllm.model_executor.layers.quantization import QuantizationMethods
23
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)
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 (
@@ -53,11 +56,6 @@ ACTIVATION_SCHEMES = ["static", "dynamic"]
56
logger = init_logger(__name__)
57
58
56
-def _swap_w13_to_w31(x: torch.Tensor) -> torch.Tensor:
57
- return x.reshape(-1, 2, x.shape[-2] // 2,
58
- x.shape[-1]).flip(dims=[1]).reshape(x.shape)
59
-
60
-
59
def _is_col_major(x: torch.Tensor) -> bool:
60
assert x.dim() == 3
61
b, m, n = x.shape
@@ -695,11 +693,13 @@ class Fp8MoEMethod(FusedMoEMethodBase):
693
elif self.flashinfer_moe_enabled:
694
# NOTE: weights have to be swapped since the activation is
695
# applied on different half for flashinfer vs vllm
698
- w13_weight = _swap_w13_to_w31(layer.w13_weight.data)
699
- w13_weight_scale_inv = _swap_w13_to_w31(
696
+ w13_weight = swap_w13_to_w31(layer.w13_weight.data)
697
+ w13_weight_scale_inv = swap_w13_to_w31(
698
layer.w13_weight_scale_inv.data)
699
w2_weight = layer.w2_weight.data
700
w2_weight_scale_inv = layer.w2_weight_scale_inv.data
701
+ if not self.block_quant:
702
+ rotate_flashinfer_fp8_moe_weights(w13_weight, w2_weight)
703
else:
704
w13_weight = layer.w13_weight.data
705
w13_weight_scale_inv = layer.w13_weight_scale_inv.data
@@ -998,30 +998,43 @@ class Fp8MoEMethod(FusedMoEMethodBase):
998
global_num_experts=global_num_experts,
999
expert_map=expert_map)
1000
elif self.flashinfer_moe_enabled:
1001
- # Currently only work with DS models
1002
- assert self.block_quant
1003
- assert (renormalize and use_grouped_topk
1004
- and scoring_func == 'sigmoid'
1005
- and custom_routing_function is None)
1006
- assert activation == "silu"
1007
- return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(
1008
- routing_logits=router_logits.to(torch.float32),
1009
- routing_bias=e_score_correction_bias,
1010
- x=x,
1011
- w13_weight=layer.w13_weight,
1012
- w13_weight_scale_inv=layer.w13_weight_scale_inv,
1013
- w2_weight=layer.w2_weight,
1014
- w2_weight_scale_inv=layer.w2_weight_scale_inv,
1015
- global_num_experts=global_num_experts,
1016
- top_k=top_k,
1017
- num_expert_group=num_expert_group,
1018
- topk_group=topk_group,
1019
- intermediate_size=layer.intermediate_size_per_partition,
1020
- expert_offset=layer.ep_rank * layer.local_num_experts,
1021
- local_num_experts=layer.local_num_experts,
1022
- block_shape=self.quant_config.weight_block_size,
1023
- routed_scaling=1.0,
1024
- )
1001
+ assert activation == 'silu'
1002
+ assert scoring_func == 'sigmoid'
1003
+ if self.block_quant:
1004
+ assert (renormalize and use_grouped_topk
1005
+ and custom_routing_function is None)
1006
+
1007
+ return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(
1008
+ routing_logits=router_logits.to(torch.float32),
1009
+ routing_bias=e_score_correction_bias,
1010
+ x=x,
1011
+ w13_weight=layer.w13_weight,
1012
+ w13_weight_scale_inv=layer.w13_weight_scale_inv,
1013
+ w2_weight=layer.w2_weight,
1014
+ w2_weight_scale_inv=layer.w2_weight_scale_inv,
1015
+ global_num_experts=global_num_experts,
1016
+ top_k=top_k,
1017
+ num_expert_group=num_expert_group,
1018
+ topk_group=topk_group,
1019
+ intermediate_size=layer.intermediate_size_per_partition,
1020
+ expert_offset=layer.ep_rank * layer.local_num_experts,
1021
+ local_num_experts=layer.local_num_experts,
1022
+ block_shape=self.quant_config.weight_block_size,
1023
+ routed_scaling=1.0,
1024
+ )
1025
+ else:
1026
+ assert (not renormalize
1027
+ and custom_routing_function is not None)
1028
+ return apply_flashinfer_per_tensor_scale_fp8(
1029
+ layer=layer,
1030
+ hidden_states=x,
1031
+ router_logits=router_logits,
1032
+ routing_bias=e_score_correction_bias,
1033
+ global_num_experts=global_num_experts,
1034
+ top_k=top_k,
1035
+ num_expert_group=num_expert_group,
1036
+ topk_group=topk_group,
1037
+ apply_router_weight_on_input=apply_router_weight_on_input)
1038
else:
1039
return self.fused_experts(
1040
hidden_states=x,
vllm/model_executor/layers/quantization/modelopt.py
+28
@@ -23,6 +23,9 @@ from vllm.model_executor.layers.quantization import QuantizationMethods
23
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)
29
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
30
apply_fp4_marlin_linear, is_fp4_marlin_supported,
31
prepare_fp4_layer_for_marlin, prepare_moe_fp4_layer_for_marlin)
@@ -34,6 +37,7 @@ from vllm.model_executor.parameter import (ModelWeightParameter,
37
PerTensorScaleParameter)
38
from vllm.platforms import current_platform
39
from vllm.scalar_type import scalar_types
40
+from vllm.utils.flashinfer import has_flashinfer_moe
41
42
logger = init_logger(__name__)
43
@@ -267,6 +271,11 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
271
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
272
cutlass_fp8_supported)
273
self.cutlass_fp8_supported = cutlass_fp8_supported()
274
+ self.flashinfer_moe_enabled = False
275
+ if envs.VLLM_USE_FLASHINFER_MOE_FP8 and has_flashinfer_moe():
276
+ logger.info_once(
277
+ "Using FlashInfer MoE FP8 kernels for ModelOptFp8MoEMethod.")
278
+ self.flashinfer_moe_enabled = True
279
280
def create_weights(
281
self,
@@ -410,6 +419,11 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
419
layer.w2_input_scale = Parameter(layer.w2_input_scale.max(),
420
requires_grad=False)
421
422
+ if self.flashinfer_moe_enabled:
423
+ layer.w13_weight.data = swap_w13_to_w31(layer.w13_weight.data)
424
+ rotate_flashinfer_fp8_moe_weights(layer.w13_weight,
425
+ layer.w2_weight)
426
+
427
def apply(
428
self,
429
layer: torch.nn.Module,
@@ -436,6 +450,20 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
450
raise NotImplementedError(
451
"EPLB not supported for `ModelOptFp8MoEMethod` yet.")
452
453
+ if self.flashinfer_moe_enabled:
454
+ assert activation == 'silu'
455
+ assert not renormalize
456
+ return apply_flashinfer_per_tensor_scale_fp8(
457
+ layer=layer,
458
+ hidden_states=x,
459
+ router_logits=router_logits,
460
+ routing_bias=e_score_correction_bias,
461
+ global_num_experts=global_num_experts,
462
+ top_k=top_k,
463
+ num_expert_group=num_expert_group,
464
+ topk_group=topk_group,
465
+ apply_router_weight_on_input=apply_router_weight_on_input)
466
+
467
# Expert selection
468
topk_weights, topk_ids = FusedMoE.select_experts(
469
hidden_states=x,
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
new
+100
@@ -0,0 +1,100 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import Optional
4
+
5
+import torch
6
+
7
+
8
+def calculate_tile_tokens_dim(num_tokens, top_k, num_experts):
9
+ from flashinfer import next_positive_power_of_2
10
+
11
+ # Guess tokens per expert assuming perfect expert distribution first.
12
+ num_tokens_per_expert = (num_tokens * top_k) // num_experts
13
+ # And pad the number to the next power of 2.
14
+ tile_tokens_dim = next_positive_power_of_2(num_tokens_per_expert)
15
+ # Cap to 8-64 tokens per CTA tile as it's the range supported by the kernel.
16
+ tile_tokens_dim = min(max(tile_tokens_dim, 8), 64)
17
+ return tile_tokens_dim
18
+
19
+
20
+def swap_w13_to_w31(x: torch.Tensor) -> torch.Tensor:
21
+ return x.reshape(-1, 2, x.shape[-2] // 2,
22
+ x.shape[-1]).flip(dims=[1]).reshape(x.shape)
23
+
24
+
25
+def rotate_flashinfer_fp8_moe_weights(gemm1_weights: torch.Tensor,
26
+ gemm2_weights: torch.Tensor):
27
+ from flashinfer import reorder_rows_for_gated_act_gemm, shuffle_matrix_a
28
+ epilogue_tile_m = 128
29
+ num_experts = gemm1_weights.shape[0]
30
+ hidden_size = gemm1_weights.shape[-1]
31
+ intermediate_size = gemm1_weights.shape[1] // 2
32
+
33
+ # Reorder rows of W1 for fused gated activation
34
+ gemm1_weights_fp8_interleaved = []
35
+ for i in range(num_experts):
36
+ gemm1_weights_fp8_interleaved.append(
37
+ reorder_rows_for_gated_act_gemm(gemm1_weights[i]))
38
+
39
+ # Stack weights and scales for all experts
40
+ gemm1_weights_fp8_interleaved = torch.stack(
41
+ gemm1_weights_fp8_interleaved).reshape(num_experts,
42
+ 2 * intermediate_size,
43
+ hidden_size)
44
+
45
+ # Shuffle weights and scaling factors for transposed mma output
46
+ gemm1_weights_fp8_shuffled = []
47
+ gemm2_weights_fp8_shuffled = []
48
+ for i in range(num_experts):
49
+ gemm1_weights_fp8_shuffled.append(
50
+ shuffle_matrix_a(
51
+ gemm1_weights_fp8_interleaved[i].view(torch.uint8),
52
+ epilogue_tile_m))
53
+
54
+ gemm2_weights_fp8_shuffled.append(
55
+ shuffle_matrix_a(gemm2_weights[i].view(torch.uint8),
56
+ epilogue_tile_m))
57
+
58
+ # Stack weights for all experts
59
+ gemm1_weights.data = torch.stack(gemm1_weights_fp8_shuffled).view(
60
+ torch.float8_e4m3fn)
61
+ gemm2_weights.data = torch.stack(gemm2_weights_fp8_shuffled).view(
62
+ torch.float8_e4m3fn)
63
+
64
+
65
+def apply_flashinfer_per_tensor_scale_fp8(
66
+ layer: torch.nn.Module,
67
+ hidden_states: torch.Tensor,
68
+ router_logits: torch.Tensor,
69
+ routing_bias: Optional[torch.Tensor],
70
+ top_k: int,
71
+ num_expert_group: Optional[int],
72
+ topk_group: Optional[int],
73
+ global_num_experts: int,
74
+ apply_router_weight_on_input: bool,
75
+) -> torch.Tensor:
76
+ from flashinfer.fused_moe import RoutingMethodType
77
+
78
+ from vllm.model_executor.models.llama4 import Llama4MoE
79
+ assert layer.custom_routing_function == Llama4MoE.custom_routing_function, \
80
+ "FusedMoE flashinfer kernels are only supported for Llama4"
81
+ return torch.ops.vllm.flashinfer_fused_moe_per_tensor_scale_fp8(
82
+ routing_logits=router_logits,
83
+ routing_bias=routing_bias,
84
+ hidden_states=hidden_states,
85
+ input_scale=layer.w13_input_scale,
86
+ gemm1_weights=layer.w13_weight,
87
+ gemm1_weights_scale=layer.w13_weight_scale,
88
+ gemm2_weights=layer.w2_weight,
89
+ gemm2_weights_scale=layer.w2_weight_scale,
90
+ activation_scale=layer.w2_input_scale,
91
+ num_experts=global_num_experts,
92
+ top_k=top_k,
93
+ num_expert_group=num_expert_group,
94
+ topk_group=topk_group,
95
+ intermediate_size=layer.intermediate_size_per_partition,
96
+ local_expert_offset=layer.ep_rank * layer.local_num_experts,
97
+ local_num_experts=layer.local_num_experts,
98
+ use_routing_scales_on_input=apply_router_weight_on_input,
99
+ routing_method_type=RoutingMethodType.Llama4,
100
+ )
\ No newline at end of file
vllm/utils/flashinfer.py
+2
@@ -66,6 +66,8 @@ def _lazy_import_wrapper(module_name: str,
66
# Create lazy wrappers for each function
67
flashinfer_trtllm_fp8_block_scale_moe = _lazy_import_wrapper(
68
"flashinfer.fused_moe", "trtllm_fp8_block_scale_moe")
69
+flashinfer_trtllm_fp8_per_tensor_scale_moe = _lazy_import_wrapper(
70
+ "flashinfer.fused_moe", "trtllm_fp8_per_tensor_scale_moe")
71
flashinfer_cutlass_fused_moe = _lazy_import_wrapper("flashinfer.fused_moe",
72
"cutlass_fused_moe")
73
fp4_quantize = _lazy_import_wrapper("flashinfer", "fp4_quantize")