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, 05:00 UTC
fix: Add SM120 (RTX Blackwell) support for FlashInfer CUTLASS NVFP4 MoE kernels (#33417)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: mgoin <mgoin64@gmail.com>
René Honig committed
Jan 31, 2026 at 23:06 UTC
079781177ae4c9fba429bf093cae73cf4cfae7a8
6 files changed
+22
-110
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+6
-1
@@ -657,7 +657,12 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):
657
658
@staticmethod
659
def _supports_current_device() -> bool:
660
- return current_platform.has_device_capability((10, 0))
660
+ p = current_platform
661
+ return p.is_cuda() and (
662
+ p.is_device_capability_family(100)
663
+ or p.is_device_capability_family(110)
664
+ or p.is_device_capability_family(120)
665
+ )
666
667
@staticmethod
668
def _supports_no_act_and_mul() -> bool:
vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py
+2
-1
@@ -54,7 +54,8 @@ class FlashInferCuteDSLExperts(mk.FusedMoEPermuteExpertsUnpermute):
54
55
@staticmethod
56
def _supports_current_device() -> bool:
57
- return current_platform.is_device_capability_family(100)
57
+ p = current_platform
58
+ return p.is_cuda() and p.is_device_capability_family(100)
59
60
@staticmethod
61
def _supports_no_act_and_mul() -> bool:
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
+14
-13
@@ -84,11 +84,14 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
84
85
@staticmethod
86
def _supports_current_device() -> bool:
87
+ p = current_platform
88
return (
88
- current_platform.is_cuda()
89
+ p.is_cuda()
90
and (
90
- current_platform.is_device_capability((9, 0))
91
- or current_platform.is_device_capability_family(100)
91
+ p.is_device_capability(90)
92
+ or p.is_device_capability_family(100)
93
+ or p.is_device_capability_family(110)
94
+ or p.is_device_capability_family(120)
95
)
96
and has_flashinfer_cutlass_fused_moe()
97
)
@@ -102,29 +105,27 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
105
weight_key: QuantKey | None,
106
activation_key: QuantKey | None,
107
) -> bool:
105
- # The following are supported by FlashInferExperts:
106
- # * unquantized
107
- # * fp8 static per-tensor on 9.0+
108
- # * fp8 block on 9.0
109
- # * nvfp4 on 10.0+
110
-
108
p = current_platform
109
scheme = (weight_key, activation_key)
110
+ # The following are supported by FlashInferExperts:
111
return (
112
+ # unquantized and fp8 static per-tensor on 9.0+
113
(
114
scheme
115
in [
116
(None, None),
117
(kFp8StaticTensorSym, kFp8StaticTensorSym),
118
]
119
+ and p.has_device_capability(90)
120
)
121
+ # fp8 block-scale on 9.0
122
or (
122
- (scheme == (kFp8Static128BlockSym, kFp8Dynamic128Sym))
123
- and (p.is_device_capability((9, 0)))
123
+ scheme == (kFp8Static128BlockSym, kFp8Dynamic128Sym)
124
+ and p.is_device_capability(90)
125
)
126
+ # nvfp4 on 10.0+
127
or (
126
- (scheme == (kNvfp4Static, kNvfp4Dynamic))
127
- and (p.is_device_capability_family(100))
128
+ scheme == (kNvfp4Static, kNvfp4Dynamic) and p.has_device_capability(100)
129
)
130
)
131
vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py
-1
@@ -30,7 +30,6 @@ from vllm.utils.torch_utils import direct_register_custom_op
30
def _supports_current_device() -> bool:
31
"""Supports only Blackwell-family GPUs."""
32
p = current_platform
33
- # Add check flashinfer trtllm is available
33
return p.is_cuda() and p.is_device_capability_family(100)
34
35
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
-27
@@ -6,7 +6,6 @@ from typing import TYPE_CHECKING
6
7
import torch
8
9
-import vllm.envs as envs
9
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
10
from vllm import _custom_ops as ops
11
from vllm.logger import init_logger
@@ -24,10 +23,6 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
23
kNvfp4Static,
24
)
25
from vllm.platforms import current_platform
27
-from vllm.utils.flashinfer import (
28
- has_flashinfer_cutedsl_grouped_gemm_nt_masked,
29
- has_flashinfer_cutlass_fused_moe,
30
-)
26
27
if TYPE_CHECKING:
28
from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
@@ -38,8 +33,6 @@ logger = init_logger(__name__)
33
34
35
__all__ = [
41
- "is_flashinfer_fp4_cutlass_moe_available",
42
- "is_flashinfer_fp4_cutedsl_moe_available",
36
"reorder_w1w3_to_w3w1",
37
]
38
@@ -124,26 +117,6 @@ def is_supported_config_trtllm(
117
return True, None
118
119
127
-def is_flashinfer_fp4_cutlass_moe_available() -> bool:
128
- """Return `True` when FlashInfer CUTLASS NV-FP4 kernels can be used."""
129
- return (
130
- envs.VLLM_USE_FLASHINFER_MOE_FP4
131
- and has_flashinfer_cutlass_fused_moe()
132
- and current_platform.is_cuda()
133
- and current_platform.has_device_capability(100)
134
- )
135
-
136
-
137
-def is_flashinfer_fp4_cutedsl_moe_available() -> bool:
138
- """Return ``True`` when FlashInfer CUTEDSL NV-FP4 kernels can be used."""
139
- return (
140
- envs.VLLM_USE_FLASHINFER_MOE_FP4
141
- and has_flashinfer_cutedsl_grouped_gemm_nt_masked()
142
- and current_platform.is_cuda()
143
- and current_platform.is_device_capability_family(100)
144
- )
145
-
146
-
120
def reorder_w1w3_to_w3w1(
121
weight: torch.Tensor, scale: torch.Tensor, dim: int = -2
122
) -> tuple[torch.Tensor, torch.Tensor]:
vllm/model_executor/layers/quantization/utils/nvfp4_moe_support.py
deleted
-67
@@ -1,67 +0,0 @@
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_cutedsl_moe_available,
9
- is_flashinfer_fp4_cutlass_moe_available,
10
-)
11
-from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
12
- is_fp4_marlin_supported,
13
-)
14
-from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
15
- cutlass_fp4_supported,
16
-)
17
-
18
-__all__ = ["detect_nvfp4_moe_support", "NvFp4Support"]
19
-
20
-_logger = init_logger(__name__)
21
-
22
-
23
-@dataclass(frozen=True)
24
-class NvFp4Support:
25
- """Result container for NV-FP4 capability probing."""
26
-
27
- cutlass_supported: bool
28
- allow_flashinfer: bool
29
- use_marlin: bool
30
-
31
-
32
-def detect_nvfp4_moe_support(class_name: str = "") -> NvFp4Support:
33
- """Detect platform support for NV-FP4 fused-MoE path"""
34
- cutlass_supported = cutlass_fp4_supported()
35
-
36
- allow_flashinfer = cutlass_supported and (
37
- is_flashinfer_fp4_cutlass_moe_available()
38
- or is_flashinfer_fp4_cutedsl_moe_available()
39
- )
40
-
41
- if allow_flashinfer:
42
- _logger.info_once(
43
- "Using FlashInfer kernels for %s.", class_name or "NVFP4 path"
44
- )
45
- else:
46
- if envs.VLLM_USE_FLASHINFER_MOE_FP4:
47
- _logger.warning_once(
48
- "FlashInfer kernels unavailable for %s on current platform.",
49
- class_name or "NVFP4 path",
50
- )
51
-
52
- use_marlin = False
53
- if not cutlass_supported:
54
- if is_fp4_marlin_supported():
55
- use_marlin = True
56
- _logger.info_once("Falling back to Marlin FP4 MoE kernel.")
57
- else:
58
- raise ValueError(
59
- "Current platform does not support NVFP4 quantization. "
60
- "Please use Blackwell GPUs or enable FlashInfer."
61
- )
62
-
63
- return NvFp4Support(
64
- cutlass_supported=cutlass_supported,
65
- allow_flashinfer=allow_flashinfer,
66
- use_marlin=use_marlin,
67
- )