| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from unittest.mock import patch |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | from vllm.model_executor.layers.quantization.utils import fp8_utils, int8_utils |
| 9 | from vllm.model_executor.layers.quantization.utils.quant_utils import get_fp8_min_max |
| 10 | from vllm.platforms import current_platform |
| 11 | |
| 12 | |
| 13 | @pytest.mark.parametrize( |
| 14 | "shape", |
| 15 | [ |
| 16 | (0, 128), |
| 17 | (1, 128), |
| 18 | (3, 384), |
| 19 | (5, 384), |
| 20 | (7, 384), |
| 21 | (15, 384), |
| 22 | (31, 128), |
| 23 | (32, 128), |
| 24 | (63, 256), |
| 25 | (64, 256), |
| 26 | (16, 512), |
| 27 | (3, 2688), |
| 28 | (3, 5376), |
| 29 | (333, 128), |
| 30 | (512, 6144), |
| 31 | ], |
| 32 | ) |
| 33 | @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) |
| 34 | @pytest.mark.parametrize("column_major", [False, True]) |
| 35 | @pytest.mark.parametrize("tma_aligned", [False, True]) |
| 36 | @pytest.mark.parametrize("scale_ue8m0", [False, True]) |
| 37 | @pytest.mark.parametrize("group_size", [64, 128]) |
| 38 | @pytest.mark.skipif( |
| 39 | not (current_platform.is_cuda_alike() or current_platform.is_xpu()), |
| 40 | reason="Only test on CUDA/ROCm/XPU.", |
| 41 | ) |
| 42 | def test_per_token_group_quant_fp8( |
| 43 | shape, |
| 44 | dtype: torch.dtype, |
| 45 | column_major: bool, |
| 46 | tma_aligned: bool, |
| 47 | scale_ue8m0: bool, |
| 48 | group_size: int, |
| 49 | ): |
| 50 | """Tail groups preserve FP8 values and scales in every supported layout.""" |
| 51 | device = current_platform.device_type |
| 52 | |
| 53 | torch.manual_seed(42) |
| 54 | num_tokens, hidden_dim = shape |
| 55 | |
| 56 | x = torch.randn((num_tokens, hidden_dim), device=device, dtype=dtype) * 8 |
| 57 | |
| 58 | # native kernel path |
| 59 | out_q, scale = fp8_utils.per_token_group_quant_fp8( |
| 60 | x, |
| 61 | group_size, |
| 62 | column_major_scales=column_major, |
| 63 | tma_aligned_scales=tma_aligned, |
| 64 | use_ue8m0=scale_ue8m0, |
| 65 | ) |
| 66 | |
| 67 | # triton ref |
| 68 | with ( |
| 69 | patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False), |
| 70 | patch("vllm.platforms.current_platform.is_xpu", return_value=False), |
| 71 | ): |
| 72 | ref_q, ref_s = fp8_utils.per_token_group_quant_fp8( |
| 73 | x, |
| 74 | group_size, |
| 75 | column_major_scales=column_major, |
| 76 | use_ue8m0=scale_ue8m0, |
| 77 | ) |
| 78 | |
| 79 | assert torch.allclose(out_q.float(), ref_q.float(), atol=0.15, rtol=0.15) |
| 80 | assert torch.allclose(scale, ref_s, atol=0.01, rtol=0.01) |
| 81 | |
| 82 | |
| 83 | def _misaligned_clone(x: torch.Tensor) -> torch.Tensor: |
| 84 | """Contiguous copy of ``x`` whose ``data_ptr`` is not 16B aligned.""" |
| 85 | numel = x.numel() |
| 86 | base = torch.empty(numel + 16 // x.element_size(), device=x.device, dtype=x.dtype) |
| 87 | view = base[1 : 1 + numel] |
| 88 | if view.data_ptr() % 16 == 0: |
| 89 | pytest.skip("allocator did not yield a misaligned offset") |
| 90 | return view.view(x.shape).copy_(x) |
| 91 | |
| 92 | |
| 93 | @pytest.mark.parametrize("shape", [(1, 128), (333, 128), (1024, 6144)]) |
| 94 | @pytest.mark.parametrize("column_major", [False, True]) |
| 95 | @pytest.mark.parametrize("scale_ue8m0", [False, True]) |
| 96 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) |
| 97 | @pytest.mark.skipif( |
| 98 | not current_platform.is_cuda_alike(), |
| 99 | reason="register-resident schedule is CUDA/ROCm only", |
| 100 | ) |
| 101 | def test_per_token_group_quant_fp8_register_matches_smem( |
| 102 | shape, column_major: bool, scale_ue8m0: bool, dtype: torch.dtype |
| 103 | ): |
| 104 | """group_size=128 with 16-bit input takes a register-resident schedule that |
| 105 | needs 16B-aligned buffers; a misaligned input falls back to the |
| 106 | shared-memory schedule. The two must agree bit-for-bit, which the |
| 107 | tolerance-based comparison against Triton above cannot establish. |
| 108 | """ |
| 109 | device = current_platform.device_type |
| 110 | torch.manual_seed(42) |
| 111 | x = torch.randn(shape, device=device, dtype=dtype) * 8 |
| 112 | |
| 113 | reg_q, reg_s = fp8_utils.per_token_group_quant_fp8( |
| 114 | x, 128, column_major_scales=column_major, use_ue8m0=scale_ue8m0 |
| 115 | ) |
| 116 | smem_q, smem_s = fp8_utils.per_token_group_quant_fp8( |
| 117 | _misaligned_clone(x), |
| 118 | 128, |
| 119 | column_major_scales=column_major, |
| 120 | use_ue8m0=scale_ue8m0, |
| 121 | ) |
| 122 | |
| 123 | assert torch.equal(reg_q.view(torch.uint8), smem_q.view(torch.uint8)) |
| 124 | assert torch.equal(reg_s, smem_s) |
| 125 | |
| 126 | |
| 127 | @pytest.mark.parametrize( |
| 128 | "num_tokens,hidden_dim,group_size", |
| 129 | [ |
| 130 | # No padding: mn=4 (mult of 4), groups_per_row=56 (mult of 4) |
| 131 | (4, 7168, 128), |
| 132 | # MN padding only: mn=1, tma_aligned_mn=4 |
| 133 | (1, 7168, 128), |
| 134 | # MN padding only: mn=3, tma_aligned_mn=4 |
| 135 | (3, 7168, 128), |
| 136 | # K padding only: groups_per_row=5 (5%4=1) |
| 137 | (4, 640, 128), |
| 138 | # K padding only: groups_per_row=6 (6%4=2) |
| 139 | (4, 768, 128), |
| 140 | # Single packed column, no padding: k_num_packed=1, mn%4=0 |
| 141 | (4, 384, 128), |
| 142 | # Both MN and K padding |
| 143 | (1, 384, 128), |
| 144 | (3, 640, 128), |
| 145 | # Larger shapes with no padding |
| 146 | (64, 7168, 128), |
| 147 | (128, 14336, 128), |
| 148 | # Larger shapes with padding |
| 149 | (127, 7168, 128), |
| 150 | (253, 640, 128), |
| 151 | (1, 32, 32), |
| 152 | (3, 64, 32), |
| 153 | (7, 160, 32), |
| 154 | (32, 512, 32), |
| 155 | (127, 4096, 32), |
| 156 | ], |
| 157 | ) |
| 158 | @pytest.mark.parametrize("poisoned_scales", [False, True]) |
| 159 | @pytest.mark.skipif( |
| 160 | not current_platform.is_cuda_alike(), |
| 161 | reason="DeepGEMM not available on this platform", |
| 162 | ) |
| 163 | def test_per_token_group_quant_fp8_packed( |
| 164 | num_tokens, hidden_dim, group_size, poisoned_scales |
| 165 | ): |
| 166 | """Test the packed DeepGEMM quantization kernel against the Triton |
| 167 | reference (row-major, UE8M0 scales).""" |
| 168 | device = "cuda" |
| 169 | torch.manual_seed(42) |
| 170 | |
| 171 | x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8 |
| 172 | |
| 173 | mn = num_tokens |
| 174 | groups_per_row = hidden_dim // group_size |
| 175 | k_num_packed = (groups_per_row + 3) // 4 |
| 176 | tma_aligned_mn = ((mn + 3) // 4) * 4 |
| 177 | num_scale_elems = mn + (k_num_packed - 1) * tma_aligned_mn |
| 178 | |
| 179 | if poisoned_scales: |
| 180 | # Call the kernel with poisoned scale buffer to |
| 181 | # ensure padded indices are correctly zeroed. |
| 182 | fp8_dtype = current_platform.fp8_dtype() |
| 183 | fp8_min, fp8_max = get_fp8_min_max() |
| 184 | out_q = torch.empty_like(x, dtype=fp8_dtype) |
| 185 | out_s_packed = torch.empty_strided( |
| 186 | (mn, k_num_packed), |
| 187 | (1, tma_aligned_mn), |
| 188 | device=device, |
| 189 | dtype=torch.int32, |
| 190 | ) |
| 191 | torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).fill_(0x7F7F7F7F) |
| 192 | torch.ops._C.per_token_group_fp8_quant_packed( |
| 193 | x, |
| 194 | out_q, |
| 195 | out_s_packed, |
| 196 | group_size, |
| 197 | 1e-10, |
| 198 | fp8_min, |
| 199 | fp8_max, |
| 200 | ) |
| 201 | else: |
| 202 | out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm( |
| 203 | x, |
| 204 | group_size=group_size, |
| 205 | use_ue8m0=True, |
| 206 | ) |
| 207 | |
| 208 | # Triton reference (row-major float32 scales, UE8M0) |
| 209 | with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False): |
| 210 | ref_q, ref_s = fp8_utils.per_token_group_quant_fp8( |
| 211 | x, |
| 212 | group_size, |
| 213 | use_ue8m0=True, |
| 214 | ) |
| 215 | |
| 216 | # Quantized values must match. |
| 217 | assert torch.equal(out_q, ref_q), "Quantized output mismatch" |
| 218 | |
| 219 | # Verify packed scales (valid exponents + padding zeros). |
| 220 | ref_s_flat = ref_s.reshape(mn, groups_per_row) |
| 221 | ref_exponents = (ref_s_flat.view(torch.int32) >> 23) & 0xFF |
| 222 | |
| 223 | expected = torch.zeros(num_scale_elems, dtype=torch.int32, device="cpu") |
| 224 | for row in range(mn): |
| 225 | for g in range(groups_per_row): |
| 226 | pack_col = g // 4 |
| 227 | pos = g % 4 |
| 228 | idx = pack_col * tma_aligned_mn + row |
| 229 | expected[idx] |= int(ref_exponents[row, g].item()) << (pos * 8) |
| 230 | |
| 231 | actual = torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).cpu() |
| 232 | assert torch.equal(actual, expected), ( |
| 233 | f"Packed scale storage mismatch.\n" |
| 234 | f"First diff at index " |
| 235 | f"{(actual != expected).nonzero(as_tuple=True)[0][0].item()}" |
| 236 | ) |
| 237 | |
| 238 | |
| 239 | @pytest.mark.skipif( |
| 240 | not current_platform.is_cuda_alike(), |
| 241 | reason="DeepGEMM not available on this platform", |
| 242 | ) |
| 243 | def test_per_token_group_quant_fp8_packed_all_zero(): |
| 244 | """All-zero input must produce well-defined UE8M0 scale bytes via the eps |
| 245 | floor in the kernel's UE8M0 path. Locks down the all-zero behavior before |
| 246 | optimization. |
| 247 | |
| 248 | The CUDA kernel computes: |
| 249 | y_s = eps / fp8_max |
| 250 | y_s = exp2(ceil(log2(fmax(y_s, 1e-10)))) |
| 251 | For all-zero input, eps/fp8_max < 1e-10, so the inner fmax clamps back to |
| 252 | 1e-10, giving exp2(ceil(log2(1e-10))) = exp2(-33) => UE8M0 byte 0x5E (94). |
| 253 | """ |
| 254 | device = "cuda" |
| 255 | num_tokens, hidden_dim, group_size = 4, 7168, 128 |
| 256 | x = torch.zeros((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) |
| 257 | |
| 258 | out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm( |
| 259 | x, |
| 260 | group_size=group_size, |
| 261 | use_ue8m0=True, |
| 262 | ) |
| 263 | |
| 264 | # Quantized values must be all zero. |
| 265 | assert torch.equal( |
| 266 | out_q.view(torch.uint8), |
| 267 | torch.zeros_like(out_q, dtype=torch.uint8), |
| 268 | ), "All-zero input should produce all-zero FP8 output" |
| 269 | |
| 270 | # UE8M0 byte produced by the kernel for all-zero input. |
| 271 | # The kernel's inner fmax(y_s, 1e-10) clamps eps/fp8_max back to 1e-10. |
| 272 | # 1e-10 as float32 has biased exponent 0x5D and a non-zero mantissa, so |
| 273 | # the kernel's bit-twiddle (exp_bits + (mant_bits != 0)) rounds up to |
| 274 | # 0x5E. This matches exp2(ceil(log2(1e-10))) = exp2(-33). |
| 275 | expected_exp_byte = 0x5E |
| 276 | |
| 277 | mn = num_tokens |
| 278 | groups_per_row = hidden_dim // group_size |
| 279 | k_num_packed = (groups_per_row + 3) // 4 |
| 280 | tma_aligned_mn = ((mn + 3) // 4) * 4 |
| 281 | num_scale_elems = mn + (k_num_packed - 1) * tma_aligned_mn |
| 282 | |
| 283 | # All valid scale slots must contain the expected packed value. |
| 284 | # Padding slots must be zero. |
| 285 | actual = torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).cpu() |
| 286 | |
| 287 | expected = torch.zeros(num_scale_elems, dtype=torch.int32, device="cpu") |
| 288 | for row in range(mn): |
| 289 | for g in range(groups_per_row): |
| 290 | pack_col = g // 4 |
| 291 | pos = g % 4 |
| 292 | idx = pack_col * tma_aligned_mn + row |
| 293 | expected[idx] |= expected_exp_byte << (pos * 8) |
| 294 | |
| 295 | assert torch.equal(actual, expected), "All-zero scale bytes mismatch" |
| 296 | |
| 297 | |
| 298 | @pytest.mark.skipif( |
| 299 | not current_platform.is_cuda_alike(), |
| 300 | reason="DeepGEMM not available on this platform", |
| 301 | ) |
| 302 | def test_per_token_group_quant_fp8_packed_mantissa_rounds_up(): |
| 303 | """Inputs whose absmax/max_8bit produces a non-power-of-2 force the |
| 304 | mantissa-rounding-up branch (exp_byte += 1). Locks down this behavior |
| 305 | before optimization.""" |
| 306 | device = "cuda" |
| 307 | num_tokens, hidden_dim, group_size = 4, 7168, 128 |
| 308 | |
| 309 | # Build a tensor whose per-group absmax = 1.5 * fp8_max * 2^k for various k. |
| 310 | # fp8_max = torch.finfo(torch.float8_e4m3fn).max = 448.0. |
| 311 | # Then absmax/fp8_max = 1.5 * 2^k -> non-zero mantissa, triggers ceil |
| 312 | # rounding to 2^(k+1). Use k=0 for simplicity; the bf16 representation of |
| 313 | # 1.5*448=672.0 is exact. |
| 314 | x = torch.full( |
| 315 | (num_tokens, hidden_dim), |
| 316 | 672.0, |
| 317 | device=device, |
| 318 | dtype=torch.bfloat16, |
| 319 | ) |
| 320 | |
| 321 | out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm( |
| 322 | x, |
| 323 | group_size=group_size, |
| 324 | use_ue8m0=True, |
| 325 | ) |
| 326 | |
| 327 | with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False): |
| 328 | ref_q, ref_s = fp8_utils.per_token_group_quant_fp8( |
| 329 | x, |
| 330 | group_size, |
| 331 | use_ue8m0=True, |
| 332 | ) |
| 333 | |
| 334 | assert torch.equal(out_q, ref_q), "Quantized output mismatch" |
| 335 | |
| 336 | mn = num_tokens |
| 337 | groups_per_row = hidden_dim // group_size |
| 338 | k_num_packed = (groups_per_row + 3) // 4 |
| 339 | tma_aligned_mn = ((mn + 3) // 4) * 4 |
| 340 | num_scale_elems = mn + (k_num_packed - 1) * tma_aligned_mn |
| 341 | |
| 342 | ref_s_flat = ref_s.reshape(mn, groups_per_row) |
| 343 | ref_exponents = (ref_s_flat.view(torch.int32) >> 23) & 0xFF |
| 344 | expected = torch.zeros(num_scale_elems, dtype=torch.int32, device="cpu") |
| 345 | for row in range(mn): |
| 346 | for g in range(groups_per_row): |
| 347 | pack_col = g // 4 |
| 348 | pos = g % 4 |
| 349 | idx = pack_col * tma_aligned_mn + row |
| 350 | expected[idx] |= int(ref_exponents[row, g].item()) << (pos * 8) |
| 351 | |
| 352 | actual = torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).cpu() |
| 353 | assert torch.equal(actual, expected), "Scale bytes mismatch" |
| 354 | |
| 355 | |
| 356 | @pytest.mark.parametrize( |
| 357 | "num_tokens,hidden_dim", |
| 358 | [ |
| 359 | (1, 7168), # mn padded 1 -> 4 |
| 360 | (2, 7168), # mn padded 2 -> 4 |
| 361 | (3, 7168), # mn padded 3 -> 4 |
| 362 | (5, 7168), # mn padded 5 -> 8 |
| 363 | (127, 7168), # mn padded 127 -> 128 |
| 364 | (253, 640), # both mn and groups padded |
| 365 | (1, 384), # extreme: 1 group, 1 mn row -> both axes padded |
| 366 | ], |
| 367 | ) |
| 368 | @pytest.mark.skipif( |
| 369 | not current_platform.is_cuda_alike(), |
| 370 | reason="DeepGEMM not available on this platform", |
| 371 | ) |
| 372 | def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q( |
| 373 | num_tokens, hidden_dim |
| 374 | ): |
| 375 | """When output_q is allocated with shape (tma_aligned_mn, k) instead of |
| 376 | (mn, k), the kernel must overwrite the padded mn rows with zeros so |
| 377 | callers can use ``torch.empty`` instead of ``torch.zeros``.""" |
| 378 | device = "cuda" |
| 379 | group_size = 128 |
| 380 | torch.manual_seed(42) |
| 381 | x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8 |
| 382 | |
| 383 | mn = num_tokens |
| 384 | groups_per_row = hidden_dim // group_size |
| 385 | k_num_packed = (groups_per_row + 3) // 4 |
| 386 | tma_aligned_mn = ((mn + 3) // 4) * 4 |
| 387 | |
| 388 | fp8_dtype = current_platform.fp8_dtype() |
| 389 | fp8_min, fp8_max = get_fp8_min_max() |
| 390 | # Allocate output_q with the padded mn extent and pre-fill with 0xFF |
| 391 | # so the kernel cannot rely on a clean buffer. |
| 392 | out_q = torch.empty((tma_aligned_mn, hidden_dim), device=device, dtype=fp8_dtype) |
| 393 | out_q.view(torch.uint8).fill_(0xFF) |
| 394 | |
| 395 | out_s_packed = torch.empty_strided( |
| 396 | (mn, k_num_packed), |
| 397 | (1, tma_aligned_mn), |
| 398 | device=device, |
| 399 | dtype=torch.int32, |
| 400 | ) |
| 401 | |
| 402 | torch.ops._C.per_token_group_fp8_quant_packed( |
| 403 | x, out_q, out_s_packed, group_size, 1e-10, fp8_min, fp8_max |
| 404 | ) |
| 405 | |
| 406 | # Live rows must match the Triton reference. |
| 407 | with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False): |
| 408 | ref_q, _ = fp8_utils.per_token_group_quant_fp8(x, group_size, use_ue8m0=True) |
| 409 | assert torch.equal(out_q[:mn], ref_q), "Live region mismatch" |
| 410 | |
| 411 | # Padded rows must be all-zero; without this, downstream TMA loads would |
| 412 | # see uninitialised data. |
| 413 | if tma_aligned_mn > mn: |
| 414 | padded_bytes = out_q[mn:tma_aligned_mn].view(torch.uint8) |
| 415 | assert padded_bytes.eq(0).all(), ( |
| 416 | f"Padded rows [{mn}, {tma_aligned_mn}) not zeroed; " |
| 417 | f"{padded_bytes.ne(0).sum().item()} non-zero bytes" |
| 418 | ) |
| 419 | |
| 420 | |
| 421 | @pytest.mark.skipif( |
| 422 | not current_platform.is_cuda_alike(), |
| 423 | reason="packed FP8 per-token-group quant kernel requires a CUDA-alike GPU", |
| 424 | ) |
| 425 | def test_per_token_group_quant_fp8_packed_large_mn(): |
| 426 | """Regression test for https://github.com/vllm-project/vllm/issues/45099. |
| 427 | |
| 428 | Some background: gridDim.x and gridDim.y have different limits of 2^31 - 1 and |
| 429 | 2^16 - 1, respectively. |
| 430 | Prior code introduced a bug where it incorrectly assumed grid.x and y both have |
| 431 | 2^31 - 1 limits and mixed them up, which doesn't surface until the kernel is |
| 432 | launched with a large mn that exceeds grid.y limit (2^16 - 1). |
| 433 | |
| 434 | This issue doesn't surface often because each forward pass only processes a |
| 435 | bounded token batch, not the full context. |
| 436 | Quantizing tensors with more rows than that will fail at launch with |
| 437 | "CUDA error: invalid argument". |
| 438 | This is a differential test that compares fp8 output against Triton output |
| 439 | reference when token size sits just above the gridDim.y 2^16 - 1 limit. |
| 440 | """ |
| 441 | device = "cuda" |
| 442 | group_size = 128 |
| 443 | # hidden 2048 -> 2048/128 = 16 groups per row -> kx=16, ry=1: one grid row per mn |
| 444 | # row, so any mn > 65535 overflowed grid.y before the fix. |
| 445 | num_tokens, hidden_dim = 65537, 2048 |
| 446 | torch.manual_seed(42) |
| 447 | x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8 |
| 448 | |
| 449 | out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm( |
| 450 | x, |
| 451 | group_size=group_size, |
| 452 | use_ue8m0=True, |
| 453 | ) |
| 454 | |
| 455 | with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False): |
| 456 | ref_q, ref_s = fp8_utils.per_token_group_quant_fp8( |
| 457 | x, group_size, use_ue8m0=True |
| 458 | ) |
| 459 | |
| 460 | assert torch.equal(out_q, ref_q), "Quantized output mismatch" |
| 461 | |
| 462 | # Vectorized packed-scale check; the per-element loop used by the smaller |
| 463 | # tests is too slow at this size. groups_per_row is a multiple of 4 here, |
| 464 | # so there is no K padding and the packed view lines up. |
| 465 | mn = num_tokens |
| 466 | groups_per_row = hidden_dim // group_size |
| 467 | k_num_packed = (groups_per_row + 3) // 4 |
| 468 | assert groups_per_row % 4 == 0 |
| 469 | ref_exponents = (ref_s.reshape(mn, groups_per_row).view(torch.int32) >> 23) & 0xFF |
| 470 | exp = ref_exponents.view(mn, k_num_packed, 4) |
| 471 | expected = ( |
| 472 | exp[..., 0] | (exp[..., 1] << 8) | (exp[..., 2] << 16) | (exp[..., 3] << 24) |
| 473 | ) |
| 474 | assert torch.equal(out_s_packed.cpu(), expected.cpu()), "Packed scale mismatch" |
| 475 | |
| 476 | |
| 477 | @pytest.mark.parametrize( |
| 478 | "shape", |
| 479 | [ |
| 480 | (0, 128), |
| 481 | (1, 128), |
| 482 | (3, 384), |
| 483 | (5, 384), |
| 484 | (7, 384), |
| 485 | (15, 384), |
| 486 | (32, 128), |
| 487 | (64, 256), |
| 488 | (16, 512), |
| 489 | ], |
| 490 | ) |
| 491 | @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) |
| 492 | @pytest.mark.parametrize("group_size", [64, 128]) |
| 493 | @pytest.mark.skipif( |
| 494 | not (current_platform.is_cuda_alike() or current_platform.is_xpu()), |
| 495 | reason="Only test on CUDA/ROCm/XPU.", |
| 496 | ) |
| 497 | def test_per_token_group_quant_int8(shape, dtype, group_size: int): |
| 498 | """The shared launch policy also handles INT8 tail groups.""" |
| 499 | device = current_platform.device_type |
| 500 | |
| 501 | torch.manual_seed(42) |
| 502 | num_tokens, hidden_dim = shape |
| 503 | |
| 504 | x = torch.randn((num_tokens, hidden_dim), device=device, dtype=dtype) * 8 |
| 505 | |
| 506 | # cuda path |
| 507 | out_q, scale = int8_utils.per_token_group_quant_int8( |
| 508 | x, |
| 509 | group_size, |
| 510 | ) |
| 511 | |
| 512 | # triton ref |
| 513 | with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False): |
| 514 | ref_q, ref_s = int8_utils.per_token_group_quant_int8( |
| 515 | x, |
| 516 | group_size, |
| 517 | ) |
| 518 | |
| 519 | assert torch.allclose(out_q.float(), ref_q.float(), atol=0.15, rtol=0.15) |
| 520 | assert torch.allclose(scale, ref_s, atol=0.01, rtol=0.01) |
| 521 | |
| 522 | |
| 523 | @pytest.mark.parametrize("shape", [(1, 128), (333, 128), (1024, 6144)]) |
| 524 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) |
| 525 | @pytest.mark.skipif( |
| 526 | not current_platform.is_cuda_alike(), |
| 527 | reason="register-resident schedule is CUDA/ROCm only", |
| 528 | ) |
| 529 | def test_per_token_group_quant_int8_register_matches_smem(shape, dtype: torch.dtype): |
| 530 | """int8 output shares the register-resident schedule with fp8, so it needs |
| 531 | the same bit-exactness guard against the shared-memory schedule. |
| 532 | """ |
| 533 | device = current_platform.device_type |
| 534 | torch.manual_seed(42) |
| 535 | x = torch.randn(shape, device=device, dtype=dtype) * 8 |
| 536 | |
| 537 | reg_q, reg_s = int8_utils.per_token_group_quant_int8(x, 128) |
| 538 | smem_q, smem_s = int8_utils.per_token_group_quant_int8(_misaligned_clone(x), 128) |
| 539 | |
| 540 | assert torch.equal(reg_q, smem_q) |
| 541 | assert torch.equal(reg_s, smem_s) |