@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 00d7b497b3

Fork 0

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 5, 12:15 UTC

[NVFP4] Support NVFP4 dense models from `modelopt` and `compressed-tensors` on AMD Instinct MI300, MI355X and Hopper through emulation (#35733)

Signed-off-by: Felix Marty <Felix.Marty@amd.com> Signed-off-by: fxmarty-amd <felmarty@amd.com> Co-authored-by: Kyle Sayers <kylesayrs@gmail.com>

fxmarty-amd committed Apr 7, 2026 at 00:18 UTC 00d7b497b33679a3ce641db9119eb09dd4ed5e30

10 files changed +191 -58

tests/models/quantization/test_nvfp4.py

+15 -4

@@ -89,22 +89,33 @@ def test_models(example_prompts, model_name) -> None:

90 EAGER = [True, False]

92 +SM_100_NVFP4_BACKENDS = [

93 + "flashinfer-cudnn",

94 + "flashinfer-trtllm",

95 + "flashinfer-cutlass",

96 +]

97 +

93 -@pytest.mark.skipif(

94 - not current_platform.has_device_capability(100),

95 - reason="modelopt_fp4 is not supported on this GPU type.",

96 -)

99 @pytest.mark.parametrize("model", ["nvidia/Llama-3.1-8B-Instruct-NVFP4"])

100 @pytest.mark.parametrize("eager", EAGER)

101 @pytest.mark.parametrize(

102 "backend",

103 [

104 + "emulation",

105 "flashinfer-cudnn",

106 "flashinfer-trtllm", # the small seq_len ensures trtllm_8x4_layout backend is used

107 "flashinfer-cutlass",

108 ],

109 )

110 def test_nvfp4(vllm_runner, model, eager, backend, monkeypatch):

111 + if (

112 + not current_platform.has_device_capability(100)

113 + and backend in SM_100_NVFP4_BACKENDS

114 + ):

115 + pytest.skip(

116 + f"The backend {backend} is not supported with current_platform.has_device_capability(100) == False"

117 + )

118 +

119 monkeypatch.setenv("VLLM_NVFP4_GEMM_BACKEND", backend)

120 with vllm_runner(model, enforce_eager=eager) as llm:

121 output = llm.generate_greedy(["1 2 3 4 5"], max_tokens=2)

tests/quantization/test_compressed_tensors.py

+1 -4

@@ -366,9 +366,6 @@ def test_compressed_tensors_kv_cache_fp8_per_attn_head(vllm_runner):

366 assert output

367

368

369 -@pytest.mark.skipif(

370 - not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."

371 -)

369 @pytest.mark.parametrize(

370 "args",

371 [

@@ -398,7 +395,7 @@ def test_compressed_tensors_nvfp4(vllm_runner, args):

395 assert qkv_proj.scheme.group_size == 16

396

397 llm.apply_model(check_model)

401 - output = llm.generate_greedy("Hello my name is", max_tokens=4)

398 + output = llm.generate_greedy(["Hello my name is"], max_tokens=4)

399 print(output)

400 assert output

401

vllm/envs.py

@@ -1464,6 +1464,10 @@ environment_variables: dict[str, Callable[[], Any]] = {

1464 # - "flashinfer-trtllm": use flashinfer trtllm GEMM backend

1465 # - "flashinfer-cutlass": use flashinfer cutlass GEMM backend

1466 # - "marlin": use marlin GEMM backend (for GPUs without native FP4 support)

1467 + # - "emulation":

1468 + # use BF16/FP16 GEMM, dequantizing weights and running QDQ on activations.

1469 + # This is only meant for research purposes to run on devices where NVFP4

1470 + # GEMM kernels are not available.

1471 # - <none>: automatically pick an available backend

1472 "VLLM_NVFP4_GEMM_BACKEND": env_with_choices(

1473 "VLLM_NVFP4_GEMM_BACKEND",

@@ -1474,6 +1478,7 @@ environment_variables: dict[str, Callable[[], Any]] = {

1478 "flashinfer-cutlass",

1479 "cutlass",

1480 "marlin",

1481 + "emulation",

1482 ],

1483 ),

1484 # Controls garbage collection during CUDA graph capture.

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py

+23

@@ -5,10 +5,12 @@ from collections.abc import Callable

5 import torch

6 from torch.nn.parameter import Parameter

8 +from vllm.logger import init_logger

9 from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (

10 CompressedTensorsScheme,

11 )

12 from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (

13 + NvFp4LinearBackend,

14 apply_nvfp4_linear,

15 convert_to_nvfp4_linear_kernel_format,

16 select_nvfp4_linear_backend,

@@ -19,6 +21,9 @@ from vllm.model_executor.parameter import (

21 PerTensorScaleParameter,

22 )

24 +logger = init_logger(__name__)

25 +

26 +

27 __all__ = ["CompressedTensorsW4A4Fp4"]

@@ -27,6 +32,10 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):

32 self.backend = select_nvfp4_linear_backend()

33 self.group_size = 16

35 + self.swizzle = None

36 + if self.backend == NvFp4LinearBackend.EMULATION:

37 + self.swizzle = False

38 +

39 @classmethod

40 def get_min_capability(cls) -> int:

41 return 75

@@ -89,6 +98,19 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):

98 # Rename CT checkpoint names to standardized names

99 layer.weight = layer.weight_packed

100 del layer.weight_packed

101 +

102 + if (

103 + torch.unique(layer.input_global_scale).numel() != 1

104 + or torch.unique(layer.weight_global_scale).numel() != 1

105 + ):

106 + logger.warning_once(

107 + "In NVFP4 linear, the global scale for input or weight are different"

108 + " for parallel layers (e.g. q_proj, k_proj, v_proj). This "

109 + " will likely result in reduced accuracy. Please verify the model"

110 + " accuracy. Consider using a checkpoint with a shared global NVFP4"

111 + " scale for fused layers."

112 + )

113 +

114 # Process global scales (CT stores as divisors, i.e. 1/scale)

115 input_global_scale_inv = layer.input_global_scale.max().to(torch.float32)

116 layer.input_global_scale = Parameter(

@@ -121,4 +143,5 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):

143 layer=layer,

144 x=x,

145 bias=bias,

146 + swizzle=self.swizzle,

147 )

vllm/model_executor/layers/quantization/modelopt.py

+19

@@ -71,6 +71,7 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (

71 mxfp8_e4m3_quantize,

72 )

73 from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (

74 + NvFp4LinearBackend,

75 apply_nvfp4_linear,

76 convert_to_nvfp4_linear_kernel_format,

77 select_nvfp4_linear_backend,

@@ -1074,6 +1075,10 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):

1075 self.marlin_input_dtype = None

1076 self.backend = select_nvfp4_linear_backend()

1077

1078 + self.swizzle = None

1079 + if self.backend == NvFp4LinearBackend.EMULATION:

1080 + self.swizzle = False

1081 +

1082 def create_weights(

1083 self,

1084 layer: torch.nn.Module,

@@ -1149,10 +1154,23 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):

1154 layer.register_parameter("weight_scale", weight_scale)

1155

1156 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

1157 + if (

1158 + torch.unique(layer.input_scale).numel() != 1

1159 + or torch.unique(layer.weight_scale_2).numel() != 1

1160 + ):

1161 + logger.warning_once(

1162 + "In NVFP4 linear, the global scale for input or weight are different"

1163 + " for parallel layers (e.g. q_proj, k_proj, v_proj). This "

1164 + " will likely results in reduce accuracy. Please verify the model"

1165 + " accuracy. Consider using a checkpoint with a shared global NVFP4"

1166 + " scale for parallel layers."

1167 + )

1168 +

1169 # Rename ModelOpt checkpoint names to standardized names

1170 input_global_scale = layer.input_scale.max().to(torch.float32)

1171 layer.input_global_scale = Parameter(input_global_scale, requires_grad=False)

1172 del layer.input_scale

1173 +

1174 weight_global_scale = layer.weight_scale_2.max().to(torch.float32)

1175 layer.weight_global_scale = Parameter(weight_global_scale, requires_grad=False)

1176 del layer.weight_scale_2

@@ -1179,6 +1197,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):

1197 layer=layer,

1198 x=x,

1199 bias=bias,

1200 + swizzle=self.swizzle,

1201 )

