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
[NVIDIA] Add support for cudnn fp4 gemm via flashinfer (#26107)
Signed-off-by: kaixih <kaixih@nvidia.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Kaixi Hou committed
Oct 15, 2025 at 10:53 UTC
de92d916fe8a897b00a8adb0aab9ed9ec99f2b6c
3 files changed
+57
-38
vllm/envs.py
+11
-6
@@ -191,6 +191,7 @@ if TYPE_CHECKING:
191
VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE: bool = False
192
VLLM_ENABLE_RESPONSES_API_STORE: bool = False
193
VLLM_USE_TRTLLM_ATTENTION: str | None = None
194
+ VLLM_NVFP4_GEMM_BACKEND: str | None = None
195
VLLM_FLASHINFER_DISABLE_Q_QUANTIZATION: bool = False
196
VLLM_HAS_FLASHINFER_CUBIN: bool = False
197
VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8: bool = False
@@ -1292,11 +1293,15 @@ environment_variables: dict[str, Callable[[], Any]] = {
1293
# If set, it means we pre-downloaded cubin files and flashinfer will
1294
# read the cubin files directly.
1295
"VLLM_HAS_FLASHINFER_CUBIN": lambda: os.getenv("VLLM_HAS_FLASHINFER_CUBIN", False),
1295
- # If set to 1, force the use of TRTLLM FP4 GEMM backend in flashinfer.
1296
- # Otherwise, uses the first available of: flashinfer cutlass GEMM,
1297
- # vllm cutlass GEMM, marlin GEMM.
1298
- "VLLM_USE_TRTLLM_FP4_GEMM": lambda: bool(
1299
- int(os.getenv("VLLM_USE_TRTLLM_FP4_GEMM", "0"))
1296
+ # Supported options:
1297
+ # - "flashinfer-cudnn": use flashinfer cudnn GEMM backend
1298
+ # - "flashinfer-trtllm": use flashinfer trtllm GEMM backend
1299
+ # - "flashinfer-cutlass": use flashinfer cutlass GEMM backend
1300
+ # - <none>: automatically pick an available backend
1301
+ "VLLM_NVFP4_GEMM_BACKEND": env_with_choices(
1302
+ "VLLM_NVFP4_GEMM_BACKEND",
1303
+ None,
1304
+ ["flashinfer-cudnn", "flashinfer-trtllm", "flashinfer-cutlass"],
1305
),
1306
# Controls garbage collection during CUDA graph capture.
1307
# If set to 0 (default), enables GC freezing to speed up capture time.
@@ -1492,7 +1497,6 @@ def compute_hash() -> str:
1497
"VLLM_DISABLED_KERNELS",
1498
"VLLM_USE_DEEP_GEMM",
1499
"VLLM_USE_DEEP_GEMM_E8M0",
1495
- "VLLM_USE_TRTLLM_FP4_GEMM",
1500
"VLLM_USE_FUSED_MOE_GROUPED_TOPK",
1501
"VLLM_USE_FLASHINFER_MOE_FP16",
1502
"VLLM_USE_FLASHINFER_MOE_FP8",
@@ -1524,6 +1528,7 @@ def compute_hash() -> str:
1528
"VLLM_ROCM_FP8_MFMA_PAGE_ATTN",
1529
"VLLM_ENABLE_INDUCTOR_MAX_AUTOTUNE",
1530
"VLLM_ENABLE_INDUCTOR_COORDINATE_DESCENT_TUNING",
1531
+ "VLLM_NVFP4_GEMM_BACKEND",
1532
"VLLM_USE_FBGEMM",
1533
]
1534
for key in environment_variables_to_hash:
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
+25
-15
@@ -14,7 +14,10 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
14
from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( # noqa: E501
15
run_nvfp4_emulations,
16
)
17
-from vllm.model_executor.layers.quantization.utils.quant_utils import swizzle_blockscale
17
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
18
+ cutlass_fp4_supported,
19
+ swizzle_blockscale,
20
+)
21
from vllm.model_executor.parameter import (
22
GroupQuantScaleParameter,
23
ModelWeightParameter,
@@ -29,10 +32,12 @@ __all__ = ["CompressedTensorsW4A4Fp4"]
32
33
class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
34
def __init__(self):
32
- if envs.VLLM_USE_TRTLLM_FP4_GEMM:
33
- assert has_flashinfer(), "TRTLLM FP4 GEMM requires FlashInfer"
34
- self.backend = "flashinfer-trtllm"
35
- logger.info_once("Using flashinfer-trtllm for FP4")
35
+ self.backend = "none"
36
+ if envs.VLLM_NVFP4_GEMM_BACKEND is None:
37
+ if has_flashinfer():
38
+ self.backend = "flashinfer-cutlass"
39
+ elif cutlass_fp4_supported():
40
+ self.backend = "cutlass"
41
elif envs.VLLM_USE_FBGEMM:
42
self.backend = "fbgemm"
43
try:
@@ -42,12 +47,17 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
47
"Backend fbgemm requires fbgemm.f4f4bf16 operator, "
48
"Please install with: pip install fbgemm-gpu-genai"
49
) from exc
45
- logger.info_once("Using FGBEMM-GPU-GENAI for FP4")
46
- elif has_flashinfer():
47
- self.backend = "flashinfer-cutlass"
48
- logger.info_once("Using flashinfer-cutlass for FP4")
49
- else:
50
- self.backend = "cutlass"
50
+ elif envs.VLLM_NVFP4_GEMM_BACKEND.startswith("flashinfer-"):
51
+ self.backend = envs.VLLM_NVFP4_GEMM_BACKEND
52
+ assert has_flashinfer(), f"FlashInfer is required for {self.backend}"
53
+
54
+ if self.backend == "none":
55
+ raise ValueError(
56
+ "No valid NVFP4 GEMM backend found. "
57
+ "Please check your platform capability."
58
+ )
59
+
60
+ logger.info_once(f"Using {self.backend} for NVFP4 GEMM")
61
self.group_size = 16
62
63
@classmethod
@@ -184,10 +194,9 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
194
layer.alpha,
195
output_dtype,
196
)
187
- if self.backend == "flashinfer-trtllm":
188
- out = flashinfer_scaled_fp4_mm(*mm_args, backend="trtllm")
189
- elif self.backend == "flashinfer-cutlass":
190
- out = flashinfer_scaled_fp4_mm(*mm_args, backend="cutlass")
197
+ if self.backend.startswith("flashinfer-"):
198
+ backend_name = self.backend[len("flashinfer-") :]
199
+ out = flashinfer_scaled_fp4_mm(*mm_args, backend=backend_name)
200
elif self.backend == "fbgemm":
201
out = torch.ops.fbgemm.f4f4bf16(
202
x_fp4,
@@ -198,6 +207,7 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
207
use_mx=False,
208
).to(output_dtype)
209
else:
210
+ assert self.backend == "cutlass"
211
out = cutlass_scaled_fp4_mm(*mm_args)
212
213
if bias is not None:
vllm/model_executor/layers/quantization/modelopt.py
+21
-17
@@ -926,22 +926,26 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
926
def __init__(self, quant_config: ModelOptNvFp4Config) -> None:
927
self.quant_config = quant_config
928
929
- if envs.VLLM_USE_TRTLLM_FP4_GEMM:
930
- assert has_flashinfer(), "TRTLLM FP4 GEMM requires FlashInfer"
931
- self.backend = "flashinfer-trtllm"
932
- elif has_flashinfer():
933
- self.backend = "flashinfer-cutlass"
934
- elif cutlass_fp4_supported():
935
- self.backend = "cutlass"
936
- elif is_fp4_marlin_supported():
937
- self.backend = "marlin"
938
- else:
929
+ self.backend = "none"
930
+ if envs.VLLM_NVFP4_GEMM_BACKEND is None:
931
+ if has_flashinfer():
932
+ self.backend = "flashinfer-cutlass"
933
+ elif cutlass_fp4_supported():
934
+ self.backend = "cutlass"
935
+ elif is_fp4_marlin_supported():
936
+ self.backend = "marlin"
937
+ elif envs.VLLM_NVFP4_GEMM_BACKEND.startswith("flashinfer-"):
938
+ self.backend = envs.VLLM_NVFP4_GEMM_BACKEND
939
+ assert has_flashinfer(), f"FlashInfer is required for {self.backend}"
940
+
941
+ if self.backend == "none":
942
raise ValueError(
940
- "Current platform does not support NVFP4"
941
- " quantization. Please use Blackwell and"
942
- " above."
943
+ "No valid NVFP4 GEMM backend found. "
944
+ "Please check your platform capability."
945
)
946
947
+ logger.info_once(f"Using {self.backend} for NVFP4 GEMM")
948
+
949
def create_weights(
950
self,
951
layer: torch.nn.Module,
@@ -1109,11 +1113,11 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
1113
layer.alpha,
1114
output_dtype,
1115
)
1112
- if self.backend == "flashinfer-trtllm":
1113
- out = flashinfer_scaled_fp4_mm(*mm_args, backend="trtllm")
1114
- elif self.backend == "flashinfer-cutlass":
1115
- out = flashinfer_scaled_fp4_mm(*mm_args, backend="cutlass")
1116
+ if self.backend.startswith("flashinfer-"):
1117
+ backend_name = self.backend[len("flashinfer-") :]
1118
+ out = flashinfer_scaled_fp4_mm(*mm_args, backend=backend_name)
1119
else:
1120
+ assert self.backend == "cutlass"
1121
out = cutlass_scaled_fp4_mm(*mm_args)
1122
1123
if bias is not None: