| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import pytest |
| 4 | import torch |
| 5 | |
| 6 | from vllm import _custom_ops as ops |
| 7 | from vllm.platforms import current_platform |
| 8 | from vllm.scalar_type import scalar_types |
| 9 | from vllm.utils.torch_utils import set_random_seed |
| 10 | |
| 11 | if not current_platform.has_device_capability(100): |
| 12 | pytest.skip( |
| 13 | reason="Nvfp4 Requires compute capability of 10 or above.", |
| 14 | allow_module_level=True, |
| 15 | ) |
| 16 | |
| 17 | DTYPES = [torch.float16, torch.bfloat16] |
| 18 | SHAPES = [(128, 64), (128, 128), (256, 64), (256, 128)] |
| 19 | PAD_SHAPES = [ |
| 20 | (90, 64), |
| 21 | (150, 64), |
| 22 | (128, 48), |
| 23 | (128, 80), |
| 24 | (150, 80), |
| 25 | (90, 48), |
| 26 | (90, 128), |
| 27 | (150, 128), |
| 28 | (150, 48), |
| 29 | (90, 80), |
| 30 | (128, 512), |
| 31 | (128, 1024), |
| 32 | (128, 2048), |
| 33 | (64, 7168), |
| 34 | (64, 7152), |
| 35 | (32, 14336), |
| 36 | ] |
| 37 | PADDED_OUTPUT_SHAPES = [(128, 48), (128, 80), (150, 48), (150, 80), (64, 7152)] |
| 38 | SEEDS = [42] |
| 39 | CUDA_DEVICES = ["cuda:0"] |
| 40 | |
| 41 | FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max() |
| 42 | FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max |
| 43 | |
| 44 | # E2M1 to float |
| 45 | # 0111 -> 6 |
| 46 | # 0110 -> 4 |
| 47 | # 0101 -> 3 |
| 48 | # 0100 -> 2 |
| 49 | # 0011 -> 1.5 |
| 50 | # 0010 -> 1 |
| 51 | # 0001 -> 0.5 |
| 52 | # 0000 -> 0 |
| 53 | E2M1_TO_FLOAT32 = [ |
| 54 | 0.0, |
| 55 | 0.5, |
| 56 | 1.0, |
| 57 | 1.5, |
| 58 | 2.0, |
| 59 | 3.0, |
| 60 | 4.0, |
| 61 | 6.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 | BLOCK_SIZE = 16 |
| 72 | |
| 73 | |
| 74 | def cast_from_fp4(x, m, n): |
| 75 | # The fp4 values are packed in uint8 as [v_1st | v_2nd] |
| 76 | v_2nd = x & 0xF |
| 77 | v_1st = (x >> 4) & 0xF |
| 78 | c = torch.stack((v_2nd, v_1st), dim=-1) |
| 79 | out = torch.tensor([E2M1_TO_FLOAT32[x] for x in c.flatten()]) |
| 80 | out = out.reshape(m, n).to(torch.float32) |
| 81 | return out |
| 82 | |
| 83 | |
| 84 | def cast_to_fp4(x): |
| 85 | sign = torch.sign(x) |
| 86 | x = torch.abs(x) |
| 87 | x[(x >= 0.0) & (x <= 0.25)] = 0.0 |
| 88 | x[(x > 0.25) & (x < 0.75)] = 0.5 |
| 89 | x[(x >= 0.75) & (x <= 1.25)] = 1.0 |
| 90 | x[(x > 1.25) & (x < 1.75)] = 1.5 |
| 91 | x[(x >= 1.75) & (x <= 2.5)] = 2.0 |
| 92 | x[(x > 2.5) & (x < 3.5)] = 3.0 |
| 93 | x[(x >= 3.5) & (x <= 5.0)] = 4.0 |
| 94 | x[x > 5.0] = 6.0 |
| 95 | return x * sign |
| 96 | |
| 97 | |
| 98 | def get_reciprocal(x): |
| 99 | if isinstance(x, torch.Tensor): |
| 100 | return torch.where(x == 0, torch.tensor(0.0, dtype=x.dtype), 1.0 / x) |
| 101 | elif isinstance(x, (float, int)): |
| 102 | return 0.0 if x == 0 else 1.0 / x |
| 103 | else: |
| 104 | raise TypeError("Input must be a float, int, or a torch.Tensor.") |
| 105 | |
| 106 | |
| 107 | def ref_nvfp4_quant(x, global_scale): |
| 108 | assert global_scale.dtype == torch.float32 |
| 109 | assert x.ndim == 2 |
| 110 | m, n = x.shape |
| 111 | x = torch.reshape(x, (m, n // BLOCK_SIZE, BLOCK_SIZE)) |
| 112 | vec_max = torch.max(torch.abs(x), dim=-1, keepdim=True)[0].to(torch.float32) |
| 113 | scale = global_scale * (vec_max * get_reciprocal(FLOAT4_E2M1_MAX)) |
| 114 | scale = scale.to(torch.float8_e4m3fn).to(torch.float32) |
| 115 | output_scale = get_reciprocal(scale * get_reciprocal(global_scale)) |
| 116 | |
| 117 | scaled_x = x.to(torch.float32) * output_scale |
| 118 | clipped_x = torch.clamp(scaled_x, -6.0, 6.0).reshape(m, n) |
| 119 | return cast_to_fp4(clipped_x), scale.squeeze(-1) |
| 120 | |
| 121 | |
| 122 | def recover_swizzled_scales(scale, m, n, include_padding=False): |
| 123 | round_up = lambda x, y: (x + y - 1) // y * y |
| 124 | rounded_m = round_up(m, 128) |
| 125 | scale_n = n // BLOCK_SIZE |
| 126 | rounded_n = round_up(scale_n, 4) |
| 127 | # Recover the swizzled scaling factor to linear layout |
| 128 | tmp = torch.reshape(scale, (1, rounded_m // 128, rounded_n // 4, 32, 4, 4)) |
| 129 | tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5)) |
| 130 | result = torch.reshape(tmp, (rounded_m, rounded_n)).to(torch.float32) |
| 131 | if include_padding: |
| 132 | return result |
| 133 | return result[:m, :scale_n] |
| 134 | |
| 135 | |
| 136 | def round_up(x: int, y: int) -> int: |
| 137 | return (x + y - 1) // y * y |
| 138 | |
| 139 | |
| 140 | @pytest.mark.parametrize("dtype", DTYPES) |
| 141 | @pytest.mark.parametrize("shape", SHAPES) |
| 142 | @pytest.mark.parametrize("seed", SEEDS) |
| 143 | @pytest.mark.parametrize("device", CUDA_DEVICES) |
| 144 | @torch.inference_mode() |
| 145 | def test_quantize_to_fp4( |
| 146 | dtype: torch.dtype, |
| 147 | shape: tuple[int, int], |
| 148 | seed: int, |
| 149 | device: str, |
| 150 | ) -> None: |
| 151 | set_random_seed(seed) |
| 152 | torch.set_default_device(device) |
| 153 | |
| 154 | m, n = shape |
| 155 | |
| 156 | x = torch.randn((m, n), dtype=dtype) |
| 157 | tensor_amax = torch.abs(x).max().to(torch.float32) |
| 158 | global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax |
| 159 | out_ref, scale_ref = ref_nvfp4_quant(x, global_scale) |
| 160 | |
| 161 | out, out_scale = ops.scaled_fp4_quant(x, global_scale) |
| 162 | scale_ans = recover_swizzled_scales(out_scale, m, n) |
| 163 | out_ans = cast_from_fp4(out, m, n) |
| 164 | |
| 165 | torch.testing.assert_close(out_ans, out_ref) |
| 166 | torch.testing.assert_close(scale_ans, scale_ref) |
| 167 | |
| 168 | |
| 169 | @pytest.mark.parametrize( |
| 170 | "shape", |
| 171 | [(32, 4096), (128, 4096), (1, 64), (127, 1024), (256, 16384)], |
| 172 | ) |
| 173 | @pytest.mark.parametrize("is_sf_swizzled_layout", [True, False]) |
| 174 | @torch.inference_mode() |
| 175 | def test_python_util_matches_cpp_allocation( |
| 176 | shape: tuple[int, int], |
| 177 | is_sf_swizzled_layout: bool, |
| 178 | ) -> None: |
| 179 | """Verify that the Python utility (create_fp4_output_tensors) allocates |
| 180 | tensors with the same shapes and dtypes as the C++ functional variant |
| 181 | (scaled_fp4_quant_func). |
| 182 | """ |
| 183 | from vllm._custom_ops import create_fp4_output_tensors |
| 184 | |
| 185 | torch.set_default_device("cuda:0") |
| 186 | m, n = shape |
| 187 | input_tensor = torch.randn((m, n), dtype=torch.bfloat16) |
| 188 | input_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda:0") |
| 189 | |
| 190 | # C++ functional variant allocates internally |
| 191 | cpp_out, cpp_scale = torch.ops._C.scaled_fp4_quant( |
| 192 | input_tensor, input_scale, is_sf_swizzled_layout |
| 193 | ) |
| 194 | |
| 195 | # Python utility |
| 196 | py_out, py_scale = create_fp4_output_tensors( |
| 197 | m, n, torch.device("cuda:0"), is_sf_swizzled_layout |
| 198 | ) |
| 199 | |
| 200 | assert py_out.shape == cpp_out.shape, ( |
| 201 | f"Output shape mismatch: Python {py_out.shape} vs C++ {cpp_out.shape}" |
| 202 | ) |
| 203 | assert py_out.dtype == cpp_out.dtype, ( |
| 204 | f"Output dtype mismatch: Python {py_out.dtype} vs C++ {cpp_out.dtype}" |
| 205 | ) |
| 206 | assert py_scale.shape == cpp_scale.shape, ( |
| 207 | f"Scale shape mismatch: Python {py_scale.shape} vs C++ {cpp_scale.shape}" |
| 208 | ) |
| 209 | assert py_scale.dtype == cpp_scale.dtype, ( |
| 210 | f"Scale dtype mismatch: Python {py_scale.dtype} vs C++ {cpp_scale.dtype}" |
| 211 | ) |
| 212 | |
| 213 | |
| 214 | @pytest.mark.parametrize("shape", PADDED_OUTPUT_SHAPES) |
| 215 | @pytest.mark.parametrize("is_sf_swizzled_layout", [True, False]) |
| 216 | @torch.inference_mode() |
| 217 | def test_quantize_to_fp4_with_padded_output( |
| 218 | shape: tuple[int, int], |
| 219 | is_sf_swizzled_layout: bool, |
| 220 | ) -> None: |
| 221 | from vllm._custom_ops import create_fp4_output_tensors |
| 222 | |
| 223 | dtype = torch.float16 |
| 224 | set_random_seed(42) |
| 225 | torch.set_default_device("cuda:0") |
| 226 | |
| 227 | m, n = shape |
| 228 | padded_n = round_up(n, 32) |
| 229 | assert padded_n > n |
| 230 | |
| 231 | x = torch.randn((m, n), dtype=dtype) |
| 232 | tensor_amax = torch.abs(x).max().to(torch.float32) |
| 233 | global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax |
| 234 | out_ref, scale_ref = ref_nvfp4_quant(x, global_scale) |
| 235 | |
| 236 | out, out_scale = ops.scaled_fp4_quant( |
| 237 | x, |
| 238 | global_scale, |
| 239 | is_sf_swizzled_layout=is_sf_swizzled_layout, |
| 240 | padded_n=padded_n, |
| 241 | ) |
| 242 | py_out, py_scale = create_fp4_output_tensors( |
| 243 | m, |
| 244 | n, |
| 245 | torch.device("cuda:0"), |
| 246 | is_sf_swizzled_layout, |
| 247 | padded_n=padded_n, |
| 248 | ) |
| 249 | |
| 250 | assert out.shape == (m, padded_n // 2) |
| 251 | assert out.shape == py_out.shape |
| 252 | assert out_scale.shape == py_scale.view(torch.float8_e4m3fn).shape |
| 253 | |
| 254 | out_ans = cast_from_fp4(out[:, : n // 2], m, n) |
| 255 | torch.testing.assert_close(out_ans, out_ref) |
| 256 | assert torch.count_nonzero(out[:, n // 2 :]) == 0 |
| 257 | |
| 258 | if is_sf_swizzled_layout: |
| 259 | scale_ans = recover_swizzled_scales(out_scale, m, padded_n) |
| 260 | torch.testing.assert_close(scale_ans[:, : n // BLOCK_SIZE], scale_ref) |
| 261 | assert torch.count_nonzero(scale_ans[:, n // BLOCK_SIZE :]) == 0 |
| 262 | else: |
| 263 | scale_ans = out_scale.to(torch.float32) |
| 264 | torch.testing.assert_close(scale_ans[:, : n // BLOCK_SIZE], scale_ref) |
| 265 | assert torch.count_nonzero(scale_ans[:, n // BLOCK_SIZE :]) == 0 |
| 266 | |
| 267 | |
| 268 | @pytest.mark.parametrize( |
| 269 | ("shape", "padded_n"), |
| 270 | [((90, 80), None), ((90, 48), 64)], |
| 271 | ) |
| 272 | @torch.inference_mode() |
| 273 | def test_swizzled_scale_padding_initialized_in_kernel( |
| 274 | shape: tuple[int, int], padded_n: int | None |
| 275 | ) -> None: |
| 276 | """The out variant must overwrite dirty scale padding without a memset.""" |
| 277 | torch.set_default_device("cuda:0") |
| 278 | m, n = shape |
| 279 | physical_n = padded_n if padded_n is not None else n |
| 280 | rounded_m = round_up(m, 128) |
| 281 | rounded_scale_n = round_up(physical_n // BLOCK_SIZE, 4) |
| 282 | |
| 283 | x = torch.randn((m, n), dtype=torch.float16) |
| 284 | global_scale = torch.tensor(1.0, dtype=torch.float32) |
| 285 | out = torch.full((m, physical_n // 2), 0xFF, dtype=torch.uint8) |
| 286 | out_scale = torch.full( |
| 287 | (rounded_m, rounded_scale_n // 4), 0x7F7F7F7F, dtype=torch.int32 |
| 288 | ) |
| 289 | |
| 290 | torch.ops._C.scaled_fp4_quant.out( |
| 291 | x, |
| 292 | global_scale, |
| 293 | True, |
| 294 | output=out, |
| 295 | output_scale=out_scale, |
| 296 | ) |
| 297 | |
| 298 | full_scale = recover_swizzled_scales( |
| 299 | out_scale.view(torch.float8_e4m3fn), m, physical_n, include_padding=True |
| 300 | ) |
| 301 | assert torch.count_nonzero(full_scale[m:, :]) == 0 |
| 302 | assert torch.count_nonzero(full_scale[:m, n // BLOCK_SIZE :]) == 0 |
| 303 | |
| 304 | |
| 305 | @pytest.mark.parametrize("pad_shape", PAD_SHAPES) |
| 306 | @torch.inference_mode() |
| 307 | def test_quantize_to_fp4_padded(pad_shape: tuple[int, int]) -> None: |
| 308 | dtype = torch.float16 |
| 309 | set_random_seed(42) |
| 310 | torch.set_default_device("cuda:0") |
| 311 | |
| 312 | m, n = pad_shape |
| 313 | |
| 314 | x = torch.randn((m, n), dtype=dtype) |
| 315 | |
| 316 | tensor_amax = torch.abs(x).max().to(torch.float32) |
| 317 | global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax |
| 318 | out_ref, scale_ref = ref_nvfp4_quant(x, global_scale) |
| 319 | |
| 320 | out, out_scale = ops.scaled_fp4_quant(x, global_scale) |
| 321 | scale_ans = recover_swizzled_scales(out_scale, m, n) |
| 322 | out_ans = cast_from_fp4(out, m, n) |
| 323 | torch.testing.assert_close(out_ans, out_ref) |
| 324 | torch.testing.assert_close(scale_ans, scale_ref) |
| 325 | |
| 326 | |
| 327 | @pytest.mark.parametrize("pad_shape", PAD_SHAPES) |
| 328 | @torch.inference_mode() |
| 329 | def test_quantize_to_fp4_padded_no_sf_swizzled(pad_shape: tuple[int, int]) -> None: |
| 330 | dtype = torch.float16 |
| 331 | set_random_seed(42) |
| 332 | torch.set_default_device("cuda:0") |
| 333 | |
| 334 | m, n = pad_shape |
| 335 | |
| 336 | x = torch.randn((m, n), dtype=dtype) |
| 337 | |
| 338 | tensor_amax = torch.abs(x).max().to(torch.float32) |
| 339 | global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax |
| 340 | out_ref, scale_ref = ref_nvfp4_quant(x, global_scale) |
| 341 | |
| 342 | out, out_scale = ops.scaled_fp4_quant(x, global_scale, is_sf_swizzled_layout=False) |
| 343 | scale_ans = out_scale.to(torch.float32) |
| 344 | out_ans = cast_from_fp4(out, m, n) |
| 345 | torch.testing.assert_close(out_ans, out_ref) |
| 346 | torch.testing.assert_close(scale_ans, scale_ref) |