1202

1203

vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py

+1 -1

@@ -24,7 +24,7 @@ logger = init_logger(__name__)

26 def is_fp4_marlin_supported():

27 - return current_platform.has_device_capability(75)

27 + return current_platform.is_cuda() and current_platform.has_device_capability(75)

30 def _nvfp4_compute_scale_factor(

vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py

+22 -9

@@ -1,5 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from types import SimpleNamespace

4 +

5 import torch

7 from vllm.scalar_type import scalar_types

@@ -11,9 +13,10 @@ __all__ = [

13 ]

15 FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()

16 +FLOAT4_E2M1_MAX_RECIPROCAL = 1 / FLOAT4_E2M1_MAX

15 -kE2M1ToFloat = torch.tensor(

16 - [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32

18 +kE2M1ToFloat_handle = SimpleNamespace(

19 + val=torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32)

20 )

@@ -29,8 +32,9 @@ def break_fp4_bytes(a, dtype):

32 # Vectorized sign and magnitude extraction

33 signs = (combined & 0x08).to(torch.bool) # Sign bits

34 abs_vals = (combined & 0x07).to(torch.long)

35 +

36 + kE2M1 = kE2M1ToFloat_handle.val

37 # Device-aware lookup and sign application

33 - kE2M1 = kE2M1ToFloat.to(device=a.device)

38 values = kE2M1[abs_vals] * torch.where(signs, -1.0, 1.0)

39 # Reshape to final form

40 return values.reshape(m, n * 2).to(dtype=dtype)

@@ -47,7 +51,12 @@ def convert_swizzled_to_linear(a_sf_swizzled: torch.Tensor, m, k, block_size):

53 def dequantize_to_dtype(

50 - tensor_fp4, tensor_sf, global_scale, dtype, device, block_size=16

54 + tensor_fp4: torch.Tensor,

55 + tensor_sf: torch.Tensor,

56 + global_scale: torch.Tensor | float,

57 + dtype: torch.dtype,

58 + block_size: int = 16,

59 + swizzle: bool | None = True,

60 ):

61 """Dequantize the fp4 tensor back to high precision."""

62 # Two fp4 values are packed into one uint8.

@@ -57,8 +66,10 @@ def dequantize_to_dtype(

66 tensor_f32 = break_fp4_bytes(tensor_fp4, torch.float32)

67 tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size)

68 tensor_sf = tensor_sf.view(torch.float8_e4m3fn)

60 - tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size)

61 - tensor_sf_dtype = tensor_sf.to(torch.float32) / global_scale

69 +

70 + if swizzle:

71 + tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size)

