@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 302 lines 9.21 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import importlib
5
6 import pytest
7 import torch
8 import torch.nn.functional as F
9
10 from vllm.models.kimi_k3.common.mtp import fused_mtp_input
11 from vllm.models.kimi_k3.nvidia.ops import attn_res
12 from vllm.platforms import current_platform
13
14 HIDDEN_SIZE = 7168
15 MAX_BLOCKS = 8
16 EPS = 1e-5
17 attn_res_module = importlib.import_module("vllm.models.kimi_k3.nvidia.ops.attn_res")
18
19
20 def _on_rocm_below_10() -> bool:
21 if not current_platform.is_rocm():
22 return False
23 from vllm.platforms.rocm import get_rocm_version
24
25 return (get_rocm_version() or (0,)) < (10,)
26
27
28 # The Triton bundled with ROCm < 10 (3.7.x) crashes in the AMD
29 # CanonicalizePointers pass on the kernel's tl.where over pointer tensors. Only
30 # the num_blocks > 0 loop contains it. ROCm serves AttnRes from
31 # vllm/models/kimi_k3/amd/ops/attn_res.py instead, so nothing real is lost.
32 _OLD_ROCM = _on_rocm_below_10()
33
34
35 def _skip_on_old_rocm(num_blocks: int) -> None:
36 if _OLD_ROCM and num_blocks > 0:
37 pytest.skip("Triton on ROCm < 10 cannot compile this kernel's pointer select")
38
39
40 def _randn_with_row_padding(*shape: int, padding: int = 0) -> torch.Tensor:
41 storage = torch.randn(
42 *shape[:-1],
43 shape[-1] + padding,
44 device="cuda",
45 dtype=torch.bfloat16,
46 )
47 return storage[..., : shape[-1]]
48
49
50 def _reference(
51 prefix: torch.Tensor,
52 delta: torch.Tensor | None,
53 blocks: torch.Tensor,
54 norm_weight: torch.Tensor,
55 qk_weight: torch.Tensor,
56 output_norm_weight: torch.Tensor | None,
57 num_blocks: int,
58 ) -> tuple[torch.Tensor, torch.Tensor]:
59 if delta is not None:
60 prefix = prefix + delta
61 values = torch.cat((blocks[:, :num_blocks], prefix.unsqueeze(1)), dim=1)
62 keys = F.rms_norm(values, (HIDDEN_SIZE,), norm_weight, EPS)
63 probs = (keys @ qk_weight).softmax(dim=-1)
64 output = torch.matmul(probs.unsqueeze(1), values).squeeze(1)
65 if output_norm_weight is not None:
66 output = F.rms_norm(output, (HIDDEN_SIZE,), output_norm_weight, EPS)
67 return output, prefix
68
69
70 @pytest.mark.parametrize(
71 (
72 "num_tokens",
73 "num_blocks",
74 "row_padding",
75 "write_block",
76 "has_delta",
77 "backend",
78 ),
79 [
80 pytest.param(1, 0, 0, True, False, "triton", id="triton-empty"),
81 pytest.param(1, 0, 0, True, True, "triton", id="triton-empty-add"),
82 pytest.param(17, 5, 7, True, False, "triton", id="triton-write"),
83 pytest.param(17, 5, 7, True, True, "triton", id="triton-write-add"),
84 pytest.param(3, 8, 0, False, False, "triton", id="triton-full"),
85 pytest.param(3, 8, 0, False, True, "triton", id="triton-full-add"),
86 pytest.param(1, 4, 0, False, True, "nvidia", id="nvidia-single-token"),
87 pytest.param(320, 1, 0, False, True, "nvidia", id="nvidia-1"),
88 pytest.param(320, 4, 0, False, True, "nvidia", id="nvidia-4"),
89 pytest.param(320, 8, 0, False, True, "nvidia", id="nvidia-8"),
90 ],
91 )
92 def test_attn_res(
93 num_tokens: int,
94 num_blocks: int,
95 row_padding: int,
96 write_block: bool,
97 has_delta: bool,
98 backend: str,
99 ):
100 if backend == "nvidia" and not current_platform.is_device_capability_family(100):
101 pytest.skip("NVIDIA AttnRes requires the SM100 family")
102 _skip_on_old_rocm(num_blocks)
103
104 prefix = _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=row_padding)
105 delta = (
106 _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=row_padding)
107 if has_delta
108 else None
109 )
110 blocks = _randn_with_row_padding(
111 num_tokens, MAX_BLOCKS, HIDDEN_SIZE, padding=row_padding
112 )
113 norm_weight = 1 + 0.1 * torch.randn(
114 HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
115 )
116 qk_weight = (
117 torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5
118 )
119 output_norm_weight = 1 + 0.1 * torch.randn(
120 HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
121 )
122 original_blocks = blocks.clone()
123 expected, expected_prefix = _reference(
124 prefix.clone(),
125 delta,
126 blocks,
127 norm_weight,
128 qk_weight,
129 output_norm_weight,
130 num_blocks,
131 )
132 block_write_idx = num_blocks if write_block else -1
133
134 actual = attn_res(
135 prefix,
136 delta,
137 blocks,
138 norm_weight,
139 qk_weight,
140 output_norm_weight,
141 num_blocks,
142 block_write_idx,
143 EPS,
144 EPS,
145 )
146
147 torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
148 torch.testing.assert_close(prefix, expected_prefix, atol=0, rtol=0)
149 if write_block:
150 original_blocks[:, block_write_idx].copy_(expected_prefix)
151 torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0)
152 assert actual.is_contiguous()
153
154
155 @pytest.mark.parametrize("num_blocks", range(MAX_BLOCKS + 1))
156 def test_attn_res_block_counts(num_blocks: int):
157 _skip_on_old_rocm(num_blocks)
158 prefix = torch.randn(1, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
159 blocks = torch.randn(
160 1, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
161 )
162 norm_weight = torch.ones(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
163 qk_weight = (
164 torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5
165 )
166 output_norm_weight = torch.ones_like(norm_weight)
167 expected, _ = _reference(
168 prefix.clone(),
169 None,
170 blocks,
171 norm_weight,
172 qk_weight,
173 output_norm_weight,
174 num_blocks,
175 )
176
177 actual = attn_res(
178 prefix,
179 None,
180 blocks,
181 norm_weight,
182 qk_weight,
183 output_norm_weight,
184 num_blocks,
185 -1,
186 EPS,
187 EPS,
188 )
189
190 torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
191
192
193 def test_attn_res_without_output_norm():
194 _skip_on_old_rocm(MAX_BLOCKS)
195 prefix = torch.randn(7, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
196 delta = torch.randn_like(prefix)
197 blocks = torch.randn(
198 7, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
199 )
200 norm_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
201 qk_weight = (
202 torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16) / HIDDEN_SIZE**0.5
203 )
204 expected, _ = _reference(
205 prefix.clone(), delta, blocks, norm_weight, qk_weight, None, MAX_BLOCKS
206 )
207
208 actual = attn_res(
209 prefix,
210 delta,
211 blocks,
212 norm_weight,
213 qk_weight,
214 None,
215 MAX_BLOCKS,
216 -1,
217 EPS,
218 0.0,
219 )
220
221 torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
222
223
224 @pytest.mark.parametrize(
225 ("num_blocks", "has_delta", "has_output_norm", "write_block"),
226 [
227 (0, False, True, True),
228 (4, True, True, True),
229 (4, False, True, False),
230 (8, True, False, False),
231 ],
232 )
233 def test_sm100_variants_do_not_fall_back_to_triton(
234 monkeypatch,
235 num_blocks: int,
236 has_delta: bool,
237 has_output_norm: bool,
238 write_block: bool,
239 ):
240 if not current_platform.is_device_capability_family(100) or not hasattr(
241 torch.ops._C, "kimi_k3_attn_res"
242 ):
243 pytest.skip("native AttnRes requires an SM100 build")
244
245 class FailingKernel:
246 def __getitem__(self, grid):
247 raise AssertionError(f"unexpected Triton launch with grid {grid}")
248
249 monkeypatch.setattr(attn_res_module, "_attn_res_kernel", FailingKernel())
250 prefix = torch.randn(1, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
251 delta = torch.zeros_like(prefix) if has_delta else None
252 blocks = torch.randn(
253 1, MAX_BLOCKS, HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16
254 )
255 norm_weight = torch.ones(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
256 qk_weight = torch.randn_like(norm_weight) / HIDDEN_SIZE**0.5
257 output_norm_weight = torch.ones_like(norm_weight) if has_output_norm else None
258
259 attn_res(
260 prefix,
261 delta,
262 blocks,
263 norm_weight,
264 qk_weight,
265 output_norm_weight,
266 num_blocks,
267 num_blocks if write_block else -1,
268 EPS,
269 EPS,
270 )
271
272
273 @pytest.mark.parametrize("num_tokens", [0, 1, 17])
274 def test_fused_mtp_input(num_tokens: int):
275 positions = torch.arange(num_tokens, device="cuda")
276 inputs_embeds = _randn_with_row_padding(num_tokens, HIDDEN_SIZE, padding=7)
277 previous_hidden_states = _randn_with_row_padding(
278 num_tokens, HIDDEN_SIZE, padding=11
279 )
280 enorm_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
281 hnorm_weight = torch.randn(HIDDEN_SIZE, device="cuda", dtype=torch.bfloat16)
282
283 masked_inputs_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, inputs_embeds)
284 expected = torch.cat(
285 (
286 F.rms_norm(masked_inputs_embeds, (HIDDEN_SIZE,), enorm_weight, EPS),
287 F.rms_norm(previous_hidden_states, (HIDDEN_SIZE,), hnorm_weight, EPS),
288 ),
289 dim=-1,
290 )
291 actual = fused_mtp_input(
292 positions,
293 inputs_embeds,
294 previous_hidden_states,
295 enorm_weight,
296 hnorm_weight,
297 EPS,
298 )
299
300 torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)
301 assert actual.shape == (num_tokens, 2 * HIDDEN_SIZE)
302 assert actual.is_contiguous()