| 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 | from vllm.utils.torch_utils import set_random_seed |
| 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_e8m0: 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 | scales_dq = x_e8m0.to(torch.float64) |
| 121 | |
| 122 | x_dq = (x_fp4_dq.unflatten(dim=-1, sizes=(-1, 32)) * scales_dq[..., None]).flatten( |
| 123 | start_dim=-2 |
| 124 | ) / alpha |
| 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( |
| 141 | x: torch.Tensor, h: torch.Tensor, rot_size: int, quest: bool = True |
| 142 | ): |
| 143 | device = x.device |
| 144 | xh_ref64 = ( |
| 145 | x.unflatten(dim=-1, sizes=(-1, rot_size)).to(dtype=torch.float64) |
| 146 | @ h.reshape(rot_size, rot_size).to(dtype=torch.float64) |
| 147 | ).flatten(start_dim=-2) |
| 148 | |
| 149 | if quest: |
| 150 | scales_ref64_ = ( |
| 151 | xh_ref64.unflatten(dim=-1, sizes=(-1, 32)).std(dim=-1, correction=0) |
| 152 | * (2.92247856 / 6.0) |
| 153 | + 1e-8 |
| 154 | ) |
| 155 | else: |
| 156 | abs_max = xh_ref64.unflatten(dim=-1, sizes=(-1, 32)).abs().amax(dim=-1) |
| 157 | scales_ref64_ = abs_max + 1e-8 |
| 158 | |
| 159 | xh_e8m0_ref = scales_ref64_.log2().floor().exp2().to(dtype=torch.float8_e8m0fnu) |
| 160 | scales_ref64 = xh_e8m0_ref.to(dtype=torch.float64) |
| 161 | |
| 162 | xh_scaled_ref64 = ( |
| 163 | xh_ref64.unflatten(dim=-1, sizes=(-1, 32)) / scales_ref64[..., None] |
| 164 | ).flatten(start_dim=-2) |
| 165 | if not quest: |
| 166 | xh_scaled_ref64 *= 3 |
| 167 | |
| 168 | clip_mask_unpacked_ref = xh_scaled_ref64.abs() < 6.0 |
| 169 | clip_mask_ref = torch.zeros( |
| 170 | *x.shape[:-1], x.size(-1) // 8, dtype=torch.uint8, device=device |
| 171 | ) |
| 172 | for i in range(8): |
| 173 | clip_mask_ref |= clip_mask_unpacked_ref[..., i::8].to(dtype=torch.uint8) << i |
| 174 | |
| 175 | xh_fp4_ref, xh_e2m1_ref = _rtne_fp4(xh_scaled_ref64) |
| 176 | xh_dq, xh_fp4_dq, scales_dq = _dq_fp4( |
| 177 | xh_e2m1_ref, xh_e8m0_ref, alpha=1.0 if quest else 3.0 |
| 178 | ) |
| 179 | clip_mask_unpacked_dq = _unpack_mask(clip_mask_ref) |
| 180 | |
| 181 | assert xh_fp4_dq.equal(xh_fp4_ref) |
| 182 | assert scales_dq.equal(scales_ref64) |
| 183 | assert clip_mask_unpacked_dq.equal(clip_mask_unpacked_ref) |
| 184 | |
| 185 | return ( |
| 186 | xh_dq, |
| 187 | clip_mask_unpacked_ref, |
| 188 | (xh_e2m1_ref, xh_e8m0_ref, clip_mask_ref), |
| 189 | ) |
| 190 | |
| 191 | |
| 192 | DTYPE = torch.bfloat16 |
| 193 | DEVICE = torch.device("cuda:0") |
| 194 | |
| 195 | ROT_SIZES = [32, 64, 128] |
| 196 | SEEDS = [0] |
| 197 | BATCHES = [1, 16] |
| 198 | |
| 199 | LLAMA_MODELS = { |
| 200 | "7B": [(4096, 3 * 4096), (4096, 4096), (4096, 2 * 10752), (10752, 4096)], |
| 201 | "13B": [(5120, 3 * 5120), (5120, 5120), (5120, 2 * 13568), (13568, 5120)], |
| 202 | "33B": [(6656, 3 * 6656), (6656, 6656), (6656, 2 * 17664), (17664, 6656)], |
| 203 | "70B": [(8192, 3 * 8192), (8192, 8192), (8192, 2 * 21760), (21760, 8192)], |
| 204 | } |
| 205 | |
| 206 | |
| 207 | @pytest.fixture(autouse=True) |
| 208 | def _seed_each_test(): |
| 209 | set_random_seed(0) |
| 210 | np.random.seed(0) |
| 211 | torch.random.manual_seed(0) |
| 212 | |
| 213 | |
| 214 | @pytest.mark.parametrize("rot_size", ROT_SIZES) |
| 215 | @torch.inference_mode() |
| 216 | def test_fused_quantization_absmax(rot_size: int): |
| 217 | dtype, device = DTYPE, DEVICE |
| 218 | h = get_hadamard_matrix(rot_size, dtype, device) |
| 219 | x = torch.randn(2, 4096, 4096, dtype=dtype, device=device) * 25.0 |
| 220 | |
| 221 | xh_dq_ref, _, _ = _forward_quantize_ref(x, h, rot_size, quest=False) |
| 222 | xh_e2m1, xh_e8m0 = fusedQuantizeMx(x, h, method="abs_max") |
| 223 | xh_e8m0 = xh_e8m0.reshape(2, 4096, 4096 // 32) |
| 224 | xh_dq, *_ = _dq_fp4(xh_e2m1, xh_e8m0, alpha=3.0) |
| 225 | |
| 226 | torch.testing.assert_close(xh_dq, xh_dq_ref, rtol=0.34, atol=100) |
| 227 | assert (xh_dq != xh_dq_ref).float().mean() <= 1e-4 |
| 228 | |
| 229 | m, n, k = 1, 504, 4096 |
| 230 | a = torch.randn(m, k, dtype=dtype, device=device) * 25.0 |
| 231 | b = torch.randn(n, k, dtype=dtype, device=device) * 25.0 |
| 232 | |
| 233 | a_e2m1, a_e8m0 = fusedQuantizeMx(a, h, method="abs_max") |
| 234 | b_e2m1, b_e8m0 = fusedQuantizeMx(b, h, method="abs_max") |
| 235 | a_dq, *_ = _dq_fp4(a_e2m1, a_e8m0[:m, :k], alpha=1.0) |
| 236 | b_dq, *_ = _dq_fp4(b_e2m1, b_e8m0[:n, :k], alpha=1.0) |
| 237 | out_ref = a_dq @ b_dq.transpose(-2, -1) |
| 238 | |
| 239 | a_scale_block = to_blocked(a_e8m0, backend="triton") |
| 240 | b_scale_block = to_blocked(b_e8m0, backend="triton") |
| 241 | alpha = torch.tensor([1.0], device=device) |
| 242 | out = matmul_mxf4_bf16_tn(a_e2m1, b_e2m1, a_scale_block, b_scale_block, alpha) |
| 243 | assert out.equal(out_ref.to(dtype=out.dtype)) |
| 244 | |
| 245 | |
| 246 | @pytest.mark.parametrize("rot_size", ROT_SIZES) |
| 247 | @torch.inference_mode() |
| 248 | def test_fused_quantization_quest(rot_size: int): |
| 249 | dtype, device = DTYPE, DEVICE |
| 250 | h = get_hadamard_matrix(rot_size, dtype, device) |
| 251 | x = torch.randn(2, 4096, 4096, dtype=dtype, device=device) * 25.0 |
| 252 | |
| 253 | xh_dq_ref, _, _ = _forward_quantize_ref(x, h, rot_size, quest=True) |
| 254 | xh_e2m1, xh_e8m0 = fusedQuantizeMx(x, h, method="quest") |
| 255 | xh_e8m0 = xh_e8m0.reshape(2, 4096, 4096 // 32) |
| 256 | xh_dq, *_ = _dq_fp4(xh_e2m1, xh_e8m0, alpha=1.0) |
| 257 | |
| 258 | torch.testing.assert_close(xh_dq, xh_dq_ref, rtol=0.34, atol=100) |
| 259 | assert (xh_dq != xh_dq_ref).float().mean() <= 1e-4 |
| 260 | |
| 261 | m, n, k = 504, 504, 2048 |
| 262 | a = torch.randn(m, k, dtype=dtype, device=device) * 25.0 |
| 263 | b = torch.randn(n, k, dtype=dtype, device=device) * 25.0 |
| 264 | |
| 265 | a_e2m1, a_e8m0 = fusedQuantizeMx(a, h, method="quest") |
| 266 | b_e2m1, b_e8m0 = fusedQuantizeMx(b, h, method="quest") |
| 267 | a_dq, *_ = _dq_fp4(a_e2m1, a_e8m0[:m, :k], alpha=1.0) |
| 268 | b_dq, *_ = _dq_fp4(b_e2m1, b_e8m0[:n, :k], alpha=1.0) |
| 269 | out_ref = a_dq @ b_dq.transpose(-2, -1) |
| 270 | |
| 271 | a_scale_block = to_blocked(a_e8m0, backend="triton") |
| 272 | b_scale_block = to_blocked(b_e8m0, backend="triton") |
| 273 | alpha = torch.tensor([1.0], device=device) |
| 274 | out = matmul_mxf4_bf16_tn(a_e2m1, b_e2m1, a_scale_block, b_scale_block, alpha) |
| 275 | assert out.equal(out_ref.to(dtype=out.dtype)) |
| 276 | |
| 277 | |
| 278 | @pytest.mark.parametrize("model", list(LLAMA_MODELS.keys())) |
| 279 | @pytest.mark.parametrize("layer_idx", [0, 1, 2, 3]) |
| 280 | @pytest.mark.parametrize("batch", [1, 16]) |
| 281 | @pytest.mark.parametrize("had_size", ROT_SIZES) |
| 282 | @torch.inference_mode() |
| 283 | def test_llama_shapes(model: str, layer_idx: int, batch: int, had_size: int): |
| 284 | dtype, device = DTYPE, DEVICE |
| 285 | m = batch |
| 286 | k, n = LLAMA_MODELS[model][layer_idx] |
| 287 | |
| 288 | h = get_hadamard_matrix(had_size, dtype, device) |
| 289 | |
| 290 | a = torch.rand(m, k, dtype=dtype, device=device) * 25.0 |
| 291 | b = torch.rand(n, k, dtype=dtype, device=device) * 25.0 |
| 292 | |
| 293 | a_e2m1, a_e8m0 = fusedQuantizeMx(a, h, method="quest") |
| 294 | b_e2m1, b_e8m0 = fusedQuantizeMx(b, h, method="quest") |
| 295 | |
| 296 | a_dq, *_ = _dq_fp4(a_e2m1, a_e8m0[:m, :k], alpha=1.0) |
| 297 | b_dq, *_ = _dq_fp4(b_e2m1, b_e8m0[:n, :k], alpha=1.0) |
| 298 | out_ref = a_dq @ b_dq.transpose(-2, -1) |
| 299 | |
| 300 | a_scale_block = to_blocked(a_e8m0, backend="triton") |
| 301 | b_scale_block = to_blocked(b_e8m0, backend="triton") |
| 302 | alpha = torch.tensor([1.0], device=device) |
| 303 | out = matmul_mxf4_bf16_tn(a_e2m1, b_e2m1, a_scale_block, b_scale_block, alpha) |
| 304 | assert out.equal(out_ref.to(dtype=out.dtype)) |