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
[Feature] Add Flashinfer MoE Support for Compressed Tensor NVFP4 (#21639)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 31, 2025 at 18:26 UTC
c3e0e9337ef0af04d2d18b263a6a0f7deed75856
4 files changed
+287
-129
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+50
-3
@@ -17,9 +17,14 @@ from vllm.model_executor.layers.fused_moe import (
17
FusedMoE, FusedMoEActivationFormat, FusedMoEConfig, FusedMoEMethodBase,
18
FusedMoEPermuteExpertsUnpermute, FusedMoEPrepareAndFinalize,
19
FusedMoeWeightScaleSupported)
20
+from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa
21
+ FlashInferCutlassMoEPrepareAndFinalize)
22
from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa
23
WNA16_SUPPORTED_BITS, WNA16_SUPPORTED_TYPES_MAP)
24
from vllm.model_executor.layers.quantization.utils import replace_parameter
25
+from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
26
+ build_flashinfer_fp4_cutlass_moe_kernel,
27
+ flashinfer_fp4_cutlass_moe_forward, reorder_w1w3_to_w3w1)
28
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
29
check_moe_marlin_supports_layer, marlin_make_workspace_new,
30
marlin_moe_permute_scales)
@@ -28,7 +33,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
33
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
34
prepare_moe_fp8_layer_for_marlin)
35
from vllm.model_executor.layers.quantization.utils.quant_utils import (
31
- cutlass_fp4_supported, swizzle_blockscale)
36
+ swizzle_blockscale)
37
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
38
all_close_1d, normalize_e4m3fn_to_e4m3fnuz, per_tensor_dequantize)
39
from vllm.model_executor.utils import set_weight_attrs
@@ -96,8 +101,14 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
101
class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
102
103
def __init__(self):
99
- self.use_marlin = not cutlass_fp4_supported()
104
+ from vllm.model_executor.layers.quantization.utils.nvfp4_moe_support import ( # noqa: E501
105
+ detect_nvfp4_moe_support)
106
+ _nvfp4 = detect_nvfp4_moe_support(self.__class__.__name__)
107
+ self.cutlass_nvfp4_supported = _nvfp4.cutlass_supported
108
+ self.allow_flashinfer_cutlass = _nvfp4.allow_flashinfer_cutlass
109
+ self.use_marlin = _nvfp4.use_marlin
110
self.group_size = 16
111
+ self.fused_experts = None # type: ignore[assignment]
112
113
def create_weights(self, layer: torch.nn.Module, num_experts: int,
114
hidden_size: int, intermediate_size_per_partition: int,
@@ -200,6 +211,14 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
211
layer.w2_weight = torch.nn.Parameter(layer.w2_weight_packed.data,
212
requires_grad=False)
213
214
+ # reorder GEMM1 weights and block scales for FlashInfer CUTLASS kernel.
215
+ if self.allow_flashinfer_cutlass:
216
+ w, s = reorder_w1w3_to_w3w1(layer.w13_weight.data,
217
+ layer.w13_weight_scale.data,
218
+ dim=-2)
219
+ layer.w13_weight = torch.nn.Parameter(w, requires_grad=False)
220
+ layer.w13_weight_scale = torch.nn.Parameter(s, requires_grad=False)
221
+
222
if not torch.allclose(layer.w13_weight_global_scale[:, 0],
223
layer.w13_weight_global_scale[:, 1]):
224
logger.warning_once(
@@ -246,6 +265,21 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
265
layer.w2_input_scale_quant = torch.nn.Parameter(
266
(layer.w2_input_global_scale), requires_grad=False)
267
268
+ def maybe_swap_experts_impl(self, moe_parallel_config):
269
+ if not self.allow_flashinfer_cutlass:
270
+ return
271
+ self.fused_experts = build_flashinfer_fp4_cutlass_moe_kernel(
272
+ moe_parallel_config)
273
+
274
+ def select_gemm_impl(self, prepare_finalize, moe):
275
+ """Return the appropriate GEMM experts implementation."""
276
+ assert moe is not None and prepare_finalize is not None
277
+ from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( # noqa: E501
278
+ select_nvfp4_gemm_impl)
279
+
280
+ return select_nvfp4_gemm_impl(self.allow_flashinfer_cutlass, moe,
281
+ logger)
282
+
283
def apply(
284
self,
285
layer: torch.nn.Module,
@@ -303,10 +337,23 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
337
global_num_experts=global_num_experts,
338
expert_map=expert_map)
339
340
+ # FlashInfer fused experts path
341
+ if self.fused_experts is not None:
342
+ return flashinfer_fp4_cutlass_moe_forward(
343
+ self.fused_experts,
344
+ layer,
345
+ x,
346
+ topk_weights,
347
+ topk_ids,
348
+ activation=activation,
349
+ global_num_experts=global_num_experts,
350
+ expert_map=expert_map,
351
+ apply_router_weight_on_input=apply_router_weight_on_input,
352
+ )
353
+
354
assert expert_map is None, ("Expert Parallelism / expert_map "
355
"is currently not supported for "
356
"CompressedTensorsW4A4MoeMethod.")
309
-
357
from vllm.model_executor.layers.fused_moe.cutlass_moe import (
358
cutlass_moe_fp4)
359
vllm/model_executor/layers/quantization/modelopt.py
+24
-126
@@ -10,11 +10,8 @@ from torch.nn.parameter import Parameter
10
import vllm.envs as envs
11
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12
from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
13
-from vllm.distributed import get_ep_group
13
from vllm.logger import init_logger
14
from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
16
-from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa: E501
17
- FlashInferCutlassMoEPrepareAndFinalize)
15
from vllm.model_executor.layers.fused_moe.layer import (
16
FusedMoE, FusedMoEMethodBase, FusedMoeWeightScaleSupported)
17
from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
@@ -23,6 +20,9 @@ from vllm.model_executor.layers.quantization import QuantizationMethods
20
from vllm.model_executor.layers.quantization.base_config import (
21
QuantizationConfig, QuantizeMethodBase)
22
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
23
+from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
24
+ build_flashinfer_fp4_cutlass_moe_kernel,
25
+ flashinfer_fp4_cutlass_moe_forward, reorder_w1w3_to_w3w1)
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)
@@ -35,7 +35,6 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
35
Fp8LinearOp, requantize_with_max_scale)
36
from vllm.model_executor.parameter import (ModelWeightParameter,
37
PerTensorScaleParameter)
38
-from vllm.platforms import current_platform
38
from vllm.scalar_type import scalar_types
39
from vllm.utils.flashinfer import has_flashinfer_moe
40
@@ -869,28 +868,12 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
868
869
def __init__(self, quant_config: ModelOptNvFp4Config):
870
self.quant_config = quant_config
872
- self.cutlass_nvfp4_supported = cutlass_fp4_supported()
873
- self.use_marlin = False
874
- self.allow_flashinfer_cutlass = False
875
-
876
- if envs.VLLM_USE_FLASHINFER_MOE_FP4:
877
- if self.cutlass_nvfp4_supported and current_platform.is_cuda() \
878
- and current_platform.is_device_capability(100):
879
- logger.info_once(
880
- "Using FlashInfer kernels for ModelOptNvFp4FusedMoE.")
881
- self.allow_flashinfer_cutlass = True
882
- else:
883
- logger.warning_once(
884
- "Flashinfer CUTLASS Fused MoE not supported "
885
- "or found on the current platform.")
886
-
887
- if not self.cutlass_nvfp4_supported:
888
- if is_fp4_marlin_supported():
889
- self.use_marlin = True
890
- else:
891
- raise ValueError("Current platform does not support NVFP4"
892
- " quantization. Please use Blackwell and"
893
- " above.")
871
+ from vllm.model_executor.layers.quantization.utils.nvfp4_moe_support import ( # noqa: E501
872
+ detect_nvfp4_moe_support)
873
+ _nvfp4 = detect_nvfp4_moe_support(self.__class__.__name__)
874
+ self.cutlass_nvfp4_supported = _nvfp4.cutlass_supported
875
+ self.allow_flashinfer_cutlass = _nvfp4.allow_flashinfer_cutlass
876
+ self.use_marlin = _nvfp4.use_marlin
877
878
self.fused_experts = None # type: ignore
879
@@ -900,29 +883,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
883
):
884
if not self.allow_flashinfer_cutlass:
885
return
903
-
904
- logger.debug_once("FlashInferExperts")
905
- # default to TP/EP case only
906
-
907
- experts_kwargs: dict[str, Any] = {
908
- "use_nvfp4_w4a4": True,
909
- "use_dp": moe_parallel_config.dp_size > 1,
910
- "ep_rank": moe_parallel_config.ep_rank,
911
- "ep_size": moe_parallel_config.ep_size,
912
- "tp_rank": moe_parallel_config.tp_rank,
913
- "tp_size": moe_parallel_config.tp_size,
914
- }
915
-
916
- from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501
917
- FlashInferExperts)
918
- experts = FlashInferExperts(**experts_kwargs)
919
- self.fused_experts = mk.FusedMoEModularKernel(
920
- FlashInferCutlassMoEPrepareAndFinalize(
921
- quant_dtype=torch.uint8,
922
- #meaning 2x e2m1 packed in one, kernel requirement
923
- ),
924
- experts,
925
- )
886
+ self.fused_experts = build_flashinfer_fp4_cutlass_moe_kernel(
887
+ moe_parallel_config)
888
889
# This method update self.fused_experts
890
# only prepare_finalize is not None call select_gemm_impl
@@ -931,32 +893,12 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
893
def select_gemm_impl(self, prepare_finalize,
894
moe) -> mk.FusedMoEPermuteExpertsUnpermute:
895
934
- assert moe is not None
935
- assert prepare_finalize is not None
936
- experts = None
937
- all2all_manager = get_ep_group().device_communicator.all2all_manager
938
- assert all2all_manager is not None
939
- if self.allow_flashinfer_cutlass:
940
- from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501
941
- FlashInferExperts)
942
- logger.debug_once("Using FlashInferExperts")
943
- experts = FlashInferExperts(
944
- use_nvfp4_w4a4=True,
945
- use_dp=moe.moe_parallel_config.dp_size > 1,
946
- ep_rank=moe.moe_parallel_config.ep_rank,
947
- ep_size=moe.moe_parallel_config.ep_size,
948
- tp_rank=moe.moe_parallel_config.tp_rank,
949
- tp_size=moe.moe_parallel_config.tp_size,
950
- )
951
- else:
952
- assert moe.dp_size > 1
953
- logger.debug_once("Using CutlassExpertsFp4")
954
- # Currently CutlassExpertsFp4 doesn't support DP
955
- raise ValueError("CutlassExpertsFp4 doesn't support DP. "
956
- "Use flashinfer CUTLASS FusedMoE backend instead "
957
- "(set VLLM_USE_FLASHINFER_MOE_FP4=1)")
896
+ assert moe is not None and prepare_finalize is not None
897
+ from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( # noqa: E501
898
+ select_nvfp4_gemm_impl)
899
959
- return experts
900
+ return select_nvfp4_gemm_impl(self.allow_flashinfer_cutlass, moe,
901
+ logger)
902
903
def uses_weight_scale_2_pattern(self) -> bool:
904
"""
@@ -1062,18 +1004,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1004
gemm1_weight_scale = layer.w13_weight_scale.data
1005
1006
if self.allow_flashinfer_cutlass:
1065
- dim = -2
1066
- size = gemm1_weight.size(dim)
1067
- assert size % 2 == 0, f"Expected even size in dim {dim}, got {size}"
1068
- half = size // 2
1069
-
1070
- # Reorder weight
1071
- w1, w3 = gemm1_weight.split(half, dim=dim)
1072
- gemm1_weight = torch.cat([w3, w1], dim=dim).contiguous()
1073
-
1074
- # Reorder scale
1075
- s1, s3 = gemm1_weight_scale.split(half, dim=dim)
1076
- gemm1_weight_scale = torch.cat([s3, s1], dim=dim).contiguous()
1007
+ gemm1_weight, gemm1_weight_scale = reorder_w1w3_to_w3w1(
1008
+ gemm1_weight, gemm1_weight_scale, dim=-2)
1009
1010
layer.w13_weight = Parameter(gemm1_weight, requires_grad=False)
1011
layer.w13_weight_scale = Parameter(gemm1_weight_scale,
@@ -1217,49 +1149,15 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1149
expert_map=expert_map,
1150
apply_router_weight_on_input=apply_router_weight_on_input)
1151
else:
1220
- # TP or DP case
1221
- from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501
1222
- is_valid_flashinfer_cutlass_fused_moe)
1223
- assert is_valid_flashinfer_cutlass_fused_moe(
1224
- x, layer.w13_weight, layer.w2_weight), (
1225
- "Flashinfer CUTLASS Fused MoE not applicable!")
1226
-
1227
- a1_gscale = layer.w13_input_scale_quant
1228
- a2_gscale = layer.w2_input_scale_quant
1229
- extra_expert_args = {
1230
- 'g1_alphas': layer.g1_alphas,
1231
- 'g2_alphas': layer.g2_alphas,
1232
- 'out_dtype': x.dtype,
1233
- # Avoid confusion with a1_scale and a2_scale
1234
- # where are batch size related.
1235
- 'a1_gscale': a1_gscale,
1236
- 'a2_gscale': a2_gscale,
1237
- }
1238
- extra_prepare_args = {
1239
- 'use_dp': layer.dp_size > 1,
1240
- 'local_tokens': x.shape[0],
1241
- 'a1_gscale': a1_gscale,
1242
- }
1243
- extra_finalize_args = {
1244
- 'use_dp': layer.dp_size > 1,
1245
- 'local_tokens': x.shape[0],
1246
- }
1247
-
1248
- out = self.fused_experts(
1249
- hidden_states=x,
1250
- w1=layer.w13_weight,
1251
- w2=layer.w2_weight,
1252
- topk_weights=topk_weights,
1253
- topk_ids=topk_ids,
1254
- inplace=False, # TODO(shuw): fix later, now output is high prec
1152
+ out = flashinfer_fp4_cutlass_moe_forward(
1153
+ self.fused_experts,
1154
+ layer,
1155
+ x,
1156
+ topk_weights,
1157
+ topk_ids,
1158
activation=activation,
1159
global_num_experts=global_num_experts,
1160
expert_map=expert_map,
1258
- w1_scale=layer.w13_blockscale_swizzled,
1259
- w2_scale=layer.w2_blockscale_swizzled,
1161
apply_router_weight_on_input=apply_router_weight_on_input,
1261
- extra_expert_args=extra_expert_args,
1262
- extra_prepare_args=extra_prepare_args,
1263
- extra_finalize_args=extra_finalize_args,
1162
)
1163
return out
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
new
+154
@@ -0,0 +1,154 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Utility helpers for NVFP4 + FlashInfer fused-MoE path"""
4
+from __future__ import annotations
5
+
6
+from typing import Optional
7
+
8
+import torch
9
+
10
+import vllm.envs as envs
11
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12
+from vllm.logger import init_logger
13
+from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
14
+from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
15
+ FlashInferExperts, is_valid_flashinfer_cutlass_fused_moe)
16
+from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa: E501
17
+ FlashInferCutlassMoEPrepareAndFinalize)
18
+from vllm.platforms import current_platform
19
+
20
+logger = init_logger(__name__)
21
+
22
+__all__ = [
23
+ "is_flashinfer_fp4_cutlass_moe_available",
24
+ "reorder_w1w3_to_w3w1",
25
+ "build_flashinfer_fp4_cutlass_moe_kernel",
26
+ "flashinfer_fp4_cutlass_moe_forward",
27
+]
28
+
29
+
30
+def is_flashinfer_fp4_cutlass_moe_available() -> bool:
31
+ """Return ``True`` when FlashInfer CUTLASS NV-FP4 kernels can be used."""
32
+ return (envs.VLLM_USE_FLASHINFER_MOE_FP4 and current_platform.is_cuda()
33
+ and current_platform.is_device_capability(100))
34
+
35
+
36
+def reorder_w1w3_to_w3w1(weight: torch.Tensor,
37
+ scale: torch.Tensor,
38
+ dim: int = -2) -> tuple[torch.Tensor, torch.Tensor]:
39
+ """Re-order the concatenated `[w1, w3]` tensors to `[w3, w1]`"""
40
+ size = weight.size(dim)
41
+ assert size % 2 == 0, f"Expected even size in dim {dim}, got {size}"
42
+ half = size // 2
43
+
44
+ w1, w3 = weight.split(half, dim=dim)
45
+ s1, s3 = scale.split(half, dim=dim)
46
+
47
+ return (torch.cat([w3, w1],
48
+ dim=dim).contiguous(), torch.cat([s3, s1],
49
+ dim=dim).contiguous())
50
+
51
+
52
+def build_flashinfer_fp4_cutlass_moe_kernel(
53
+ moe_parallel_config: FusedMoEParallelConfig, ) -> mk.FusedMoEModularKernel:
54
+ """Create *and return* a FlashInfer CUTLASS fused-MoE modular kernel"""
55
+ experts = FlashInferExperts(
56
+ use_nvfp4_w4a4=True,
57
+ use_dp=moe_parallel_config.dp_size > 1,
58
+ ep_rank=moe_parallel_config.ep_rank,
59
+ ep_size=moe_parallel_config.ep_size,
60
+ tp_rank=moe_parallel_config.tp_rank,
61
+ tp_size=moe_parallel_config.tp_size,
62
+ )
63
+ logger.debug_once("FlashInferExperts (util)")
64
+ return mk.FusedMoEModularKernel(
65
+ FlashInferCutlassMoEPrepareAndFinalize(quant_dtype=torch.uint8),
66
+ experts,
67
+ )
68
+
69
+
70
+def flashinfer_fp4_cutlass_moe_forward(
71
+ fused_experts: mk.FusedMoEModularKernel,
72
+ layer: torch.nn.Module,
73
+ x: torch.Tensor,
74
+ topk_weights: torch.Tensor,
75
+ topk_ids: torch.Tensor,
76
+ activation: str,
77
+ global_num_experts: int,
78
+ expert_map: Optional[torch.Tensor],
79
+ apply_router_weight_on_input: bool,
80
+) -> torch.Tensor:
81
+ """Common forward wrapper for FlashInfer NV-FP4 fused-MoE"""
82
+
83
+ assert is_valid_flashinfer_cutlass_fused_moe(
84
+ x, layer.w13_weight,
85
+ layer.w2_weight), ("FlashInfer CUTLASS fused-MoE not applicable!")
86
+
87
+ a1_gscale = layer.w13_input_scale_quant
88
+ a2_gscale = layer.w2_input_scale_quant
89
+
90
+ extra_expert_args = {
91
+ "g1_alphas": layer.g1_alphas,
92
+ "g2_alphas": layer.g2_alphas,
93
+ # Avoid confusion with a1_scale and a2_scale
94
+ # where are batch size related.
95
+ "a1_gscale": a1_gscale,
96
+ "a2_gscale": a2_gscale,
97
+ "out_dtype": x.dtype,
98
+ }
99
+ extra_prepare_args = {
100
+ "use_dp": layer.dp_size > 1,
101
+ "local_tokens": x.shape[0],
102
+ "a1_gscale": a1_gscale,
103
+ }
104
+ extra_finalize_args = {
105
+ "use_dp": layer.dp_size > 1,
106
+ "local_tokens": x.shape[0],
107
+ }
108
+
109
+ return fused_experts(
110
+ hidden_states=x,
111
+ w1=layer.w13_weight,
112
+ w2=layer.w2_weight,
113
+ topk_weights=topk_weights,
114
+ topk_ids=topk_ids,
115
+ inplace=False, # TODO(shuw): fix later, now output is high prec
116
+ activation=activation,
117
+ global_num_experts=global_num_experts,
118
+ expert_map=expert_map,
119
+ w1_scale=layer.w13_blockscale_swizzled,
120
+ w2_scale=layer.w2_blockscale_swizzled,
121
+ apply_router_weight_on_input=apply_router_weight_on_input,
122
+ extra_expert_args=extra_expert_args,
123
+ extra_prepare_args=extra_prepare_args,
124
+ extra_finalize_args=extra_finalize_args,
125
+ )
126
+
127
+
128
+def select_nvfp4_gemm_impl(
129
+ allow_flashinfer_cutlass: bool,
130
+ moe, # FusedMoEConfig
131
+ logger):
132
+ """Return a GEMM *experts* implementation for NV-FP4 fused-MoE layers"""
133
+
134
+ # lazy import
135
+ from vllm.distributed import get_ep_group
136
+
137
+ all2all_manager = get_ep_group().device_communicator.all2all_manager
138
+ assert all2all_manager is not None
139
+
140
+ if allow_flashinfer_cutlass:
141
+ logger.debug_once("Using FlashInferExperts")
142
+ return FlashInferExperts(
143
+ use_nvfp4_w4a4=True,
144
+ use_dp=moe.moe_parallel_config.dp_size > 1,
145
+ ep_rank=moe.moe_parallel_config.ep_rank,
146
+ ep_size=moe.moe_parallel_config.ep_size,
147
+ tp_rank=moe.moe_parallel_config.tp_rank,
148
+ tp_size=moe.moe_parallel_config.tp_size,
149
+ )
150
+
151
+ # native cutlass experts currently don't support DP; TP case won't call this
152
+ raise ValueError(
153
+ "CutlassExpertsFp4 doesn't support DP. Use flashinfer CUTLASS "
154
+ "Fused MoE backend instead (set VLLM_USE_FLASHINFER_MOE_FP4=1)")
vllm/model_executor/layers/quantization/utils/nvfp4_moe_support.py
new
+59
@@ -0,0 +1,59 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import dataclass
4
+
5
+import vllm.envs as envs
6
+from vllm.logger import init_logger
7
+from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
8
+ is_flashinfer_fp4_cutlass_moe_available)
9
+from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
10
+ is_fp4_marlin_supported)
11
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
12
+ cutlass_fp4_supported)
13
+
14
+__all__ = ["detect_nvfp4_moe_support", "NvFp4Support"]
15
+
16
+_logger = init_logger(__name__)
17
+
18
+
19
+@dataclass(frozen=True)
20
+class NvFp4Support:
21
+ """Result container for NV-FP4 capability probing."""
22
+
23
+ cutlass_supported: bool
24
+ allow_flashinfer_cutlass: bool
25
+ use_marlin: bool
26
+
27
+
28
+def detect_nvfp4_moe_support(class_name: str = "") -> NvFp4Support:
29
+ """Detect platform support for NV-FP4 fused-MoE path"""
30
+ cutlass_supported = cutlass_fp4_supported()
31
+
32
+ allow_flashinfer = (cutlass_supported
33
+ and is_flashinfer_fp4_cutlass_moe_available())
34
+
35
+ if allow_flashinfer:
36
+ _logger.info_once("Using FlashInfer kernels for %s.", class_name
37
+ or "NVFP4 path")
38
+ else:
39
+ if envs.VLLM_USE_FLASHINFER_MOE_FP4:
40
+ _logger.warning_once(
41
+ "FlashInfer kernels unavailable for %s on current platform.",
42
+ class_name or "NVFP4 path",
43
+ )
44
+
45
+ use_marlin = False
46
+ if not cutlass_supported:
47
+ if is_fp4_marlin_supported():
48
+ use_marlin = True
49
+ _logger.info_once("Falling back to Marlin FP4 MoE kernel.")
50
+ else:
51
+ raise ValueError(
52
+ "Current platform does not support NVFP4 quantization. "
53
+ "Please use Blackwell GPUs or enable FlashInfer.")
54
+
55
+ return NvFp4Support(
56
+ cutlass_supported=cutlass_supported,
57
+ allow_flashinfer_cutlass=allow_flashinfer,
58
+ use_marlin=use_marlin,
59
+ )