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 9, 10:00 UTC
[NVIDIA] [feat] Integrate flashinfer Trtllmgen bf16 moe (#32954)
Signed-off-by: Linda-Stadter <57756729+Linda-Stadter@users.noreply.github.com>
Linda committed
Jan 29, 2026 at 19:00 UTC
0493d897c4b7ed645b506bcd0452b7f3326f85f6
5 files changed
+290
-17
vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py
+103
@@ -77,6 +77,18 @@ def _supports_routing_method(
77
raise ValueError("Unsupported quantization scheme.")
78
79
80
+def _supports_routing_method_bf16(
81
+ routing_method: RoutingMethodType,
82
+) -> bool:
83
+ return routing_method in [
84
+ RoutingMethodType.Default,
85
+ RoutingMethodType.Renormalize,
86
+ RoutingMethodType.DeepSeekV3,
87
+ RoutingMethodType.Llama4,
88
+ RoutingMethodType.RenormalizeNaive,
89
+ ]
90
+
91
+
92
def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
93
"""Supports TRTLLM Kernel does not support EPLB."""
94
return not moe_parallel_config.enable_eplb
@@ -115,6 +127,34 @@ def is_supported_config_trtllm(
127
return True, None
128
129
130
+def is_supported_config_trtllm_bf16(
131
+ moe_config: FusedMoEConfig,
132
+ activation_format: mk.FusedMoEActivationFormat,
133
+) -> tuple[bool, str | None]:
134
+ """
135
+ This method mirrors mk.FusedMoEPermuteExpertsUnpermute.is_supported_config
136
+ for BF16 unquantized kernels.
137
+ """
138
+
139
+ def _make_reason(reason: str) -> str:
140
+ return f"kernel does not support {reason}"
141
+
142
+ if not _supports_current_device():
143
+ return False, _make_reason("current device")
144
+ elif not (moe_config.is_act_and_mul or _supports_no_act_and_mul()):
145
+ return False, _make_reason("no act_and_mul MLP layer")
146
+ elif not _supports_activation(moe_config.activation):
147
+ return False, _make_reason(f"{moe_config.activation} activation")
148
+ elif not _supports_parallel_config(moe_config.moe_parallel_config):
149
+ return False, _make_reason("parallel config")
150
+ elif not _supports_routing_method_bf16(moe_config.routing_method):
151
+ return False, _make_reason("routing method")
152
+ elif activation_format != mk.FusedMoEActivationFormat.Standard:
153
+ return False, _make_reason("activation format")
154
+
155
+ return True, None
156
+
157
+
158
def flashinfer_fused_moe_blockscale_fp8(
159
routing_logits: torch.Tensor,
160
routing_bias: torch.Tensor,
@@ -287,3 +327,66 @@ direct_register_custom_op(
327
fake_impl=fi_trtllm_fp8_per_tensor_moe_fake,
328
tags=(torch.Tag.needs_fixed_stride_order,),
329
)
330
+
331
+
332
+def flashinfer_fused_moe_bf16(
333
+ routing_logits: torch.Tensor,
334
+ routing_bias: torch.Tensor | None,
335
+ hidden_states: torch.Tensor,
336
+ gemm1_weights: torch.Tensor,
337
+ gemm2_weights: torch.Tensor,
338
+ num_experts: int,
339
+ top_k: int,
340
+ n_group: int | None,
341
+ topk_group: int | None,
342
+ intermediate_size: int,
343
+ local_expert_offset: int,
344
+ local_num_experts: int,
345
+ routing_method_type: int,
346
+ tune_max_num_tokens: int = 8192,
347
+) -> torch.Tensor:
348
+ from vllm.utils.flashinfer import flashinfer_trtllm_bf16_moe
349
+
350
+ return flashinfer_trtllm_bf16_moe(
351
+ routing_logits=routing_logits,
352
+ routing_bias=routing_bias,
353
+ hidden_states=hidden_states,
354
+ gemm1_weights=gemm1_weights,
355
+ gemm2_weights=gemm2_weights,
356
+ num_experts=num_experts,
357
+ top_k=top_k,
358
+ n_group=n_group,
359
+ topk_group=topk_group,
360
+ intermediate_size=intermediate_size,
361
+ local_expert_offset=local_expert_offset,
362
+ local_num_experts=local_num_experts,
363
+ routing_method_type=routing_method_type,
364
+ tune_max_num_tokens=tune_max_num_tokens,
365
+ )
366
+
367
+
368
+def flashinfer_fused_moe_bf16_fake(
369
+ routing_logits: torch.Tensor,
370
+ routing_bias: torch.Tensor | None,
371
+ hidden_states: torch.Tensor,
372
+ gemm1_weights: torch.Tensor,
373
+ gemm2_weights: torch.Tensor,
374
+ num_experts: int,
375
+ top_k: int,
376
+ n_group: int | None,
377
+ topk_group: int | None,
378
+ intermediate_size: int,
379
+ local_expert_offset: int,
380
+ local_num_experts: int,
381
+ routing_method_type: int = RoutingMethodType.Renormalize,
382
+ tune_max_num_tokens: int = 8192,
383
+) -> torch.Tensor:
384
+ return torch.empty_like(hidden_states)
385
+
386
+
387
+direct_register_custom_op(
388
+ op_name="flashinfer_fused_moe_bf16",
389
+ op_func=flashinfer_fused_moe_bf16,
390
+ fake_impl=flashinfer_fused_moe_bf16_fake,
391
+ tags=(torch.Tag.needs_fixed_stride_order,),
392
+)
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+35
-6
@@ -14,6 +14,9 @@ from vllm.model_executor.layers.fused_moe.config import (
14
FusedMoEConfig,
15
FusedMoEQuantConfig,
16
)
17
+from vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe import (
18
+ is_supported_config_trtllm_bf16,
19
+)
20
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
21
MoEPrepareAndFinalizeNoEP,
22
)
@@ -21,12 +24,13 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
24
swap_w13_to_w31,
25
)
26
from vllm.platforms import current_platform
24
-from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
27
+from vllm.utils.flashinfer import has_flashinfer, has_flashinfer_cutlass_fused_moe
28
29
logger = init_logger(__name__)
30
31
32
class UnquantizedMoeBackend(Enum):
33
+ FLASHINFER_TRTLLM = "FlashInfer TRTLLM"
34
FLASHINFER_CUTLASS = "FlashInfer CUTLASS"
35
AITER = "ROCm AITER"
36
TRITON = "TRITON"
@@ -40,6 +44,7 @@ class UnquantizedMoeBackend(Enum):
44
# that is not conform with Modular kernel format.
45
# We will directly call the kernel for those backend
46
UNSUPPORTED_BACKEND = [
47
+ UnquantizedMoeBackend.FLASHINFER_TRTLLM,
48
UnquantizedMoeBackend.CPU,
49
UnquantizedMoeBackend.XPU,
50
UnquantizedMoeBackend.TPU,
@@ -48,11 +53,12 @@ UNSUPPORTED_BACKEND = [
53
54
55
def select_unquantized_moe_backend(
56
+ moe_config: FusedMoEConfig,
57
use_ep: bool,
58
use_dp: bool,
59
) -> UnquantizedMoeBackend:
60
"""
55
- Select the primary FP8 MoE backend
61
+ Select the primary Unquantized MoE backend
62
Note: Shape-specific fallbacks may still occur at runtime.
63
"""
64
@@ -61,13 +67,27 @@ def select_unquantized_moe_backend(
67
68
rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
69
70
+ activation_format = (
71
+ mk.FusedMoEActivationFormat.BatchedExperts
72
+ if moe_config.moe_parallel_config.use_batched_activation_format
73
+ else mk.FusedMoEActivationFormat.Standard
74
+ )
75
+
76
+ # Check if FlashInfer TRTLLM BF16 MoE is supported
77
+ trtllm_supported, _ = is_supported_config_trtllm_bf16(
78
+ moe_config=moe_config,
79
+ activation_format=activation_format,
80
+ )
81
+ flashinfer_trtllm_moe_enabled = (
82
+ has_flashinfer() and envs.VLLM_USE_FLASHINFER_MOE_FP16 and trtllm_supported
83
+ )
84
# FlashInfer CUTLASS MoE is only supported on Hopper and later GPUS
85
flashinfer_cutlass_moe_enabled = (
86
has_flashinfer_cutlass_fused_moe()
87
and envs.VLLM_USE_FLASHINFER_MOE_FP16
88
and use_ep
89
and (not use_dp)
70
- and current_platform.get_device_capability()[0] >= 9
90
+ and current_platform.has_device_capability(90)
91
)
92
if current_platform.is_rocm():
93
if rocm_aiter_moe_enabled:
@@ -75,12 +95,21 @@ def select_unquantized_moe_backend(
95
else:
96
backend = UnquantizedMoeBackend.TRITON
97
if current_platform.is_cuda():
78
- if flashinfer_cutlass_moe_enabled:
98
+ if flashinfer_trtllm_moe_enabled:
99
+ backend = UnquantizedMoeBackend.FLASHINFER_TRTLLM
100
+ elif flashinfer_cutlass_moe_enabled:
101
backend = UnquantizedMoeBackend.FLASHINFER_CUTLASS
102
else:
81
- if use_ep and (not use_dp):
103
+ if not envs.VLLM_USE_FLASHINFER_MOE_FP16 and trtllm_supported:
104
+ logger.info_once(
105
+ "FlashInfer TRTLLM MoE is available but not enabled, "
106
+ "consider setting VLLM_USE_FLASHINFER_MOE_FP16=1 "
107
+ "to enable it for better performance.",
108
+ scope="local",
109
+ )
110
+ elif use_ep and (not use_dp):
111
logger.info_once(
83
- "FlashInfer CUTLASS MoE is available for EP"
112
+ "FlashInfer MoE is available for EP"
113
" but not enabled, consider setting"
114
" VLLM_USE_FLASHINFER_MOE_FP16=1 to enable it.",
115
scope="local",
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+74
-11
@@ -6,6 +6,7 @@ from collections.abc import Callable
6
import torch
7
import torch.nn.functional as F
8
from torch.nn import Module
9
+from torch.nn.parameter import Parameter
10
11
import vllm.envs as envs
12
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
@@ -32,6 +33,9 @@ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
33
make_unquantized_moe_kernel,
34
select_unquantized_moe_backend,
35
)
36
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
37
+ convert_moe_weights_to_flashinfer_trtllm_block_layout,
38
+)
39
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
40
from vllm.platforms import current_platform
41
from vllm.platforms.interface import CpuArchEnum
@@ -54,6 +58,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
58
def __init__(self, moe: FusedMoEConfig):
59
super().__init__(moe)
60
self.unquantized_backend = select_unquantized_moe_backend(
61
+ moe_config=self.moe,
62
use_ep=self.moe.moe_parallel_config.use_ep,
63
use_dp=self.moe.moe_parallel_config.dp_size > 1,
64
)
@@ -64,7 +69,32 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
69
rocm_aiter_ops.is_fused_moe_enabled() and moe.is_act_and_mul
70
)
71
self.kernel: mk.FusedMoEModularKernel | None = None
67
- self._is_monolithic = current_platform.is_cpu() or current_platform.is_xpu()
72
+ self._is_monolithic = (
73
+ current_platform.is_cpu()
74
+ or current_platform.is_xpu()
75
+ or self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM
76
+ )
77
+
78
+ if self.is_monolithic:
79
+ self.apply_monolithic: Callable = self._select_monolithic()
80
+
81
+ def _select_monolithic(self) -> Callable:
82
+ """Select the monolithic implementation based on platform."""
83
+ if current_platform.is_cpu():
84
+ return self.forward_monolithic_cpu
85
+ elif current_platform.is_xpu():
86
+ return self.forward_monolithic_xpu
87
+ else:
88
+ return self.forward_monolithic_cuda
89
+
90
+ def forward_native(
91
+ self,
92
+ layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
93
+ x: torch.Tensor,
94
+ topk_weights: torch.Tensor,
95
+ topk_ids: torch.Tensor,
96
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
97
+ return self.forward_cuda(layer, x, topk_weights, topk_ids)
98
99
@property
100
def is_monolithic(self) -> bool:
@@ -211,7 +241,22 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
241
layer.w13_weight.data = self._maybe_pad_weight(layer.w13_weight.data)
242
layer.w2_weight.data = self._maybe_pad_weight(layer.w2_weight.data)
243
214
- if self.unquantized_backend == UnquantizedMoeBackend.XPU:
244
+ if self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:
245
+ _cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
246
+ # Swap halves to arrange as [w3; w1] (kernel expectation)
247
+ w1_w, w3_w = torch.chunk(layer.w13_weight.data, 2, dim=1)
248
+ w13_weight_swapped = torch.cat([w3_w, w1_w], dim=1)
249
+ layer.w13_weight.data = w13_weight_swapped.contiguous()
250
+ w13_weights_shuffled, w2_weights_shuffled = (
251
+ convert_moe_weights_to_flashinfer_trtllm_block_layout(
252
+ _cache_permute_indices,
253
+ layer.w13_weight.data,
254
+ layer.w2_weight.data,
255
+ )
256
+ )
257
+ layer.w13_weight = Parameter(w13_weights_shuffled, requires_grad=False)
258
+ layer.w2_weight = Parameter(w2_weights_shuffled, requires_grad=False)
259
+ elif self.unquantized_backend == UnquantizedMoeBackend.XPU:
260
import intel_extension_for_pytorch as ipex
261
262
ep_rank_start = self.moe.ep_rank * self.moe.num_local_experts
@@ -290,6 +335,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
335
topk_ids: torch.Tensor,
336
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
337
assert self.kernel is not None
338
+
339
return self.kernel(
340
hidden_states=x,
341
w1=layer.w13_weight,
@@ -303,6 +349,32 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
349
expert_map=layer.expert_map,
350
)
351
352
+ def forward_monolithic_cuda(
353
+ self,
354
+ layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
355
+ x: torch.Tensor,
356
+ router_logits: torch.Tensor,
357
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
358
+ import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: F401
359
+
360
+ assert self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM
361
+
362
+ return torch.ops.vllm.flashinfer_fused_moe_bf16(
363
+ routing_logits=router_logits,
364
+ routing_bias=layer.e_score_correction_bias,
365
+ hidden_states=x,
366
+ gemm1_weights=layer.w13_weight,
367
+ gemm2_weights=layer.w2_weight,
368
+ num_experts=layer.global_num_experts,
369
+ top_k=layer.top_k,
370
+ n_group=layer.num_expert_group,
371
+ topk_group=layer.topk_group,
372
+ intermediate_size=layer.intermediate_size_per_partition,
373
+ local_expert_offset=layer.ep_rank * layer.local_num_experts,
374
+ local_num_experts=layer.local_num_experts,
375
+ routing_method_type=layer.routing_method_type,
376
+ )
377
+
378
def forward_monolithic_cpu(
379
self,
380
layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
@@ -344,12 +416,3 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
416
layer.num_expert_group,
417
custom_routing_function=layer.custom_routing_function,
418
)
347
-
348
- if current_platform.is_cpu():
349
- forward_native: Callable = forward_monolithic_cpu
350
- apply_monolithic = forward_monolithic_cpu
351
- elif current_platform.is_xpu():
352
- forward_native = forward_monolithic_xpu
353
- apply_monolithic = forward_monolithic_xpu
354
- else:
355
- forward_native = forward_cuda
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+75
@@ -195,6 +195,81 @@ def is_flashinfer_supporting_global_sf(backend: FlashinferMoeBackend | None) ->
195
return backend in backends_supporting_global_sf
196
197
198
+def convert_moe_weights_to_flashinfer_trtllm_block_layout(
199
+ cache_permute_indices: dict[torch.Size, torch.Tensor],
200
+ w13_weight: torch.Tensor,
201
+ w2_weight: torch.Tensor,
202
+) -> tuple[torch.Tensor, torch.Tensor]:
203
+ """Convert expert weights to FlashInfer's block layout.
204
+
205
+ This reorders W13 and W2 into the expected epilogue-tiled block layout and
206
+ returns the shuffled weight tensors.
207
+ """
208
+ if w13_weight.dtype != torch.bfloat16 or w2_weight.dtype != torch.bfloat16:
209
+ raise ValueError(
210
+ "Unquantized Moe Backend FlashInfer TRTLLM requires bfloat16 weights"
211
+ )
212
+
213
+ from flashinfer.fused_moe.core import (
214
+ _maybe_get_cached_w3_w1_permute_indices,
215
+ convert_to_block_layout,
216
+ get_w2_permute_indices_with_cache,
217
+ )
218
+
219
+ epilogue_tile_m = 128
220
+ block_k = 128
221
+
222
+ # Reorder rows of W13 and W2 for fused gated activation and convert to the
223
+ # block layout expected by the FlashInfer kernel.
224
+ num_experts = w13_weight.shape[0]
225
+ device_w13 = w13_weight.device
226
+ device_w2 = w2_weight.device
227
+
228
+ w13_weights_shuffled: list[torch.Tensor] = []
229
+ w2_weights_shuffled: list[torch.Tensor] = []
230
+
231
+ for i in range(num_experts):
232
+ permute_indices = _maybe_get_cached_w3_w1_permute_indices(
233
+ cache_permute_indices,
234
+ w13_weight[i].view(torch.uint8),
235
+ epilogue_tile_m,
236
+ )
237
+ tmp_weights1 = (
238
+ w13_weight[i]
239
+ .clone()
240
+ .view(torch.uint8)[permute_indices.to(device_w13)]
241
+ .contiguous()
242
+ )
243
+
244
+ permute_indices = get_w2_permute_indices_with_cache(
245
+ cache_permute_indices,
246
+ w2_weight[i].view(torch.uint8),
247
+ epilogue_tile_m,
248
+ )
249
+ tmp_weights2 = (
250
+ w2_weight[i]
251
+ .clone()
252
+ .view(torch.uint8)[permute_indices.to(device_w2)]
253
+ .contiguous()
254
+ )
255
+
256
+ tmp_weights1 = convert_to_block_layout(tmp_weights1.view(torch.uint8), block_k)
257
+ tmp_weights2 = convert_to_block_layout(tmp_weights2.view(torch.uint8), block_k)
258
+
259
+ w13_weights_shuffled.append(tmp_weights1.view(torch.bfloat16))
260
+ w2_weights_shuffled.append(tmp_weights2.view(torch.bfloat16))
261
+
262
+ # Stack weights for all experts and return as BF16 tensors.
263
+ w13_weights_shuffled_tensor = (
264
+ torch.stack(w13_weights_shuffled).view(torch.bfloat16).contiguous()
265
+ )
266
+ w2_weights_shuffled_tensor = (
267
+ torch.stack(w2_weights_shuffled).view(torch.bfloat16).contiguous()
268
+ )
269
+
270
+ return w13_weights_shuffled_tensor, w2_weights_shuffled_tensor
271
+
272
+
273
def align_fp8_moe_weights_for_fi(
274
w13: torch.Tensor, w2: torch.Tensor, is_act_and_mul: bool
275
) -> tuple[torch.Tensor, torch.Tensor, int]:
vllm/utils/flashinfer.py
+3
@@ -105,6 +105,9 @@ def _lazy_import_wrapper(
105
106
107
# Create lazy wrappers for each function
108
+flashinfer_trtllm_bf16_moe = _lazy_import_wrapper(
109
+ "flashinfer.fused_moe", "trtllm_bf16_moe"
110
+)
111
flashinfer_trtllm_fp8_block_scale_moe = _lazy_import_wrapper(
112
"flashinfer.fused_moe", "trtllm_fp8_block_scale_moe"
113
)