[CompressedTensors] FP4 Qutlass Integration (#43229)
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com> Signed-off-by: Brian Dellabetta <bdellabe@redhat.com> Signed-off-by: Brian Dellabetta <brian-dellabetta@users.noreply.github.com> Co-authored-by: Brian Dellabetta <bdellabe@redhat.com> Co-authored-by: Brian Dellabetta <brian-dellabetta@users.noreply.github.com> Co-authored-by: Dipika Sikka <dipikasikka1@gmail.com>
Kyle Sayers committed
Jul 28, 2026 at 22:34 UTC
30c2718eaafc230927f0d8ac47439d040b46a7c3
7 files changed
+169
-34
tests/quantization/test_compressed_tensors.py
+5
@@ -473,6 +473,11 @@ def test_compressed_tensors_w4a8_fp8(vllm_runner, args):
473
"Flat is better than nested.\nSparse is better than dense.",
474
150.0,
475
),
476
+ (
477
+ "nm-testing/Llama-3.2-1B-Instruct-quipv16-nvfp4",
478
+ "Flat is better than nested.\nSparse is better than dense.",
479
+ 150.0,
480
+ ),
481
],
482
)
483
def test_compressed_tensors_transforms_perplexity(
vllm/_custom_ops.py
+34
-3
@@ -3994,9 +3994,40 @@ def fusedQuantizeNv(
3994
padded_rows, padded_cols, dtype=torch.float8_e4m3fn, device=a.device
3995
)
3996
3997
- return torch.ops._qutlass_C.fusedQuantizeNvAbsMax(
3998
- a, b, xh_e2m1, xh_e4m3, global_scale
3999
- )
3997
+ safeFusedQuantizeNv(a, b, xh_e2m1, xh_e4m3, global_scale)
3998
+ return xh_e2m1, xh_e4m3
3999
+
4000
+
4001
+@torch.library.custom_op(
4002
+ "vllm::safeFusedQuantizeNv", mutates_args=("xh_e2m1", "xh_e4m3")
4003
+)
4004
+def safeFusedQuantizeNv(
4005
+ a: torch.Tensor,
4006
+ b: torch.Tensor,
4007
+ xh_e2m1: torch.Tensor,
4008
+ xh_e4m3: torch.Tensor,
4009
+ global_scale: torch.Tensor,
4010
+) -> None:
4011
+ """
4012
+ Wrapper for QUTLASS fusedQuantizeNv method that operates on tensors in-place
4013
+ rather than returning them, to prevent torch 2.12+ errors that outputs of custom
4014
+ operators may not alias any inputs to the custom operator.
4015
+ """
4016
+ torch.ops._qutlass_C.fusedQuantizeNvAbsMax(a, b, xh_e2m1, xh_e4m3, global_scale)
4017
+ return
4018
+
4019
+
4020
+if hasattr(torch.ops._qutlass_C, "fusedQuantizeNv"):
4021
+
4022
+ @register_fake("vllm::safeFusedQuantizeNv")
4023
+ def _fake_fused_quantize_nv(
4024
+ a: torch.Tensor,
4025
+ b: torch.Tensor,
4026
+ xh_e2m1: torch.Tensor,
4027
+ xh_e4m3: torch.Tensor,
4028
+ global_scale: torch.Tensor,
4029
+ ) -> None:
4030
+ return
4031
4032
4033
def hadacore_transform(x: torch.Tensor, inplace: bool = True) -> torch.Tensor:
vllm/model_executor/layers/linear.py
+1
@@ -50,6 +50,7 @@ WEIGHT_LOADER_V2_SUPPORTED = [
50
"UnquantizedLinearMethod",
51
"CompressedTensorsLinearMethod",
52
"CompressedTensorsLinearTransformMethod",
53
+ "QutlassNvFP4LinearMethod",
54
"AutoAWQMarlinLinearMethod",
55
"AutoAWQLinearMethod",
56
"AutoGPTQLinearMethod",
vllm/model_executor/layers/quantization/compressed_tensors/transform/linear.py
+4
-2
@@ -25,6 +25,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.transform.module
25
from vllm.model_executor.layers.quantization.compressed_tensors.transform.utils import ( # noqa: E501
26
TransformTuple,
27
)
28
+from vllm.platforms import current_platform
29
30
31
class CompressedTensorsLinearTransformMethod(LinearMethodBase):
@@ -48,11 +49,12 @@ class CompressedTensorsLinearTransformMethod(LinearMethodBase):
49
50
assert input_tfms or output_tfms
51
51
- if is_qutlass_fp4_scheme(quant_scheme, input_tfms):
52
+ if is_qutlass_fp4_scheme(
53
+ quant_scheme, input_tfms
54
+ ) and current_platform.has_device_capability(100):
55
return QutlassNvFP4LinearMethod(quant_method, input_tfms, output_tfms)
56
57
# hadacore or dense gemm is selected by Transform module
55
-
58
return cls(quant_method, input_tfms, output_tfms)
59
60
def __init__(
vllm/model_executor/layers/quantization/compressed_tensors/transform/module.py
+14
-22
@@ -32,7 +32,7 @@ class HadamardTransform(torch.nn.Module):
32
transforms: dict[int, TransformTuple] # info parsed from transforms config
33
weight: SharedWeightParameter # container for shared tensors
34
35
- scales: dict[int, float] # hadamard scale, usually sqrt(matrix.size(0))
35
+ scaled_data_ptrs: set[int] = set()
36
37
def __init__(
38
self,
@@ -44,7 +44,6 @@ class HadamardTransform(torch.nn.Module):
44
):
45
super().__init__()
46
self.transforms = transforms
47
- self.scales = {}
47
48
if get_tensor_model_parallel_world_size() > 1:
49
raise NotImplementedError(
@@ -64,26 +63,26 @@ class HadamardTransform(torch.nn.Module):
63
)
64
65
data_key = self._get_data_key(scheme, weight_size)
66
+ # load up in model's default precision, rather than using scheme.precision
67
self.weight.add_partition(
68
part_index,
69
data_key,
70
size=(weight_size, weight_size),
71
- dtype=scheme.precision,
71
)
72
73
# validate that shared tensors and schemes are correct
74
self._validate_input_transforms()
75
76
def process_weights_after_loading(self):
78
- for part_id in self.weight.partitions:
79
- data = self.weight.partitions[part_id].data
80
-
77
+ for part_id, partition in self.weight.partitions.items():
78
# required by torch.compile
79
self.weight.process_weights_after_loading()
80
84
- # precompute scale as a runtime multiply, not division
85
- # do not fold into weight in order to utilize FWHT
86
- self.scales[part_id] = 1 / math.sqrt(data.size(0))
81
+ # Merge normalization scale directly into weight, must be done only once
82
+ data_ptr = partition.data.data_ptr()
83
+ if data_ptr not in HadamardTransform.scaled_data_ptrs:
84
+ partition.data.div_(math.sqrt(partition.data.size(0)))
85
+ HadamardTransform.scaled_data_ptrs.add(data_ptr)
86
87
# FUTURE: avoid runtime transpose by processing weights
88
# prior to apply
@@ -111,26 +110,19 @@ class HadamardTransform(torch.nn.Module):
110
weight = (
111
weight if self.transforms[part_id].args.inverse else weight.T
112
) # linear := x(W.T)
114
- scale = self.scales[part_id]
113
114
if self.transforms[part_id].scheme.head_dim is not None:
115
value = value.unflatten(-1, (-1, weight.size(0)))
118
- value = (
119
- dispatch_unquantized_gemm()(
120
- self, value.to(weight.dtype), weight, None
121
- ).to(value.dtype)
122
- * scale
123
- )
116
+ value = dispatch_unquantized_gemm()(
117
+ self, value.to(weight.dtype), weight, None
118
+ ).to(value.dtype)
119
value = value.flatten(-2, -1)
120
121
return value
122
128
- return (
129
- dispatch_unquantized_gemm()(
130
- self, value.to(weight.dtype), weight, None
131
- ).to(value.dtype)
132
- * scale
133
- )
123
+ return dispatch_unquantized_gemm()(
124
+ self, value.to(weight.dtype), weight, None
125
+ ).to(value.dtype)
126
127
def _get_data_key(self, scheme: TransformScheme, weight_size: int) -> Hashable:
128
return (id(scheme), weight_size)
vllm/model_executor/layers/quantization/compressed_tensors/transform/schemes/linear_qutlass_nvfp4.py
+101
-6
@@ -2,7 +2,13 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import torch
5
+from torch.nn.parameter import Parameter
6
7
+from vllm._custom_ops import fusedQuantizeNv
8
+from vllm.model_executor.kernels.linear import (
9
+ _LINEAR_BACKEND_KERNEL_MAP,
10
+ NvFp4LinearKernel,
11
+)
12
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
13
CompressedTensorsScheme,
14
CompressedTensorsW4A4Fp4,
@@ -11,18 +17,32 @@ from vllm.model_executor.layers.quantization.compressed_tensors.transform.linear
17
CompressedTensorsLinearTransformMethod,
18
TransformTuple,
19
)
20
+from vllm.model_executor.layers.quantization.qutlass_utils import to_blocked
21
+from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
22
+ slice_nvfp4_output,
23
+)
24
+from vllm.utils.flashinfer import (
25
+ flashinfer_scaled_fp4_mm,
26
+)
27
28
__all__ = ["is_qutlass_fp4_scheme", "QutlassNvFP4LinearMethod"]
29
30
+NVFP4_MAX = 6.0
31
+
32
33
+# QUTLASS supports transform block sizes (16, 32, 64, 128) for NVFP4
34
+# https://github.com/IST-DASLab/qutlass/blob/v0.2.0/qutlass/csrc/bindings.cpp#L413-L414
35
def is_qutlass_fp4_scheme(
36
quant_scheme: CompressedTensorsScheme | None,
37
input_tfms: dict[int, TransformTuple],
38
) -> bool:
39
return (
23
- isinstance(quant_scheme, (CompressedTensorsW4A4Fp4,))
24
- and len(input_tfms) == 1
25
- and input_tfms[0].scheme.head_dim == quant_scheme.group_size
40
+ isinstance(quant_scheme, CompressedTensorsW4A4Fp4)
41
+ and len(input_tfms) >= 1
42
+ and all(
43
+ input_tfm.scheme.head_dim in (16, 32, 64, 128)
44
+ for input_tfm in input_tfms.values()
45
+ )
46
)
47
48
@@ -50,15 +70,90 @@ class QutlassNvFP4LinearMethod(CompressedTensorsLinearTransformMethod):
70
)
71
72
assert self.input_transform is not None
53
- assert len(self.input_transform.weight) == 1
54
- assert self.input_transform.weight[0].size(0) == layer.scheme.group_size
73
+ assert len(self.input_transform.weight.partitions) >= 1
74
75
return ret
76
77
+ @staticmethod
78
+ def _get_flashinfer_gemm_backend(kernel: NvFp4LinearKernel) -> str:
79
+ """
80
+ Given a kernel, find the string that is needed to be passed into
81
+ `flashinfer_scaled_fp4_mm`, using
82
+ vllm.model_executor.kernels.linear._LINEAR_BACKEND_KERNEL_MAP as source of truth
83
+ """
84
+ kernel_type = type(kernel)
85
+ for key, kernels in _LINEAR_BACKEND_KERNEL_MAP.items():
86
+ if not key.startswith("flashinfer_") or kernel_type not in kernels:
87
+ continue
88
+ backend = key.removeprefix("flashinfer_")
89
+ # flashinfer GEMM backend uses "cute-dsl", not "cutedsl"
90
+ return backend.replace("cutedsl", "cute-dsl")
91
+ raise ValueError(
92
+ f"QutlassNvFP4 transform requires a FlashInfer kernel, "
93
+ f"got {kernel_type.__name__}"
94
+ )
95
+
96
+ def process_weights_after_loading(self, layer):
97
+ super().process_weights_after_loading(layer)
98
+
99
+ assert self.input_transform is not None
100
+ layer.hadamard_matrix = self.input_transform.weight.partitions[0].data
101
+
102
+ # fusedQuantizeNv stores raw absmax as block scales (sf = absmax),
103
+ # while CT weights use sf = absmax * SFScaleVal / 6.0. The GEMM
104
+ # computes alpha * sum(fp4_a * sf_a * fp4_w * sf_w), so alpha must
105
+ # compensate: alpha = weight_global_scale / 6.0
106
+ layer.fused_alpha = Parameter(
107
+ layer.weight_global_scale / NVFP4_MAX, requires_grad=False
108
+ )
109
+
110
+ layer.fused_global_scale = Parameter(
111
+ torch.tensor(
112
+ [NVFP4_MAX],
113
+ dtype=torch.float32,
114
+ device=layer.weight_global_scale.device,
115
+ ),
116
+ requires_grad=False,
117
+ )
118
+
119
+ layer.flashinfer_gemm_backend = self._get_flashinfer_gemm_backend(
120
+ layer.scheme.kernel
121
+ )
122
+
123
def apply(
124
self,
125
layer: torch.nn.Module,
126
x: torch.Tensor,
127
bias: torch.Tensor | None = None,
128
) -> torch.Tensor:
64
- raise NotImplementedError()
129
+ assert bias is None
130
+ output_size = layer.output_size_per_partition
131
+ output_shape = [*x.shape[:-1], output_size]
132
+
133
+ x_flat = x.contiguous().flatten(end_dim=-2)
134
+
135
+ x_fp4, x_scales = fusedQuantizeNv(
136
+ x_flat, layer.hadamard_matrix, layer.fused_global_scale
137
+ )
138
+
139
+ x_scales_blocked = to_blocked(x_scales, backend="triton").view(x_scales.shape)
140
+
141
+ out = flashinfer_scaled_fp4_mm(
142
+ x_fp4,
143
+ layer.weight,
144
+ x_scales_blocked,
145
+ layer.weight_scale,
146
+ layer.fused_alpha,
147
+ x.dtype,
148
+ backend=layer.flashinfer_gemm_backend,
149
+ )
150
+
151
+ out = slice_nvfp4_output(out, output_size)
152
+
153
+ if self.output_transform is not None:
154
+ for part_id, (start, length) in enumerate(self.partition_ranges):
155
+ out[:, start : start + length] = self.output_transform(
156
+ out[:, start : start + length].clone(), part_id=part_id
157
+ )
158
+
159
+ return out.view(*output_shape)
vllm/model_executor/layers/quantization/qutlass_utils.py
+10
-1
@@ -84,6 +84,7 @@ def triton_scale_swizzle(
84
)
85
86
87
+@torch.library.custom_op("vllm::triton_mx_block_rearrange", mutates_args=())
88
def triton_mx_block_rearrange(scale_tensor: torch.Tensor) -> torch.Tensor:
89
"""
90
Rearranges an E8M0 tensor scale from row-major format to
@@ -142,6 +143,14 @@ def triton_mx_block_rearrange(scale_tensor: torch.Tensor) -> torch.Tensor:
143
return out
144
145
146
+@triton_mx_block_rearrange.register_fake
147
+def _triton_mx_block_rearrange_fake(scale_tensor: torch.Tensor) -> torch.Tensor:
148
+ rows, cols = scale_tensor.shape
149
+ padded_rows = cdiv(rows, 128) * 128
150
+ padded_cols = cdiv(cols, 4) * 4
151
+ return scale_tensor.new_empty((padded_rows, padded_cols))
152
+
153
+
154
def to_blocked(
155
input_matrix: torch.Tensor, backend: Literal["torch", "triton"] = "triton"
156
) -> torch.Tensor:
@@ -157,7 +166,7 @@ def to_blocked(
166
backend: "torch" (PyTorch path) or "triton" (Triton kernel)
167
168
Returns:
160
- Rearranged tensor of shape (32*cdiv(H,128), 16*cdiv(W,4))
169
+ Rearranged flattened tensor of size (32*cdiv(H,128) * 16*cdiv(W,4))
170
"""
171
if backend == "triton":
172
return triton_mx_block_rearrange(input_matrix).flatten()