@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 250 lines 7.97 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for the fused_qk_norm_rope helion kernel.
4
5 Run `pytest tests/kernels/helion/test_fused_qk_norm_rope.py`.
6 """
7
8 from typing import Any
9
10 import pytest
11 import torch
12 from torch._subclasses.fake_tensor import FakeTensorMode
13
14 from tests.kernels.helion.utils import skip_if_platform_unsupported
15 from vllm.benchmarks.lib.utils import default_vllm_config
16 from vllm.kernels.helion.case_key import CaseKey
17 from vllm.kernels.helion.config_manager import ConfigManager
18 from vllm.kernels.helion.ops.fused_qk_norm_rope import (
19 _pick_cache,
20 baseline,
21 fused_qk_norm_rope,
22 pick_config,
23 )
24 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
25 from vllm.utils.import_utils import has_helion
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 @default_vllm_config()
35 def _generate_fake_input(
36 num_tokens: int, num_q_heads: int, num_kv_heads: int
37 ) -> tuple[Any, ...]:
38 with FakeTensorMode():
39 head_dim = 128
40 eps = 1e-6
41 is_neox = True
42 rotary_ratio = 1.0
43 device = "cuda"
44 dtype = torch.bfloat16
45 total_dim = (num_q_heads + 2 * num_kv_heads) * head_dim
46 qkv = torch.randn(num_tokens, total_dim, dtype=dtype, device=device)
47 positions = torch.arange(num_tokens, dtype=torch.long, device=device)
48 q_weight = torch.normal(
49 mean=1.0,
50 std=1.0,
51 size=(head_dim,),
52 dtype=qkv.dtype,
53 device=device,
54 )
55 k_weight = torch.normal(
56 mean=1.0,
57 std=1.0,
58 size=(head_dim,),
59 dtype=qkv.dtype,
60 device=device,
61 )
62 rotary_dim = int(head_dim * rotary_ratio)
63 rope = RotaryEmbedding(
64 head_size=head_dim,
65 rotary_dim=rotary_dim,
66 max_position_embeddings=4096,
67 base=10000.0,
68 is_neox_style=is_neox,
69 dtype=dtype,
70 ).to(device)
71 args = (
72 qkv,
73 num_q_heads,
74 num_kv_heads,
75 num_kv_heads,
76 head_dim,
77 eps,
78 q_weight,
79 k_weight,
80 rope.cos_sin_cache,
81 is_neox,
82 positions.view(-1),
83 )
84 return args
85
86
87 @pytest.fixture(autouse=True)
88 def reset_config_manager_singleton():
89 ConfigManager.reset_instance()
90 ConfigManager()
91 yield
92 ConfigManager.reset_instance()
93
94
95 class TestFusedQkNormRopeConfigPicker:
96 def setup_method(self):
97 _pick_cache.clear()
98
99 def test_config_picker_exact_match(self):
100 config_keys = [
101 CaseKey({"q_heads": 2048, "kv_heads": 64, "num_tokens": 16}),
102 CaseKey({"q_heads": 4096, "kv_heads": 128, "num_tokens": 16}),
103 ]
104
105 args = _generate_fake_input(16, 4096, 128)
106 selected_key = pick_config(args, config_keys)
107 assert selected_key == CaseKey(
108 {"q_heads": 4096, "kv_heads": 128, "num_tokens": 16}
109 )
110
111 def test_config_picker_closest_match(self):
112 config_keys = [
113 CaseKey({"q_heads": 2048, "kv_heads": 64, "num_tokens": 16}),
114 CaseKey({"q_heads": 2048, "kv_heads": 64, "num_tokens": 32}),
115 CaseKey({"q_heads": 2048, "kv_heads": 128, "num_tokens": 16}),
116 CaseKey({"q_heads": 2048, "kv_heads": 128, "num_tokens": 32}),
117 CaseKey({"q_heads": 4096, "kv_heads": 64, "num_tokens": 16}),
118 CaseKey({"q_heads": 4096, "kv_heads": 64, "num_tokens": 32}),
119 CaseKey({"q_heads": 4096, "kv_heads": 128, "num_tokens": 16}),
120 CaseKey({"q_heads": 4096, "kv_heads": 128, "num_tokens": 32}),
121 ]
122
123 args = _generate_fake_input(20, 3000, 70)
124 selected_key = pick_config(args, config_keys)
125 assert selected_key == CaseKey(
126 {"q_heads": 2048, "kv_heads": 64, "num_tokens": 32}
127 )
128
129 def test_config_picker_no_configs(self):
130 config_keys: list[CaseKey] = []
131
132 args = _generate_fake_input(16, 4096, 128)
133 selected_key = pick_config(args, config_keys)
134 assert selected_key is None
135
136 def test_config_picker_fallback_to_largest(self):
137 config_keys = [
138 CaseKey({"q_heads": 2048, "kv_heads": 64, "num_tokens": 16}),
139 CaseKey({"q_heads": 2048, "kv_heads": 64, "num_tokens": 32}),
140 CaseKey({"q_heads": 2048, "kv_heads": 128, "num_tokens": 16}),
141 CaseKey({"q_heads": 2048, "kv_heads": 128, "num_tokens": 32}),
142 CaseKey({"q_heads": 4096, "kv_heads": 64, "num_tokens": 16}),
143 CaseKey({"q_heads": 4096, "kv_heads": 64, "num_tokens": 32}),
144 CaseKey({"q_heads": 4096, "kv_heads": 128, "num_tokens": 16}),
145 CaseKey({"q_heads": 4096, "kv_heads": 128, "num_tokens": 32}),
146 ]
147
148 args = _generate_fake_input(64, 8192, 256)
149 selected_key = pick_config(args, config_keys)
150 assert selected_key == CaseKey(
151 {"q_heads": 4096, "kv_heads": 128, "num_tokens": 32}
152 )
153
154
155 class TestFusedQkNormRopeCorrectness:
156 @pytest.mark.parametrize(
157 "num_heads, num_kv_heads, head_dim", [(16, 4, 128), (64, 8, 128)]
158 )
159 @pytest.mark.parametrize("num_tokens", [1, 7, 1024, 1025])
160 @pytest.mark.parametrize("is_neox", [False, True])
161 @pytest.mark.parametrize("rotary_ratio", [1.0, 0.5, 0.25])
162 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
163 @default_vllm_config()
164 def test_fused_qk_norm_rope(
165 self,
166 num_heads: int,
167 num_kv_heads: int,
168 head_dim: int,
169 num_tokens: int,
170 is_neox: bool,
171 rotary_ratio: float,
172 dtype: torch.dtype,
173 ):
174 skip_if_platform_unsupported("fused_qk_norm_rope")
175
176 torch.manual_seed(42)
177 eps = 1e-6
178 device = "cuda"
179 total_dim = (num_heads + 2 * num_kv_heads) * head_dim
180 ref_qkv = torch.empty(
181 num_tokens, total_dim, dtype=dtype, device=device
182 ).uniform_(-0.1, 0.1)
183 ops_qkv = ref_qkv.clone()
184 positions = torch.arange(num_tokens, dtype=torch.long, device=device)
185 q_weight = torch.empty(head_dim, dtype=dtype, device=device).uniform_(0.8, 1.2)
186 k_weight = torch.empty(head_dim, dtype=dtype, device=device).uniform_(0.8, 1.2)
187 rotary_dim = int(head_dim * rotary_ratio)
188 rope = RotaryEmbedding(
189 head_size=head_dim,
190 rotary_dim=rotary_dim,
191 max_position_embeddings=40960,
192 base=10000.0,
193 is_neox_style=is_neox,
194 dtype=dtype,
195 ).to(device)
196
197 baseline(
198 ref_qkv,
199 num_heads,
200 num_kv_heads,
201 num_kv_heads,
202 head_dim,
203 eps,
204 q_weight,
205 k_weight,
206 rope.cos_sin_cache,
207 is_neox,
208 positions.view(-1),
209 )
210
211 fused_qk_norm_rope(
212 ops_qkv,
213 num_heads,
214 num_kv_heads,
215 num_kv_heads,
216 head_dim,
217 eps,
218 q_weight,
219 k_weight,
220 rope.cos_sin_cache,
221 is_neox,
222 positions.view(-1),
223 )
224
225 if dtype == torch.bfloat16:
226 atol = 5e-2
227 rtol = 5e-2
228 else:
229 atol = 1e-2
230 rtol = 1e-2
231
232 torch.testing.assert_close(
233 ref_qkv,
234 ops_qkv,
235 atol=atol,
236 rtol=rtol,
237 )
238
239
240 class TestFusedQkNormRopeIntegration:
241 def test_kernel_registration_integration(self):
242 from vllm.kernels.helion.register import get_registered_kernels
243
244 registered_kernels = get_registered_kernels()
245 assert "fused_qk_norm_rope" in registered_kernels
246
247 kernel_wrapper = registered_kernels["fused_qk_norm_rope"]
248 assert kernel_wrapper.op_name == "fused_qk_norm_rope"
249 assert kernel_wrapper._config_picker is not None
250 assert kernel_wrapper._mutates_args == ["qkv"]