@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 213 lines 7.83 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for the silu_and_mul_per_block_quant helion kernel
4 Run `pytest tests/kernels/helion/test_silu_and_mul_per_block_quant.py`.
5 """
6
7 from typing import Any
8
9 import pytest
10 import torch
11 from torch._subclasses.fake_tensor import FakeTensorMode
12
13 from tests.kernels.helion.utils import skip_if_platform_unsupported
14 from tests.kernels.quant_utils import FP8_DTYPE
15 from vllm.kernels.helion.case_key import CaseKey
16 from vllm.kernels.helion.config_manager import ConfigManager
17 from vllm.kernels.helion.ops.silu_and_mul_per_block_quant import (
18 _pick_cache,
19 baseline,
20 pick_config,
21 silu_and_mul_per_block_quant,
22 )
23 from vllm.platforms import current_platform
24 from vllm.utils.import_utils import has_helion
25 from vllm.utils.torch_utils import set_random_seed
26
27 if not has_helion():
28 pytest.skip(
29 "Helion is not installed. Install with: pip install vllm[helion]",
30 allow_module_level=True,
31 )
32
33
34 def _generate_fake_input(
35 num_tokens: int, intermediate_size: int, group_size: int
36 ) -> tuple[Any, ...]:
37 with FakeTensorMode():
38 in_dtype: torch.dtype = torch.bfloat16
39 out_dtype: torch.dtype = current_platform.fp8_dtype()
40 scale_dtype: torch.dtype = torch.float32
41 input = torch.randn(
42 num_tokens, 2 * intermediate_size, device="cuda", dtype=in_dtype
43 )
44 result = torch.empty(
45 num_tokens, intermediate_size, device=input.device, dtype=out_dtype
46 )
47 scale = torch.empty(
48 (num_tokens, intermediate_size // group_size),
49 device=input.device,
50 dtype=scale_dtype,
51 )
52 scale_ub = torch.mean(input).to(scale_dtype)
53 args = (
54 result,
55 input,
56 scale,
57 group_size,
58 scale_ub,
59 False,
60 )
61 return args
62
63
64 class TestSiluAndMulPerBlockQuantConfigPicker:
65 def setup_method(self):
66 _pick_cache.clear()
67
68 def test_config_picker_exact_match(self):
69 config_keys = [
70 CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}),
71 CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}),
72 ]
73
74 args = _generate_fake_input(16, 4096, 128)
75 selected_key = pick_config(args, config_keys)
76 assert selected_key == CaseKey(
77 {"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}
78 )
79
80 def test_config_picker_closest_match(self):
81 config_keys = [
82 CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}),
83 CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 32}),
84 CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 16}),
85 CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 32}),
86 CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 16}),
87 CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 32}),
88 CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}),
89 CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 32}),
90 ]
91
92 args = _generate_fake_input(20, 3000, 70)
93 selected_key = pick_config(args, config_keys)
94 assert selected_key == CaseKey(
95 {"intermediate_size": 2048, "group_size": 64, "num_tokens": 32}
96 )
97
98 def test_config_picker_no_configs(self):
99 config_keys: list[CaseKey] = []
100
101 args = _generate_fake_input(16, 4096, 128)
102 selected_key = pick_config(args, config_keys)
103 assert selected_key is None
104
105 def test_config_picker_fallback_to_largest(self):
106 config_keys = [
107 CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}),
108 CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 32}),
109 CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 16}),
110 CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 32}),
111 CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 16}),
112 CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 32}),
113 CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}),
114 CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 32}),
115 ]
116
117 args = _generate_fake_input(64, 8192, 256)
118 selected_key = pick_config(args, config_keys)
119 assert selected_key == CaseKey(
120 {"intermediate_size": 4096, "group_size": 128, "num_tokens": 32}
121 )
122
123
124 @pytest.fixture(autouse=True)
125 def reset_config_manager_singleton():
126 ConfigManager.reset_instance()
127 ConfigManager()
128 yield
129 ConfigManager.reset_instance()
130
131
132 class TestSiluAndMulPerBlockQuantCorrectness:
133 @pytest.mark.parametrize("num_tokens", [1, 7, 4096])
134 @pytest.mark.parametrize("hidden_size", [1024, 2048, 5120])
135 @pytest.mark.parametrize("group_size", [64, 128])
136 @pytest.mark.parametrize("is_scale_transposed", [False, True])
137 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
138 @pytest.mark.parametrize("quant_dtype", [current_platform.fp8_dtype(), torch.int8])
139 @pytest.mark.parametrize("has_scale_ub", [True, False])
140 @pytest.mark.parametrize("seed", [0])
141 def test_silu_and_mul_per_block_quant(
142 self,
143 num_tokens: int,
144 hidden_size: int,
145 group_size: int,
146 is_scale_transposed: bool,
147 dtype: torch.dtype,
148 quant_dtype: torch.dtype,
149 has_scale_ub: bool,
150 seed: int,
151 ) -> None:
152 skip_if_platform_unsupported("silu_and_mul_per_block_quant")
153 set_random_seed(seed)
154
155 if hidden_size % group_size != 0:
156 return
157
158 if has_scale_ub and quant_dtype != FP8_DTYPE:
159 # skip
160 return
161
162 scale = 1 / hidden_size
163 x = torch.randn(num_tokens, 2 * hidden_size, dtype=dtype, device="cuda") * scale
164
165 if has_scale_ub:
166 act = torch.nn.functional.silu(x[:, :hidden_size]) * x[:, hidden_size:]
167 act_abs = act.abs().float()
168 scale_ub = 0.5 * (act_abs.mean() + act_abs.amax())
169 else:
170 scale_ub = None
171
172 ref_out = torch.empty(num_tokens, hidden_size, device="cuda", dtype=quant_dtype)
173
174 if is_scale_transposed:
175 ref_scales = torch.empty(
176 (hidden_size // group_size, x.shape[0]),
177 device="cuda",
178 dtype=torch.float32,
179 ).t()
180 else:
181 ref_scales = torch.empty(
182 (x.shape[0], hidden_size // group_size),
183 device="cuda",
184 dtype=torch.float32,
185 )
186
187 ops_out = ref_out.clone()
188 ops_scales = ref_scales.clone()
189
190 baseline(ref_out, x, ref_scales, group_size, scale_ub, is_scale_transposed)
191 silu_and_mul_per_block_quant(
192 ops_out, x, ops_scales, group_size, scale_ub, is_scale_transposed
193 )
194
195 torch.testing.assert_close(ref_scales, ops_scales)
196 # allow 1 ULP difference
197 assert (
198 ref_out.view(torch.uint8).to(torch.int16)
199 - ops_out.view(torch.uint8).to(torch.int16)
200 ).abs().max() <= 1
201
202
203 class TestSiluAndMulPerBlockQuantIntegration:
204 def test_kernel_registration_integration(self):
205 from vllm.kernels.helion.register import get_registered_kernels
206
207 registered_kernels = get_registered_kernels()
208 assert "silu_and_mul_per_block_quant" in registered_kernels
209
210 kernel_wrapper = registered_kernels["silu_and_mul_per_block_quant"]
211 assert kernel_wrapper.op_name == "silu_and_mul_per_block_quant"
212 assert kernel_wrapper._config_picker is not None
213 assert kernel_wrapper._mutates_args == ["out", "scales"]