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 13, 13:15 UTC
[Performance] Support FP8 flashinfer TRTLLM MOE on Qwen3 and Qwen-3next (#27492)
Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com>
jiahanc committed
Nov 10, 2025 at 09:34 UTC
34553b9d2702dd2a27a578fec819e88e76dcbfb4
7 files changed
+78
-30
vllm/model_executor/layers/fused_moe/config.py
+21
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from dataclasses import dataclass
4
+from enum import IntEnum
5
from typing import Optional, Union
6
7
import torch
@@ -91,6 +92,26 @@ def _quant_flags_to_group_shape(
92
return a_shape, w_shape
93
94
95
+# The type of method in top-K routing
96
+# Please keep this in sync with the counterpart defined in https://github.com/flashinfer-ai/flashinfer/blob/main/include/flashinfer/trtllm/fused_moe/runner.h
97
+class RoutingMethodType(IntEnum):
98
+ # Default: Softmax -> TopK
99
+ Default = (0,)
100
+ # Renormalize: TopK -> Softmax
101
+ Renormalize = (1,)
102
+ # DeepSeekV3: Sigmoid -> RoutingBiasAdd -> Top2 in group -> Top4 groups
103
+ # -> Top8 experts from the Top4 groups
104
+ DeepSeekV3 = (2,)
105
+ # Llama4: Top1 -> Sigmoid
106
+ Llama4 = (3,)
107
+ # RenormalizeNaive: Softmax -> TopK -> Renormalize
108
+ RenormalizeNaive = (4,)
109
+ # TopK: TopK (no softmax)
110
+ TopK = (5,)
111
+ # Unspecified
112
+ Unspecified = 6.0
113
+
114
+
115
@dataclass
116
class FusedMoEQuantDesc:
117
"""
vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py
+12
-14
@@ -3,6 +3,7 @@
3
4
import torch
5
6
+from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
7
from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
8
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
9
calculate_tile_tokens_dim,
@@ -23,26 +24,24 @@ def flashinfer_fused_moe_blockscale_fp8(
24
w2_weight_scale_inv: torch.Tensor,
25
global_num_experts: int,
26
top_k: int,
26
- num_expert_group: int,
27
- topk_group: int,
27
+ num_expert_group: int | None,
28
+ topk_group: int | None,
29
intermediate_size: int,
30
expert_offset: int,
31
local_num_experts: int,
32
block_shape: list[int],
32
- routed_scaling: float = 1.0,
33
+ routing_method_type: int = RoutingMethodType.DeepSeekV3,
34
+ routed_scaling: float | None = 1.0,
35
) -> torch.Tensor:
36
from vllm.utils.flashinfer import flashinfer_trtllm_fp8_block_scale_moe
37
38
+ topk_group = topk_group if topk_group is not None else 0
39
assert top_k <= global_num_experts
37
- assert top_k <= 8
38
- assert topk_group <= 4
39
- assert global_num_experts > num_expert_group
40
- assert global_num_experts % num_expert_group == 0
40
+ assert top_k <= 10
41
assert global_num_experts % 4 == 0
42
- assert top_k < (topk_group * global_num_experts / num_expert_group)
42
assert block_shape == [128, 128]
44
- # Routing kernel expects #experts <= #threads 256
45
- assert global_num_experts <= 256
43
+ # Routing kernel expects #experts <= #threads 512
44
+ assert global_num_experts <= 512
45
46
a_q, a_sf = per_token_group_quant_fp8(x, block_shape[1])
47
# NOTE: scales of hidden states have to be transposed!
@@ -64,10 +63,8 @@ def flashinfer_fused_moe_blockscale_fp8(
63
local_expert_offset=expert_offset,
64
local_num_experts=local_num_experts,
65
routed_scaling_factor=routed_scaling,
67
- tile_tokens_dim=calculate_tile_tokens_dim(
68
- x.shape[0], top_k, global_num_experts
69
- ),
70
- routing_method_type=2, # DeepSeek-styled routing method
66
+ tile_tokens_dim=None,
67
+ routing_method_type=routing_method_type,
68
use_shuffled_weight=False,
69
)
70
@@ -88,6 +85,7 @@ def flashinfer_fused_moe_blockscale_fp8_fake(
85
expert_offset: int,
86
local_num_experts: int,
87
block_shape: list[int],
88
+ routing_method_type: int,
89
routed_scaling: float = 1.0,
90
) -> torch.Tensor:
91
return torch.empty_like(x)
vllm/model_executor/layers/fused_moe/layer.py
+20
@@ -31,6 +31,7 @@ from vllm.model_executor.layers.fused_moe.config import (
31
FusedMoEConfig,
32
FusedMoEParallelConfig,
33
FusedMoEQuantConfig,
34
+ RoutingMethodType,
35
biased_moe_quant_config,
36
)
37
from vllm.model_executor.layers.fused_moe.fused_moe import zero_experts_compute_triton
@@ -1213,6 +1214,7 @@ class FusedMoE(CustomOp):
1214
zero_expert_type: str | None = None,
1215
expert_mapping: list[tuple[str, str, int, str]] | None = None,
1216
n_shared_experts: int | None = None,
1217
+ routing_method_type: int | None = None,
1218
):
1219
super().__init__()
1220
@@ -1397,6 +1399,24 @@ class FusedMoE(CustomOp):
1399
"Only softmax scoring function is supported for non-grouped topk."
1400
)
1401
1402
+ # ToDo: Better logic to determine the routing method type
1403
+ if routing_method_type is not None:
1404
+ self.routing_method_type = routing_method_type
1405
+ else:
1406
+ if scoring_func == "sigmoid":
1407
+ if self.use_grouped_topk:
1408
+ self.routing_method_type = RoutingMethodType.DeepSeekV3
1409
+ elif self.top_k == 1:
1410
+ self.routing_method_type = RoutingMethodType.Llama4
1411
+ elif self.scoring_func == "softmax":
1412
+ self.routing_method_type = (
1413
+ RoutingMethodType.Renormalize
1414
+ if not self.renormalize
1415
+ else RoutingMethodType.RenormalizeNaive
1416
+ )
1417
+ else:
1418
+ self.routing_method_type = RoutingMethodType.TopK
1419
+
1420
self.moe_config: FusedMoEConfig = FusedMoEConfig(
1421
num_experts=self.global_num_experts,
1422
experts_per_token=top_k,
vllm/model_executor/layers/quantization/fp8.py
+7
-7
@@ -28,6 +28,7 @@ from vllm.model_executor.layers.fused_moe import (
28
)
29
from vllm.model_executor.layers.fused_moe.config import (
30
FusedMoEQuantConfig,
31
+ RoutingMethodType,
32
fp8_w8a8_moe_quant_config,
33
)
34
from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe
@@ -1222,22 +1223,20 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1223
assert activation == "silu", (
1224
f"Expected 'silu' activation but got {activation}"
1225
)
1225
- assert scoring_func == "sigmoid", (
1226
- f"Expected 'sigmoid' scoring func but got {scoring_func}"
1227
- )
1226
+
1227
if self.block_quant:
1228
import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401
1229
1231
- assert (
1232
- renormalize and use_grouped_topk and custom_routing_function is None
1233
- )
1230
e_score_correction_bias = (
1231
e_score_correction_bias.to(x.dtype)
1232
if e_score_correction_bias is not None
1233
else None
1234
)
1235
+ routing_method_type = layer.routing_method_type
1236
return torch.ops.vllm.flashinfer_fused_moe_blockscale_fp8(
1240
- routing_logits=router_logits.to(torch.float32),
1237
+ routing_logits=router_logits.to(torch.float32)
1238
+ if routing_method_type == RoutingMethodType.DeepSeekV3
1239
+ else router_logits,
1240
routing_bias=e_score_correction_bias,
1241
x=x,
1242
w13_weight=layer.w13_weight,
@@ -1252,6 +1251,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1251
expert_offset=layer.ep_rank * layer.local_num_experts,
1252
local_num_experts=layer.local_num_experts,
1253
block_shape=self.weight_block_size,
1254
+ routing_method_type=routing_method_type,
1255
routed_scaling=routed_scaling_factor,
1256
)
1257
else:
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+14
-9
@@ -27,20 +27,25 @@ class FlashinferMoeBackend(Enum):
27
28
29
def calculate_tile_tokens_dim(num_tokens, top_k, num_experts):
30
+ from flashinfer import next_positive_power_of_2
31
+
32
# FlashInfer 0.2.10 has issues with larger tile sizes. Set to 8 for now.
33
# TODO: Revert this to dynamic calculation once a new version of FlashInfer
34
# with the necessary kernels is released.
35
tile_tokens_dim = 8
36
35
- # from flashinfer import next_positive_power_of_2
36
-
37
- # # Guess tokens per expert assuming perfect expert distribution first.
38
- # num_tokens_per_expert = (num_tokens * top_k) // num_experts
39
- # # And pad the number to the next power of 2.
40
- # tile_tokens_dim = next_positive_power_of_2(num_tokens_per_expert)
41
- # # Cap to 8-64 tokens per CTA tile as it's the range supported by the
42
- # # kernel.
43
- # tile_tokens_dim = min(max(tile_tokens_dim, 8), 64)
37
+ # A factor considering tokens are not perfectly balanced among experts.
38
+ imbalance_factor = 1.3
39
+ # Calculate the number of tokens per expert
40
+ # assuming perfect distribution.
41
+ num_tokens_per_expert = (num_tokens * top_k) // num_experts
42
+ # Apply the imbalance factor.
43
+ num_tokens_per_expert = int(num_tokens_per_expert * imbalance_factor)
44
+ # And pad the number to the next power of 2.
45
+ tile_tokens_dim = next_positive_power_of_2(num_tokens_per_expert)
46
+ # Cap to 8-max_tile_tokens_dim tokens per CTA tile
47
+ # as it's the range supported by the kernel.
48
+ tile_tokens_dim = min(max(tile_tokens_dim, 8), 64)
49
50
return tile_tokens_dim
51
vllm/model_executor/models/qwen3_moe.py
+2
@@ -43,6 +43,7 @@ from vllm.distributed import (
43
from vllm.logger import init_logger
44
from vllm.model_executor.layers.activation import SiluAndMul
45
from vllm.model_executor.layers.fused_moe import FusedMoE
46
+from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
47
from vllm.model_executor.layers.layernorm import RMSNorm
48
from vllm.model_executor.layers.linear import (
49
MergedColumnParallelLinear,
@@ -171,6 +172,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
172
enable_eplb=self.enable_eplb,
173
num_redundant_experts=self.n_redundant_experts,
174
is_sequence_parallel=self.is_sequence_parallel,
175
+ routing_method_type=RoutingMethodType.Renormalize,
176
)
177
178
self.gate = ReplicatedLinear(
vllm/model_executor/models/qwen3_next.py
+2
@@ -34,6 +34,7 @@ from vllm.model_executor.layers.fla.ops import (
34
fused_recurrent_gated_delta_rule,
35
)
36
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
37
+from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
38
from vllm.model_executor.layers.layernorm import (
39
GemmaRMSNorm as Qwen3NextRMSNorm,
40
)
@@ -173,6 +174,7 @@ class Qwen3NextSparseMoeBlock(nn.Module):
174
enable_eplb=self.enable_eplb,
175
num_redundant_experts=self.n_redundant_experts,
176
is_sequence_parallel=self.is_sequence_parallel,
177
+ routing_method_type=RoutingMethodType.Renormalize,
178
)
179
180
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: