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
Enable Fbgemm NVFP4 on Dense models (#25609)
Signed-off-by: Saman Keon <samanamp@outlook.com>
Saman A. Pour committed
Sep 24, 2025 at 21:12 UTC
90b139cfffd52ee2563260a89409fee60f7329f2
3 files changed
+89
-5
benchmarks/kernels/bench_nvfp4_gemm.py
+61
-4
@@ -3,6 +3,7 @@
3
import argparse
4
import copy
5
import itertools
6
+import os
7
8
import torch
9
from weight_shapes import WEIGHT_SHAPES
@@ -23,21 +24,45 @@ PROVIDER_CFGS = {
24
"torch-bf16": dict(enabled=True),
25
"nvfp4": dict(no_a_quant=False, enabled=True),
26
"nvfp4-noquant": dict(no_a_quant=True, enabled=True),
27
+ "fbgemm-nvfp4": dict(fbgemm=True, no_a_quant=False, enabled=True),
28
+ "fbgemm-nvfp4-noquant": dict(fbgemm=True, no_a_quant=True, enabled=True),
29
}
30
31
+_needs_fbgemm = any(
32
+ v.get("fbgemm", False) for v in PROVIDER_CFGS.values() if v.get("enabled", False)
33
+)
34
+if _needs_fbgemm:
35
+ try:
36
+ from fbgemm_gpu.experimental.gemm.triton_gemm.fp4_quantize import (
37
+ triton_scale_nvfp4_quant,
38
+ )
39
+ except ImportError:
40
+ print(
41
+ "WARNING: FBGEMM providers are enabled but fbgemm_gpu is not installed. "
42
+ "These providers will be skipped. Please install fbgemm_gpu with: "
43
+ "'pip install fbgemm-gpu-genai' to run them."
44
+ )
45
+ # Disable FBGEMM providers so the benchmark can run.
46
+ for cfg in PROVIDER_CFGS.values():
47
+ if cfg.get("fbgemm"):
48
+ cfg["enabled"] = False
49
+
50
_enabled = [k for k, v in PROVIDER_CFGS.items() if v["enabled"]]
51
52
31
-def _quant_weight_nvfp4(b: torch.Tensor, device: str):
53
+def _quant_weight_nvfp4(b: torch.Tensor, device: str, cfg):
54
# Compute global scale for weight
55
b_amax = torch.abs(b).max().to(torch.float32)
56
b_global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / b_amax
35
- b_fp4, scale_b_fp4 = ops.scaled_fp4_quant(b, b_global_scale)
57
+ if "fbgemm" in cfg and cfg["fbgemm"]:
58
+ b_fp4, scale_b_fp4 = triton_scale_nvfp4_quant(b, b_global_scale)
59
+ else:
60
+ b_fp4, scale_b_fp4 = ops.scaled_fp4_quant(b, b_global_scale)
61
return b_fp4, scale_b_fp4, b_global_scale
62
63
64
def build_nvfp4_runner(cfg, a, b, dtype, device):
40
- b_fp4, scale_b_fp4, b_global_scale = _quant_weight_nvfp4(b, device)
65
+ b_fp4, scale_b_fp4, b_global_scale = _quant_weight_nvfp4(b, device, cfg)
66
67
# Compute global scale for activation
68
# NOTE: This is generally provided ahead-of-time by the model checkpoint.
@@ -46,6 +71,35 @@ def build_nvfp4_runner(cfg, a, b, dtype, device):
71
72
# Alpha for the GEMM operation
73
alpha = 1.0 / (a_global_scale * b_global_scale)
74
+ if "fbgemm" in cfg and cfg["fbgemm"]:
75
+ if cfg["no_a_quant"]:
76
+ a_fp4, scale_a_fp4 = triton_scale_nvfp4_quant(a, a_global_scale)
77
+
78
+ def run():
79
+ return torch.ops.fbgemm.f4f4bf16(
80
+ a_fp4,
81
+ b_fp4,
82
+ scale_a_fp4,
83
+ scale_b_fp4,
84
+ global_scale=alpha,
85
+ use_mx=False,
86
+ )
87
+
88
+ return run
89
+ else:
90
+
91
+ def run():
92
+ a_fp4, scale_a_fp4 = triton_scale_nvfp4_quant(a, a_global_scale)
93
+ return torch.ops.fbgemm.f4f4bf16(
94
+ a_fp4,
95
+ b_fp4,
96
+ scale_a_fp4,
97
+ scale_b_fp4,
98
+ global_scale=alpha,
99
+ use_mx=False,
100
+ )
101
+
102
+ return run
103
104
if cfg["no_a_quant"]:
105
# Pre-quantize activation
@@ -130,10 +184,13 @@ if __name__ == "__main__":
184
185
for K, N, model in prepare_shapes(args):
186
print(f"{model}, N={N} K={K}, BF16 vs NVFP4 GEMMs TFLOP/s:")
187
+ save_dir = f"bench_nvfp4_res_n{N}_k{K}"
188
+ os.makedirs(save_dir, exist_ok=True)
189
+
190
benchmark.run(
191
print_data=True,
192
show_plots=True,
136
- save_path=f"bench_nvfp4_res_n{N}_k{K}",
193
+ save_path=save_dir,
194
N=N,
195
K=K,
196
)
vllm/envs.py
+4
-1
@@ -201,6 +201,7 @@ if TYPE_CHECKING:
201
VLLM_ENABLE_INDUCTOR_COORDINATE_DESCENT_TUNING: bool = True
202
VLLM_USE_NCCL_SYMM_MEM: bool = False
203
VLLM_NCCL_INCLUDE_PATH: Optional[str] = None
204
+ VLLM_USE_FBGEMM: bool = False
205
206
207
def get_default_cache_root():
@@ -1452,7 +1453,8 @@ environment_variables: dict[str, Callable[[], Any]] = {
1453
# NCCL header path
1454
"VLLM_NCCL_INCLUDE_PATH":
1455
lambda: os.environ.get("VLLM_NCCL_INCLUDE_PATH", None),
1455
-
1456
+ # Flag to enable FBGemm kernels on model execution
1457
+ "VLLM_USE_FBGEMM": lambda: bool(int(os.getenv("VLLM_USE_FBGEMM", "0"))),
1458
}
1459
1460
# --8<-- [end:env-vars-definition]
@@ -1548,6 +1550,7 @@ def compute_hash() -> str:
1550
"VLLM_ROCM_FP8_MFMA_PAGE_ATTN",
1551
"VLLM_ENABLE_INDUCTOR_MAX_AUTOTUNE",
1552
"VLLM_ENABLE_INDUCTOR_COORDINATE_DESCENT_TUNING",
1553
+ "VLLM_USE_FBGEMM",
1554
]
1555
for key in environment_variables_to_hash:
1556
# if this goes out of sync with environment_variables,
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
+24
@@ -30,8 +30,20 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
30
if envs.VLLM_USE_TRTLLM_FP4_GEMM:
31
assert has_flashinfer(), "TRTLLM FP4 GEMM requires FlashInfer"
32
self.backend = "flashinfer-trtllm"
33
+ logger.info_once("Using flashinfer-trtllm for FP4")
34
+ elif envs.VLLM_USE_FBGEMM:
35
+ self.backend = "fbgemm"
36
+ try:
37
+ import fbgemm_gpu # noqa: F401
38
+ except ImportError as exc:
39
+ raise ImportError(
40
+ "Backend fbgemm requires fbgemm.f4f4bf16 operator, "
41
+ "Please install with: pip install fbgemm-gpu-genai"
42
+ ) from exc
43
+ logger.info_once("Using FGBEMM-GPU-GENAI for FP4")
44
elif has_flashinfer():
45
self.backend = "flashinfer-cutlass"
46
+ logger.info_once("Using flashinfer-cutlass for FP4")
47
else:
48
self.backend = "cutlass"
49
self.group_size = 16
@@ -116,6 +128,9 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
128
layer.weight_packed = Parameter(weight, requires_grad=False)
129
else:
130
swizzled_weight_scale = swizzle_blockscale(layer.weight_scale)
131
+ if self.backend == "fbgemm":
132
+ swizzled_weight_scale = swizzled_weight_scale.view(-1).view(
133
+ torch.uint8)
134
layer.weight_scale = Parameter(swizzled_weight_scale,
135
requires_grad=False)
136
layer.weight_packed = Parameter(layer.weight_packed.data,
@@ -153,6 +168,15 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
168
out = flashinfer_scaled_fp4_mm(*mm_args, backend="trtllm")
169
elif self.backend == "flashinfer-cutlass":
170
out = flashinfer_scaled_fp4_mm(*mm_args, backend="cutlass")
171
+ elif self.backend == "fbgemm":
172
+ out = torch.ops.fbgemm.f4f4bf16(
173
+ x_fp4,
174
+ layer.weight_packed,
175
+ x_blockscale.view(-1).view(torch.uint8),
176
+ layer.weight_scale,
177
+ layer.alpha,
178
+ use_mx=False,
179
+ ).to(output_dtype)
180
else:
181
out = cutlass_scaled_fp4_mm(*mm_args)
182