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, 07:30 UTC
Refactor NVFP4 Linear utils for ModelOpt and CT (#33201)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jan 30, 2026 at 19:37 UTC
67ebaff528be7faef14fa4eb86cfd8c4b5714e38
12 files changed
+462
-483
tests/kernels/moe/modular_kernel_tools/mk_objects.py
+1
-1
@@ -25,7 +25,7 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize import (
25
from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
26
TritonOrDeepGemmExperts,
27
)
28
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
28
+from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
29
cutlass_fp4_supported,
30
)
31
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
tests/quantization/test_compressed_tensors.py
+1
-1
@@ -24,7 +24,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
24
)
25
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
26
from vllm.model_executor.layers.quantization.utils.fp8_utils import W8A8BlockFp8LinearOp
27
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
27
+from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
28
cutlass_fp4_supported,
29
)
30
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+1
-12
@@ -18,7 +18,6 @@ from compressed_tensors.quantization import (
18
)
19
from compressed_tensors.transform import TransformConfig
20
21
-import vllm.envs as envs
21
from vllm.distributed import (
22
get_tensor_model_parallel_rank,
23
get_tensor_model_parallel_world_size,
@@ -63,9 +62,6 @@ from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
62
should_ignore_layer,
63
)
64
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
66
-from vllm.model_executor.layers.quantization.utils.quant_utils import (
67
- cutlass_fp4_supported,
68
-)
65
from vllm.platforms import current_platform
66
67
if TYPE_CHECKING:
@@ -627,14 +623,7 @@ class CompressedTensorsConfig(QuantizationConfig):
623
if self._is_nvfp4_format(weight_quant) and self._is_nvfp4_format(
624
input_quant
625
):
630
- if cutlass_fp4_supported() or envs.VLLM_USE_NVFP4_CT_EMULATIONS:
631
- return CompressedTensorsW4A4Fp4()
632
- else:
633
- logger.warning_once(
634
- "Current platform does not support cutlass NVFP4."
635
- " Running CompressedTensorsW4A16Fp4."
636
- )
637
- return CompressedTensorsW4A16Fp4(has_input_global_scale=True)
626
+ return CompressedTensorsW4A4Fp4()
627
628
if self._is_fp8_w8a8(weight_quant, input_quant):
629
is_fp8_w8a8_supported = self._check_scheme_supported(
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a16_mxfp4.py
+2
-2
@@ -81,7 +81,7 @@ class CompressedTensorsW4A16Mxfp4(CompressedTensorsScheme):
81
)
82
layer.register_parameter("weight_scale", weight_scale)
83
84
- def process_weights_after_loading(self, layer) -> None:
84
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
85
# Rename weight_packed to weight that marlin expects
86
layer.weight = Parameter(layer.weight_packed.data, requires_grad=False)
87
del layer.weight_packed
@@ -98,7 +98,7 @@ class CompressedTensorsW4A16Mxfp4(CompressedTensorsScheme):
98
input=x,
99
weight=layer.weight,
100
weight_scale=layer.weight_scale,
101
- weight_scale_2=None,
101
+ weight_global_scale=None,
102
workspace=layer.workspace,
103
size_n=layer.output_size_per_partition,
104
size_k=layer.input_size_per_partition,
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a16_nvfp4.py
+6
-21
@@ -22,8 +22,7 @@ __all__ = ["CompressedTensorsW4A16Fp4"]
22
23
24
class CompressedTensorsW4A16Fp4(CompressedTensorsScheme):
25
- def __init__(self, has_input_global_scale: bool = False):
26
- self.has_input_global_scale = has_input_global_scale
25
+ def __init__(self):
26
self.group_size = 16
27
28
@classmethod
@@ -79,30 +78,16 @@ class CompressedTensorsW4A16Fp4(CompressedTensorsScheme):
78
79
layer.register_parameter("weight_scale", weight_scale)
80
82
- if self.has_input_global_scale:
83
- input_global_scale = PerTensorScaleParameter(
84
- data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
85
- weight_loader=weight_loader,
86
- )
87
- layer.register_parameter("input_global_scale", input_global_scale)
88
-
89
- def process_weights_after_loading(self, layer) -> None:
81
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
82
# Process parameters for marlin repacking
83
84
# Rename weight_packed to weight that marlin expects
85
layer.weight = Parameter(layer.weight_packed.data, requires_grad=False)
86
del layer.weight_packed
95
- # Rename weight_global_scale to weight_scale_2 that marlin expects
96
- # Note: ct stores the inverse of what is expected by the marlin kernel
97
- layer.weight_scale_2 = Parameter(
98
- 1 / layer.weight_global_scale.max().to(torch.float32), requires_grad=False
87
+ # ct stores the inverse of what is expected by the marlin kernel
88
+ layer.weight_global_scale = Parameter(
89
+ 1.0 / layer.weight_global_scale.max().to(torch.float32), requires_grad=False
90
)
100
- del layer.weight_global_scale
101
-
102
- if self.has_input_global_scale:
103
- layer.input_global_scale = torch.nn.Parameter(
104
- layer.input_global_scale.data, requires_grad=False
105
- )
91
92
prepare_fp4_layer_for_marlin(layer)
93
@@ -116,7 +101,7 @@ class CompressedTensorsW4A16Fp4(CompressedTensorsScheme):
101
input=x,
102
weight=layer.weight,
103
weight_scale=layer.weight_scale,
119
- weight_scale_2=layer.weight_scale_2,
104
+ weight_global_scale=layer.weight_global_scale,
105
workspace=layer.workspace,
106
size_n=layer.output_size_per_partition,
107
size_k=layer.input_size_per_partition,
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
+29
-153
@@ -5,75 +5,31 @@ from collections.abc import Callable
5
import torch
6
from torch.nn.parameter import Parameter
7
8
-import vllm.envs as envs
9
-from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
10
-from vllm.logger import init_logger
8
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
9
CompressedTensorsScheme,
10
)
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 (
18
- cutlass_fp4_supported,
19
- pad_nvfp4_activation_for_cutlass,
20
- pad_nvfp4_weight_for_cutlass,
21
- slice_nvfp4_output,
22
- swizzle_blockscale,
11
+from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
12
+ apply_nvfp4_linear,
13
+ convert_to_nvfp4_linear_kernel_format,
14
+ select_nvfp4_linear_backend,
15
)
16
from vllm.model_executor.parameter import (
17
GroupQuantScaleParameter,
18
ModelWeightParameter,
19
PerTensorScaleParameter,
20
)
29
-from vllm.utils.flashinfer import (
30
- flashinfer_scaled_fp4_mm,
31
- has_flashinfer,
32
-)
33
-
34
-logger = init_logger(__name__)
21
22
__all__ = ["CompressedTensorsW4A4Fp4"]
23
24
25
class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
26
def __init__(self):
41
- self.backend = "none"
42
- if envs.VLLM_NVFP4_GEMM_BACKEND is None:
43
- if has_flashinfer():
44
- self.backend = "flashinfer-cutlass"
45
- elif cutlass_fp4_supported():
46
- self.backend = "cutlass"
47
- elif envs.VLLM_USE_FBGEMM:
48
- self.backend = "fbgemm"
49
- try:
50
- import fbgemm_gpu # noqa: F401
51
- except ImportError as exc:
52
- raise ImportError(
53
- "Backend fbgemm requires fbgemm.f4f4bf16 operator, "
54
- "Please install with: pip install fbgemm-gpu-genai"
55
- ) from exc
56
- elif envs.VLLM_NVFP4_GEMM_BACKEND.startswith("flashinfer-"):
57
- self.backend = envs.VLLM_NVFP4_GEMM_BACKEND
58
- assert has_flashinfer(), f"FlashInfer is required for {self.backend}"
59
- elif envs.VLLM_NVFP4_GEMM_BACKEND == "cutlass":
60
- self.backend = "cutlass"
61
- assert cutlass_fp4_supported(), f"Cutlass is required for {self.backend}"
62
-
63
- if self.backend == "none":
64
- raise ValueError(
65
- "No valid NVFP4 GEMM backend found. "
66
- "Please check your platform capability."
67
- )
68
-
69
- logger.info_once(f"Using {self.backend} for NVFP4 GEMM")
27
+ self.backend = select_nvfp4_linear_backend()
28
self.group_size = 16
29
30
@classmethod
31
def get_min_capability(cls) -> int:
74
- if envs.VLLM_USE_NVFP4_CT_EMULATIONS:
75
- return 80
76
- return 100
32
+ return 75
33
34
def create_weights(
35
self,
@@ -129,120 +85,40 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
85
)
86
layer.register_parameter("input_global_scale", input_global_scale)
87
132
- def process_weights_after_loading(self, layer) -> None:
133
- global_input_scale = layer.input_global_scale.max().to(torch.float32)
134
- layer.input_global_scale = Parameter(global_input_scale, requires_grad=False)
135
-
88
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
89
+ # Rename CT checkpoint names to standardized names
90
+ layer.weight = layer.weight_packed
91
+ del layer.weight_packed
92
+ # Process global scales (CT stores as divisors, i.e. 1/scale)
93
+ input_global_scale_inv = layer.input_global_scale.max().to(torch.float32)
94
+ layer.input_global_scale = Parameter(
95
+ (1.0 / input_global_scale_inv).to(torch.float32), requires_grad=False
96
+ )
97
+ weight_global_scale = layer.weight_global_scale.max().to(torch.float32)
98
layer.weight_global_scale = Parameter(
137
- layer.weight_global_scale.max().to(torch.float32), requires_grad=False
99
+ 1.0 / weight_global_scale, requires_grad=False
100
)
101
140
- if self.backend == "flashinfer-trtllm":
141
- # FlashInfer TRTLLM FP4 GEMM requires a different weight layout.
142
- # FlashInfer provides nvfp4_quantize to quantize + shuffle the
143
- # layout but we use our own quantization so we have to call
144
- # shuffles ourselves.
145
- from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a
146
-
147
- weight = layer.weight_packed.data
148
- weight_scale = layer.weight_scale.data
149
-
150
- epilogue_tile_m = 128
151
- weight = shuffle_matrix_a(weight.view(torch.uint8), epilogue_tile_m)
152
- weight_scale = (
153
- shuffle_matrix_sf_a(weight_scale.view(torch.uint8), epilogue_tile_m)
154
- .reshape(weight_scale.shape)
155
- .view(torch.float8_e4m3fn)
156
- )
157
-
158
- layer.weight_scale = Parameter(weight_scale, requires_grad=False)
159
- layer.weight_packed = Parameter(weight, requires_grad=False)
160
- else:
161
- swizzled_weight_scale = swizzle_blockscale(layer.weight_scale)
162
- if self.backend == "fbgemm":
163
- swizzled_weight_scale = swizzled_weight_scale.view(-1).view(torch.uint8)
164
- layer.weight_scale = Parameter(swizzled_weight_scale, requires_grad=False)
165
-
166
- # Pad weights for CUTLASS/FlashInfer kernel alignment (K and N
167
- # divisible by 32). fbgemm has its own layout requirements.
168
- if self.backend in ("cutlass", "flashinfer-cutlass"):
169
- weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
170
- layer.weight_packed.data
171
- )
172
- layer.weights_padding_cols = weights_padding_cols
173
- layer.weight_packed = Parameter(weight, requires_grad=False)
174
- else:
175
- layer.weights_padding_cols = 0
176
- layer.weight_packed = Parameter(
177
- layer.weight_packed.data, requires_grad=False
178
- )
179
-
102
+ # Pre-compute alpha and inverse for runtime quantization
103
+ layer.input_global_scale_inv = Parameter(
104
+ input_global_scale_inv, requires_grad=False
105
+ )
106
layer.alpha = Parameter(
181
- 1 / (layer.input_global_scale * layer.weight_global_scale),
182
- requires_grad=False,
107
+ layer.input_global_scale * layer.weight_global_scale, requires_grad=False
108
)
109
110
+ # Convert layer to NVFP4 linear kernel format
111
+ convert_to_nvfp4_linear_kernel_format(self.backend, layer)
112
+
113
def apply_weights(
114
self,
115
layer: torch.nn.Module,
116
x: torch.Tensor,
117
bias: torch.Tensor | None = None,
118
) -> torch.Tensor:
191
- if envs.VLLM_USE_NVFP4_CT_EMULATIONS:
192
- out = run_nvfp4_emulations(
193
- x=x,
194
- input_global_scale=layer.input_global_scale,
195
- weight=layer.weight_packed,
196
- weight_scale_swizzled=layer.weight_scale,
197
- weight_global_scale=layer.weight_global_scale,
198
- )
199
- if bias is not None:
200
- out = out + bias
201
- return out
202
-
203
- output_dtype = x.dtype
204
- output_size = layer.output_size_per_partition
205
- output_shape = [*x.shape[:-1], output_size]
206
-
207
- # quantize BF16 or FP16 to (FP4 and interleaved block scale)
208
- x_fp4, x_blockscale = scaled_fp4_quant(
209
- x,
210
- layer.input_global_scale,
211
- is_sf_swizzled_layout=True,
119
+ return apply_nvfp4_linear(
120
backend=self.backend,
121
+ layer=layer,
122
+ x=x,
123
+ bias=bias,
124
)
214
-
215
- # Pad activations to match weight K-dimension padding
216
- weights_padding_cols = getattr(layer, "weights_padding_cols", 0)
217
- x_fp4 = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_cols)
218
-
219
- mm_args = (
220
- x_fp4,
221
- layer.weight_packed,
222
- x_blockscale,
223
- layer.weight_scale,
224
- layer.alpha,
225
- output_dtype,
226
- )
227
- if self.backend.startswith("flashinfer-"):
228
- backend_name = self.backend[len("flashinfer-") :]
229
- out = flashinfer_scaled_fp4_mm(*mm_args, backend=backend_name)
230
- elif self.backend == "fbgemm":
231
- out = torch.ops.fbgemm.f4f4bf16(
232
- x_fp4,
233
- layer.weight_packed,
234
- x_blockscale.view(-1).view(torch.uint8),
235
- layer.weight_scale,
236
- layer.alpha,
237
- use_mx=False,
238
- ).to(output_dtype)
239
- else:
240
- assert self.backend == "cutlass"
241
- out = cutlass_scaled_fp4_mm(*mm_args)
242
-
243
- # Slice output to remove N-dimension padding
244
- out = slice_nvfp4_output(out, output_size)
245
-
246
- if bias is not None:
247
- out = out + bias
248
- return out.view(*output_shape)
vllm/model_executor/layers/quantization/modelopt.py
+34
-162
@@ -5,12 +5,9 @@ from fnmatch import fnmatch
5
from typing import TYPE_CHECKING, Any
6
7
import torch
8
-from torch.nn import Module
8
from torch.nn.parameter import Parameter
9
11
-import vllm.envs as envs
10
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
13
-from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
11
from vllm.logger import init_logger
12
from vllm.model_executor.layers.attention import Attention
13
from vllm.model_executor.layers.fused_moe.config import (
@@ -66,24 +63,19 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
63
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
64
get_marlin_input_dtype,
65
)
69
-from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
70
- apply_fp4_marlin_linear,
71
- is_fp4_marlin_supported,
72
- prepare_fp4_layer_for_marlin,
66
+from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
67
+ apply_nvfp4_linear,
68
+ convert_to_nvfp4_linear_kernel_format,
69
+ select_nvfp4_linear_backend,
70
)
71
from vllm.model_executor.layers.quantization.utils.quant_utils import (
72
GroupShape,
76
- cutlass_fp4_supported,
73
is_layer_skipped,
74
kFp8DynamicTokenSym,
75
kFp8StaticTensorSym,
76
kFp8StaticTokenSym,
77
kNvfp4Dynamic,
78
kNvfp4Static,
83
- pad_nvfp4_activation_for_cutlass,
84
- pad_nvfp4_weight_for_cutlass,
85
- slice_nvfp4_output,
86
- swizzle_blockscale,
79
)
80
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
81
cutlass_block_fp8_supported,
@@ -96,11 +88,6 @@ from vllm.model_executor.parameter import (
88
PerTensorScaleParameter,
89
)
90
from vllm.model_executor.utils import replace_parameter
99
-from vllm.platforms import current_platform
100
-from vllm.utils.flashinfer import (
101
- flashinfer_scaled_fp4_mm,
102
- has_flashinfer,
103
-)
91
92
if TYPE_CHECKING:
93
from vllm.model_executor.models.utils import WeightsMapper
@@ -498,7 +485,7 @@ class ModelOptFp8LinearMethod(LinearMethodBase):
485
scale[:] = torch.finfo(torch.float32).min
486
layer.register_parameter("input_scale", scale)
487
501
- def process_weights_after_loading(self, layer: Module) -> None:
488
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
489
weight = layer.weight
490
max_w_scale = layer.weight_scale.max()
491
if not (layer.weight_scale == layer.weight_scale[0]).all():
@@ -580,7 +567,7 @@ class ModelOptFp8PcPtLinearMethod(LinearMethodBase):
567
weight_scale[:] = torch.finfo(torch.float32).min
568
layer.register_parameter("weight_scale", weight_scale)
569
583
- def process_weights_after_loading(self, layer: Module) -> None:
570
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
571
layer.weight = Parameter(layer.weight.t(), requires_grad=False)
572
layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False)
573
@@ -681,7 +668,7 @@ class ModelOptFp8PbWoLinearMethod(LinearMethodBase):
668
weight_scale[:] = torch.finfo(torch.float32).min
669
layer.register_parameter("weight_scale", weight_scale)
670
684
- def process_weights_after_loading(self, layer: Module) -> None:
671
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
672
# Keep weight in [out, in] layout for W8A8BlockFp8LinearOp.
673
layer.weight = Parameter(layer.weight.data, requires_grad=False)
674
@@ -1108,32 +1095,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
1095
def __init__(self, quant_config: ModelOptNvFp4Config) -> None:
1096
self.quant_config = quant_config
1097
self.marlin_input_dtype = None
1111
-
1112
- self.backend = "none"
1113
- if envs.VLLM_NVFP4_GEMM_BACKEND is None:
1114
- if current_platform.has_device_capability(100) and has_flashinfer():
1115
- self.backend = "flashinfer-cutlass"
1116
- elif cutlass_fp4_supported():
1117
- self.backend = "cutlass"
1118
- elif is_fp4_marlin_supported():
1119
- self.backend = "marlin"
1120
- elif envs.VLLM_NVFP4_GEMM_BACKEND.startswith("flashinfer-"):
1121
- self.backend = envs.VLLM_NVFP4_GEMM_BACKEND
1122
- assert has_flashinfer(), f"FlashInfer is required for {self.backend}"
1123
- elif envs.VLLM_NVFP4_GEMM_BACKEND == "cutlass":
1124
- self.backend = "cutlass"
1125
- assert cutlass_fp4_supported(), f"Cutlass is required for {self.backend}"
1126
- elif envs.VLLM_NVFP4_GEMM_BACKEND == "marlin":
1127
- self.backend = "marlin"
1128
- assert is_fp4_marlin_supported(), f"Marlin is required for {self.backend}"
1129
-
1130
- if self.backend == "none":
1131
- raise ValueError(
1132
- "No valid NVFP4 GEMM backend found. "
1133
- "Please check your platform capability."
1134
- )
1135
-
1136
- logger.info_once(f"Using {self.backend} for NVFP4 GEMM")
1098
+ self.backend = select_nvfp4_linear_backend()
1099
1100
def create_weights(
1101
self,
@@ -1181,19 +1143,19 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
1143
)
1144
layer.register_parameter("weight", weight)
1145
1184
- # Input Weight Scale
1185
- input_scale = PerTensorScaleParameter(
1146
+ # Input Global Scale
1147
+ input_global_scale = PerTensorScaleParameter(
1148
data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
1149
weight_loader=weight_loader,
1150
)
1189
- layer.register_parameter("input_scale", input_scale)
1151
+ layer.register_parameter("input_scale", input_global_scale)
1152
1191
- # Global Weight Scale
1192
- weight_scale_2 = PerTensorScaleParameter(
1153
+ # Weight Global Scale
1154
+ weight_global_scale = PerTensorScaleParameter(
1155
data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
1156
weight_loader=weight_loader,
1157
)
1196
- layer.register_parameter("weight_scale_2", weight_scale_2)
1158
+ layer.register_parameter("weight_scale_2", weight_global_scale)
1159
1160
# Per Block Weight Scale
1161
weight_scale = ModelWeightParameter(
@@ -1209,65 +1171,25 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
1171
1172
layer.register_parameter("weight_scale", weight_scale)
1173
1212
- def process_weights_after_loading(self, layer: Module) -> None:
1213
- # global scales:
1214
- input_scale_2 = layer.input_scale.max().to(torch.float32)
1215
- layer.input_scale = Parameter(input_scale_2, requires_grad=False)
1216
-
1217
- weight_scale_2 = layer.weight_scale_2.max().to(torch.float32)
1218
- layer.weight_scale_2 = Parameter(weight_scale_2, requires_grad=False)
1219
-
1174
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1175
+ # Rename ModelOpt checkpoint names to standardized names
1176
+ input_global_scale = layer.input_scale.max().to(torch.float32)
1177
+ layer.input_global_scale = Parameter(input_global_scale, requires_grad=False)
1178
+ del layer.input_scale
1179
+ weight_global_scale = layer.weight_scale_2.max().to(torch.float32)
1180
+ layer.weight_global_scale = Parameter(weight_global_scale, requires_grad=False)
1181
+ del layer.weight_scale_2
1182
+
1183
+ # Pre-compute alpha and inverse for runtime quantization
1184
layer.alpha = Parameter(
1221
- layer.input_scale * layer.weight_scale_2, requires_grad=False
1222
- )
1223
-
1224
- # Calculate `1 / input_scale` so that we don't need to do so at runtime
1225
- layer.input_scale_inv = Parameter(
1226
- (1 / layer.input_scale).to(torch.float32), requires_grad=False
1185
+ layer.input_global_scale * layer.weight_global_scale, requires_grad=False
1186
)
1228
-
1229
- # Swizzle the weight blockscale.
1230
- # contracting dimension is input dimension
1231
- # block_size = 16;
1232
- assert layer.weight_scale.dtype == torch.float8_e4m3fn, (
1233
- "Weight Block scale must be represented as FP8-E4M3"
1187
+ layer.input_global_scale_inv = Parameter(
1188
+ (1.0 / layer.input_global_scale).to(torch.float32), requires_grad=False
1189
)
1190
1236
- if self.backend == "marlin":
1237
- prepare_fp4_layer_for_marlin(layer)
1238
- del layer.alpha
1239
- del layer.input_scale
1240
- elif self.backend == "flashinfer-trtllm":
1241
- # FlashInfer TRTLLM FP4 GEMM requires a different weight layout.
1242
- # FlashInfer provides nvfp4_quantize to quantize + shuffle the
1243
- # layout but we use our own quantization so we have to call
1244
- # shuffles ourselves.
1245
- from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a
1246
-
1247
- weight = layer.weight.data
1248
- weight_scale = layer.weight_scale.data
1249
-
1250
- epilogue_tile_m = 128
1251
- weight = shuffle_matrix_a(weight.view(torch.uint8), epilogue_tile_m)
1252
- weight_scale = (
1253
- shuffle_matrix_sf_a(weight_scale.view(torch.uint8), epilogue_tile_m)
1254
- .reshape(weight_scale.shape)
1255
- .view(torch.float8_e4m3fn)
1256
- )
1257
-
1258
- layer.weight_scale = Parameter(weight_scale, requires_grad=False)
1259
- layer.weight = Parameter(weight, requires_grad=False)
1260
- else:
1261
- # Swizzle block scales and pad the packed NVFP4 weights for kernel
1262
- # alignment (CUTLASS/FlashInfer require K and N divisible by 32).
1263
- swizzled_weight_scale = swizzle_blockscale(layer.weight_scale)
1264
- layer.weight_scale = Parameter(swizzled_weight_scale, requires_grad=False)
1265
-
1266
- weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
1267
- layer.weight.data
1268
- )
1269
- layer.weights_padding_cols = weights_padding_cols
1270
- layer.weight = Parameter(weight, requires_grad=False)
1191
+ # Convert layer to NVFP4 linear kernel format
1192
+ convert_to_nvfp4_linear_kernel_format(self.backend, layer)
1193
1194
def apply(
1195
self,
@@ -1275,63 +1197,13 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
1197
x: torch.Tensor,
1198
bias: torch.Tensor | None = None,
1199
) -> torch.Tensor:
1278
- if self.backend == "marlin":
1279
- return apply_fp4_marlin_linear(
1280
- input=x,
1281
- weight=layer.weight,
1282
- weight_scale=layer.weight_scale,
1283
- weight_scale_2=layer.weight_scale_2,
1284
- workspace=layer.workspace,
1285
- size_n=layer.output_size_per_partition,
1286
- size_k=layer.input_size_per_partition,
1287
- bias=bias,
1288
- input_dtype=self.marlin_input_dtype,
1289
- )
1290
-
1291
- output_dtype = x.dtype
1292
-
1293
- # quantize BF16 or FP16 to (FP4 and interleaved block scale)
1294
- x_fp4, x_blockscale = scaled_fp4_quant(
1295
- x, layer.input_scale_inv, is_sf_swizzled_layout=True, backend=self.backend
1296
- )
1297
-
1298
- # validate dtypes of quantized input, input block scale,
1299
- # weight and weight_blockscale
1300
- assert x_fp4.dtype == torch.uint8
1301
- assert layer.weight.dtype == torch.uint8
1302
- assert x_blockscale.dtype == torch.float8_e4m3fn
1303
- assert layer.weight_scale.dtype == torch.float8_e4m3fn
1304
- assert layer.alpha.dtype == torch.float32
1305
-
1306
- # Pad activations to match weight K-dimension padding
1307
- weights_padding_cols = getattr(layer, "weights_padding_cols", 0)
1308
- output_size = layer.output_size_per_partition
1309
- output_shape = [x.shape[0], output_size]
1310
- x_fp4 = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_cols)
1311
-
1312
- mm_args = (
1313
- x_fp4,
1314
- layer.weight,
1315
- x_blockscale,
1316
- layer.weight_scale,
1317
- layer.alpha,
1318
- output_dtype,
1200
+ return apply_nvfp4_linear(
1201
+ backend=self.backend,
1202
+ layer=layer,
1203
+ x=x,
1204
+ bias=bias,
1205
)
1206
1321
- if self.backend.startswith("flashinfer-"):
1322
- backend_name = self.backend[len("flashinfer-") :]
1323
- out = flashinfer_scaled_fp4_mm(*mm_args, backend=backend_name)
1324
- else:
1325
- assert self.backend == "cutlass"
1326
- out = cutlass_scaled_fp4_mm(*mm_args)
1327
-
1328
- # Slice output to remove N-dimension padding
1329
- out = slice_nvfp4_output(out, output_size)
1330
-
1331
- if bias is not None:
1332
- out = out + bias
1333
- return out.view(*output_shape)
1334
-
1207
1208
class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1209
"""
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+3
-1
@@ -15,11 +15,13 @@ from vllm.model_executor.layers.fused_moe.config import (
15
FusedMoEParallelConfig,
16
RoutingMethodType,
17
)
18
+from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
19
+ swizzle_blockscale,
20
+)
21
from vllm.model_executor.layers.quantization.utils.quant_utils import (
22
QuantKey,
23
kNvfp4Dynamic,
24
kNvfp4Static,
22
- swizzle_blockscale,
25
)
26
from vllm.platforms import current_platform
27
from vllm.utils.flashinfer import (
vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py
+9
-7
@@ -92,7 +92,7 @@ def apply_fp4_marlin_linear(
92
input: torch.Tensor,
93
weight: torch.Tensor,
94
weight_scale: torch.Tensor,
95
- weight_scale_2: torch.Tensor | None,
95
+ weight_global_scale: torch.Tensor | None,
96
workspace: torch.Tensor,
97
size_n: int,
98
size_k: int,
@@ -112,7 +112,7 @@ def apply_fp4_marlin_linear(
112
113
inputs = reshaped_x
114
a_scales = None
115
- is_nvfp4 = weight_scale_2 is not None
115
+ is_nvfp4 = weight_global_scale is not None
116
if input_dtype is not None and input_dtype.itemsize == 1:
117
if is_nvfp4:
118
raise RuntimeError("NVFP4 weight + INT8/FP8 activation is not supported.")
@@ -128,7 +128,7 @@ def apply_fp4_marlin_linear(
128
b_bias=bias,
129
b_scales=weight_scale,
130
a_scales=a_scales,
131
- global_scale=weight_scale_2,
131
+ global_scale=weight_global_scale,
132
b_zeros=None,
133
g_idx=None,
134
perm=None,
@@ -154,7 +154,7 @@ def prepare_fp4_layer_for_marlin(
154
"performance for compute-heavy workloads."
155
)
156
157
- is_nvfp4 = hasattr(layer, "weight_scale_2")
157
+ is_nvfp4 = hasattr(layer, "weight_global_scale")
158
if input_dtype is not None and input_dtype.itemsize == 1:
159
if is_nvfp4:
160
raise RuntimeError("NVFP4 weight + INT8/FP8 activation is not supported.")
@@ -210,9 +210,11 @@ def prepare_fp4_layer_for_marlin(
210
weight_scale = nvfp4_marlin_process_scales(weight_scale)
211
layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
212
213
- weight_scale_2 = layer.weight_scale_2.to(param_dtype)
214
- weight_scale_2 = nvfp4_marlin_process_global_scale(weight_scale_2)
215
- layer.weight_scale_2 = torch.nn.Parameter(weight_scale_2, requires_grad=False)
213
+ weight_global_scale = layer.weight_global_scale.to(param_dtype)
214
+ weight_global_scale = nvfp4_marlin_process_global_scale(weight_global_scale)
215
+ layer.weight_global_scale = torch.nn.Parameter(
216
+ weight_global_scale, requires_grad=False
217
+ )
218
else:
219
weight_scale = mxfp4_marlin_process_scales(
220
weight_scale, input_dtype=input_dtype
vllm/model_executor/layers/quantization/utils/nvfp4_moe_support.py
+1
-1
@@ -11,7 +11,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
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.quant_utils import (
14
+from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
15
cutlass_fp4_supported,
16
)
17
vllm/model_executor/layers/quantization/utils/nvfp4_utils.py
new
+375
@@ -0,0 +1,375 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from enum import Enum
4
+
5
+import torch
6
+
7
+import vllm.envs as envs
8
+from vllm._custom_ops import (
9
+ cutlass_scaled_fp4_mm,
10
+ cutlass_scaled_mm_supports_fp4,
11
+ scaled_fp4_quant,
12
+)
13
+from vllm.logger import init_logger
14
+from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
15
+ apply_fp4_marlin_linear,
16
+ is_fp4_marlin_supported,
17
+ prepare_fp4_layer_for_marlin,
18
+)
19
+from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
20
+ run_nvfp4_emulations,
21
+)
22
+from vllm.platforms import current_platform
23
+from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm, has_flashinfer
24
+from vllm.utils.math_utils import round_up
25
+
26
+logger = init_logger(__name__)
27
+
28
+
29
+class NvFp4LinearBackend(Enum):
30
+ VLLM_CUTLASS = "cutlass"
31
+ FLASHINFER_CUTLASS = "flashinfer-cutlass"
32
+ FLASHINFER_TRTLLM = "flashinfer-trtllm"
33
+ FLASHINFER_CUDNN = "flashinfer-cudnn"
34
+ FBGEMM = "fbgemm"
35
+ MARLIN = "marlin"
36
+ EMULATION = "emulation"
37
+
38
+
39
+def select_nvfp4_linear_backend() -> NvFp4LinearBackend:
40
+ """
41
+ Select the best available NVFP4 GEMM backend based on environment
42
+ configuration and platform capabilities.
43
+ """
44
+ backend: NvFp4LinearBackend | None = None
45
+
46
+ if envs.VLLM_USE_FBGEMM:
47
+ try:
48
+ import fbgemm_gpu # noqa: F401
49
+ except ImportError as exc:
50
+ raise ImportError(
51
+ "Backend fbgemm requires fbgemm.f4f4bf16 operator, "
52
+ "Please install with: pip install fbgemm-gpu-genai"
53
+ ) from exc
54
+ backend = NvFp4LinearBackend.FBGEMM
55
+ elif envs.VLLM_USE_NVFP4_CT_EMULATIONS:
56
+ backend = NvFp4LinearBackend.EMULATION
57
+ elif envs.VLLM_NVFP4_GEMM_BACKEND is None:
58
+ # Auto-select best available backend
59
+ if current_platform.has_device_capability(100) and has_flashinfer():
60
+ backend = NvFp4LinearBackend.FLASHINFER_CUTLASS
61
+ elif cutlass_fp4_supported():
62
+ backend = NvFp4LinearBackend.VLLM_CUTLASS
63
+ elif is_fp4_marlin_supported():
64
+ backend = NvFp4LinearBackend.MARLIN
65
+ else:
66
+ backend = NvFp4LinearBackend(envs.VLLM_NVFP4_GEMM_BACKEND)
67
+
68
+ # Validate that the backend is supported
69
+ if backend in (
70
+ NvFp4LinearBackend.FLASHINFER_CUTLASS,
71
+ NvFp4LinearBackend.FLASHINFER_TRTLLM,
72
+ NvFp4LinearBackend.FLASHINFER_CUDNN,
73
+ ):
74
+ assert has_flashinfer(), f"FlashInfer is required for {backend}"
75
+ elif backend == NvFp4LinearBackend.VLLM_CUTLASS:
76
+ assert cutlass_fp4_supported(), f"Cutlass is required for {backend}"
77
+ elif backend == NvFp4LinearBackend.MARLIN:
78
+ assert is_fp4_marlin_supported(), f"Marlin is required for {backend}"
79
+ elif backend is None:
80
+ raise ValueError(
81
+ f"No NVFP4 GEMM backend selected, "
82
+ f"available backends: {list(NvFp4LinearBackend)}"
83
+ )
84
+
85
+ logger.info_once(f"Using {backend} for NVFP4 GEMM")
86
+ return backend
87
+
88
+
89
+def prepare_weights_for_nvfp4_flashinfer_trtllm(
90
+ weight: torch.Tensor,
91
+ weight_scale: torch.Tensor,
92
+) -> tuple[torch.Tensor, torch.Tensor]:
93
+ """Prepare weights and scales for FlashInfer TRTLLM FP4 GEMM."""
94
+ from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a
95
+
96
+ epilogue_tile_m = 128
97
+ shuffled_weight = shuffle_matrix_a(weight.view(torch.uint8), epilogue_tile_m)
98
+ shuffled_weight_scale = (
99
+ shuffle_matrix_sf_a(weight_scale.view(torch.uint8), epilogue_tile_m)
100
+ .reshape(weight_scale.shape)
101
+ .view(torch.float8_e4m3fn)
102
+ )
103
+
104
+ return shuffled_weight, shuffled_weight_scale
105
+
106
+
107
+def prepare_weights_for_nvfp4_cutlass(
108
+ weight: torch.Tensor,
109
+ weight_scale: torch.Tensor,
110
+) -> tuple[torch.Tensor, torch.Tensor, int]:
111
+ """
112
+ Prepare weights and scales for CUTLASS/FlashInfer-CUTLASS FP4 GEMM.
113
+ This involves padding weights for alignment (K and N divisible by 32)
114
+ """
115
+ swizzled_weight_scale = swizzle_blockscale(weight_scale)
116
+ padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(weight)
117
+ return padded_weight, swizzled_weight_scale, weights_padding_cols
118
+
119
+
120
+def prepare_weights_for_nvfp4_fbgemm(
121
+ weight: torch.Tensor,
122
+ weight_scale: torch.Tensor,
123
+) -> tuple[torch.Tensor, torch.Tensor]:
124
+ """Prepare weights and scales for FBGEMM FP4 GEMM."""
125
+ swizzled_weight_scale = swizzle_blockscale(weight_scale)
126
+ swizzled_weight_scale = swizzled_weight_scale.view(-1).view(torch.uint8)
127
+ return weight, swizzled_weight_scale
128
+
129
+
130
+def convert_to_nvfp4_linear_kernel_format(
131
+ backend: NvFp4LinearBackend,
132
+ layer: torch.nn.Module,
133
+) -> None:
134
+ """Convert layer to NVFP4 linear kernel format."""
135
+
136
+ assert layer.weight_scale.dtype == torch.float8_e4m3fn, (
137
+ "Weight Block scale must be represented as FP8-E4M3"
138
+ )
139
+
140
+ # Default to no padding
141
+ layer.weights_padding_cols = 0
142
+
143
+ if backend == NvFp4LinearBackend.MARLIN:
144
+ prepare_fp4_layer_for_marlin(layer)
145
+ elif backend == NvFp4LinearBackend.FLASHINFER_TRTLLM:
146
+ weight, weight_scale = prepare_weights_for_nvfp4_flashinfer_trtllm(
147
+ layer.weight.data, layer.weight_scale.data
148
+ )
149
+ layer.weight = torch.nn.Parameter(weight, requires_grad=False)
150
+ layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
151
+ elif backend == NvFp4LinearBackend.FBGEMM:
152
+ weight, weight_scale = prepare_weights_for_nvfp4_fbgemm(
153
+ layer.weight.data, layer.weight_scale.data
154
+ )
155
+ layer.weight = torch.nn.Parameter(weight, requires_grad=False)
156
+ layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
157
+ elif (
158
+ backend == NvFp4LinearBackend.VLLM_CUTLASS
159
+ or backend == NvFp4LinearBackend.FLASHINFER_CUTLASS
160
+ ):
161
+ weight, weight_scale, weights_padding_cols = prepare_weights_for_nvfp4_cutlass(
162
+ layer.weight.data, layer.weight_scale.data
163
+ )
164
+ layer.weight = torch.nn.Parameter(weight, requires_grad=False)
165
+ layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
166
+ layer.weights_padding_cols = weights_padding_cols
167
+
168
+
169
+def apply_nvfp4_linear(
170
+ backend: NvFp4LinearBackend,
171
+ layer: torch.nn.Module,
172
+ x: torch.Tensor,
173
+ bias: torch.Tensor | None = None,
174
+) -> torch.Tensor:
175
+ """
176
+ Apply NVFP4 linear transformation using the specified backend.
177
+ """
178
+ weight = layer.weight
179
+ weight_scale = layer.weight_scale
180
+ weight_global_scale = layer.weight_global_scale
181
+ input_global_scale_inv = layer.input_global_scale_inv
182
+ alpha = layer.alpha
183
+ output_size = layer.output_size_per_partition
184
+ input_size = layer.input_size_per_partition
185
+
186
+ if backend == NvFp4LinearBackend.MARLIN:
187
+ return apply_fp4_marlin_linear(
188
+ input=x,
189
+ weight=weight,
190
+ weight_scale=weight_scale,
191
+ weight_global_scale=weight_global_scale,
192
+ workspace=layer.workspace,
193
+ size_n=output_size,
194
+ size_k=input_size,
195
+ bias=bias,
196
+ )
197
+ elif backend == NvFp4LinearBackend.EMULATION:
198
+ out = run_nvfp4_emulations(
199
+ x=x,
200
+ input_global_scale=input_global_scale_inv,
201
+ weight=weight,
202
+ weight_scale_swizzled=weight_scale,
203
+ weight_global_scale=weight_global_scale,
204
+ )
205
+ if bias is not None:
206
+ out = out + bias
207
+ return out
208
+
209
+ output_dtype = x.dtype
210
+ output_shape = [*x.shape[:-1], output_size]
211
+
212
+ # Quantize BF16 or FP16 to (FP4 and interleaved block scale)
213
+ x_fp4, x_blockscale = scaled_fp4_quant(
214
+ x, input_global_scale_inv, is_sf_swizzled_layout=True, backend=backend.value
215
+ )
216
+
217
+ # Validate dtypes
218
+ assert x_fp4.dtype == torch.uint8
219
+ assert weight.dtype == torch.uint8
220
+ assert x_blockscale.dtype == torch.float8_e4m3fn
221
+ # weight_scale is fp8 for most backends, but uint8 for fbgemm
222
+ assert weight_scale.dtype in (torch.float8_e4m3fn, torch.uint8)
223
+ assert alpha.dtype == torch.float32
224
+
225
+ # Pad activations to match weight K-dimension padding
226
+ weights_padding_cols = getattr(layer, "weights_padding_cols", 0)
227
+ x_fp4 = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_cols)
228
+
229
+ # Prepare args for the matmul
230
+ mm_args = (
231
+ x_fp4,
232
+ weight,
233
+ x_blockscale,
234
+ weight_scale,
235
+ alpha,
236
+ output_dtype,
237
+ )
238
+
239
+ # Call the appropriate backend
240
+ if backend.value.startswith("flashinfer-"):
241
+ backend_name = backend.value[len("flashinfer-") :]
242
+ out = flashinfer_scaled_fp4_mm(*mm_args, backend=backend_name)
243
+ elif backend == NvFp4LinearBackend.FBGEMM:
244
+ out = torch.ops.fbgemm.f4f4bf16(
245
+ x_fp4,
246
+ weight,
247
+ x_blockscale.view(-1).view(torch.uint8),
248
+ weight_scale,
249
+ alpha,
250
+ use_mx=False,
251
+ ).to(output_dtype)
252
+ else:
253
+ assert backend == NvFp4LinearBackend.VLLM_CUTLASS
254
+ out = cutlass_scaled_fp4_mm(*mm_args)
255
+
256
+ # Slice output to remove N-dimension padding
257
+ out = slice_nvfp4_output(out, output_size)
258
+
259
+ if bias is not None:
260
+ out = out + bias
261
+
262
+ return out.view(*output_shape)
263
+
264
+
265
+def swizzle_blockscale(scale: torch.Tensor) -> torch.Tensor:
266
+ """
267
+ Pad and block-interleave the FP4 block-scales so that they match the data
268
+ layout expected by the CUTLASS / FlashInfer kernels.
269
+
270
+ Parameters
271
+ ----------
272
+ scale: torch.Tensor
273
+
274
+ Returns
275
+ -------
276
+ torch.Tensor
277
+ The swizzled tensor with the same logical shape as *scale*.
278
+ """
279
+ assert scale.dtype == torch.float8_e4m3fn, (
280
+ "swizzle_blockscale expects the input tensor to be in "
281
+ "torch.float8_e4m3fn format."
282
+ )
283
+
284
+ scale_ndim = scale.ndim
285
+ if scale_ndim == 2:
286
+ scale = scale.unsqueeze(0) # (1, M, K)
287
+ assert scale.ndim == 3, "Expected a 2-D or 3-D tensor for block scales."
288
+
289
+ B, M, K = scale.shape
290
+
291
+ M_padded = round_up(M, 128)
292
+ K_padded = round_up(K, 4)
293
+
294
+ padded = torch.zeros(
295
+ (B, M_padded, K_padded), dtype=scale.dtype, device=scale.device
296
+ )
297
+ padded[:B, :M, :K] = scale
298
+
299
+ # Reshape / permute to the layout required by the kernel.
300
+ padded = padded.reshape(B, M_padded // 128, 4, 32, K_padded // 4, 4)
301
+ swizzled = padded.permute(0, 1, 4, 3, 2, 5).contiguous().cuda()
302
+
303
+ if scale_ndim == 2:
304
+ return swizzled.reshape(M_padded, K_padded)
305
+ return swizzled.reshape(B, M_padded, K_padded)
306
+
307
+
308
+def cutlass_fp4_supported() -> bool:
309
+ if not current_platform.is_cuda():
310
+ return False
311
+ capability_tuple = current_platform.get_device_capability()
312
+ capability = -1 if capability_tuple is None else capability_tuple.to_int()
313
+ return cutlass_scaled_mm_supports_fp4(capability)
314
+
315
+
316
+def pad_nvfp4_weight_for_cutlass(
317
+ weight: torch.Tensor,
318
+ alignment: int = 32,
319
+) -> tuple[torch.Tensor, int]:
320
+ """
321
+ Pad packed NVFP4 weights so that both N (rows) and K (columns) satisfy
322
+ the alignment constraints required by CUTLASS / FlashInfer FP4 kernels.
323
+
324
+ CUTLASS FP4 kernel requires both K and N matrix dimensions to be divisible
325
+ by 32 for aligned memory access and efficient tensor core operations.
326
+ """
327
+ weight_current_rows = weight.shape[0]
328
+
329
+ # Pad N dimension (rows) if not aligned
330
+ if weight_current_rows % alignment != 0:
331
+ total_rows = round_up(weight_current_rows, alignment)
332
+ pad_rows = total_rows - weight_current_rows
333
+ weight = torch.nn.functional.pad(weight, (0, 0, 0, pad_rows)).contiguous()
334
+
335
+ # Check K dimension alignment
336
+ # 2 FP4 items are packed per byte in the input dimension
337
+ weight_current_col_bytes = weight.shape[1]
338
+ weight_current_col_elements = weight_current_col_bytes * 2
339
+
340
+ weights_padding_bytes = 0
341
+ if weight_current_col_elements % alignment != 0:
342
+ total_cols = round_up(weight_current_col_elements, alignment)
343
+ pad_cols = total_cols - weight_current_col_elements
344
+ # Convert from FP4 element count to bytes (2 FP4 values per byte)
345
+ # pad_cols is always even since alignment=32 and current elements are even
346
+ pad_bytes = pad_cols // 2
347
+ weight = torch.nn.functional.pad(weight, (0, pad_bytes, 0, 0)).contiguous()
348
+ weights_padding_bytes = pad_bytes
349
+
350
+ return weight, weights_padding_bytes
351
+
352
+
353
+def pad_nvfp4_activation_for_cutlass(
354
+ x_fp4: torch.Tensor,
355
+ weights_padding_bytes: int,
356
+) -> torch.Tensor:
357
+ """
358
+ Pad packed FP4 activations to match the K-dimension padding applied to weights.
359
+ The padding is in bytes (tensor dimension), not FP4 elements.
360
+ """
361
+ if weights_padding_bytes > 0:
362
+ return torch.nn.functional.pad(x_fp4, (0, weights_padding_bytes)).contiguous()
363
+ return x_fp4
364
+
365
+
366
+def slice_nvfp4_output(
367
+ out: torch.Tensor,
368
+ output_size: int,
369
+) -> torch.Tensor:
370
+ """
371
+ Slice the output tensor to remove padding in N dimension if weight was padded.
372
+ """
373
+ if out.shape[-1] != output_size:
374
+ return out[..., :output_size].contiguous()
375
+ return out
vllm/model_executor/layers/quantization/utils/quant_utils.py
-122
@@ -11,7 +11,6 @@ import numpy
11
import torch
12
from torch import fx
13
14
-from vllm._custom_ops import cutlass_scaled_mm_supports_fp4
14
from vllm.platforms import current_platform
15
from vllm.scalar_type import ScalarType, scalar_types
16
@@ -768,60 +767,6 @@ def awq_pack(
767
return pack_cols(q_w, num_bits, size_k, size_n)
768
769
771
-def swizzle_blockscale(scale: torch.Tensor) -> torch.Tensor:
772
- """
773
- Pad and block-interleave the FP4 block-scales so that they match the data
774
- layout expected by the CUTLASS / FlashInfer kernels.
775
-
776
- Parameters
777
- ----------
778
- scale: torch.Tensor
779
-
780
- Returns
781
- -------
782
- torch.Tensor
783
- The swizzled tensor with the same logical shape as *scale*.
784
- """
785
- assert scale.dtype == torch.float8_e4m3fn, (
786
- "swizzle_blockscale expects the input tensor to be in "
787
- "torch.float8_e4m3fn format."
788
- )
789
-
790
- scale_ndim = scale.ndim
791
- if scale_ndim == 2:
792
- scale = scale.unsqueeze(0) # (1, M, K)
793
- assert scale.ndim == 3, "Expected a 2-D or 3-D tensor for block scales."
794
-
795
- B, M, K = scale.shape
796
-
797
- def _round_up(x: int, m: int) -> int:
798
- return (x + m - 1) // m * m
799
-
800
- M_padded = _round_up(M, 128)
801
- K_padded = _round_up(K, 4)
802
-
803
- padded = torch.zeros(
804
- (B, M_padded, K_padded), dtype=scale.dtype, device=scale.device
805
- )
806
- padded[:B, :M, :K] = scale
807
-
808
- # Reshape / permute to the layout required by the kernel.
809
- padded = padded.reshape(B, M_padded // 128, 4, 32, K_padded // 4, 4)
810
- swizzled = padded.permute(0, 1, 4, 3, 2, 5).contiguous().cuda()
811
-
812
- if scale_ndim == 2:
813
- return swizzled.reshape(M_padded, K_padded)
814
- return swizzled.reshape(B, M_padded, K_padded)
815
-
816
-
817
-def cutlass_fp4_supported() -> bool:
818
- if not current_platform.is_cuda():
819
- return False
820
- capability_tuple = current_platform.get_device_capability()
821
- capability = -1 if capability_tuple is None else capability_tuple.to_int()
822
- return cutlass_scaled_mm_supports_fp4(capability)
823
-
824
-
770
def convert_bf16_scales_to_fp8(
771
quant_fp8: Callable, scales: torch.Tensor
772
) -> tuple[torch.Tensor, torch.Tensor]:
@@ -868,70 +813,3 @@ def convert_packed_uint4b8_to_signed_int4_inplace(t: torch.Tensor) -> torch.Tens
813
t |= ((nib - 8) & 0xF) << shift
814
815
return t
871
-
872
-
873
-def round_up(x: int, m: int) -> int:
874
- """Round up x to the nearest multiple of m."""
875
- return (x + m - 1) // m * m
876
-
877
-
878
-def pad_nvfp4_weight_for_cutlass(
879
- weight: torch.Tensor,
880
- alignment: int = 32,
881
-) -> tuple[torch.Tensor, int]:
882
- """
883
- Pad packed NVFP4 weights so that both N (rows) and K (columns) satisfy
884
- the alignment constraints required by CUTLASS / FlashInfer FP4 kernels.
885
-
886
- CUTLASS FP4 kernel requires both K and N matrix dimensions to be divisible
887
- by 32 for aligned memory access and efficient tensor core operations.
888
- """
889
- weight_current_rows = weight.shape[0]
890
-
891
- # Pad N dimension (rows) if not aligned
892
- if weight_current_rows % alignment != 0:
893
- total_rows = round_up(weight_current_rows, alignment)
894
- pad_rows = total_rows - weight_current_rows
895
- weight = torch.nn.functional.pad(weight, (0, 0, 0, pad_rows)).contiguous()
896
-
897
- # Check K dimension alignment
898
- # 2 FP4 items are packed per byte in the input dimension
899
- weight_current_col_bytes = weight.shape[1]
900
- weight_current_col_elements = weight_current_col_bytes * 2
901
-
902
- weights_padding_bytes = 0
903
- if weight_current_col_elements % alignment != 0:
904
- total_cols = round_up(weight_current_col_elements, alignment)
905
- pad_cols = total_cols - weight_current_col_elements
906
- # Convert from FP4 element count to bytes (2 FP4 values per byte)
907
- # pad_cols is always even since alignment=32 and current elements are even
908
- pad_bytes = pad_cols // 2
909
- weight = torch.nn.functional.pad(weight, (0, pad_bytes, 0, 0)).contiguous()
910
- weights_padding_bytes = pad_bytes
911
-
912
- return weight, weights_padding_bytes
913
-
914
-
915
-def pad_nvfp4_activation_for_cutlass(
916
- x_fp4: torch.Tensor,
917
- weights_padding_bytes: int,
918
-) -> torch.Tensor:
919
- """
920
- Pad packed FP4 activations to match the K-dimension padding applied to weights.
921
- The padding is in bytes (tensor dimension), not FP4 elements.
922
- """
923
- if weights_padding_bytes > 0:
924
- return torch.nn.functional.pad(x_fp4, (0, weights_padding_bytes)).contiguous()
925
- return x_fp4
926
-
927
-
928
-def slice_nvfp4_output(
929
- out: torch.Tensor,
930
- output_size: int,
931
-) -> torch.Tensor:
932
- """
933
- Slice the output tensor to remove padding in N dimension if weight was padded.
934
- """
935
- if out.shape[-1] != output_size:
936
- return out[..., :output_size].contiguous()
937
- return out