| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | |
| 7 | import vllm._custom_ops as ops |
| 8 | from tests.kernels.quant_utils import ( |
| 9 | FP8_DTYPE, |
| 10 | ref_dynamic_per_tensor_fp8_quant, |
| 11 | ref_dynamic_per_token_quant, |
| 12 | ) |
| 13 | from tests.kernels.utils import opcheck |
| 14 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 15 | GroupShape, |
| 16 | scaled_quantize, |
| 17 | ) |
| 18 | from vllm.platforms import current_platform |
| 19 | from vllm.utils.torch_utils import set_random_seed |
| 20 | |
| 21 | DTYPES = [torch.bfloat16, torch.float] |
| 22 | HIDDEN_SIZES = [17, 1024, 1025, 1026, 5137, 8193] |
| 23 | NUM_TOKENS = [1, 7, 4096] |
| 24 | SCALE_UBS = [True, False] |
| 25 | SEEDS = [0] |
| 26 | |
| 27 | |
| 28 | def opcheck_fp8_quant( |
| 29 | output, |
| 30 | input, |
| 31 | scale=None, |
| 32 | scale_ub=None, |
| 33 | use_per_token_if_dynamic=False, |
| 34 | group_shape=None, |
| 35 | ): |
| 36 | if scale is not None: |
| 37 | opcheck( |
| 38 | torch.ops._C.static_scaled_fp8_quant, |
| 39 | (output, input, scale, group_shape), |
| 40 | ) |
| 41 | elif use_per_token_if_dynamic: |
| 42 | scale = torch.empty( |
| 43 | (input.shape[0], 1), device=input.device, dtype=torch.float32 |
| 44 | ) |
| 45 | opcheck( |
| 46 | torch.ops._C.dynamic_per_token_scaled_fp8_quant, |
| 47 | (output, input, scale, scale_ub), |
| 48 | ) |
| 49 | else: |
| 50 | scale = torch.empty( |
| 51 | (input.numel() // input.shape[-1], 1), |
| 52 | device=input.device, |
| 53 | dtype=torch.float32, |
| 54 | ) |
| 55 | opcheck(torch.ops._C.dynamic_scaled_fp8_quant, (output, input, scale)) |
| 56 | |
| 57 | |
| 58 | @pytest.mark.parametrize("num_tokens", NUM_TOKENS) |
| 59 | @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) |
| 60 | @pytest.mark.parametrize("dtype", DTYPES) |
| 61 | @pytest.mark.parametrize("do_scale_ub", SCALE_UBS) |
| 62 | @pytest.mark.parametrize("seed", SEEDS) |
| 63 | @torch.inference_mode() |
| 64 | def test_dynamic_per_token_fp8_quant( |
| 65 | num_tokens: int, hidden_size: int, dtype: torch.dtype, do_scale_ub: bool, seed: int |
| 66 | ) -> None: |
| 67 | set_random_seed(seed) |
| 68 | |
| 69 | x = ( |
| 70 | torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") + 1e-6 |
| 71 | ) # avoid nans |
| 72 | |
| 73 | scale_ub = ( |
| 74 | torch.mean(x).to(dtype=torch.float32, device="cuda") if do_scale_ub else None |
| 75 | ) |
| 76 | ref_out, ref_scales = ref_dynamic_per_token_quant(x, FP8_DTYPE, scale_ub) |
| 77 | ops_out, ops_scales = ops.scaled_fp8_quant( |
| 78 | x, scale_ub=scale_ub, use_per_token_if_dynamic=True |
| 79 | ) |
| 80 | |
| 81 | torch.testing.assert_close(ref_scales, ops_scales) |
| 82 | torch.testing.assert_close( |
| 83 | ref_out.to(dtype=torch.float32), ops_out.to(dtype=torch.float32) |
| 84 | ) |
| 85 | |
| 86 | opcheck_fp8_quant(ops_out, x, None, scale_ub, use_per_token_if_dynamic=True) |
| 87 | |
| 88 | |
| 89 | @pytest.mark.parametrize("num_tokens", NUM_TOKENS) |
| 90 | @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES) |
| 91 | @pytest.mark.parametrize("dtype", DTYPES) |
| 92 | @pytest.mark.parametrize("seed", SEEDS) |
| 93 | @torch.inference_mode() |
| 94 | def test_dynamic_per_tensor_fp8_quant( |
| 95 | num_tokens: int, hidden_size: int, dtype: torch.dtype, seed: int |
| 96 | ) -> None: |
| 97 | set_random_seed(seed) |
| 98 | |
| 99 | x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") |
| 100 | |
| 101 | ref_out, ref_scale = ref_dynamic_per_tensor_fp8_quant(x) |
| 102 | ops_out, ops_scale = ops.scaled_fp8_quant(x) |
| 103 | |
| 104 | torch.testing.assert_close(ref_scale, ops_scale) |
| 105 | torch.testing.assert_close( |
| 106 | ref_out.to(dtype=torch.float32), ops_out.to(dtype=torch.float32) |
| 107 | ) |
| 108 | |
| 109 | opcheck_fp8_quant(ops_out, x) |
| 110 | |
| 111 | |
| 112 | @pytest.mark.parametrize("pad", [8, 16]) |
| 113 | @pytest.mark.parametrize("dtype", DTYPES) |
| 114 | @pytest.mark.parametrize("seed", SEEDS) |
| 115 | @torch.inference_mode() |
| 116 | def test_dynamic_per_tensor_fp8_quant_strided_input( |
| 117 | pad: int, dtype: torch.dtype, seed: int |
| 118 | ) -> None: |
| 119 | """Non-contiguous rows take the strided abs-max kernel; contiguous rows take |
| 120 | the flat one. Both must produce exactly the same scale and bytes. |
| 121 | |
| 122 | pad=8 (16 B for bf16) leaves rows misaligned for 32 B vector loads, pad=16 |
| 123 | keeps them aligned but non-contiguous. |
| 124 | """ |
| 125 | set_random_seed(seed) |
| 126 | num_tokens, hidden_size = 64, 4096 |
| 127 | |
| 128 | full = torch.rand(num_tokens, hidden_size + pad, dtype=dtype, device="cuda") |
| 129 | x = full[:, :hidden_size] |
| 130 | assert not x.is_contiguous() |
| 131 | |
| 132 | ops_out, ops_scale = ops.scaled_fp8_quant(x) |
| 133 | flat_out, flat_scale = ops.scaled_fp8_quant(x.contiguous()) |
| 134 | assert torch.equal(ops_scale, flat_scale) |
| 135 | assert torch.equal(ops_out.view(torch.uint8), flat_out.view(torch.uint8)) |
| 136 | |
| 137 | ref_out, ref_scale = ref_dynamic_per_tensor_fp8_quant(x) |
| 138 | assert torch.equal(ref_scale, ops_scale) |
| 139 | assert torch.equal(ref_out.view(torch.uint8), ops_out.view(torch.uint8)) |
| 140 | |
| 141 | |
| 142 | # Regression test for a case with large activations where an int32 index cannot |
| 143 | # represent the number of elements. |
| 144 | @torch.inference_mode() |
| 145 | @pytest.mark.parametrize("seed", SEEDS) |
| 146 | def test_fp8_quant_large(seed: int) -> None: |
| 147 | set_random_seed(seed) |
| 148 | |
| 149 | num_tokens = 1024000 # Mistral-Nemo's max_position_embeddings |
| 150 | hidden_size = 1152 # Smallest hidden_size to reproduce the error |
| 151 | dtype = torch.bfloat16 |
| 152 | |
| 153 | x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") |
| 154 | ref_out, scale = ref_dynamic_per_tensor_fp8_quant(x) |
| 155 | ops_out, _ = ops.scaled_fp8_quant(x, scale) |
| 156 | |
| 157 | # Minimize memory footprint in this test by freeing x and upconverting |
| 158 | # the outputs in place. (torch.allclose does not support fp8) |
| 159 | del x |
| 160 | ref_out = ref_out.to(dtype=dtype) |
| 161 | ops_out = ops_out.to(dtype=dtype) |
| 162 | |
| 163 | torch.testing.assert_close(ref_out, ops_out) |
| 164 | |
| 165 | |
| 166 | # Test static FP8 quantization with 2D group scales |
| 167 | GROUP_SHAPES_2D = [ |
| 168 | (-1, -1), # Per-tensor |
| 169 | (-1, 1), # Per-channel |
| 170 | (1, -1), # Per-token |
| 171 | (-1, 128), # Per-head quantization |
| 172 | (1, 128), # DeepSeek-style per-token-per-group (group_m=1, group_n=128) |
| 173 | (128, 128), # DeepSeek-style block quantization |
| 174 | (1, 64), # Smaller group size |
| 175 | (1, 16), # Small group (scalar path in kernel) |
| 176 | (4, 256), # Non-trivial both dimensions |
| 177 | ] |
| 178 | # Use sizes divisible by all group shapes |
| 179 | NUM_TOKENS_GROUP = [128, 512] |
| 180 | HIDDEN_SIZES_GROUP = [256, 1024, 2048] |
| 181 | |
| 182 | |
| 183 | @pytest.mark.parametrize("num_tokens", NUM_TOKENS_GROUP) |
| 184 | @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES_GROUP) |
| 185 | @pytest.mark.parametrize("group_shape", GROUP_SHAPES_2D) |
| 186 | @pytest.mark.parametrize("dtype", DTYPES) |
| 187 | @pytest.mark.parametrize("seed", SEEDS) |
| 188 | @torch.inference_mode() |
| 189 | def test_static_fp8_quant_group_2d( |
| 190 | num_tokens: int, |
| 191 | hidden_size: int, |
| 192 | group_shape: tuple[int, int], |
| 193 | dtype: torch.dtype, |
| 194 | seed: int, |
| 195 | ) -> None: |
| 196 | """Test static FP8 quantization with 2D group scales using scaled_quantize.""" |
| 197 | # Normalize group_shape (-1 means full extent) |
| 198 | norm_group_m = num_tokens if group_shape[0] == -1 else group_shape[0] |
| 199 | norm_group_n = hidden_size if group_shape[1] == -1 else group_shape[1] |
| 200 | |
| 201 | # Skip if sizes are not divisible by group shape |
| 202 | if num_tokens % norm_group_m != 0 or hidden_size % norm_group_n != 0: |
| 203 | pytest.skip( |
| 204 | f"Skipping: ({num_tokens}, {hidden_size}) not divisible by " |
| 205 | f"group_shape ({group_shape[0]}, {group_shape[1]})" |
| 206 | ) |
| 207 | |
| 208 | set_random_seed(seed) |
| 209 | |
| 210 | x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") |
| 211 | ref_out, scale = scaled_quantize( |
| 212 | x, |
| 213 | GroupShape(*group_shape), |
| 214 | current_platform.fp8_dtype(), |
| 215 | compute_dtype=torch.float32, |
| 216 | ) |
| 217 | ops_out, ops_scale = ops.scaled_fp8_quant(x, scale=scale, group_shape=group_shape) |
| 218 | |
| 219 | torch.testing.assert_close(scale, ops_scale) |
| 220 | torch.testing.assert_close(ref_out.float(), ops_out.float(), rtol=1.2e-1, atol=1e-5) |
| 221 | |
| 222 | opcheck_fp8_quant(ops_out, x, scale=scale) |
| 223 | |
| 224 | |
| 225 | @pytest.mark.parametrize("num_tokens", NUM_TOKENS_GROUP) |
| 226 | @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES_GROUP) |
| 227 | @pytest.mark.parametrize("dtype", DTYPES) |
| 228 | @pytest.mark.parametrize("seed", SEEDS) |
| 229 | @pytest.mark.parametrize("group_shape", [(1, -1), (-1, 1)]) # per-token, per-channel |
| 230 | @torch.inference_mode() |
| 231 | def test_static_fp8_quant_1d_scale( |
| 232 | num_tokens: int, |
| 233 | hidden_size: int, |
| 234 | dtype: torch.dtype, |
| 235 | seed: int, |
| 236 | group_shape: tuple[int, int], |
| 237 | ) -> None: |
| 238 | """Test static FP8 quantization with 1D scale (per-token or per-channel).""" |
| 239 | set_random_seed(seed) |
| 240 | |
| 241 | x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") |
| 242 | ref_out, scale_2d = scaled_quantize( |
| 243 | x, GroupShape(*group_shape), FP8_DTYPE, compute_dtype=torch.float32 |
| 244 | ) |
| 245 | |
| 246 | # Flatten scale to 1D for testing 1D scale path |
| 247 | scale_1d = scale_2d.flatten() |
| 248 | ops_out, ops_scale = ops.scaled_fp8_quant( |
| 249 | x, scale=scale_1d, group_shape=group_shape |
| 250 | ) |
| 251 | |
| 252 | torch.testing.assert_close(scale_1d, ops_scale) |
| 253 | torch.testing.assert_close(ref_out.float(), ops_out.float(), rtol=0.12, atol=0.0) |
| 254 | |
| 255 | opcheck_fp8_quant(ops_out, x, scale=scale_1d, group_shape=group_shape) |