72 + tensor_sf_dtype = tensor_sf.to(torch.float32) * global_scale

74 # scale the tensor

75 out = (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k)

@@ -67,7 +78,8 @@ def dequantize_to_dtype(

79 def get_reciprocal(x):

80 if isinstance(x, torch.Tensor):

70 - return torch.where(x == 0, torch.tensor(0.0, dtype=x.dtype), 1.0 / x)

81 + # torch.where yields operation not permitted when stream is capturing.

82 + return 1.0 / (x + (x == 0) * 1e8)

83 elif isinstance(x, (float, int)):

84 return 0.0 if x == 0 else 1.0 / x

85 else:

@@ -94,7 +106,7 @@ def ref_nvfp4_quant(x, global_scale, block_size):

106 m, n = x.shape

107 x = torch.reshape(x, (m, n // block_size, block_size))

108 vec_max = torch.max(torch.abs(x), dim=-1, keepdim=True)[0].to(torch.float32)

97 - scale = global_scale * (vec_max * get_reciprocal(FLOAT4_E2M1_MAX))

109 + scale = global_scale * (vec_max * FLOAT4_E2M1_MAX_RECIPROCAL)

110 scale = torch.clamp(scale, max=448, min=-448)

111 scale = scale.to(torch.float8_e4m3fn).to(torch.float32)

112 output_scale = get_reciprocal(scale * get_reciprocal(global_scale))

@@ -111,6 +123,7 @@ def run_nvfp4_emulations(

123 weight: torch.Tensor,

124 weight_scale_swizzled: torch.Tensor,

125 weight_global_scale: torch.Tensor,

126 + swizzle: bool | None = True,

127 ):

128 group_size = 16

129 x_m, x_k = x.shape

@@ -132,8 +145,8 @@ def run_nvfp4_emulations(

145 weight_scale_swizzled.data,

146 weight_global_scale,

147 output_dtype,

135 - x.device,

148 group_size,

149 + swizzle=swizzle,

150 )

151

152 # matmul

vllm/model_executor/layers/quantization/utils/nvfp4_utils.py

+99 -40

@@ -17,31 +17,99 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (

17 prepare_fp4_layer_for_marlin,

18 )

19 from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (

20 + kE2M1ToFloat_handle,

21 run_nvfp4_emulations,

22 )

23 from vllm.platforms import current_platform

24 from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm, has_flashinfer

25 +from vllm.utils.import_utils import has_fbgemm_gpu

26 from vllm.utils.math_utils import round_up

28 logger = init_logger(__name__)

31 +# NOTE: This is ordered by preferred backend.

32 +# Example: if both are available, FLASHINFER_CUTLASS is preferred to VLLM_CUTLASS.

33 class NvFp4LinearBackend(Enum):

30 - VLLM_CUTLASS = "cutlass"

34 FLASHINFER_CUTLASS = "flashinfer-cutlass"

35 + VLLM_CUTLASS = "cutlass"

36 + MARLIN = "marlin"

37 FLASHINFER_TRTLLM = "flashinfer-trtllm"

38 FLASHINFER_CUDNN = "flashinfer-cudnn"

39 FBGEMM = "fbgemm"

35 - MARLIN = "marlin"

40 EMULATION = "emulation"

43 +NVFP4_LINEAR_BACKENDS = list(NvFp4LinearBackend)

44 +

45 +

46 +def is_backend_supported(backend: NvFp4LinearBackend) -> tuple[bool, str | None]:

47 + reason = None

48 + supported = True

49 +

50 + if backend == NvFp4LinearBackend.FLASHINFER_CUTLASS:

51 + # cutlass_fp4_supported() checks that the vLLM NVFP4 kernels (both

52 + # quantization and GEMM) were compiled for the current SM version.

53 + # FlashInfer backends still rely on the vLLM quantization kernels,

54 + # so we gate them on the same check.

55 + supported = (

56 + cutlass_fp4_supported()

57 + and current_platform.has_device_capability(100)

58 + and has_flashinfer()

59 + )

60 +

61 + if not supported:

62 + reason = "FlashInfer is required, >=sm_100 is required"

63 + elif backend == NvFp4LinearBackend.VLLM_CUTLASS:

64 + supported = cutlass_fp4_supported()

65 + if not supported:

66 + reason = "Cutlass is required"

67 + elif backend == NvFp4LinearBackend.MARLIN:

68 + supported = is_fp4_marlin_supported()

69 + if not supported:

70 + reason = "Marlin is required"

71 + elif backend in [

72 + NvFp4LinearBackend.FLASHINFER_TRTLLM,

73 + NvFp4LinearBackend.FLASHINFER_CUDNN,

74 + ]:

75 + supported = has_flashinfer()

76 + if not supported:

77 + reason = "FlashInfer is required"

78 + elif backend == NvFp4LinearBackend.FBGEMM:

79 + supported = has_fbgemm_gpu()

80 + if not supported:

81 + reason = "fbgemm_gpu is required"

82 + elif backend == NvFp4LinearBackend.EMULATION:

83 + # e.g. AMD Instinct does not support native NVFP4.

84 + unsupported_reasons = {}

85 + for other_backend in NVFP4_LINEAR_BACKENDS:

86 + if other_backend == NvFp4LinearBackend.EMULATION:

87 + continue

88 + other_supported, other_reason = is_backend_supported(other_backend)

89 + if not other_supported:

90 + unsupported_reasons[other_backend] = other_reason

91 +

92 + if unsupported_reasons:

93 + unsupported_reasons_str = "\n - ".join(

94 + [f"{b.value}: {r}" for b, r in unsupported_reasons.items()]

95 + )

96 + logger.warning_once(

97 + f"NVFP4 linear falling back to the slow and unoptimized "

98 + f"backend=NvFp4LinearBackend.EMULATION as no optimized backend is "

99 + f"available (unavailable reasons:\n - {unsupported_reasons_str}\n). "

100 + "In case you expect one of these backend to be used, "

101 + "please verify your environment."

102 + )

103 +

104 + return supported, reason

105 +

106 +

107 def select_nvfp4_linear_backend() -> NvFp4LinearBackend:

108 """

109 Select the best available NVFP4 GEMM backend based on environment

110 configuration and platform capabilities.

111 """

44 - backend: NvFp4LinearBackend | None = None

112 + selected_backend: NvFp4LinearBackend | None = None

113

114 if envs.VLLM_USE_FBGEMM:

115 try:

@@ -51,51 +119,36 @@ def select_nvfp4_linear_backend() -> NvFp4LinearBackend:

119 "Backend fbgemm requires fbgemm.f4f4bf16 operator, "

120 "Please install with: pip install fbgemm-gpu-genai"

121 ) from exc

54 - backend = NvFp4LinearBackend.FBGEMM

122 + selected_backend = NvFp4LinearBackend.FBGEMM

123 elif envs.VLLM_USE_NVFP4_CT_EMULATIONS:

56 - backend = NvFp4LinearBackend.EMULATION

124 + selected_backend = NvFp4LinearBackend.EMULATION

125 elif envs.VLLM_NVFP4_GEMM_BACKEND is None:

58 - # Auto-select best available backend.

59 - # cutlass_fp4_supported() checks that the vLLM NVFP4 kernels (both

60 - # quantization and GEMM) were compiled for the current SM version.

61 - # FlashInfer backends still rely on the vLLM quantization kernels,

62 - # so we gate them on the same check.

63 - if (

64 - cutlass_fp4_supported()

65 - and current_platform.has_device_capability(100)

66 - and has_flashinfer()

67 - ):

68 - backend = NvFp4LinearBackend.FLASHINFER_CUTLASS

69 - elif cutlass_fp4_supported():

70 - backend = NvFp4LinearBackend.VLLM_CUTLASS

71 - elif is_fp4_marlin_supported():

72 - backend = NvFp4LinearBackend.MARLIN

126 + for backend in NVFP4_LINEAR_BACKENDS:

127 + supported, reason = is_backend_supported(backend)

128 + if supported:

129 + selected_backend = backend

130 + break

131 else:

74 - backend = NvFp4LinearBackend(envs.VLLM_NVFP4_GEMM_BACKEND)

132 + selected_backend = NvFp4LinearBackend(envs.VLLM_NVFP4_GEMM_BACKEND)

133

76 - # Validate that the backend is supported

77 - if backend in (

78 - NvFp4LinearBackend.FLASHINFER_CUTLASS,

79 - NvFp4LinearBackend.FLASHINFER_TRTLLM,

80 - NvFp4LinearBackend.FLASHINFER_CUDNN,

81 - ):

82 - assert has_flashinfer(), f"FlashInfer is required for {backend}"

83 - assert cutlass_fp4_supported(), (

84 - f"{backend} requires vLLM NVFP4 quantization kernels compiled "

85 - f"for the current GPU (SM {current_platform.get_device_capability()})"

86 - )

87 - elif backend == NvFp4LinearBackend.VLLM_CUTLASS:

88 - assert cutlass_fp4_supported(), f"Cutlass is required for {backend}"

89 - elif backend == NvFp4LinearBackend.MARLIN:

90 - assert is_fp4_marlin_supported(), f"Marlin is required for {backend}"

91 - elif backend is None:

134 + if selected_backend is None:

135 raise ValueError(

136 f"No NVFP4 GEMM backend selected, "

94 - f"available backends: {list(NvFp4LinearBackend)}"

137 + f"available backends: {NVFP4_LINEAR_BACKENDS}"

138 )

139

97 - logger.info_once(f"Using {backend} for NVFP4 GEMM")

98 - return backend

140 + supported, reason = is_backend_supported(selected_backend)

141 +

142 + if not supported:

143 + raise ValueError(

144 + f"The selected backend={selected_backend} is not supported in current "

145 + f"environment. Reason: {reason}. Current environment: "

146 + f"{envs.VLLM_USE_FBGEMM=}, {envs.VLLM_USE_NVFP4_CT_EMULATIONS=}, "

147 + f"{envs.VLLM_NVFP4_GEMM_BACKEND}."

148 + )

149 +

150 + logger.info_once(f"Using {selected_backend} for NVFP4 GEMM")

151 + return selected_backend

152

153

154 def prepare_weights_for_nvfp4_flashinfer_trtllm(

@@ -183,6 +236,10 @@ def convert_to_nvfp4_linear_kernel_format(

236 layer.weight = torch.nn.Parameter(weight, requires_grad=False)

237 layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)

238 layer.weights_padding_cols = weights_padding_cols

239 + elif backend == NvFp4LinearBackend.EMULATION:

240 + # We can not call `.to(device)` during cuda graph capture - do it here instead.

241 + # (operation not permitted when stream is capturing)

242 + kE2M1ToFloat_handle.val = kE2M1ToFloat_handle.val.to(layer.weight.device)

243

244

245 def apply_nvfp4_linear(

@@ -190,6 +247,7 @@ def apply_nvfp4_linear(

247 layer: torch.nn.Module,

248 x: torch.Tensor,

249 bias: torch.Tensor | None = None,

250 + swizzle: bool | None = None,

251 ) -> torch.Tensor:

252 """

253 Apply NVFP4 linear transformation using the specified backend.

@@ -220,6 +278,7 @@ def apply_nvfp4_linear(

278 weight=weight,

279 weight_scale_swizzled=weight_scale,

280 weight_global_scale=weight_global_scale,

281 + swizzle=swizzle,

282 )

283 if bias is not None:

284 out = out + bias

vllm/platforms/rocm.py

@@ -409,6 +409,7 @@ class RocmPlatform(Platform):

409 "mxfp4",

410 "torchao",

411 "bitsandbytes",

412 + "modelopt_fp4",

413 ]

414

415 @classmethod

vllm/utils/import_utils.py

@@ -461,3 +461,8 @@ def has_aiter() -> bool:

461 def has_mori() -> bool:

462 """Whether the optional `mori` package is available."""

463 return _has_module("mori")

464 +

465 +

466 +def has_fbgemm_gpu() -> bool:

467 + """Whether the optional `fbgemm_gpu` package is available."""

468 + return _has_module("fbgemm_gpu")