@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 300 lines 10.2 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for the rms_norm_per_block_quant helion kernel.
4
5 Run `pytest tests/kernels/helion/test_rms_norm_per_block_quant.py`.
6 """
7
8 import itertools
9 from typing import Any
10
11 import pytest
12 import torch
13 from torch._subclasses.fake_tensor import FakeTensorMode
14
15 from tests.kernels.helion.utils import skip_if_platform_unsupported
16 from tests.kernels.quant_utils import FP8_DTYPE
17 from vllm.kernels.helion.case_key import CaseKey
18 from vllm.kernels.helion.config_manager import ConfigManager
19 from vllm.kernels.helion.ops.rms_norm_per_block_quant import (
20 _pick_cache,
21 baseline,
22 pick_config,
23 rms_norm_per_block_quant,
24 )
25 from vllm.utils.import_utils import has_helion
26 from vllm.utils.torch_utils import set_random_seed
27
28 if not has_helion():
29 pytest.skip(
30 "Helion is not installed. Install with: pip install vllm[helion]",
31 allow_module_level=True,
32 )
33
34
35 def _generate_fake_input(
36 num_tokens: int, hidden_size: int, group_size: int
37 ) -> tuple[Any, ...]:
38 with FakeTensorMode():
39 input = torch.randn(
40 (num_tokens, hidden_size), device="cuda", dtype=torch.bfloat16
41 )
42 result = torch.empty(input.shape, device=input.device, dtype=FP8_DTYPE)
43 scale = torch.empty(
44 (num_tokens, hidden_size // group_size),
45 device=input.device,
46 dtype=torch.float32,
47 )
48 scale_ub = torch.mean(input).to(scale.dtype)
49 residual = torch.randn_like(input)
50 weight = torch.normal(
51 mean=1.0,
52 std=1.0,
53 size=(hidden_size,),
54 dtype=input.dtype,
55 device=input.device,
56 )
57 epsilon = 1e-6
58 args = (
59 result,
60 input,
61 weight,
62 scale,
63 epsilon,
64 scale_ub,
65 residual,
66 group_size,
67 False,
68 )
69 return args
70
71
72 @pytest.fixture(autouse=True)
73 def reset_config_manager_singleton():
74 ConfigManager.reset_instance()
75 ConfigManager()
76 yield
77 ConfigManager.reset_instance()
78
79
80 class TestRmsNormPerBlockQuantConfigPicker:
81 def setup_method(self):
82 _pick_cache.clear()
83
84 def test_config_picker_exact_match(self):
85 config_keys = [
86 CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}),
87 CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}),
88 ]
89
90 args = _generate_fake_input(16, 4096, 128)
91 selected_key = pick_config(args, config_keys)
92 assert selected_key == CaseKey(
93 {"hidden_size": 4096, "group_size": 128, "num_tokens": 16}
94 )
95
96 def test_config_picker_closest_match(self):
97 config_keys = [
98 CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}),
99 CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 32}),
100 CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 16}),
101 CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 32}),
102 CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 16}),
103 CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 32}),
104 CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}),
105 CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 32}),
106 ]
107
108 args = _generate_fake_input(20, 3000, 70)
109 selected_key = pick_config(args, config_keys)
110 assert selected_key == CaseKey(
111 {"hidden_size": 2048, "group_size": 64, "num_tokens": 32}
112 )
113
114 def test_config_picker_no_configs(self):
115 config_keys: list[CaseKey] = []
116
117 args = _generate_fake_input(16, 4096, 128)
118 selected_key = pick_config(args, config_keys)
119 assert selected_key is None
120
121 def test_config_picker_fallback_to_largest(self):
122 config_keys = [
123 CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 16}),
124 CaseKey({"hidden_size": 2048, "group_size": 64, "num_tokens": 32}),
125 CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 16}),
126 CaseKey({"hidden_size": 2048, "group_size": 128, "num_tokens": 32}),
127 CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 16}),
128 CaseKey({"hidden_size": 4096, "group_size": 64, "num_tokens": 32}),
129 CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 16}),
130 CaseKey({"hidden_size": 4096, "group_size": 128, "num_tokens": 32}),
131 ]
132
133 args = _generate_fake_input(64, 8192, 256)
134 selected_key = pick_config(args, config_keys)
135 assert selected_key == CaseKey(
136 {"hidden_size": 4096, "group_size": 128, "num_tokens": 32}
137 )
138
139 def test_b200_configs_disable_reduction_warp_specialization(self):
140 config_set = ConfigManager.get_instance().load_config_set(
141 "rms_norm_per_block_quant"
142 )
143 configs = config_set.to_dict()["nvidia_b200"].values()
144
145 # TODO: Remove once the Triton pin includes
146 # https://github.com/triton-lang/triton/pull/9716. Tracked by
147 # https://github.com/triton-lang/triton/issues/10901.
148 assert all(
149 config["range_warp_specializes"][1] is not True for config in configs
150 )
151
152
153 DTYPES = [torch.bfloat16, torch.float]
154 QUANT_DTYPES = [torch.int8, FP8_DTYPE]
155 VEC_HIDDEN_SIZES = [64, 1024]
156 # Avoid combinatorial explosion with full Cartesian product
157 NUM_TOKENS_HIDDEN_SIZES = [
158 *[(1, i) for i in [64, 128, 1024, 5120]],
159 *[(2048, i) for i in [64, 1024]],
160 *[(4096, i) for i in [64]],
161 ]
162
163 ADD_RESIDUAL = [False, True]
164 SCALE_UBS = [True, False]
165 GROUP_SIZES = [64, 128]
166 TMA_ALIGNMENTS = [0, 4]
167 SEEDS = [0]
168 EPS = 1e-6
169
170
171 class TestRmsNormPerBlockQuantCorrectness:
172 @pytest.mark.parametrize("num_tokens, hidden_size", NUM_TOKENS_HIDDEN_SIZES)
173 @pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
174 @pytest.mark.parametrize("has_scale_ub", SCALE_UBS)
175 @pytest.mark.parametrize("dtype", DTYPES)
176 @pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
177 @pytest.mark.parametrize("is_scale_transposed", [False, True])
178 @pytest.mark.parametrize(
179 "group_size, tma_alignment",
180 [*itertools.product(GROUP_SIZES, TMA_ALIGNMENTS)],
181 )
182 @pytest.mark.parametrize("seed", SEEDS)
183 def test_rms_norm_per_block_quant(
184 self,
185 num_tokens: int,
186 hidden_size: int,
187 add_residual: bool,
188 has_scale_ub: bool,
189 dtype: torch.dtype,
190 quant_dtype: torch.dtype,
191 is_scale_transposed: bool,
192 group_size: int,
193 tma_alignment: int,
194 seed: int,
195 ) -> None:
196 skip_if_platform_unsupported("rms_norm_per_block_quant")
197
198 set_random_seed(seed)
199
200 if hidden_size % group_size != 0:
201 # skip
202 return
203
204 if tma_alignment != 0 and hidden_size // group_size % tma_alignment == 0:
205 # Skip tests where TMA alignment doesn't create extra padding to save time
206 return
207
208 if has_scale_ub and quant_dtype != FP8_DTYPE:
209 # skip
210 return
211
212 scale = 1 / (hidden_size)
213 input = torch.randn(num_tokens, hidden_size, dtype=dtype, device="cuda") * scale
214 weight = torch.normal(
215 mean=1.0, std=1.0, size=(hidden_size,), dtype=dtype, device=input.device
216 )
217 residual = torch.randn_like(input) * scale if add_residual else None
218 scale_ub = (
219 torch.mean(input).to(dtype=torch.float32, device="cuda")
220 if has_scale_ub
221 else None
222 )
223 groups_per_row = hidden_size // group_size
224
225 ref_residual = residual.clone() if residual is not None else None
226 ops_residual = residual.clone() if residual is not None else None
227 ref_out = torch.empty(input.shape, device=input.device, dtype=quant_dtype)
228 ops_out = ref_out.clone()
229
230 if is_scale_transposed:
231 if tma_alignment == 0:
232 ref_scales = torch.empty(
233 (groups_per_row, num_tokens),
234 device=input.device,
235 dtype=torch.float32,
236 ).transpose(0, 1)
237 else:
238 tma_aligned_m = (
239 (num_tokens + tma_alignment - 1) // tma_alignment * tma_alignment
240 )
241 shape = (num_tokens, groups_per_row)
242 stride = (1, tma_aligned_m)
243 ref_scales = torch.empty_strided(
244 shape, stride, device=input.device, dtype=torch.float32
245 )
246 else:
247 ref_scales = torch.empty(
248 (num_tokens, groups_per_row), device=input.device, dtype=torch.float32
249 )
250
251 ops_scales = ref_scales.clone()
252
253 baseline(
254 ref_out,
255 input,
256 weight,
257 ref_scales,
258 EPS,
259 scale_ub,
260 ref_residual,
261 group_size,
262 is_scale_transposed,
263 )
264 ref_scales = ref_scales.contiguous()
265
266 rms_norm_per_block_quant(
267 ops_out,
268 input,
269 weight,
270 ops_scales,
271 EPS,
272 scale_ub,
273 ops_residual,
274 group_size,
275 is_scale_transposed,
276 )
277 ops_scales = ops_scales.contiguous()
278
279 torch.testing.assert_close(ref_scales, ops_scales)
280 # allow 1 ULP difference
281 assert (
282 ref_out.view(torch.uint8).to(torch.int16)
283 - ops_out.view(torch.uint8).to(torch.int16)
284 ).abs().max() <= 1
285
286 if add_residual:
287 torch.testing.assert_close(ref_residual, ops_residual)
288
289
290 class TestRmsNormPerBlockQuantIntegration:
291 def test_kernel_registration_integration(self):
292 from vllm.kernels.helion.register import get_registered_kernels
293
294 registered_kernels = get_registered_kernels()
295 assert "rms_norm_per_block_quant" in registered_kernels
296
297 kernel_wrapper = registered_kernels["rms_norm_per_block_quant"]
298 assert kernel_wrapper.op_name == "rms_norm_per_block_quant"
299 assert kernel_wrapper._config_picker is not None
300 assert kernel_wrapper._mutates_args == ["result", "scale", "residual"]