| 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 | from vllm.utils.torch_utils import set_random_seed |
| 29 | |
| 30 | if not torch.cuda.is_available(): |
| 31 | pytest.skip("CUDA required for these tests.", allow_module_level=True) |
| 32 | |
| 33 | if not ( |
| 34 | current_platform.has_device_capability(100) |
| 35 | or current_platform.has_device_capability(120) |
| 36 | ): |
| 37 | pytest.skip( |
| 38 | reason="Tests require compute capability 10.0 (100) or 12.0 (120).", |
| 39 | allow_module_level=True, |
| 40 | ) |
| 41 | |
| 42 | |
| 43 | # ----- Helpers ----- |
| 44 | def get_hadamard_matrix(group_size: int, dtype: torch.dtype, device: torch.device): |
| 45 | return ( |
| 46 | deterministic_hadamard_matrix(group_size, dtype=dtype, device=device) |
| 47 | * group_size**-0.5 |
| 48 | ) |
| 49 | |
| 50 | |
| 51 | def _rtne_fp4(x: torch.Tensor): |
| 52 | device = x.device |
| 53 | grid = torch.tensor( |
| 54 | [ |
| 55 | -6.0, |
| 56 | -4.0, |
| 57 | -3.0, |
| 58 | -2.0, |
| 59 | -1.5, |
| 60 | -1.0, |
| 61 | -0.5, |
| 62 | -0.0, |
| 63 | 0.0, |
| 64 | 0.5, |
| 65 | 1.0, |
| 66 | 1.5, |
| 67 | 2.0, |
| 68 | 3.0, |
| 69 | 4.0, |
| 70 | 6.0, |
| 71 | ], |
| 72 | dtype=x.dtype, |
| 73 | device=x.device, |
| 74 | ) |
| 75 | grid_int = torch.tensor( |
| 76 | [-1, -2, -3, -4, -5, -6, -7, -8, 0, 1, 2, 3, 4, 5, 6, 7], |
| 77 | dtype=torch.uint8, |
| 78 | device=device, |
| 79 | ) |
| 80 | inds = torch.bucketize(x, grid) |
| 81 | lo, hi = (inds - 1).clamp(min=0, max=15), inds.clamp(min=0, max=15) |
| 82 | g_lo, g_hi = grid[lo], grid[hi] |
| 83 | pick_hi = (g_hi - x < x - g_lo) | (g_hi - x == x - g_lo) & (grid_int[hi] % 2 == 0) |
| 84 | y = torch.where(pick_hi, g_hi, g_lo) |
| 85 | y_int = torch.where(pick_hi, grid_int[hi], grid_int[lo]) |
| 86 | y_int_packed = (y_int[..., 1::2] & 0xF) << 4 | y_int[..., ::2] & 0xF |
| 87 | return y, y_int_packed |
| 88 | |
| 89 | |
| 90 | def _dq_fp4(x_e2m1: torch.Tensor, x_e4m3: torch.Tensor, alpha: float): |
| 91 | device = x_e2m1.device |
| 92 | |
| 93 | x_e2m1_i32 = x_e2m1.view(dtype=torch.uint8).to(dtype=torch.int32) |
| 94 | x_e2m1_unpacked = torch.stack( |
| 95 | [x_e2m1_i32 & 0xF, (x_e2m1_i32 >> 4) & 0xF], dim=-1 |
| 96 | ).flatten(start_dim=-2) |
| 97 | |
| 98 | grid_dq = torch.tensor( |
| 99 | [ |
| 100 | 0.0, |
| 101 | 0.5, |
| 102 | 1.0, |
| 103 | 1.5, |
| 104 | 2.0, |
| 105 | 3.0, |
| 106 | 4.0, |
| 107 | 6.0, |
| 108 | -0.0, |
| 109 | -0.5, |
| 110 | -1.0, |
| 111 | -1.5, |
| 112 | -2.0, |
| 113 | -3.0, |
| 114 | -4.0, |
| 115 | -6.0, |
| 116 | ], |
| 117 | dtype=torch.float64, |
| 118 | device=device, |
| 119 | ) |
| 120 | x_fp4_dq = grid_dq[x_e2m1_unpacked] |
| 121 | |
| 122 | scales_dq = x_e4m3.to(torch.float64) |
| 123 | x_dq = (x_fp4_dq.unflatten(dim=-1, sizes=(-1, 16)) * scales_dq[..., None]).flatten( |
| 124 | start_dim=-2 |
| 125 | ) / alpha # * (4. / 3.) |
| 126 | return x_dq, x_fp4_dq, scales_dq |
| 127 | |
| 128 | |
| 129 | def _unpack_mask(clip_mask: torch.Tensor) -> torch.Tensor: |
| 130 | clip_mask_unpacked_dq = torch.zeros( |
| 131 | *clip_mask.shape[:-1], |
| 132 | clip_mask.size(-1) * 8, |
| 133 | dtype=torch.bool, |
| 134 | device=clip_mask.device, |
| 135 | ) |
| 136 | for i in range(8): |
| 137 | clip_mask_unpacked_dq[..., i::8] = (clip_mask >> i) & 1 |
| 138 | return clip_mask_unpacked_dq |
| 139 | |
| 140 | |
| 141 | def _forward_quantize_ref(x: torch.Tensor, h: torch.Tensor, rot_size: int): |
| 142 | device = x.device |
| 143 | |
| 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 | abs_max = xh_ref64.unflatten(dim=-1, sizes=(-1, 16)).abs().amax(dim=-1) |
| 150 | scales_ref64_ = abs_max + 1e-8 |
| 151 | |
| 152 | xh_e4m3_ref = scales_ref64_.to(dtype=torch.float8_e4m3fn) |
| 153 | scales_ref64 = xh_e4m3_ref.to(dtype=torch.float64) |
| 154 | xh_scaled_ref64 = ( |
| 155 | xh_ref64.unflatten(dim=-1, sizes=(-1, 16)) / scales_ref64[..., None] |
| 156 | ).flatten(start_dim=-2) |
| 157 | |
| 158 | xh_scaled_ref64 *= 6.0 |
| 159 | |
| 160 | clip_mask_unpacked_ref = xh_scaled_ref64.abs() < 6.0 |
| 161 | clip_mask_ref = torch.zeros( |
| 162 | *x.shape[:-1], x.size(-1) // 8, dtype=torch.uint8, device=device |
| 163 | ) |
| 164 | for i in range(8): |
| 165 | clip_mask_ref |= clip_mask_unpacked_ref[..., i::8].to(dtype=torch.uint8) << i |
| 166 | |
| 167 | xh_fp4_ref, xh_e2m1_ref = _rtne_fp4(xh_scaled_ref64) |
| 168 | xh_dq, xh_fp4_dq, scales_dq = _dq_fp4(xh_e2m1_ref, xh_e4m3_ref, 6.0) |
| 169 | clip_mask_unpacked_dq = _unpack_mask(clip_mask_ref) |
| 170 | |
| 171 | assert xh_fp4_dq.equal(xh_fp4_ref) |
| 172 | assert scales_dq.equal(scales_ref64) |
| 173 | assert clip_mask_unpacked_dq.equal(clip_mask_unpacked_ref) |
| 174 | |
| 175 | return ( |
| 176 | xh_dq, |
| 177 | clip_mask_unpacked_ref, |
| 178 | (xh_e2m1_ref, xh_e4m3_ref, clip_mask_ref), |
| 179 | ) |
| 180 | |
| 181 | |
| 182 | DTYPE = torch.bfloat16 |
| 183 | DEVICE = torch.device("cuda:0") |
| 184 | ROT_SIZES = [16, 32, 64, 128] |
| 185 | GLOBAL_SCALES = [6.0] |
| 186 | |
| 187 | LLAMA_MODELS = { |
| 188 | "7B": [(4096, 3 * 4096), (4096, 4096), (4096, 2 * 10752), (10752, 4096)], |
| 189 | "13B": [(5120, 3 * 5120), (5120, 5120), (5120, 2 * 13568), (13568, 5120)], |
| 190 | "33B": [(6656, 3 * 6656), (6656, 6656), (6656, 2 * 17664), (17664, 6656)], |
| 191 | "70B": [(8192, 3 * 8192), (8192, 8192), (8192, 2 * 21760), (21760, 8192)], |
| 192 | } |
| 193 | |
| 194 | |
| 195 | @pytest.fixture(autouse=True) |
| 196 | def _seed_each_test(): |
| 197 | set_random_seed(0) |
| 198 | np.random.seed(0) |
| 199 | torch.random.manual_seed(0) |
| 200 | |
| 201 | |
| 202 | @pytest.mark.parametrize("rot_size", ROT_SIZES) |
| 203 | @pytest.mark.parametrize("global_scale_value", GLOBAL_SCALES) |
| 204 | @torch.inference_mode() |
| 205 | def test_fused_quantization(rot_size: int, global_scale_value: float): |
| 206 | dtype, device = DTYPE, DEVICE |
| 207 | h = get_hadamard_matrix(rot_size, dtype, device) |
| 208 | x = torch.randn(2, 4096, 4096, dtype=dtype, device=device) * 25.0 |
| 209 | global_scale = torch.tensor([global_scale_value], device=device) |
| 210 | |
| 211 | xh_dq_ref, _, _ = _forward_quantize_ref(x, h, rot_size) |
| 212 | xh_e2m1, xh_e4m3 = fusedQuantizeNv(x, h, global_scale) |
| 213 | xh_e4m3 = xh_e4m3.reshape(2, 4096, 4096 // 16) |
| 214 | xh_dq, *_ = _dq_fp4(xh_e2m1, xh_e4m3, alpha=global_scale_value) |
| 215 | |
| 216 | torch.testing.assert_close(xh_dq, xh_dq_ref, rtol=0.34, atol=100) |
| 217 | assert (xh_dq != xh_dq_ref).float().mean() <= 1e-1 |
| 218 | |
| 219 | m, n, k = 504, 4096 * 2, 4096 |
| 220 | a = torch.randn(m, k, dtype=dtype, device=device) * 25.0 |
| 221 | b = torch.randn(n, k, dtype=dtype, device=device) * 25.0 |
| 222 | |
| 223 | a_e2m1, a_e4m3 = fusedQuantizeNv(a, h, global_scale) |
| 224 | b_e2m1, b_e4m3 = fusedQuantizeNv(b, h, global_scale) |
| 225 | |
| 226 | a_dq, *_ = _dq_fp4(a_e2m1, a_e4m3[:m, :k], alpha=1.0) |
| 227 | b_dq, *_ = _dq_fp4(b_e2m1, b_e4m3[:n, :k], alpha=1.0) |
| 228 | out_ref = a_dq @ b_dq.transpose(-2, -1) |
| 229 | |
| 230 | a_scale_block = to_blocked(a_e4m3, backend="triton").view(-1, k // 16) |
| 231 | b_scale_block = to_blocked(b_e4m3, backend="triton").view(-1, k // 16) |
| 232 | alpha = torch.tensor([1.0], device=device) |
| 233 | out = ops.cutlass_scaled_fp4_mm( |
| 234 | a_e2m1, b_e2m1, a_scale_block, b_scale_block, alpha, torch.bfloat16 |
| 235 | ) |
| 236 | assert out.equal(out_ref.to(dtype=out.dtype)) |
| 237 | |
| 238 | |
| 239 | @pytest.mark.parametrize("model", list(LLAMA_MODELS.keys())) |
| 240 | @pytest.mark.parametrize("layer_idx", [0, 1, 2, 3]) |
| 241 | @pytest.mark.parametrize("batch", [1, 16]) |
| 242 | @pytest.mark.parametrize("rot_size", ROT_SIZES) |
| 243 | @torch.inference_mode() |
| 244 | def test_llama_shapes(model: str, layer_idx: int, batch: int, rot_size: int): |
| 245 | dtype, device = DTYPE, DEVICE |
| 246 | m = batch |
| 247 | k, n = LLAMA_MODELS[model][layer_idx] |
| 248 | |
| 249 | h = get_hadamard_matrix(rot_size, dtype, device) |
| 250 | |
| 251 | a = torch.randn(m, k, dtype=dtype, device=device) * 25.0 |
| 252 | b = torch.randn(n, k, dtype=dtype, device=device) * 25.0 |
| 253 | |
| 254 | global_scale = torch.tensor([1.0], device=device) |
| 255 | |
| 256 | a_e2m1, a_e4m3 = fusedQuantizeNv(a, h, global_scale) |
| 257 | b_e2m1, b_e4m3 = fusedQuantizeNv(b, h, global_scale) |
| 258 | |
| 259 | a_dq, *_ = _dq_fp4(a_e2m1, a_e4m3[:m, :k], alpha=1.0) |
| 260 | b_dq, *_ = _dq_fp4(b_e2m1, b_e4m3[:n, :k], alpha=1.0) |
| 261 | out_ref = a_dq @ b_dq.transpose(-2, -1) |
| 262 | |
| 263 | a_scale_block = to_blocked(a_e4m3, backend="triton").view(-1, k // 16) |
| 264 | b_scale_block = to_blocked(b_e4m3, backend="triton").view(-1, k // 16) |
| 265 | alpha = torch.tensor([1.0], device=device) |
| 266 | out = ops.cutlass_scaled_fp4_mm( |
| 267 | a_e2m1, b_e2m1, a_scale_block, b_scale_block, alpha, torch.bfloat16 |
| 268 | ) |
| 269 | assert out.equal(out_ref.to(dtype=out.dtype)) |