| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for QuantFP8 Group Quantization implementation.""" |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 |
| 9 | from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape |
| 10 | from vllm.utils.torch_utils import set_random_seed |
| 11 | |
| 12 | |
| 13 | @pytest.mark.parametrize( |
| 14 | "batch_size,hidden_dim,group_size", |
| 15 | [ |
| 16 | (16, 256, 32), # Small |
| 17 | (64, 1024, 64), # Medium |
| 18 | (128, 2048, 128), # Large |
| 19 | (8, 513, 64), # Non-divisible (native only) |
| 20 | ], |
| 21 | ) |
| 22 | @pytest.mark.parametrize("seed", [42]) |
| 23 | @pytest.mark.parametrize("use_ue8m0", [True, False]) |
| 24 | @torch.inference_mode() |
| 25 | def test_quantfp8_group_functionality( |
| 26 | default_vllm_config, |
| 27 | batch_size: int, |
| 28 | hidden_dim: int, |
| 29 | group_size: int, |
| 30 | seed: int, |
| 31 | use_ue8m0: bool, |
| 32 | ) -> None: |
| 33 | """Test QuantFP8 group quantization with various configurations. |
| 34 | |
| 35 | Tests both CUDA and native implementations, column-major scales, |
| 36 | and verifies consistency between implementations. |
| 37 | """ |
| 38 | set_random_seed(seed) |
| 39 | |
| 40 | x = torch.randn((batch_size, hidden_dim), dtype=torch.bfloat16, device="cuda") * 8 |
| 41 | expected_num_groups = (hidden_dim + group_size - 1) // group_size |
| 42 | is_divisible = hidden_dim % group_size == 0 |
| 43 | |
| 44 | group_shape = GroupShape(1, group_size) |
| 45 | quant_op = QuantFP8( |
| 46 | static=False, |
| 47 | group_shape=group_shape, |
| 48 | column_major_scales=False, |
| 49 | use_ue8m0=use_ue8m0, |
| 50 | ) |
| 51 | |
| 52 | # 1. Test native implementation (always available) |
| 53 | x_quant_native, scales_native = quant_op.forward_native(x.clone()) |
| 54 | assert x_quant_native.shape == x.shape |
| 55 | assert scales_native.shape == (batch_size, expected_num_groups) |
| 56 | |
| 57 | # 2. Test column-major scales configuration |
| 58 | quant_op_col = QuantFP8( |
| 59 | static=False, |
| 60 | group_shape=group_shape, |
| 61 | column_major_scales=True, |
| 62 | use_ue8m0=use_ue8m0, |
| 63 | ) |
| 64 | _, scales_col = quant_op_col.forward_native(x.clone()) |
| 65 | assert scales_col.shape == (batch_size, expected_num_groups) |
| 66 | assert scales_col.stride(0) == 1 |
| 67 | assert scales_col.stride(1) == batch_size |
| 68 | |
| 69 | # Test column-major scales consistency |
| 70 | torch.testing.assert_close(scales_col, scales_native, rtol=1e-9, atol=1e-8) |
| 71 | |
| 72 | # 3. Test CUDA implementation (only for divisible dimensions) |
| 73 | if is_divisible: |
| 74 | x_quant_cuda, scales_cuda = quant_op.forward_cuda(x.clone()) |
| 75 | assert x_quant_cuda.shape == x.shape |
| 76 | |
| 77 | if scales_cuda.dtype == torch.int32: |
| 78 | # DeepGEMM UE8M0 fast path (group_size == 128 on Blackwell with |
| 79 | # the UE8M0 oracle): scales are packed 4-per-int32 in a |
| 80 | # TMA-aligned layout. Unpack and compare against native. |
| 81 | assert use_ue8m0 and group_size == 128 |
| 82 | packed_groups = (expected_num_groups + 3) // 4 |
| 83 | assert scales_cuda.shape == (batch_size, packed_groups) |
| 84 | unpacked = ( |
| 85 | scales_cuda.contiguous() |
| 86 | .view(torch.uint8) |
| 87 | .reshape(batch_size, -1)[:, :expected_num_groups] |
| 88 | ) |
| 89 | scales_cuda = torch.pow(2.0, unpacked.float() - 127) |
| 90 | |
| 91 | assert scales_cuda.shape == (batch_size, expected_num_groups) |
| 92 | |
| 93 | # Verify CUDA/native consistency |
| 94 | torch.testing.assert_close(scales_cuda, scales_native, rtol=2e-7, atol=2e-8) |
| 95 | |
| 96 | # Quantized values should mostly match |
| 97 | diff_count = (x_quant_cuda != x_quant_native).sum().item() |
| 98 | diff_ratio = diff_count / x_quant_cuda.numel() |
| 99 | assert diff_ratio < 0.002, f"Too many differences: {diff_ratio:.4%}" |
| 100 | |
| 101 | |
| 102 | @pytest.mark.parametrize("seed", [42]) |
| 103 | @pytest.mark.parametrize("use_ue8m0", [True, False]) |
| 104 | @torch.inference_mode() |
| 105 | def test_quantfp8_group_multidimensional( |
| 106 | default_vllm_config, seed: int, use_ue8m0: bool |
| 107 | ) -> None: |
| 108 | set_random_seed(seed) |
| 109 | |
| 110 | group_size = 64 |
| 111 | |
| 112 | # Test with 3D input |
| 113 | batch1, batch2, hidden_dim = 4, 8, 1024 |
| 114 | x_3d = ( |
| 115 | torch.randn((batch1, batch2, hidden_dim), dtype=torch.bfloat16, device="cuda") |
| 116 | * 8 |
| 117 | ) |
| 118 | |
| 119 | group_shape = GroupShape(1, group_size) |
| 120 | quant_op = QuantFP8( |
| 121 | static=False, |
| 122 | group_shape=group_shape, |
| 123 | column_major_scales=False, |
| 124 | use_ue8m0=use_ue8m0, |
| 125 | ) |
| 126 | |
| 127 | x_quant, scales = quant_op.forward_native(x_3d.clone()) |
| 128 | assert x_quant.shape == x_3d.shape |
| 129 | assert scales.shape == (batch1, batch2, hidden_dim // group_size) |
| 130 | |
| 131 | # Test column_major_scales with multi-dim |
| 132 | quant_op_col = QuantFP8( |
| 133 | static=False, |
| 134 | group_shape=group_shape, |
| 135 | column_major_scales=True, |
| 136 | use_ue8m0=use_ue8m0, |
| 137 | ) |
| 138 | _, scales_col = quant_op_col.forward_native(x_3d.clone()) |
| 139 | assert scales_col.shape == (batch1, batch2, hidden_dim // group_size) |
| 140 | |
| 141 | # Test with 4D input |
| 142 | batch1, batch2, batch3, hidden_dim = 2, 3, 4, 256 |
| 143 | x_4d = ( |
| 144 | torch.randn( |
| 145 | (batch1, batch2, batch3, hidden_dim), dtype=torch.bfloat16, device="cuda" |
| 146 | ) |
| 147 | * 8 |
| 148 | ) |
| 149 | |
| 150 | x_quant_4d, scales_4d = quant_op.forward_native(x_4d.clone()) |
| 151 | assert x_quant_4d.shape == x_4d.shape |
| 152 | assert scales_4d.shape == (batch1, batch2, batch3, hidden_dim // group_size) |
| 153 | |
| 154 | _, scales_4d_col = quant_op_col.forward_native(x_4d.clone()) |
| 155 | assert scales_4d_col.shape == (batch1, batch2, hidden_dim // group_size, batch3) |
| 156 | |
| 157 | |
| 158 | @pytest.mark.parametrize("seed", [42]) |
| 159 | @torch.inference_mode() |
| 160 | def test_quantfp8_group_edge_cases(default_vllm_config, seed: int) -> None: |
| 161 | set_random_seed(seed) |
| 162 | |
| 163 | batch_size = 16 |
| 164 | group_size = 64 |
| 165 | |
| 166 | # Test with single group (group_size >= hidden_dim) |
| 167 | x_small = torch.randn((batch_size, 32), dtype=torch.bfloat16, device="cuda") * 8 |
| 168 | group_shape = GroupShape(1, group_size) |
| 169 | quant_op = QuantFP8( |
| 170 | static=False, group_shape=group_shape, column_major_scales=False |
| 171 | ) |
| 172 | |
| 173 | x_quant_small, scales_small = quant_op.forward_native(x_small.clone()) |
| 174 | assert x_quant_small.shape == x_small.shape |
| 175 | assert scales_small.shape == (batch_size, 1) |
| 176 | |
| 177 | # Test with zero inputs |
| 178 | x_zero = torch.zeros((batch_size, 256), dtype=torch.bfloat16, device="cuda") |
| 179 | x_quant_zero, scales_zero = quant_op.forward_native(x_zero.clone()) |
| 180 | assert x_quant_zero.shape == x_zero.shape |
| 181 | assert (scales_zero > 0).all(), "Scales should be clamped to minimum" |
| 182 | |
| 183 | # Test very large values |
| 184 | x_large = torch.full((batch_size, 256), 1000.0, dtype=torch.bfloat16, device="cuda") |
| 185 | x_quant_large, scales_large = quant_op.forward_native(x_large.clone()) |
| 186 | assert x_quant_large.shape == x_large.shape |
| 187 | # FP8 max is typically 448 or 224, so scales should be > 1 |
| 188 | assert (scales_large > 1.0).all(), "Large values should have scales > 1" |