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
Use Blackwell FlashInfer MXFP4 MoE by default if available (#23008)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Aug 18, 2025 at 18:25 UTC
6d25e3fd6ea9562b8d73ecf0a527ce3da23261b2
2 files changed
+51
-19
vllm/model_executor/layers/fused_moe/layer.py
+5
-5
@@ -762,11 +762,11 @@ class FusedMoE(CustomOp):
762
self.global_num_experts = num_experts + num_redundant_experts
763
764
# we padding globally so EP buffer allocation works
765
- if (quant_config and quant_config.get_name() == "mxfp4"
766
- and (current_platform.is_rocm()
767
- or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
768
- or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16)):
769
- hidden_size = round_up(hidden_size, 256)
765
+ if quant_config and quant_config.get_name() == "mxfp4":
766
+ from vllm.model_executor.layers.quantization.mxfp4 import ( # noqa: E501
767
+ should_use_flashinfer_mxfp4)
768
+ if current_platform.is_rocm() or should_use_flashinfer_mxfp4():
769
+ hidden_size = round_up(hidden_size, 256)
770
771
# For smuggling this layer into the fused moe custom op
772
compilation_config = vllm_config.compilation_config
vllm/model_executor/layers/quantization/mxfp4.py
+46
-14
@@ -6,6 +6,7 @@ import torch
6
from torch.nn.parameter import Parameter
7
8
from vllm import envs
9
+from vllm.logger import init_logger
10
from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEConfig,
11
FusedMoEMethodBase)
12
from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
@@ -26,12 +27,38 @@ from vllm.platforms import current_platform
27
from vllm.scalar_type import scalar_types
28
from vllm.utils import (has_triton_kernels, is_torch_equal_or_newer,
29
next_power_of_2, round_up)
30
+from vllm.utils.flashinfer import has_flashinfer
31
30
-if (envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
31
- or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16):
32
- # from flashinfer.fused_moe import cutlass_fused_moe
33
- from flashinfer import (mxfp8_quantize, shuffle_matrix_a,
34
- shuffle_matrix_sf_a, trtllm_fp4_block_scale_moe)
32
+logger = init_logger(__name__)
33
+
34
+
35
+def _should_use_flashinfer_mxfp4_bf16():
36
+ """Determine if FlashInfer MXFP4 BF16 should be used."""
37
+ # If explicitly set, respect the setting
38
+ if envs.is_set("VLLM_USE_FLASHINFER_MOE_MXFP4_BF16"):
39
+ return envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16
40
+
41
+ # Enable by default on SM100 if MXFP8 is not explicitly enabled
42
+ if (current_platform.is_device_capability(100) and has_flashinfer()
43
+ and not envs.is_set("VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8")):
44
+ logger.info_once(
45
+ "Enabling FlashInfer MXFP4 BF16 backend by default for Blackwell. "
46
+ "For faster performance, consider setting "
47
+ "VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8=1, "
48
+ "though this may impact accuracy.")
49
+ return True
50
+
51
+ return False
52
+
53
+
54
+def _should_use_flashinfer_mxfp4_mxfp8():
55
+ """Determine if FlashInfer MXFP4 MXFP8 should be used."""
56
+ return envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
57
+
58
+
59
+def should_use_flashinfer_mxfp4():
60
+ return (_should_use_flashinfer_mxfp4_mxfp8()
61
+ or _should_use_flashinfer_mxfp4_bf16())
62
63
64
class Mxfp4Config(QuantizationConfig):
@@ -87,12 +114,18 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
114
self.moe = moe
115
self.use_marlin = self._should_use_marlin()
116
117
+ if current_platform.is_device_capability(100) and not has_flashinfer():
118
+ logger.warning_once(
119
+ "MXFP4 MoE is enabled on Blackwell but FlashInfer "
120
+ "is not available. This may result in degraded performance. "
121
+ "Please `pip install vllm[flashinfer]` for best results.")
122
+
123
def _should_use_marlin(self):
124
if envs.VLLM_MXFP4_USE_MARLIN is not None:
125
return envs.VLLM_MXFP4_USE_MARLIN
126
if current_platform.is_cuda() and \
94
- not current_platform.has_device_capability(100):
95
- if not current_platform.is_device_capability(90):
127
+ not current_platform.is_device_capability(100):
128
+ if not current_platform.has_device_capability(90):
129
# marlin kernel has better performance on ampere
130
return True
131
if not has_triton_kernels():
@@ -138,8 +171,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
171
layer.hidden_size = hidden_size
172
layer.intermediate_size_per_partition = \
173
intermediate_size_per_partition_after_pad
141
- elif (envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
142
- or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16):
174
+ elif should_use_flashinfer_mxfp4():
175
# pad the intermediate size to be a multiple of 2 * mxfp4_block
176
# for to hold non-uniform sharded tensor as well as swizzling
177
# other padding to increase performance
@@ -230,8 +262,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
262
def process_weights_after_loading(self, layer):
263
if self.use_marlin:
264
prepare_moe_fp4_layer_for_marlin(layer)
233
- elif (envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
234
- or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16):
265
+ elif should_use_flashinfer_mxfp4():
266
+ from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a
267
layer.gemm1_alpha = Parameter(torch.tensor(
268
[1.702] * self.num_experts, dtype=torch.float32).cuda(),
269
requires_grad=False)
@@ -478,11 +510,11 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
510
logical_replica_count), (
511
"MXFP4 are not supported with this configuration.")
512
481
- if (envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8
482
- or envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16):
513
+ if should_use_flashinfer_mxfp4():
514
+ from flashinfer import mxfp8_quantize, trtllm_fp4_block_scale_moe
515
assert not self.moe.use_ep, (
516
"EP is not supported for flashinfer mxfp4 moe backend yet.")
485
- if envs.VLLM_USE_FLASHINFER_MOE_MXFP4_BF16:
517
+ if _should_use_flashinfer_mxfp4_bf16():
518
assert x.dtype == torch.bfloat16
519
x_quant = x
520
x_scale = None