@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 195 lines 6.93 KB
Raw
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 from tests.kernels.quant_utils import ref_dynamic_per_token_quant
8 from tests.kernels.utils import opcheck
9 from vllm._custom_ops import scaled_int8_quant
10 from vllm.utils.torch_utils import set_random_seed
11
12 DTYPES = [torch.bfloat16, torch.float]
13 HIDDEN_SIZES = [17, 1024, 1025, 1026, 5137, 8193]
14 NUM_TOKENS = [1, 7, 4096]
15 SEEDS = [0]
16 SCALE = [0.1, 2.1]
17
18
19 def opcheck_int8_quant_static(output, input, scale, azp=None):
20 if azp is None:
21 opcheck(torch.ops._C.static_scaled_int8_quant, (output, input, scale, None))
22 else:
23 opcheck(torch.ops._C.static_scaled_int8_quant, (output, input, scale, azp))
24
25
26 def opcheck_int8_quant_dynamic(output, input, symmetric=True):
27 scale = torch.empty(
28 (input.numel() // input.shape[-1], 1), device=input.device, dtype=torch.float32
29 )
30 if symmetric:
31 opcheck(torch.ops._C.dynamic_scaled_int8_quant, (output, input, scale, None))
32 else:
33 azp = torch.empty(
34 (input.numel() // input.shape[-1], 1),
35 device=input.device,
36 dtype=torch.int32,
37 )
38 opcheck(torch.ops._C.dynamic_scaled_int8_quant, (output, input, scale, azp))
39
40
41 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
42 @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
43 @pytest.mark.parametrize("dtype", DTYPES)
44 @pytest.mark.parametrize("seed", SEEDS)
45 @torch.inference_mode()
46 def test_dynamic_scaled_int8_quant(
47 num_tokens: int, hidden_size: int, dtype: torch.dtype, seed: int
48 ) -> None:
49 set_random_seed(seed)
50
51 x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") * 1000
52
53 # reference
54 ref_out, ref_scales = ref_dynamic_per_token_quant(x, torch.int8)
55 # kernel
56 ops_out, ops_scales, _ = scaled_int8_quant(x)
57
58 torch.testing.assert_close(ops_scales, ref_scales)
59 # big atol to account for rounding errors
60 torch.testing.assert_close(ops_out, ref_out, atol=1, rtol=0.0)
61
62 opcheck_int8_quant_dynamic(ops_out, x)
63
64
65 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
66 @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
67 @pytest.mark.parametrize("dtype", DTYPES)
68 @pytest.mark.parametrize("seed", SEEDS)
69 @torch.inference_mode()
70 def test_dynamic_scaled_int8_azp_quant(
71 num_tokens: int, hidden_size: int, dtype: torch.dtype, seed: int
72 ) -> None:
73 set_random_seed(seed)
74 int8_traits = torch.iinfo(torch.int8)
75
76 x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") * 1000 - 300
77
78 x_token_max, _ = x.to(dtype=torch.float32).max(dim=1, keepdim=True)
79 x_token_min, _ = x.to(dtype=torch.float32).min(dim=1, keepdim=True)
80
81 # calculate scale and azp, and adjust the range
82 scales = (x_token_max - x_token_min) / torch.tensor(255.0)
83 azps = torch.round(torch.tensor(-128.0) - x_token_min / scales).to(torch.int32)
84
85 torch_out = (
86 ((x / scales).round() + azps)
87 .clamp(int8_traits.min, int8_traits.max)
88 .to(torch.int8)
89 )
90 assert torch_out.min() >= int8_traits.min and torch_out.max() <= int8_traits.max
91
92 ops_out, scales_out, azp_out = scaled_int8_quant(x, symmetric=False)
93
94 if not torch.allclose(scales_out, scales):
95 print(torch.argmax(torch.abs(scales_out - scales)))
96 torch.testing.assert_close(scales_out, scales)
97 # big atol to account for rounding errors
98 torch.testing.assert_close(azp_out, azps, atol=1, rtol=0.0)
99 # if AZP is off by 1, after rounding-to-even, the output may be off by 2
100 torch.testing.assert_close(ops_out, torch_out, atol=2, rtol=0.0)
101
102 opcheck_int8_quant_dynamic(ops_out, x, False)
103
104
105 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
106 @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
107 @pytest.mark.parametrize("dtype", DTYPES)
108 @pytest.mark.parametrize("seed", SEEDS)
109 @pytest.mark.parametrize("scale", SCALE)
110 @torch.inference_mode()
111 def test_static_scaled_int8_quant(
112 num_tokens: int, hidden_size: int, dtype: torch.dtype, seed: int, scale: float
113 ) -> None:
114 set_random_seed(seed)
115 int8_traits = torch.iinfo(torch.int8)
116
117 x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") * 1000
118 scale_arg = torch.tensor([scale], dtype=torch.float32, device="cuda")
119
120 out1 = (
121 (x / scale_arg).round().clamp(int8_traits.min, int8_traits.max).to(torch.int8)
122 )
123 out2, scale2, _ = scaled_int8_quant(x, scale_arg)
124 assert scale2 is scale_arg
125
126 # big atol to account for rounding errors
127 torch.testing.assert_close(out1, out2, atol=1, rtol=0.0)
128
129 opcheck_int8_quant_static(out2, x, scale_arg)
130
131
132 @pytest.mark.parametrize("num_tokens", NUM_TOKENS)
133 @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
134 @pytest.mark.parametrize("dtype", DTYPES)
135 @pytest.mark.parametrize("seed", SEEDS)
136 @pytest.mark.parametrize("scale", SCALE)
137 @pytest.mark.parametrize("azp", [-255, 54])
138 @torch.inference_mode()
139 def test_static_scaled_int8_azp_quant(
140 num_tokens: int,
141 hidden_size: int,
142 dtype: torch.dtype,
143 seed: int,
144 scale: float,
145 azp: int,
146 ) -> None:
147 set_random_seed(seed)
148 int8_traits = torch.iinfo(torch.int8)
149
150 x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") * 1000 - 300
151
152 out1 = (
153 ((x / scale).round() + azp)
154 .clamp(int8_traits.min, int8_traits.max)
155 .to(torch.int8)
156 )
157 scale_arg = torch.tensor([scale], dtype=torch.float32, device="cuda")
158 azp_arg = torch.tensor([azp], dtype=torch.int32, device="cuda")
159
160 out2, scale2, azp2 = scaled_int8_quant(x, scale_arg, azp_arg, symmetric=False)
161 assert scale2 is scale_arg
162 assert azp2 is azp_arg
163
164 # big atol to account for rounding errors
165 torch.testing.assert_close(out1, out2, atol=1, rtol=0.0)
166
167 opcheck_int8_quant_static(out2, x, scale_arg, azp_arg)
168
169
170 @pytest.mark.parametrize("is_max", [True, False])
171 @torch.inference_mode()
172 def test_static_scaled_int8_azp_quant_saturating_cast(is_max: bool) -> None:
173 # Test that the saturating cast works correctly for values near i32 max/min
174
175 from numpy import inf, nextafter
176
177 int32_traits = torch.iinfo(torch.int32)
178 val = float(int32_traits.max if is_max else int32_traits.min)
179
180 x_vals = [[nextafter(val, inf), val + 1, val, val - 1, nextafter(val, -inf)]]
181 x = torch.tensor(x_vals, dtype=torch.float32, device="cuda")
182
183 # The calculation in the kernel is: cast<int8>(cast<int32>(x / scale) + azp)
184 # where cast<T> is a saturating cast to type T.
185 # Scale is set to 1.0 so that the input values are the ones that are cast.
186 # AZP is set to 0 to make sure the int8 saturating cast is tested as well.
187 scale = torch.scalar_tensor(1.0, dtype=torch.float32, device="cuda")
188 azp = torch.scalar_tensor(0, dtype=torch.int32, device="cuda")
189
190 int8_traits = torch.iinfo(torch.int8)
191 val_i8 = int8_traits.max if is_max else int8_traits.min
192 expected = torch.full((1, 5), val_i8, dtype=torch.int8, device="cuda")
193
194 out, _, _ = scaled_int8_quant(x, scale, azp, symmetric=False)
195 torch.testing.assert_close(expected, out, atol=0, rtol=0)