@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 96ad65b7fe
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 14, 05:30 UTC

[Transform] [Quantization] Add QuTLASS support to vLLM (#24440)

Signed-off-by: LopezCastroRoberto <roberto.lopez.castro@udc.es> Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> Signed-off-by: Andrei Panferov <andrei@panferov.org> Co-authored-by: Andrei Panferov <andrei@panferov.org> Co-authored-by: Michael Goin <mgoin64@gmail.com>

Roberto L. Castro committed Oct 10, 2025 at 18:43 UTC 96ad65b7fe515663da8ede09a1aa7f74aa500c97
12 files changed +1848 -1
.buildkite/test-pipeline.yaml
+2
@@ -834,6 +834,8 @@ steps:
834 - pytest -v -s tests/compile/test_fusion_attn.py::test_attention_quant_pattern
835 - pytest -v -s tests/kernels/moe/test_flashinfer.py
836 - pytest -v -s tests/compile/test_silu_mul_quant_fusion.py
837 + - pytest -v -s tests/kernels/quantization/test_nvfp4_qutlass.py
838 + - pytest -v -s tests/kernels/quantization/test_mxfp4_qutlass.py
839
840 - label: Blackwell GPT-OSS Eval
841 timeout_in_minutes: 60
CMakeLists.txt
+1
@@ -1007,6 +1007,7 @@ endif()
1007 # For CUDA we also build and ship some external projects.
1008 if (VLLM_GPU_LANG STREQUAL "CUDA")
1009 include(cmake/external_projects/flashmla.cmake)
1010 + include(cmake/external_projects/qutlass.cmake)
1011
1012 # vllm-flash-attn should be last as it overwrites some CMake functions
1013 include(cmake/external_projects/vllm_flash_attn.cmake)
benchmarks/kernels/bench_mxfp4_qutlass.py new
+191
@@ -0,0 +1,191 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +#
4 +# Copyright (C) 2025 Roberto L. Castro (Roberto.LopezCastro@ist.ac.at).
5 +# All Rights Reserved.
6 +#
7 +# Licensed under the Apache License, Version 2.0 (the "License");
8 +# you may not use this file except in compliance with the License.
9 +# You may obtain a copy of the License at
10 +#
11 +# http://www.apache.org/licenses/LICENSE-2.0
12 +#
13 +# Unless required by applicable law or agreed to in writing, software
14 +# distributed under the License is distributed on an "AS IS" BASIS,
15 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16 +# See the License for the specific language governing permissions and
17 +# limitations under the License.
18 +#
19 +
20 +import argparse
21 +import copy
22 +import itertools
23 +
24 +import torch
25 +from compressed_tensors.transform.utils.hadamard import deterministic_hadamard_matrix
26 +from weight_shapes import WEIGHT_SHAPES
27 +
28 +from vllm._custom_ops import fusedQuantizeMx, matmul_mxf4_bf16_tn
29 +from vllm.model_executor.layers.quantization.qutlass_utils import to_blocked
30 +from vllm.triton_utils import triton
31 +
32 +PROVIDER_CFGS = {
33 + "torch-bf16": dict(enabled=True),
34 + "mxfp4": dict(no_a_quant=False, enabled=True),
35 + "mxfp4-noquant": dict(no_a_quant=True, enabled=True),
36 +}
37 +
38 +_enabled = [k for k, v in PROVIDER_CFGS.items() if v["enabled"]]
39 +
40 +
41 +def get_hadamard_matrix(group_size: int, dtype: torch.dtype, device: torch.device):
42 + return (
43 + deterministic_hadamard_matrix(group_size, dtype=dtype, device=device)
44 + * group_size**-0.5
45 + )
46 +
47 +
48 +def _quant_weight_mxfp4(
49 + b: torch.Tensor, forward_hadamard_matrix: torch.Tensor, device: str
50 +):
51 + weight_hf_e2m1, weight_hf_e8m0 = fusedQuantizeMx(
52 + b, forward_hadamard_matrix, method="abs_max"
53 + )
54 + weight_hf_scale_block = to_blocked(weight_hf_e8m0, backend="triton")
55 + return weight_hf_e2m1, weight_hf_scale_block
56 +
57 +
58 +def build_mxfp4_runner(cfg, a, b, forward_hadamard_matrix, dtype, device):
59 + weight_hf_e2m1, weight_hf_scale_block = _quant_weight_mxfp4(
60 + b, forward_hadamard_matrix, device
61 + )
62 + alpha = torch.tensor([1.0], device="cuda")
63 +
64 + if cfg["no_a_quant"]:
65 + # Pre-quantize activation
66 + input_hf_e2m1, input_hf_e8m0 = fusedQuantizeMx(
67 + a, forward_hadamard_matrix, method="abs_max"
68 + )
69 + input_hf_scale_block = to_blocked(input_hf_e8m0, backend="triton")
70 +
71 + def run():
72 + return matmul_mxf4_bf16_tn(
73 + input_hf_e2m1,
74 + weight_hf_e2m1,
75 + input_hf_scale_block,
76 + weight_hf_scale_block,
77 + alpha,
78 + )
79 +
80 + return run
81 +
82 + # Quantize activation on-the-fly
83 + def run():
84 + input_hf_e2m1, input_hf_e8m0 = fusedQuantizeMx(
85 + a, forward_hadamard_matrix, method="abs_max"
86 + )
87 + input_hf_scale_block = to_blocked(input_hf_e8m0, backend="triton")
88 + return matmul_mxf4_bf16_tn(
89 + input_hf_e2m1,
90 + weight_hf_e2m1,
91 + input_hf_scale_block,
92 + weight_hf_scale_block,
93 + alpha,
94 + )
95 +
96 + return run
97 +
98 +
99 +@triton.testing.perf_report(
100 + triton.testing.Benchmark(
101 + x_names=["batch_size"],
102 + x_vals=[
103 + 1,
104 + 4,
105 + 8,
106 + 16,
107 + 32,
108 + 64,
109 + 128,
110 + 256,
111 + 512,
112 + 1024,
113 + 2048,
114 + 4096,
115 + 8192,
116 + 16384,
117 + 24576,
118 + 32768,
119 + ],
120 + x_log=False,
121 + line_arg="provider",
122 + line_vals=_enabled,
123 + line_names=_enabled,
124 + ylabel="TFLOP/s (larger is better)",
125 + plot_name="BF16 vs MXFP4 GEMMs",
126 + args={},
127 + )
128 +)
129 +def benchmark(batch_size, provider, N, K, had_size):
130 + M = batch_size
131 + device = "cuda"
132 + dtype = torch.bfloat16
133 +
134 + a = torch.randn((M, K), device=device, dtype=dtype)
135 + b = torch.randn((N, K), device=device, dtype=dtype)
136 + forward_hadamard_matrix = get_hadamard_matrix(had_size, dtype, device)
137 +
138 + quantiles = [0.5, 0.2, 0.8]
139 +
140 + if provider == "torch-bf16":
141 + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
142 + lambda: torch.nn.functional.linear(a, b), rep=200, quantiles=quantiles
143 + )
144 + else:
145 + cfg = PROVIDER_CFGS[provider]
146 + run_quant = build_mxfp4_runner(
147 + cfg, a, b, forward_hadamard_matrix, dtype, device
148 + )
149 + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
150 + lambda: run_quant(), rep=200, quantiles=quantiles
151 + )
152 +
153 + to_tflops = lambda t_ms: (2 * M * N * K) * 1e-12 / (t_ms * 1e-3)
154 + return to_tflops(ms), to_tflops(max_ms), to_tflops(min_ms)
155 +
156 +
157 +def prepare_shapes(args):
158 + out = []
159 + for model, tp_size in itertools.product(args.models, args.tp_sizes):
160 + for KN, tp_dim in copy.deepcopy(WEIGHT_SHAPES[model]):
161 + KN[tp_dim] //= tp_size
162 + KN.append(model)
163 + out.append(KN)
164 + return out
165 +
166 +
167 +if __name__ == "__main__":
168 + parser = argparse.ArgumentParser()
169 + parser.add_argument(
170 + "--models",
171 + nargs="+",
172 + type=str,
173 + default=["meta-llama/Llama-3.3-70B-Instruct"],
174 + choices=list(WEIGHT_SHAPES.keys()),
175 + )
176 + parser.add_argument("--tp-sizes", nargs="+", type=int, default=[1])
177 + args = parser.parse_args()
178 +
179 + for K, N, model in prepare_shapes(args):
180 + for had_size in [32, 64, 128]:
181 + print(f"{model}, N={N} K={K}, HAD={had_size}, BF16 vs MXFP4 GEMMs TFLOP/s:")
182 + benchmark.run(
183 + print_data=True,
184 + show_plots=True,
185 + save_path=f"bench_mxfp4_res_n{N}_k{K}",
186 + N=N,
187 + K=K,
188 + had_size=had_size,
189 + )
190 +
191 + print("Benchmark finished!")
benchmarks/kernels/bench_nvfp4_qutlass.py new
+207
@@ -0,0 +1,207 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +#
4 +# Copyright (C) 2025 Roberto L. Castro (Roberto.LopezCastro@ist.ac.at).
5 +# All Rights Reserved.
6 +#
7 +# Licensed under the Apache License, Version 2.0 (the "License");
8 +# you may not use this file except in compliance with the License.
9 +# You may obtain a copy of the License at
10 +#
11 +# http://www.apache.org/licenses/LICENSE-2.0
12 +#
13 +# Unless required by applicable law or agreed to in writing, software
14 +# distributed under the License is distributed on an "AS IS" BASIS,
15 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16 +# See the License for the specific language governing permissions and
17 +# limitations under the License.
18 +#
19 +
20 +import argparse
21 +import copy
22 +import itertools
23 +
24 +import torch
25 +from compressed_tensors.transform.utils.hadamard import deterministic_hadamard_matrix
26 +from weight_shapes import WEIGHT_SHAPES
27 +
28 +from vllm import _custom_ops as ops # use existing nvfp4 gemm in vllm
29 +from vllm._custom_ops import fusedQuantizeNv
30 +from vllm.model_executor.layers.quantization.qutlass_utils import to_blocked
31 +from vllm.triton_utils import triton
32 +
33 +PROVIDER_CFGS = {
34 + "torch-bf16": dict(enabled=True),
35 + "nvfp4": dict(no_a_quant=False, enabled=True),
36 + "nvfp4-noquant": dict(no_a_quant=True, enabled=True),
37 +}
38 +
39 +_enabled = [k for k, v in PROVIDER_CFGS.items() if v["enabled"]]
40 +
41 +
42 +def get_hadamard_matrix(group_size: int, dtype: torch.dtype, device: torch.device):
43 + return (
44 + deterministic_hadamard_matrix(group_size, dtype=dtype, device=device)
45 + * group_size**-0.5
46 + )
47 +
48 +
49 +def _quant_weight_nvfp4(
50 + b: torch.Tensor,
51 + forward_hadamard_matrix: torch.Tensor,
52 + global_scale: torch.Tensor,
53 + device: str,
54 + M: int,
55 + N: int,
56 + K: int,
57 +):
58 + weight_hf_e2m1, weight_hf_e8m0 = fusedQuantizeNv(
59 + b, forward_hadamard_matrix, global_scale
60 + )
61 + weight_hf_scale_block = to_blocked(weight_hf_e8m0, backend="triton").view(
62 + -1, K // 16
63 + )
64 + return weight_hf_e2m1, weight_hf_scale_block
65 +
66 +
67 +def build_nvfp4_runner(cfg, a, b, forward_hadamard_matrix, dtype, device, M, N, K):
68 + alpha = torch.tensor([1.0], device="cuda")
69 + global_scale = torch.tensor([1.0], device="cuda")
70 + weight_hf_e2m1, weight_hf_scale_block = _quant_weight_nvfp4(
71 + b, forward_hadamard_matrix, global_scale, device, M, N, K
72 + )
73 +
74 + if cfg["no_a_quant"]:
75 + # Pre-quantize activation
76 + input_hf_e2m1, input_hf_e8m0 = fusedQuantizeNv(
77 + a, forward_hadamard_matrix, global_scale
78 + )
79 + input_hf_scale_block = to_blocked(input_hf_e8m0, backend="triton").view(
80 + -1, K // 16
81 + )
82 +
83 + def run():
84 + return ops.cutlass_scaled_fp4_mm(
85 + input_hf_e2m1,
86 + weight_hf_e2m1,
87 + input_hf_scale_block,
88 + weight_hf_scale_block,
89 + alpha,
90 + torch.bfloat16,
91 + )
92 +
93 + return run
94 +
95 + # Quantize activation on-the-fly
96 + def run():
97 + input_hf_e2m1, input_hf_e8m0 = fusedQuantizeNv(
98 + a, forward_hadamard_matrix, global_scale
99 + )
100 + input_hf_scale_block = to_blocked(input_hf_e8m0, backend="triton").view(
101 + -1, K // 16
102 + )
103 + return ops.cutlass_scaled_fp4_mm(
104 + input_hf_e2m1,
105 + weight_hf_e2m1,
106 + input_hf_scale_block,
107 + weight_hf_scale_block,
108 + alpha,
109 + torch.bfloat16,
110 + )
111 +
112 + return run
113 +
114 +
115 +@triton.testing.perf_report(
116 + triton.testing.Benchmark(
117 + x_names=["batch_size"],
118 + x_vals=[
119 + 1,
120 + 4,
121 + 8,
122 + 16,
123 + 32,
124 + 64,
125 + 128,
126 + 256,
127 + 512,
128 + 1024,
129 + 2048,
130 + 4096,
131 + 8192,
132 + 16384,
133 + 24576,
134 + 32768,
135 + ],
136 + x_log=False,
137 + line_arg="provider",
138 + line_vals=_enabled,
139 + line_names=_enabled,
140 + ylabel="TFLOP/s (larger is better)",
141 + plot_name="BF16 vs NVFP4 GEMMs",
142 + args={},
143 + )
144 +)
145 +def benchmark(batch_size, provider, N, K, had_size):
146 + M = batch_size
147 + device = "cuda"
148 + dtype = torch.bfloat16
149 +
150 + a = torch.randn((M, K), device=device, dtype=dtype)
151 + b = torch.randn((N, K), device=device, dtype=dtype)
152 + forward_hadamard_matrix = get_hadamard_matrix(had_size, dtype, device)
153 +
154 + quantiles = [0.5, 0.2, 0.8]
155 +
156 + if provider == "torch-bf16":
157 + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
158 + lambda: torch.nn.functional.linear(a, b), rep=200, quantiles=quantiles
159 + )
160 + else:
161 + cfg = PROVIDER_CFGS[provider]
162 + run_quant = build_nvfp4_runner(
163 + cfg, a, b, forward_hadamard_matrix, dtype, device, M, N, K
164 + )
165 + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
166 + lambda: run_quant(), rep=200, quantiles=quantiles
167 + )
168 +
169 + to_tflops = lambda t_ms: (2 * M * N * K) * 1e-12 / (t_ms * 1e-3)
170 + return to_tflops(ms), to_tflops(max_ms), to_tflops(min_ms)
171 +
172 +
173 +def prepare_shapes(args):
174 + out = []
175 + for model, tp_size in itertools.product(args.models, args.tp_sizes):
176 + for KN, tp_dim in copy.deepcopy(WEIGHT_SHAPES[model]):
177 + KN[tp_dim] //= tp_size
178 + KN.append(model)
179 + out.append(KN)
180 + return out
181 +
182 +
183 +if __name__ == "__main__":
184 + parser = argparse.ArgumentParser()
185 + parser.add_argument(
186 + "--models",
187 + nargs="+",
188 + type=str,
189 + default=["meta-llama/Llama-3.3-70B-Instruct"],
190 + choices=list(WEIGHT_SHAPES.keys()),
191 + )
192 + parser.add_argument("--tp-sizes", nargs="+", type=int, default=[1])
193 + args = parser.parse_args()
194 +
195 + for K, N, model in prepare_shapes(args):
196 + for had_size in [16, 32, 64, 128]:
197 + print(f"{model}, N={N} K={K}, HAD={had_size}, BF16 vs NVFP4 GEMMs TFLOP/s:")
198 + benchmark.run(
199 + print_data=True,
200 + show_plots=True,
201 + save_path=f"bench_nvfp4_res_n{N}_k{K}",
202 + N=N,
203 + K=K,
204 + had_size=had_size,
205 + )
206 +
207 + print("Benchmark finished!")
cmake/external_projects/qutlass.cmake new
+97
@@ -0,0 +1,97 @@
1 +include(FetchContent)
2 +
3 +set(CUTLASS_INCLUDE_DIR "${CUTLASS_INCLUDE_DIR}" CACHE PATH "Path to CUTLASS include/ directory")
4 +
5 +if(DEFINED ENV{QUTLASS_SRC_DIR})
6 + set(QUTLASS_SRC_DIR $ENV{QUTLASS_SRC_DIR})
7 +endif()
8 +
9 +if(QUTLASS_SRC_DIR)
10 + FetchContent_Declare(
11 + qutlass
12 + SOURCE_DIR ${QUTLASS_SRC_DIR}
13 + CONFIGURE_COMMAND ""
14 + BUILD_COMMAND ""
15 + )
16 +else()
17 + FetchContent_Declare(
18 + qutlass
19 + GIT_REPOSITORY https://github.com/IST-DASLab/qutlass.git
20 + GIT_TAG 830d2c4537c7396e14a02a46fbddd18b5d107c65
21 + GIT_PROGRESS TRUE
22 + CONFIGURE_COMMAND ""
23 + BUILD_COMMAND ""
24 + )
25 + FetchContent_Populate(qutlass)
26 + set(qutlass_SOURCE_DIR "${qutlass_SOURCE_DIR}")
27 +endif()
28 +
29 +if(NOT qutlass_SOURCE_DIR)
30 + message(FATAL_ERROR "[QUTLASS] source directory could not be resolved.")
31 +endif()
32 +message(STATUS "[QUTLASS] QuTLASS is available at ${qutlass_SOURCE_DIR}")
33 +
34 +cuda_archs_loose_intersection(QUTLASS_ARCHS "12.0a;10.0a" "${CUDA_ARCHS}")
35 +if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.8 AND QUTLASS_ARCHS)
36 +
37 + if(QUTLASS_ARCHS MATCHES "10\\.0a")
38 + set(QUTLASS_TARGET_CC 100)
39 + elseif(QUTLASS_ARCHS MATCHES "12\\.0a")
40 + set(QUTLASS_TARGET_CC 120)
41 + else()
42 + message(FATAL_ERROR "[QUTLASS] internal error parsing CUDA_ARCHS='${QUTLASS_ARCHS}'.")
43 + endif()
44 +
45 + set(QUTLASS_SOURCES
46 + ${qutlass_SOURCE_DIR}/qutlass/csrc/bindings.cpp
47 + ${qutlass_SOURCE_DIR}/qutlass/csrc/gemm.cu
48 + ${qutlass_SOURCE_DIR}/qutlass/csrc/gemm_ada.cu
49 + ${qutlass_SOURCE_DIR}/qutlass/csrc/fused_quantize_mx.cu
50 + ${qutlass_SOURCE_DIR}/qutlass/csrc/fused_quantize_nv.cu
51 + ${qutlass_SOURCE_DIR}/qutlass/csrc/fused_quantize_mx_sm100.cu
52 + ${qutlass_SOURCE_DIR}/qutlass/csrc/fused_quantize_nv_sm100.cu
53 + )
54 +
55 + set(QUTLASS_INCLUDES
56 + ${qutlass_SOURCE_DIR}
57 + ${qutlass_SOURCE_DIR}/qutlass
58 + ${qutlass_SOURCE_DIR}/qutlass/csrc/include
59 + ${qutlass_SOURCE_DIR}/qutlass/csrc/include/cutlass_extensions
60 + )
61 +
62 + if(CUTLASS_INCLUDE_DIR AND EXISTS "${CUTLASS_INCLUDE_DIR}/cutlass/cutlass.h")
63 + list(APPEND QUTLASS_INCLUDES "${CUTLASS_INCLUDE_DIR}")
64 + elseif(EXISTS "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include/cutlass/cutlass.h")
65 + list(APPEND QUTLASS_INCLUDES "${qutlass_SOURCE_DIR}/qutlass/third_party/cutlass/include")
66 + message(STATUS "[QUTLASS] Using QuTLASS vendored CUTLASS headers (no vLLM CUTLASS detected).")
67 + else()
68 + message(FATAL_ERROR "[QUTLASS] CUTLASS headers not found. "
69 + "Set -DCUTLASS_INCLUDE_DIR=/path/to/cutlass/include")
70 + endif()
71 +
72 + set_gencode_flags_for_srcs(
73 + SRCS "${QUTLASS_SOURCES}"
74 + CUDA_ARCHS "${QUTLASS_ARCHS}"
75 + )
76 +
77 + target_sources(_C PRIVATE ${QUTLASS_SOURCES})
78 + target_include_directories(_C PRIVATE ${QUTLASS_INCLUDES})
79 + target_compile_definitions(_C PRIVATE
80 + QUTLASS_DISABLE_PYBIND=1
81 + TARGET_CUDA_ARCH=${QUTLASS_TARGET_CC}
82 + )
83 +
84 + set_property(SOURCE ${QUTLASS_SOURCES} APPEND PROPERTY COMPILE_OPTIONS
85 + $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr --use_fast_math -O3>
86 + )
87 +
88 +else()
89 + if("${CMAKE_CUDA_COMPILER_VERSION}" VERSION_LESS "12.8")
90 + message(STATUS
91 + "[QUTLASS] Skipping build: CUDA 12.8 or newer is required (found ${CMAKE_CUDA_COMPILER_VERSION}).")
92 + else()
93 + message(STATUS
94 + "[QUTLASS] Skipping build: no supported arch (12.0a / 10.0a) found in "
95 + "CUDA_ARCHS='${CUDA_ARCHS}'.")
96 + endif()
97 +endif()
tests/kernels/quantization/test_mxfp4_qutlass.py new
+303
@@ -0,0 +1,303 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +#
4 +# Copyright (C) 2025 Roberto L. Castro (Roberto.LopezCastro@ist.ac.at).
5 +# All Rights Reserved.
6 +#
7 +# Licensed under the Apache License, Version 2.0 (the "License");
8 +# you may not use this file except in compliance with the License.
9 +# You may obtain a copy of the License at
10 +#
11 +# http://www.apache.org/licenses/LICENSE-2.0
12 +#
13 +# Unless required by applicable law or agreed to in writing, software
14 +# distributed under the License is distributed on an "AS IS" BASIS,
15 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16 +# See the License for the specific language governing permissions and
17 +# limitations under the License.
18 +#
19 +import numpy as np
20 +import pytest
21 +import torch
22 +from compressed_tensors.transform.utils.hadamard import deterministic_hadamard_matrix
23 +
24 +from vllm._custom_ops import fusedQuantizeMx, matmul_mxf4_bf16_tn
25 +from vllm.model_executor.layers.quantization.qutlass_utils import to_blocked
26 +from vllm.platforms import current_platform
27 +
28 +if not torch.cuda.is_available():
29 + pytest.skip("CUDA required for these tests.", allow_module_level=True)
30 +
31 +if not (
32 + current_platform.has_device_capability(100)
33 + or current_platform.has_device_capability(120)
34 +):
35 + pytest.skip(
36 + reason="Tests require compute capability 10.0 (100) or 12.0 (120).",
37 + allow_module_level=True,
38 + )
39 +
40 +
41 +# ----- Helpers -----
42 +def get_hadamard_matrix(group_size: int, dtype: torch.dtype, device: torch.device):
43 + return (
44 + deterministic_hadamard_matrix(group_size, dtype=dtype, device=device)
45 + * group_size**-0.5
46 + )
47 +
48 +
49 +def _rtne_fp4(x: torch.Tensor):
50 + device = x.device
51 + grid = torch.tensor(
52 + [
53 + -6.0,
54 + -4.0,
55 + -3.0,
56 + -2.0,
57 + -1.5,
58 + -1.0,
59 + -0.5,
60 + -0.0,
61 + 0.0,
62 + 0.5,
63 + 1.0,
64 + 1.5,
65 + 2.0,
66 + 3.0,
67 + 4.0,
68 + 6.0,
69 + ],
70 + dtype=x.dtype,
71 + device=x.device,
72 + )
73 + grid_int = torch.tensor(
74 + [-1, -2, -3, -4, -5, -6, -7, -8, 0, 1, 2, 3, 4, 5, 6, 7],
75 + dtype=torch.uint8,
76 + device=device,
77 + )
78 + inds = torch.bucketize(x, grid)
79 + lo, hi = (inds - 1).clamp(min=0, max=15), inds.clamp(min=0, max=15)
80 + g_lo, g_hi = grid[lo], grid[hi]
81 + pick_hi = (g_hi - x < x - g_lo) | (g_hi - x == x - g_lo) & (grid_int[hi] % 2 == 0)
82 + y = torch.where(pick_hi, g_hi, g_lo)
83 + y_int = torch.where(pick_hi, grid_int[hi], grid_int[lo])
84 + y_int_packed = (y_int[..., 1::2] & 0xF) << 4 | y_int[..., ::2] & 0xF
85 + return y, y_int_packed
86 +
87 +
88 +def _dq_fp4(x_e2m1: torch.Tensor, x_e8m0: torch.Tensor, alpha: float):
89 + device = x_e2m1.device
90 +
91 + x_e2m1_i32 = x_e2m1.view(dtype=torch.uint8).to(dtype=torch.int32)
92 + x_e2m1_unpacked = torch.stack(
93 + [x_e2m1_i32 & 0xF, (x_e2m1_i32 >> 4) & 0xF], dim=-1
94 + ).flatten(start_dim=-2)
95 +
96 + grid_dq = torch.tensor(
97 + [
98 + 0.0,
99 + 0.5,
100 + 1.0,
101 + 1.5,
102 + 2.0,
103 + 3.0,
104 + 4.0,
105 + 6.0,
106 + -0.0,
107 + -0.5,
108 + -1.0,
109 + -1.5,
110 + -2.0,
111 + -3.0,
112 + -4.0,
113 + -6.0,
114 + ],
115 + dtype=torch.float64,
116 + device=device,
117 + )
118 + x_fp4_dq = grid_dq[x_e2m1_unpacked]
119 + scales_dq = x_e8m0.to(torch.float64)
120 +
121 + x_dq = (x_fp4_dq.unflatten(dim=-1, sizes=(-1, 32)) * scales_dq[..., None]).flatten(
122 + start_dim=-2
123 + ) / alpha
124 + return x_dq, x_fp4_dq, scales_dq
125 +
126 +
127 +def _unpack_mask(clip_mask: torch.Tensor) -> torch.Tensor:
128 + clip_mask_unpacked_dq = torch.zeros(
129 + *clip_mask.shape[:-1],
130 + clip_mask.size(-1) * 8,
131 + dtype=torch.bool,
132 + device=clip_mask.device,
133 + )
134 + for i in range(8):
135 + clip_mask_unpacked_dq[..., i::8] = (clip_mask >> i) & 1
136 + return clip_mask_unpacked_dq
137 +
138 +
139 +def _forward_quantize_ref(
140 + x: torch.Tensor, h: torch.Tensor, rot_size: int, quest: bool = True
141 +):
142 + device = x.device
143 + xh_ref64 = (
144 + x.unflatten(dim=-1, sizes=(-1, rot_size)).to(dtype=torch.float64)
145 + @ h.reshape(rot_size, rot_size).to(dtype=torch.float64)
146 + ).flatten(start_dim=-2)
147 +
148 + if quest:
149 + scales_ref64_ = (
150 + xh_ref64.unflatten(dim=-1, sizes=(-1, 32)).std(dim=-1, correction=0)
151 + * (2.92247856 / 6.0)
152 + + 1e-8
153 + )
154 + else:
155 + abs_max = xh_ref64.unflatten(dim=-1, sizes=(-1, 32)).abs().amax(dim=-1)
156 + scales_ref64_ = abs_max + 1e-8
157 +
158 + xh_e8m0_ref = scales_ref64_.log2().floor().exp2().to(dtype=torch.float8_e8m0fnu)
159 + scales_ref64 = xh_e8m0_ref.to(dtype=torch.float64)
160 +
161 + xh_scaled_ref64 = (
162 + xh_ref64.unflatten(dim=-1, sizes=(-1, 32)) / scales_ref64[..., None]
163 + ).flatten(start_dim=-2)
164 + if not quest:
165 + xh_scaled_ref64 *= 3
166 +
167 + clip_mask_unpacked_ref = xh_scaled_ref64.abs() < 6.0
168 + clip_mask_ref = torch.zeros(
169 + *x.shape[:-1], x.size(-1) // 8, dtype=torch.uint8, device=device
170 + )
171 + for i in range(8):
172 + clip_mask_ref |= clip_mask_unpacked_ref[..., i::8].to(dtype=torch.uint8) << i
173 +
174 + xh_fp4_ref, xh_e2m1_ref = _rtne_fp4(xh_scaled_ref64)
175 + xh_dq, xh_fp4_dq, scales_dq = _dq_fp4(
176 + xh_e2m1_ref, xh_e8m0_ref, alpha=1.0 if quest else 3.0
177 + )
178 + clip_mask_unpacked_dq = _unpack_mask(clip_mask_ref)
179 +
180 + assert xh_fp4_dq.equal(xh_fp4_ref)
181 + assert scales_dq.equal(scales_ref64)
182 + assert clip_mask_unpacked_dq.equal(clip_mask_unpacked_ref)
183 +
184 + return (
185 + xh_dq,
186 + clip_mask_unpacked_ref,
187 + (xh_e2m1_ref, xh_e8m0_ref, clip_mask_ref),
188 + )
189 +
190 +
191 +DTYPE = torch.bfloat16
192 +DEVICE = torch.device("cuda:0")
193 +
194 +ROT_SIZES = [32, 64, 128]
195 +SEEDS = [0]
196 +BATCHES = [1, 16]
197 +
198 +LLAMA_MODELS = {
199 + "7B": [(4096, 3 * 4096), (4096, 4096), (4096, 2 * 10752), (10752, 4096)],
200 + "13B": [(5120, 3 * 5120), (5120, 5120), (5120, 2 * 13568), (13568, 5120)],
201 + "33B": [(6656, 3 * 6656), (6656, 6656), (6656, 2 * 17664), (17664, 6656)],
202 + "70B": [(8192, 3 * 8192), (8192, 8192), (8192, 2 * 21760), (21760, 8192)],
203 +}
204 +
205 +
206 +@pytest.fixture(autouse=True)
207 +def _seed_each_test():
208 + current_platform.seed_everything(0)
209 + np.random.seed(0)
210 + torch.random.manual_seed(0)
211 +
212 +
213 +@pytest.mark.parametrize("rot_size", ROT_SIZES)
214 +@torch.inference_mode()
215 +def test_fused_quantization_absmax(rot_size: int):
216 + dtype, device = DTYPE, DEVICE
217 + h = get_hadamard_matrix(rot_size, dtype, device)
218 + x = torch.randn(2, 4096, 4096, dtype=dtype, device=device) * 25.0
219 +
220 + xh_dq_ref, _, _ = _forward_quantize_ref(x, h, rot_size, quest=False)
221 + xh_e2m1, xh_e8m0 = fusedQuantizeMx(x, h, method="abs_max")
222 + xh_e8m0 = xh_e8m0.reshape(2, 4096, 4096 // 32)
223 + xh_dq, *_ = _dq_fp4(xh_e2m1, xh_e8m0, alpha=3.0)
224 +
225 + torch.testing.assert_close(xh_dq, xh_dq_ref, rtol=0.34, atol=100)
226 + assert (xh_dq != xh_dq_ref).float().mean() <= 1e-4
227 +
228 + m, n, k = 1, 504, 4096
229 + a = torch.randn(m, k, dtype=dtype, device=device) * 25.0
230 + b = torch.randn(n, k, dtype=dtype, device=device) * 25.0
231 +
232 + a_e2m1, a_e8m0 = fusedQuantizeMx(a, h, method="abs_max")
233 + b_e2m1, b_e8m0 = fusedQuantizeMx(b, h, method="abs_max")
234 + a_dq, *_ = _dq_fp4(a_e2m1, a_e8m0[:m, :k], alpha=1.0)
235 + b_dq, *_ = _dq_fp4(b_e2m1, b_e8m0[:n, :k], alpha=1.0)
236 + out_ref = a_dq @ b_dq.transpose(-2, -1)
237 +
238 + a_scale_block = to_blocked(a_e8m0, backend="triton")
239 + b_scale_block = to_blocked(b_e8m0, backend="triton")
240 + alpha = torch.tensor([1.0], device=device)
241 + out = matmul_mxf4_bf16_tn(a_e2m1, b_e2m1, a_scale_block, b_scale_block, alpha)
242 + assert out.equal(out_ref.to(dtype=out.dtype))
243 +
244 +
245 +@pytest.mark.parametrize("rot_size", ROT_SIZES)
246 +@torch.inference_mode()
247 +def test_fused_quantization_quest(rot_size: int):
248 + dtype, device = DTYPE, DEVICE
249 + h = get_hadamard_matrix(rot_size, dtype, device)
250 + x = torch.randn(2, 4096, 4096, dtype=dtype, device=device) * 25.0
251 +
252 + xh_dq_ref, _, _ = _forward_quantize_ref(x, h, rot_size, quest=True)
253 + xh_e2m1, xh_e8m0 = fusedQuantizeMx(x, h, method="quest")
254 + xh_e8m0 = xh_e8m0.reshape(2, 4096, 4096 // 32)
255 + xh_dq, *_ = _dq_fp4(xh_e2m1, xh_e8m0, alpha=1.0)
256 +
257 + torch.testing.assert_close(xh_dq, xh_dq_ref, rtol=0.34, atol=100)
258 + assert (xh_dq != xh_dq_ref).float().mean() <= 1e-4
259 +
260 + m, n, k = 504, 504, 2048
261 + a = torch.randn(m, k, dtype=dtype, device=device) * 25.0
262 + b = torch.randn(n, k, dtype=dtype, device=device) * 25.0
263 +
264 + a_e2m1, a_e8m0 = fusedQuantizeMx(a, h, method="quest")
265 + b_e2m1, b_e8m0 = fusedQuantizeMx(b, h, method="quest")
266 + a_dq, *_ = _dq_fp4(a_e2m1, a_e8m0[:m, :k], alpha=1.0)
267 + b_dq, *_ = _dq_fp4(b_e2m1, b_e8m0[:n, :k], alpha=1.0)
268 + out_ref = a_dq @ b_dq.transpose(-2, -1)
269 +
270 + a_scale_block = to_blocked(a_e8m0, backend="triton")
271 + b_scale_block = to_blocked(b_e8m0, backend="triton")
272 + alpha = torch.tensor([1.0], device=device)
273 + out = matmul_mxf4_bf16_tn(a_e2m1, b_e2m1, a_scale_block, b_scale_block, alpha)
274 + assert out.equal(out_ref.to(dtype=out.dtype))
275 +
276 +
277 +@pytest.mark.parametrize("model", list(LLAMA_MODELS.keys()))
278 +@pytest.mark.parametrize("layer_idx", [0, 1, 2, 3])
279 +@pytest.mark.parametrize("batch", [1, 16])
280 +@pytest.mark.parametrize("had_size", ROT_SIZES)
281 +@torch.inference_mode()
282 +def test_llama_shapes(model: str, layer_idx: int, batch: int, had_size: int):
283 + dtype, device = DTYPE, DEVICE
284 + m = batch
285 + k, n = LLAMA_MODELS[model][layer_idx]
286 +
287 + h = get_hadamard_matrix(had_size, dtype, device)
288 +
289 + a = torch.rand(m, k, dtype=dtype, device=device) * 25.0
290 + b = torch.rand(n, k, dtype=dtype, device=device) * 25.0
291 +
292 + a_e2m1, a_e8m0 = fusedQuantizeMx(a, h, method="quest")
293 + b_e2m1, b_e8m0 = fusedQuantizeMx(b, h, method="quest")
294 +
295 + a_dq, *_ = _dq_fp4(a_e2m1, a_e8m0[:m, :k], alpha=1.0)
296 + b_dq, *_ = _dq_fp4(b_e2m1, b_e8m0[:n, :k], alpha=1.0)
297 + out_ref = a_dq @ b_dq.transpose(-2, -1)
298 +
299 + a_scale_block = to_blocked(a_e8m0, backend="triton")
300 + b_scale_block = to_blocked(b_e8m0, backend="triton")
301 + alpha = torch.tensor([1.0], device=device)
302 + out = matmul_mxf4_bf16_tn(a_e2m1, b_e2m1, a_scale_block, b_scale_block, alpha)
303 + assert out.equal(out_ref.to(dtype=out.dtype))
tests/kernels/quantization/test_nvfp4_qutlass.py new
+268
@@ -0,0 +1,268 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +#
4 +# Copyright (C) 2025 Roberto L. Castro (Roberto.LopezCastro@ist.ac.at).
5 +# All Rights Reserved.
6 +#
7 +# Licensed under the Apache License, Version 2.0 (the "License");
8 +# you may not use this file except in compliance with the License.
9 +# You may obtain a copy of the License at
10 +#
11 +# http://www.apache.org/licenses/LICENSE-2.0
12 +#
13 +# Unless required by applicable law or agreed to in writing, software
14 +# distributed under the License is distributed on an "AS IS" BASIS,
15 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16 +# See the License for the specific language governing permissions and
17 +# limitations under the License.
18 +#
19 +import numpy as np
20 +import pytest
21 +import torch
22 +from compressed_tensors.transform.utils.hadamard import deterministic_hadamard_matrix
23 +
24 +from vllm import _custom_ops as ops # use existing nvfp4 gemm in vllm
25 +from vllm._custom_ops import fusedQuantizeNv
26 +from vllm.model_executor.layers.quantization.qutlass_utils import to_blocked
27 +from vllm.platforms import current_platform
28 +
29 +if not torch.cuda.is_available():
30 + pytest.skip("CUDA required for these tests.", allow_module_level=True)
31 +
32 +if not (
33 + current_platform.has_device_capability(100)
34 + or current_platform.has_device_capability(120)
35 +):
36 + pytest.skip(
37 + reason="Tests require compute capability 10.0 (100) or 12.0 (120).",
38 + allow_module_level=True,
39 + )
40 +
41 +
42 +# ----- Helpers -----
43 +def get_hadamard_matrix(group_size: int, dtype: torch.dtype, device: torch.device):
44 + return (
45 + deterministic_hadamard_matrix(group_size, dtype=dtype, device=device)
46 + * group_size**-0.5
47 + )
48 +
49 +
50 +def _rtne_fp4(x: torch.Tensor):
51 + device = x.device
52 + grid = torch.tensor(
53 + [
54 + -6.0,
55 + -4.0,
56 + -3.0,
57 + -2.0,
58 + -1.5,
59 + -1.0,
60 + -0.5,
61 + -0.0,
62 + 0.0,
63 + 0.5,
64 + 1.0,
65 + 1.5,
66 + 2.0,
67 + 3.0,
68 + 4.0,
69 + 6.0,
70 + ],
71 + dtype=x.dtype,
72 + device=x.device,
73 + )
74 + grid_int = torch.tensor(
75 + [-1, -2, -3, -4, -5, -6, -7, -8, 0, 1, 2, 3, 4, 5, 6, 7],
76 + dtype=torch.uint8,
77 + device=device,
78 + )
79 + inds = torch.bucketize(x, grid)
80 + lo, hi = (inds - 1).clamp(min=0, max=15), inds.clamp(min=0, max=15)
81 + g_lo, g_hi = grid[lo], grid[hi]
82 + pick_hi = (g_hi - x < x - g_lo) | (g_hi - x == x - g_lo) & (grid_int[hi] % 2 == 0)
83 + y = torch.where(pick_hi, g_hi, g_lo)
84 + y_int = torch.where(pick_hi, grid_int[hi], grid_int[lo])
85 + y_int_packed = (y_int[..., 1::2] & 0xF) << 4 | y_int[..., ::2] & 0xF
86 + return y, y_int_packed
87 +
88 +
89 +def _dq_fp4(x_e2m1: torch.Tensor, x_e4m3: torch.Tensor, alpha: float):
90 + device = x_e2m1.device
91 +
92 + x_e2m1_i32 = x_e2m1.view(dtype=torch.uint8).to(dtype=torch.int32)
93 + x_e2m1_unpacked = torch.stack(
94 + [x_e2m1_i32 & 0xF, (x_e2m1_i32 >> 4) & 0xF], dim=-1
95 + ).flatten(start_dim=-2)
96 +
97 + grid_dq = torch.tensor(
98 + [
99 + 0.0,
100 + 0.5,
101 + 1.0,
102 + 1.5,
103 + 2.0,
104 + 3.0,
105 + 4.0,
106 + 6.0,
107 + -0.0,
108 + -0.5,
109 + -1.0,
110 + -1.5,
111 + -2.0,
112 + -3.0,
113 + -4.0,
114 + -6.0,
115 + ],
116 + dtype=torch.float64,
117 + device=device,
118 + )
119 + x_fp4_dq = grid_dq[x_e2m1_unpacked]
120 +
121 + scales_dq = x_e4m3.to(torch.float64)
122 + x_dq = (x_fp4_dq.unflatten(dim=-1, sizes=(-1, 16)) * scales_dq[..., None]).flatten(
123 + start_dim=-2
124 + ) / alpha # * (4. / 3.)
125 + return x_dq, x_fp4_dq, scales_dq
126 +
127 +
128 +def _unpack_mask(clip_mask: torch.Tensor) -> torch.Tensor:
129 + clip_mask_unpacked_dq = torch.zeros(
130 + *clip_mask.shape[:-1],
131 + clip_mask.size(-1) * 8,
132 + dtype=torch.bool,
133 + device=clip_mask.device,
134 + )
135 + for i in range(8):
136 + clip_mask_unpacked_dq[..., i::8] = (clip_mask >> i) & 1
137 + return clip_mask_unpacked_dq
138 +
139 +
140 +def _forward_quantize_ref(x: torch.Tensor, h: torch.Tensor, rot_size: int):
141 + device = x.device
142 +
143 + xh_ref64 = (
144 + x.unflatten(dim=-1, sizes=(-1, rot_size)).to(dtype=torch.float64)
145 + @ h.reshape(rot_size, rot_size).to(dtype=torch.float64)
146 + ).flatten(start_dim=-2)
147 +
148 + abs_max = xh_ref64.unflatten(dim=-1, sizes=(-1, 16)).abs().amax(dim=-1)
149 + scales_ref64_ = abs_max + 1e-8
150 +
151 + xh_e4m3_ref = scales_ref64_.to(dtype=torch.float8_e4m3fn)
152 + scales_ref64 = xh_e4m3_ref.to(dtype=torch.float64)
153 + xh_scaled_ref64 = (
154 + xh_ref64.unflatten(dim=-1, sizes=(-1, 16)) / scales_ref64[..., None]
155 + ).flatten(start_dim=-2)
156 +
157 + xh_scaled_ref64 *= 6.0
158 +
159 + clip_mask_unpacked_ref = xh_scaled_ref64.abs() < 6.0
160 + clip_mask_ref = torch.zeros(
161 + *x.shape[:-1], x.size(-1) // 8, dtype=torch.uint8, device=device
162 + )
163 + for i in range(8):
164 + clip_mask_ref |= clip_mask_unpacked_ref[..., i::8].to(dtype=torch.uint8) << i
165 +
166 + xh_fp4_ref, xh_e2m1_ref = _rtne_fp4(xh_scaled_ref64)
167 + xh_dq, xh_fp4_dq, scales_dq = _dq_fp4(xh_e2m1_ref, xh_e4m3_ref, 6.0)
168 + clip_mask_unpacked_dq = _unpack_mask(clip_mask_ref)
169 +
170 + assert xh_fp4_dq.equal(xh_fp4_ref)
171 + assert scales_dq.equal(scales_ref64)
172 + assert clip_mask_unpacked_dq.equal(clip_mask_unpacked_ref)
173 +
174 + return (
175 + xh_dq,
176 + clip_mask_unpacked_ref,
177 + (xh_e2m1_ref, xh_e4m3_ref, clip_mask_ref),
178 + )
179 +
180 +
181 +DTYPE = torch.bfloat16
182 +DEVICE = torch.device("cuda:0")
183 +ROT_SIZES = [16, 32, 64, 128]
184 +GLOBAL_SCALES = [6.0]
185 +
186 +LLAMA_MODELS = {
187 + "7B": [(4096, 3 * 4096), (4096, 4096), (4096, 2 * 10752), (10752, 4096)],
188 + "13B": [(5120, 3 * 5120), (5120, 5120), (5120, 2 * 13568), (13568, 5120)],
189 + "33B": [(6656, 3 * 6656), (6656, 6656), (6656, 2 * 17664), (17664, 6656)],
190 + "70B": [(8192, 3 * 8192), (8192, 8192), (8192, 2 * 21760), (21760, 8192)],
191 +}
192 +
193 +
194 +@pytest.fixture(autouse=True)
195 +def _seed_each_test():
196 + current_platform.seed_everything(0)
197 + np.random.seed(0)
198 + torch.random.manual_seed(0)
199 +
200 +
201 +@pytest.mark.parametrize("rot_size", ROT_SIZES)
202 +@pytest.mark.parametrize("global_scale_value", GLOBAL_SCALES)
203 +@torch.inference_mode()
204 +def test_fused_quantization(rot_size: int, global_scale_value: float):
205 + dtype, device = DTYPE, DEVICE
206 + h = get_hadamard_matrix(rot_size, dtype, device)
207 + x = torch.randn(2, 4096, 4096, dtype=dtype, device=device) * 25.0
208 + global_scale = torch.tensor([global_scale_value], device=device)
209 +
210 + xh_dq_ref, _, _ = _forward_quantize_ref(x, h, rot_size)
211 + xh_e2m1, xh_e4m3 = fusedQuantizeNv(x, h, global_scale)
212 + xh_e4m3 = xh_e4m3.reshape(2, 4096, 4096 // 16)
213 + xh_dq, *_ = _dq_fp4(xh_e2m1, xh_e4m3, alpha=global_scale_value)
214 +
215 + torch.testing.assert_close(xh_dq, xh_dq_ref, rtol=0.34, atol=100)
216 + assert (xh_dq != xh_dq_ref).float().mean() <= 1e-1
217 +
218 + m, n, k = 504, 4096 * 2, 4096
219 + a = torch.randn(m, k, dtype=dtype, device=device) * 25.0
220 + b = torch.randn(n, k, dtype=dtype, device=device) * 25.0
221 +
222 + a_e2m1, a_e4m3 = fusedQuantizeNv(a, h, global_scale)
223 + b_e2m1, b_e4m3 = fusedQuantizeNv(b, h, global_scale)
224 +
225 + a_dq, *_ = _dq_fp4(a_e2m1, a_e4m3[:m, :k], alpha=1.0)
226 + b_dq, *_ = _dq_fp4(b_e2m1, b_e4m3[:n, :k], alpha=1.0)
227 + out_ref = a_dq @ b_dq.transpose(-2, -1)
228 +
229 + a_scale_block = to_blocked(a_e4m3, backend="triton").view(-1, k // 16)
230 + b_scale_block = to_blocked(b_e4m3, backend="triton").view(-1, k // 16)
231 + alpha = torch.tensor([1.0], device=device)
232 + out = ops.cutlass_scaled_fp4_mm(
233 + a_e2m1, b_e2m1, a_scale_block, b_scale_block, alpha, torch.bfloat16
234 + )
235 + assert out.equal(out_ref.to(dtype=out.dtype))
236 +
237 +
238 +@pytest.mark.parametrize("model", list(LLAMA_MODELS.keys()))
239 +@pytest.mark.parametrize("layer_idx", [0, 1, 2, 3])
240 +@pytest.mark.parametrize("batch", [1, 16])
241 +@pytest.mark.parametrize("rot_size", ROT_SIZES)
242 +@torch.inference_mode()
243 +def test_llama_shapes(model: str, layer_idx: int, batch: int, rot_size: int):
244 + dtype, device = DTYPE, DEVICE
245 + m = batch
246 + k, n = LLAMA_MODELS[model][layer_idx]
247 +
248 + h = get_hadamard_matrix(rot_size, dtype, device)
249 +
250 + a = torch.randn(m, k, dtype=dtype, device=device) * 25.0
251 + b = torch.randn(n, k, dtype=dtype, device=device) * 25.0
252 +
253 + global_scale = torch.tensor([1.0], device=device)
254 +
255 + a_e2m1, a_e4m3 = fusedQuantizeNv(a, h, global_scale)
256 + b_e2m1, b_e4m3 = fusedQuantizeNv(b, h, global_scale)
257 +
258 + a_dq, *_ = _dq_fp4(a_e2m1, a_e4m3[:m, :k], alpha=1.0)
259 + b_dq, *_ = _dq_fp4(b_e2m1, b_e4m3[:n, :k], alpha=1.0)
260 + out_ref = a_dq @ b_dq.transpose(-2, -1)
261 +
262 + a_scale_block = to_blocked(a_e4m3, backend="triton").view(-1, k // 16)
263 + b_scale_block = to_blocked(b_e4m3, backend="triton").view(-1, k // 16)
264 + alpha = torch.tensor([1.0], device=device)
265 + out = ops.cutlass_scaled_fp4_mm(
266 + a_e2m1, b_e2m1, a_scale_block, b_scale_block, alpha, torch.bfloat16
267 + )
268 + assert out.equal(out_ref.to(dtype=out.dtype))
tests/quantization/fp_quant.py new
+32
@@ -0,0 +1,32 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Test model set-up and inference for quantized HF models supported
4 +on the GPU backend using FPQuant.
5 +
6 +Validating the configuration and printing results for manual checking.
7 +
8 +Run `pytest tests/quantization/test_fp_quant.py`.
9 +"""
10 +
11 +import pytest
12 +
13 +from tests.quantization.utils import is_quant_method_supported
14 +
15 +MODELS = [
16 + "ISTA-DASLab/Qwen3-0.6B-RTN-NVFP4",
17 + "ISTA-DASLab/Qwen3-0.6B-RTN-MXFP4",
18 +]
19 +DTYPE = ["bfloat16"]
20 +EAGER = [True, False]
21 +
22 +
23 +@pytest.mark.skipif(
24 + not is_quant_method_supported("fp_quant"),
25 + reason="FPQuant is not supported on this GPU type.",
26 +)
27 +@pytest.mark.parametrize("model", MODELS)
28 +@pytest.mark.parametrize("eager", EAGER)
29 +def test_fpquant(vllm_runner, model, eager):
30 + with vllm_runner(model, enforce_eager=eager) as llm:
31 + output = llm.generate_greedy(["1 2 3 4 5"], max_tokens=2)
32 + assert output[0][1] == "1 2 3 4 5 6"
vllm/_custom_ops.py
+139 -1
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from typing import TYPE_CHECKING, Optional, Union
4 +from typing import TYPE_CHECKING, Literal, Optional, Union
5
6 import torch
7
@@ -2483,6 +2483,144 @@ def onednn_scaled_mm(
2483 return output
2484
2485
2486 +if hasattr(torch.ops._qutlass_C, "matmul_mxf4_bf16_tn"):
2487 +
2488 + @register_fake("_qutlass_C::matmul_mxf4_bf16_tn")
2489 + def _fake_matmul_mxf4_bf16_tn(
2490 + a: torch.Tensor,
2491 + b: torch.Tensor,
2492 + a_sf: torch.Tensor,
2493 + b_sf: torch.Tensor,
2494 + alpha: torch.Tensor,
2495 + ):
2496 + return a.new_empty(*a.shape[:-1], b.shape[0], dtype=torch.bfloat16)
2497 +
2498 +
2499 +def matmul_mxf4_bf16_tn(
2500 + a: torch.Tensor,
2501 + b: torch.Tensor,
2502 + a_sf: torch.Tensor,
2503 + b_sf: torch.Tensor,
2504 + alpha: torch.Tensor,
2505 +) -> torch.Tensor:
2506 + return torch.ops._qutlass_C.matmul_mxf4_bf16_tn(a, b, a_sf, b_sf, alpha)
2507 +
2508 +
2509 +if hasattr(torch.ops._qutlass_C, "matmul_ada_mxf4_bf16_tn"):
2510 +
2511 + @register_fake("_qutlass_C::matmul_ada_mxf4_bf16_tn")
2512 + def _fake_matmul_ada_mxf4_bf16_tn(
2513 + a: torch.Tensor,
2514 + b: torch.Tensor,
2515 + a_sf: torch.Tensor,
2516 + b_sf: torch.Tensor,
2517 + alpha: torch.Tensor,
2518 + ):
2519 + return a.new_empty(*a.shape[:-1], b.shape[0], dtype=torch.bfloat16)
2520 +
2521 +
2522 +def matmul_ada_mxf4_bf16_tn(
2523 + a: torch.Tensor,
2524 + b: torch.Tensor,
2525 + a_sf: torch.Tensor,
2526 + b_sf: torch.Tensor,
2527 + alpha: torch.Tensor,
2528 +) -> torch.Tensor:
2529 + return torch.ops._qutlass_C.matmul_ada_mxf4_bf16_tn(a, b, a_sf, b_sf, alpha)
2530 +
2531 +
2532 +def ceil_div(a, b):
2533 + return (a + b - 1) // b
2534 +
2535 +
2536 +if hasattr(torch.ops._qutlass_C, "fusedQuantizeMxQuest"):
2537 +
2538 + @register_fake("_qutlass_C::fusedQuantizeMxQuest")
2539 + def _fake_fused_quantize_mx_quest(
2540 + a: torch.Tensor, b: torch.Tensor, xh_e2m1: torch.Tensor, xh_e8m0: torch.Tensor
2541 + ):
2542 + return xh_e2m1, xh_e8m0
2543 +
2544 +
2545 +if hasattr(torch.ops._qutlass_C, "fusedQuantizeMxAbsMax"):
2546 +
2547 + @register_fake("_qutlass_C::fusedQuantizeMxAbsMax")
2548 + def _fake_fused_quantize_mx_absmax(
2549 + a: torch.Tensor, b: torch.Tensor, xh_e2m1: torch.Tensor, xh_e8m0: torch.Tensor
2550 + ):
2551 + return xh_e2m1, xh_e8m0
2552 +
2553 +
2554 +def fusedQuantizeMx(
2555 + a: torch.Tensor, b: torch.Tensor, *, method: Literal["quest", "abs_max"] = "quest"
2556 +) -> tuple[torch.Tensor, torch.Tensor]:
2557 + if a.dim() == 0:
2558 + raise ValueError("`a` must have at least 1 dimension.")
2559 + if a.size(-1) % 32 != 0:
2560 + raise ValueError(f"last dim of `a` must be divisible by 32, got {a.size(-1)}.")
2561 + if b.device != a.device:
2562 + raise ValueError("`a` and `b` must be on the same device.")
2563 +
2564 + xh_e2m1 = torch.empty(
2565 + *a.shape[:-1], a.size(-1) // 2, dtype=torch.uint8, device=a.device
2566 + )
2567 +
2568 + rows, cols = a.numel() // a.size(-1), a.size(-1) // 32
2569 + n_row_blocks = ceil_div(rows, 128)
2570 + n_col_blocks = ceil_div(cols, 4)
2571 + padded_rows = n_row_blocks * 128
2572 + padded_cols = n_col_blocks * 4
2573 +
2574 + xh_e8m0 = torch.empty(
2575 + padded_rows, padded_cols, dtype=torch.float8_e8m0fnu, device=a.device
2576 + )
2577 +
2578 + if not hasattr(torch.ops, "_qutlass_C"):
2579 + raise RuntimeError(
2580 + "The `_qutlass_C` extension is not loaded. "
2581 + "Make sure your custom op library is imported before calling fusedQuantizeMx."
2582 + )
2583 +
2584 + if method == "quest":
2585 + return torch.ops._qutlass_C.fusedQuantizeMxQuest(a, b, xh_e2m1, xh_e8m0)
2586 + elif method == "abs_max":
2587 + return torch.ops._qutlass_C.fusedQuantizeMxAbsMax(a, b, xh_e2m1, xh_e8m0)
2588 + else:
2589 + raise ValueError(f"invalid method {method!r}, must be 'quest' or 'abs_max'")
2590 +
2591 +
2592 +if hasattr(torch.ops._qutlass_C, "fusedQuantizeNv"):
2593 +
2594 + @register_fake("_qutlass_C::fusedQuantizeNv")
2595 + def _fake_fused_quantize_nv(
2596 + a: torch.Tensor,
2597 + b: torch.Tensor,
2598 + xh_e2m1: torch.Tensor,
2599 + xh_e4m3: torch.Tensor,
2600 + global_scale: torch.Tensor,
2601 + ):
2602 + return xh_e2m1, xh_e4m3
2603 +
2604 +
2605 +def fusedQuantizeNv(
2606 + a: torch.Tensor, b: torch.Tensor, global_scale: torch.Tensor
2607 +) -> tuple[torch.Tensor, torch.Tensor]:
2608 + xh_e2m1 = torch.empty(
2609 + *a.shape[:-1], a.size(-1) // 2, dtype=torch.uint8, device=a.device
2610 + )
2611 +
2612 + rows, cols = a.numel() // a.size(-1), a.size(-1) // 16
2613 + n_row_blocks = ceil_div(rows, 128)
2614 + n_col_blocks = ceil_div(cols, 4)
2615 + padded_rows = n_row_blocks * 128
2616 + padded_cols = n_col_blocks * 4
2617 + xh_e4m3 = torch.empty(
2618 + padded_rows, padded_cols, dtype=torch.float8_e4m3fn, device=a.device
2619 + )
2620 +
2621 + return torch.ops._qutlass_C.fusedQuantizeNv(a, b, xh_e2m1, xh_e4m3, global_scale)
2622 +
2623 +
2624 def hadacore_transform(x: torch.Tensor, inplace: bool = True) -> torch.Tensor:
2625 """
2626 Perform Hadamard transforms using [Hadacore](https://arxiv.org/abs/2412.08832)
vllm/model_executor/layers/quantization/__init__.py
+3
@@ -12,6 +12,7 @@ QuantizationMethods = Literal[
12 "fp8",
13 "ptpc_fp8",
14 "fbgemm_fp8",
15 + "fp_quant",
16 "modelopt",
17 "modelopt_fp4",
18 "bitblas",
@@ -102,6 +103,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
103 from .experts_int8 import ExpertsInt8Config
104 from .fbgemm_fp8 import FBGEMMFp8Config
105 from .fp8 import Fp8Config
106 + from .fp_quant import FPQuantConfig
107 from .gguf import GGUFConfig
108 from .gptq import GPTQConfig
109 from .gptq_bitblas import GPTQBitBLASConfig
@@ -125,6 +127,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]:
127 "tpu_int8": Int8TpuConfig,
128 "fp8": Fp8Config,
129 "fbgemm_fp8": FBGEMMFp8Config,
130 + "fp_quant": FPQuantConfig,
131 "modelopt": ModelOptFp8Config,
132 "modelopt_fp4": ModelOptNvFp4Config,
133 "bitblas": BitBLASConfig,
vllm/model_executor/layers/quantization/fp_quant.py new
+420
@@ -0,0 +1,420 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +# Supports FP-Quant compression, see https://arxiv.org/abs/2509.23202
5 +
6 +from typing import Any, Optional
7 +
8 +import torch
9 +from torch.nn.parameter import Parameter
10 +
11 +from vllm._custom_ops import (
12 + cutlass_scaled_fp4_mm,
13 + fusedQuantizeMx,
14 + fusedQuantizeNv,
15 + matmul_mxf4_bf16_tn,
16 +)
17 +from vllm.model_executor.layers.linear import (
18 + LinearBase,
19 + LinearMethodBase,
20 + UnquantizedLinearMethod,
21 +)
22 +from vllm.model_executor.layers.quantization import QuantizationMethods
23 +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
24 +from vllm.model_executor.layers.quantization.qutlass_utils import to_blocked
25 +from vllm.model_executor.utils import set_weight_attrs
26 +from vllm.platforms import current_platform
27 +from vllm.utils import direct_register_custom_op
28 +
29 +
30 +class FPQuantConfig(QuantizationConfig):
31 + """Config class for FPQuant."""
32 +
33 + def __init__(
34 + self,
35 + hadamard_group_size: int = 32,
36 + forward_dtype: str = "mxfp4",
37 + forward_method: str = "abs_max",
38 + pseudoquantization: bool = False,
39 + modules_to_not_convert: Optional[list[str]] = None,
40 + ) -> None:
41 + super().__init__()
42 + self.hadamard_group_size = hadamard_group_size
43 + self.forward_dtype = forward_dtype
44 + self.forward_method = forward_method
45 + self.pseudoquantization = pseudoquantization
46 + self.modules_to_not_convert = modules_to_not_convert
47 +
48 + if pseudoquantization:
49 + raise ValueError("Pseudoquantization is not supported for vLLM")
50 +
51 + def __repr__(self) -> str:
52 + return (
53 + f"FPQuantConfig(hadamard_group_size={self.hadamard_group_size}, "
54 + f"forward_dtype={self.forward_dtype}, "
55 + f"forward_method={self.forward_method}, "
56 + f"pseudoquantization={self.pseudoquantization}, "
57 + f"modules_to_not_convert={self.modules_to_not_convert})"
58 + )
59 +
60 + @classmethod
61 + def get_name(cls) -> QuantizationMethods:
62 + return "fp_quant"
63 +
64 + @classmethod
65 + def get_supported_act_dtypes(cls) -> list[torch.dtype]:
66 + return [torch.bfloat16]
67 +
68 + @classmethod
69 + def get_min_capability(cls) -> int:
70 + return 100
71 +
72 + @classmethod
73 + def get_config_filenames(cls) -> list[str]:
74 + return [] # no extra configs.
75 +
76 + @classmethod
77 + def from_config(cls, config: dict[str, Any]) -> "FPQuantConfig":
78 + hadamard_group_size = cls.get_from_keys(config, ["hadamard_group_size"])
79 + forward_dtype = cls.get_from_keys(config, ["forward_dtype"])
80 + forward_method = cls.get_from_keys(config, ["forward_method"])
81 + pseudoquantization = cls.get_from_keys(config, ["pseudoquantization"])
82 + modules_to_not_convert = cls.get_from_keys(config, ["modules_to_not_convert"])
83 + return cls(
84 + hadamard_group_size,
85 + forward_dtype,
86 + forward_method,
87 + pseudoquantization,
88 + modules_to_not_convert,
89 + )
90 +
91 + def get_quant_method(
92 + self, layer: torch.nn.Module, prefix: str
93 + ) -> Optional[LinearMethodBase]:
94 + if self.modules_to_not_convert is not None and any(
95 + prefix.endswith(module) for module in self.modules_to_not_convert
96 + ):
97 + return UnquantizedLinearMethod()
98 +
99 + if isinstance(layer, LinearBase):
100 + return FPQuantLinearMethod(self)
101 + return None
102 +
103 +
104 +class FPQuantLinearMethod(LinearMethodBase):
105 + """Linear method for FPQuant.
106 +
107 + Args:
108 + quant_config: The FPQuant quantization config.
109 + """
110 +
111 + def __init__(self, quant_config: FPQuantConfig):
112 + self.quant_config = quant_config
113 +
114 + def create_weights(
115 + self,
116 + layer: torch.nn.Module,
117 + input_size_per_partition: int,
118 + output_partition_sizes: list[int],
119 + input_size: int,
120 + output_size: int,
121 + params_dtype: torch.dtype,
122 + **extra_weight_attrs,
123 + ):
124 + del output_size # Unused.
125 + del input_size # Unused.
126 +
127 + if params_dtype != torch.bfloat16:
128 + raise ValueError("Only bfloat16 is currently supported by FPQuant")
129 + if input_size_per_partition % self.quant_config.hadamard_group_size != 0: # noqa: E501
130 + raise ValueError(
131 + "The input size is not aligned with the quantized "
132 + "weight shape. This can be caused by too large "
133 + "tensor parallel size. Or other skill issues."
134 + )
135 +
136 + assert self.quant_config.forward_dtype in ["mxfp4", "nvfp4"], (
137 + "Only mxfp4 and nvfp4 are supported for now"
138 + )
139 + if self.quant_config.forward_dtype == "mxfp4":
140 + group_size = 32
141 + elif self.quant_config.forward_dtype == "nvfp4":
142 + group_size = 16
143 + else:
144 + raise ValueError(
145 + f"Unsupported forward_dtype: {self.quant_config.forward_dtype}"
146 + )
147 +
148 + qweight = Parameter(
149 + torch.empty(
150 + sum(output_partition_sizes),
151 + input_size_per_partition // 2,
152 + dtype=torch.uint8,
153 + ),
154 + requires_grad=False,
155 + )
156 + set_weight_attrs(
157 + qweight,
158 + {
159 + "input_dim": 1,
160 + "output_dim": 0,
161 + "packed_dim": 1,
162 + "pack_factor": 2,
163 + }
164 + | extra_weight_attrs,
165 + )
166 + layer.register_parameter("qweight", qweight)
167 +
168 + scales = Parameter(
169 + torch.empty(
170 + sum(output_partition_sizes),
171 + input_size_per_partition // group_size,
172 + dtype=torch.uint8,
173 + ),
174 + requires_grad=False,
175 + )
176 + set_weight_attrs(
177 + scales,
178 + {
179 + "input_dim": 1,
180 + "output_dim": 0,
181 + "packed_dim": 1,
182 + "pack_factor": group_size,
183 + }
184 + | extra_weight_attrs,
185 + )
186 + layer.register_parameter("scales", scales)
187 +
188 + weight_global_scale = Parameter(
189 + torch.empty(1, dtype=torch.float32),
190 + requires_grad=False,
191 + )
192 + set_weight_attrs(
193 + weight_global_scale, {"ignore_warning": True} | extra_weight_attrs
194 + )
195 + layer.register_parameter("weight_global_scale", weight_global_scale)
196 +
197 + act_global_scale = Parameter(
198 + torch.empty(1, dtype=torch.float32),
199 + requires_grad=False,
200 + )
201 + set_weight_attrs(
202 + act_global_scale, {"ignore_warning": True} | extra_weight_attrs
203 + )
204 + layer.register_parameter("act_global_scale", act_global_scale)
205 +
206 + forward_hadamard_matrix = Parameter(
207 + torch.empty(
208 + self.quant_config.hadamard_group_size,
209 + self.quant_config.hadamard_group_size,
210 + dtype=params_dtype,
211 + ),
212 + requires_grad=False,
213 + )
214 + set_weight_attrs(
215 + forward_hadamard_matrix, {"ignore_warning": True} | extra_weight_attrs
216 + )
217 + layer.register_parameter("forward_hadamard_matrix", forward_hadamard_matrix)
218 +
219 + backward_hadamard_matrix = Parameter(
220 + torch.empty(
221 + self.quant_config.hadamard_group_size,
222 + self.quant_config.hadamard_group_size,
223 + dtype=params_dtype,
224 + ),
225 + requires_grad=False,
226 + )
227 + set_weight_attrs(
228 + backward_hadamard_matrix, {"ignore_warning": True} | extra_weight_attrs
229 + )
230 + layer.register_parameter("backward_hadamard_matrix", backward_hadamard_matrix)
231 +
232 + def apply(
233 + self,
234 + layer: torch.nn.Module,
235 + x: torch.Tensor,
236 + bias: Optional[torch.Tensor] = None,
237 + ) -> torch.Tensor:
238 + return quantized_forward(
239 + x,
240 + layer.qweight,
241 + layer.scales,
242 + layer.weight_global_scale,
243 + layer.act_global_scale,
244 + bias,
245 + layer.forward_hadamard_matrix,
246 + self.quant_config.forward_method,
247 + self.quant_config.forward_dtype,
248 + )
249 +
250 +
251 +def ceil_div(a, b):
252 + return (a + b - 1) // b
253 +
254 +
255 +def fused_quantize_mx(
256 + x_flat: torch.Tensor, hadamard_matrix: torch.Tensor, forward_method: str
257 +) -> tuple[torch.Tensor, torch.Tensor]:
258 + return fusedQuantizeMx(x_flat, hadamard_matrix, method=forward_method)
259 +
260 +
261 +def fused_quantize_mx_fake(x_flat, hadamard_matrix, forward_method):
262 + rows, cols = x_flat.size(0), x_flat.size(1) // 32
263 + padded_rows = ((rows + 128 - 1) // 128) * 128
264 + padded_cols = ((cols + 4 - 1) // 4) * 4
265 +
266 + xh_e2m1 = torch.empty(
267 + x_flat.size(0), x_flat.size(1) // 2, dtype=torch.uint8, device=x_flat.device
268 + )
269 + xh_e8m0 = torch.empty(
270 + padded_rows, padded_cols, dtype=torch.float8_e8m0fnu, device=x_flat.device
271 + )
272 +
273 + return xh_e2m1, xh_e8m0
274 +
275 +
276 +direct_register_custom_op(
277 + op_name="fused_quantize_mx",
278 + op_func=fused_quantize_mx,
279 + mutates_args=[],
280 + fake_impl=fused_quantize_mx_fake,
281 + dispatch_key=current_platform.dispatch_key,
282 +)
283 +
284 +
285 +def matmul_mxf4_bf16(
286 + x: torch.Tensor,
287 + w: torch.Tensor,
288 + xs: torch.Tensor,
289 + ws: torch.Tensor,
290 + alpha: torch.Tensor,
291 +) -> torch.Tensor:
292 + return matmul_mxf4_bf16_tn(
293 + x,
294 + w,
295 + to_blocked(xs, backend="triton").view(torch.float8_e8m0fnu),
296 + to_blocked(ws, backend="triton").view(torch.float8_e8m0fnu),
297 + alpha,
298 + )
299 +
300 +
301 +def matmul_mxf4_bf16_fake(x, w, xs, ws, alpha):
302 + return torch.empty(*x.shape[:-1], w.shape[0], dtype=torch.bfloat16, device=x.device)
303 +
304 +
305 +direct_register_custom_op(
306 + op_name="matmul_mxf4_bf16",
307 + op_func=matmul_mxf4_bf16,
308 + mutates_args=[],
309 + fake_impl=matmul_mxf4_bf16_fake,
310 + dispatch_key=current_platform.dispatch_key,
311 +)
312 +
313 +
314 +def fused_quantize_nv(
315 + x_flat: torch.Tensor, hadamard_matrix: torch.Tensor, global_scale: torch.Tensor
316 +) -> tuple[torch.Tensor, torch.Tensor]:
317 + return fusedQuantizeNv(x_flat, hadamard_matrix, global_scale)
318 +
319 +
320 +def fused_quantize_nv_fake(x_flat, hadamard_matrix, global_scale):
321 + rows, cols = x_flat.size(0), x_flat.size(1) // 16
322 + padded_rows = ((rows + 128 - 1) // 128) * 128
323 + padded_cols = ((cols + 4 - 1) // 4) * 4
324 +
325 + xh_e2m1 = torch.empty(
326 + x_flat.size(0), x_flat.size(1) // 2, dtype=torch.uint8, device=x_flat.device
327 + )
328 + xh_e8m0 = torch.empty(
329 + padded_rows, padded_cols, dtype=torch.float8_e4m3fn, device=x_flat.device
330 + )
331 +
332 + return xh_e2m1, xh_e8m0
333 +
334 +
335 +direct_register_custom_op(
336 + op_name="fused_quantize_nv",
337 + op_func=fused_quantize_nv,
338 + mutates_args=[],
339 + fake_impl=fused_quantize_nv_fake,
340 + dispatch_key=current_platform.dispatch_key,
341 +)
342 +
343 +
344 +def matmul_nvf4_bf16(
345 + x: torch.Tensor,
346 + w: torch.Tensor,
347 + xs: torch.Tensor,
348 + ws: torch.Tensor,
349 + alpha: torch.Tensor,
350 +) -> torch.Tensor:
351 + return cutlass_scaled_fp4_mm(
352 + x,
353 + w,
354 + to_blocked(xs, backend="triton")
355 + .view(torch.float8_e4m3fn)
356 + .view(-1, x.shape[1] // 8), # *2//16
357 + to_blocked(ws, backend="triton")
358 + .view(torch.float8_e4m3fn)
359 + .view(-1, x.shape[1] // 8),
360 + alpha,
361 + torch.bfloat16,
362 + )
363 +
364 +
365 +def matmul_nvf4_bf16_fake(x, w, xs, ws, alpha):
366 + return torch.empty(*x.shape[:-1], w.shape[0], dtype=torch.bfloat16, device=x.device)
367 +
368 +
369 +direct_register_custom_op(
370 + op_name="matmul_nvf4_bf16",
371 + op_func=matmul_nvf4_bf16,
372 + mutates_args=[],
373 + fake_impl=matmul_nvf4_bf16_fake,
374 + dispatch_key=current_platform.dispatch_key,
375 +)
376 +
377 +
378 +def quantized_forward(
379 + x: torch.Tensor,
380 + qweight: torch.Tensor,
381 + weight_scales: torch.Tensor,
382 + weight_global_scale: torch.Tensor,
383 + act_global_scale: torch.Tensor,
384 + bias: Optional[torch.Tensor],
385 + forward_hadamard_matrix: torch.Tensor,
386 + forward_method: str,
387 + forward_dtype: str,
388 +) -> torch.Tensor:
389 + x_flat = x.contiguous().flatten(end_dim=-2)
390 +
391 + if forward_dtype == "mxfp4":
392 + x_flat_q, x_flat_scales = torch.ops.vllm.fused_quantize_mx(
393 + x_flat, forward_hadamard_matrix, forward_method
394 + )
395 + y = torch.ops.vllm.matmul_mxf4_bf16(
396 + x_flat_q,
397 + qweight,
398 + x_flat_scales,
399 + weight_scales,
400 + 1 / (weight_global_scale * act_global_scale),
401 + )
402 + elif forward_dtype == "nvfp4":
403 + x_flat_q, x_flat_scales = torch.ops.vllm.fused_quantize_nv(
404 + x_flat, forward_hadamard_matrix, act_global_scale
405 + )
406 + y = torch.ops.vllm.matmul_nvf4_bf16(
407 + x_flat_q,
408 + qweight,
409 + x_flat_scales,
410 + weight_scales,
411 + 1 / (weight_global_scale * act_global_scale),
412 + )
413 + else:
414 + raise ValueError(f"Unsupported forward_dtype: {forward_dtype}")
415 +
416 + y = y.view(*x.shape[:-1], y.shape[-1])
417 + if bias is not None:
418 + y += bias
419 +
420 + return y
vllm/model_executor/layers/quantization/qutlass_utils.py new
+185
@@ -0,0 +1,185 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +#
4 +# Modified by Roberto L. Castro (Roberto.LopezCastro@ist.ac.at).
5 +#
6 +# Copied from https://github.com/pytorch/ao/tree/main/torchao/prototype/mx_formats
7 +#
8 +# Copyright (c) Meta Platforms, Inc. and affiliates.
9 +# All rights reserved.
10 +
11 +# This source code is licensed under the license found in the
12 +# LICENSE file in the root directory of this source tree.
13 +
14 +from typing import Literal
15 +
16 +import torch
17 +import triton
18 +import triton.language as tl
19 +from torch.library import wrap_triton
20 +
21 +
22 +@triton.jit
23 +def triton_scale_swizzle(
24 + scale_ptr: torch.Tensor,
25 + scale_rows: int,
26 + scale_cols: int,
27 + output_ptr: torch.Tensor,
28 + input_row_stride: int,
29 + output_block_stride: int,
30 + BLOCK_ROWS: tl.constexpr,
31 + BLOCK_COLS: tl.constexpr,
32 +):
33 + """
34 + Rearranges tensor data from row-major to block-scaled swizzle format.
35 +
36 + Args:
37 + scale_ptr: Pointer to the input scale tensor
38 + scale_rows: Number of rows in the scale tensor
39 + scale_cols: Number of columns in the scale tensor
40 + output_ptr: Pointer to the output tensor
41 + input_row_stride: Stride between rows in the input tensor
42 + output_block_stride: Stride between blocks in the output tensor
43 + BLOCK_ROWS: Number of rows in a tile (compile-time constant)
44 + BLOCK_COLS: Number of columns in a tile (compile-time constant)
45 + """
46 + pid_row = tl.program_id(0)
47 + pid_col = tl.program_id(1)
48 +
49 + rows = tl.arange(0, BLOCK_ROWS)[:, None]
50 + cols = tl.arange(0, BLOCK_COLS)[None, :]
51 +
52 + # Calculate starting row and column for this tile
53 + start_row = pid_row * BLOCK_ROWS
54 + start_col = pid_col * BLOCK_COLS
55 + global_rows = start_row + rows
56 + global_cols = start_col + cols
57 +
58 + mask = (global_rows < scale_rows) & (global_cols < scale_cols)
59 +
60 + input_scales = tl.load(
61 + scale_ptr + global_rows * input_row_stride + global_cols,
62 + mask=mask,
63 + other=0.0,
64 + )
65 +
66 + r_div_32 = rows // 32
67 + r_mod_32 = rows % 32
68 +
69 + # 2) Rearrange to (32, 4, 4) then to final (32, 16) coordinates
70 + dest_indices = r_mod_32 * 16 + r_div_32 * 4 + cols
71 +
72 + # Flatten
73 + dest_indices_flat = tl.reshape(dest_indices, (BLOCK_ROWS * BLOCK_COLS))
74 + scales_flat = tl.reshape(input_scales, (BLOCK_ROWS * BLOCK_COLS))
75 +
76 + # Calculate block offset using provided output block stride
77 + LOCAL_NUMEL = BLOCK_ROWS * BLOCK_COLS
78 + block_offset = pid_col * LOCAL_NUMEL + (pid_row * output_block_stride)
79 +
80 + tl.store(
81 + output_ptr + block_offset + dest_indices_flat,
82 + scales_flat,
83 + )
84 +
85 +
86 +def triton_mx_block_rearrange(scale_tensor: torch.Tensor) -> torch.Tensor:
87 + """
88 + Rearranges an E8M0 tensor scale from row-major format to
89 + block-scaled swizzle format.
90 +
91 + This format is suitable for Tmem as described in NVIDIA documentation:
92 + https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout
93 +
94 + Args:
95 + scale_tensor: Input tensor in row-major format with 8-bit elements
96 +
97 + Returns:
98 + Rearranged tensor in block-scaled swizzle format
99 + """
100 + assert scale_tensor.element_size() == 1, (
101 + "Expected element size to be 1 byte (8 bits)"
102 + )
103 + assert scale_tensor.is_contiguous(), "Input tensor must be contiguous"
104 +
105 + rows, cols = scale_tensor.shape
106 +
107 + # Calculate blocks needed
108 + n_row_blocks = triton.cdiv(rows, 128)
109 + n_col_blocks = triton.cdiv(cols, 4)
110 + padded_rows = n_row_blocks * 128
111 + padded_cols = n_col_blocks * 4
112 +
113 + out = scale_tensor.new_empty((padded_rows, padded_cols))
114 +
115 + # Input stride (for row-major format)
116 + input_row_stride = cols
117 +
118 + # We probably want handle multiple blocks per tile but
119 + # for now keep it simple
120 + BLOCK_ROWS, BLOCK_COLS = 128, 4
121 +
122 + # Output block stride for the rearranged format
123 + output_block_stride = BLOCK_ROWS * BLOCK_COLS * (padded_cols // BLOCK_COLS)
124 +
125 + grid = lambda META: (
126 + triton.cdiv(padded_rows, BLOCK_ROWS),
127 + triton.cdiv(padded_cols, BLOCK_COLS),
128 + )
129 +
130 + wrap_triton(triton_scale_swizzle)[grid](
131 + scale_tensor.view(torch.uint8),
132 + rows,
133 + cols,
134 + out.view(torch.uint8),
135 + input_row_stride,
136 + output_block_stride,
137 + BLOCK_ROWS=BLOCK_ROWS,
138 + BLOCK_COLS=BLOCK_COLS,
139 + )
140 +
141 + return out
142 +
143 +
144 +def ceil_div(a, b):
145 + return (a + b - 1) // b
146 +
147 +
148 +def to_blocked(
149 + input_matrix: torch.Tensor, backend: Literal["torch", "triton"] = "triton"
150 +) -> torch.Tensor:
151 + """
152 + Rearrange a large matrix by breaking it into blocks and applying
153 + the rearrangement pattern.
154 +
155 + See:
156 + https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout
157 +
158 + Args:
159 + input_matrix: Input tensor of shape (H, W)
160 + backend: "torch" (PyTorch path) or "triton" (Triton kernel)
161 +
162 + Returns:
163 + Rearranged tensor of shape (32*ceil_div(H,128), 16*ceil_div(W,4))
164 + """
165 + if backend == "triton":
166 + return triton_mx_block_rearrange(input_matrix).flatten()
167 + elif backend != "torch":
168 + raise ValueError(f'backend must be "torch" or "triton", got {backend!r}')
169 +
170 + rows, cols = input_matrix.shape
171 + n_row_blocks = ceil_div(rows, 128)
172 + n_col_blocks = ceil_div(cols, 4)
173 +
174 + # Calculate the padded shape
175 + padded_rows = n_row_blocks * 128
176 + padded_cols = n_col_blocks * 4
177 +
178 + padded = input_matrix
179 + assert (rows, cols) == (padded_rows, padded_cols)
180 +
181 + # Rearrange the blocks
182 + blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3)
183 + rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16)
184 +
185 + return rearranged.flatten()