@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 193 lines 5.49 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 import torch.nn.functional as F
7
8 from vllm.models.kimi_k3.amd.ops.attn_res import attn_res
9 from vllm.platforms import current_platform
10
11 pytestmark = pytest.mark.skipif(
12 not current_platform.is_rocm(),
13 reason="AMD AttnRes requires ROCm",
14 )
15
16
17 def _randn_with_row_padding(*shape: int, padding: int = 0) -> torch.Tensor:
18 storage = torch.randn(
19 *shape[:-1],
20 shape[-1] + padding,
21 device="cuda",
22 dtype=torch.bfloat16,
23 )
24 return storage[..., : shape[-1]]
25
26
27 def _reference(
28 prefix: torch.Tensor,
29 blocks: torch.Tensor,
30 norm_weight: torch.Tensor,
31 qk_weight: torch.Tensor,
32 num_blocks: int,
33 eps: float,
34 ) -> torch.Tensor:
35 hidden_size = prefix.shape[-1]
36 values = torch.cat((blocks[:, :num_blocks], prefix.unsqueeze(1)), dim=1)
37 keys = F.rms_norm(values, (hidden_size,), norm_weight, eps)
38 probs = (keys @ qk_weight).softmax(dim=-1)
39 return torch.matmul(probs.unsqueeze(1), values).squeeze(1)
40
41
42 @pytest.mark.parametrize(
43 (
44 "num_tokens",
45 "num_blocks",
46 "block_capacity",
47 "hidden_size",
48 "row_padding",
49 ),
50 [
51 pytest.param(0, 3, 5, 128, 0, id="empty"),
52 pytest.param(1, 1, 2, 128, 0, id="decode-single"),
53 pytest.param(17, 4, 6, 1024, 7, id="decode-padded"),
54 pytest.param(320, 8, 10, 7168, 0, id="prefill-full"),
55 ],
56 )
57 def test_amd_attn_res_matches_reference(
58 num_tokens: int,
59 num_blocks: int,
60 block_capacity: int,
61 hidden_size: int,
62 row_padding: int,
63 ) -> None:
64 eps = 1e-5
65 prefix = _randn_with_row_padding(num_tokens, hidden_size, padding=row_padding)
66 blocks = _randn_with_row_padding(
67 num_tokens,
68 block_capacity,
69 hidden_size,
70 padding=row_padding,
71 )
72 norm_weight = 1 + 0.1 * torch.randn(
73 hidden_size, device="cuda", dtype=torch.bfloat16
74 )
75 qk_weight = (
76 torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16) / hidden_size**0.5
77 )
78 expected = _reference(
79 prefix,
80 blocks,
81 norm_weight,
82 qk_weight,
83 num_blocks,
84 eps,
85 )
86 original_prefix = prefix.clone()
87 original_blocks = blocks.clone()
88
89 actual = attn_res(
90 prefix,
91 None,
92 blocks,
93 norm_weight,
94 qk_weight,
95 None,
96 num_blocks,
97 -1,
98 eps,
99 0.0,
100 )
101
102 torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
103 torch.testing.assert_close(prefix, original_prefix, atol=0, rtol=0)
104 torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0)
105 assert actual.shape == prefix.shape
106 assert actual.is_contiguous()
107
108
109 @pytest.mark.parametrize(
110 (
111 "num_tokens",
112 "num_blocks",
113 "hidden_size",
114 "has_delta",
115 "write_block",
116 "apply_output_norm",
117 ),
118 [
119 pytest.param(1, 0, 128, False, True, True, id="empty-write-norm"),
120 pytest.param(7, 1, 1024, True, False, True, id="single-add-norm"),
121 pytest.param(17, 5, 7168, True, True, True, id="padded-write-add"),
122 pytest.param(3, 8, 7168, True, False, True, id="full-add-norm"),
123 pytest.param(320, 4, 7168, True, False, False, id="prefill-add"),
124 ],
125 )
126 def test_amd_attn_res_fused_contract(
127 num_tokens: int,
128 num_blocks: int,
129 hidden_size: int,
130 has_delta: bool,
131 write_block: bool,
132 apply_output_norm: bool,
133 ) -> None:
134 torch.manual_seed(42)
135 eps = 1e-5
136 output_eps = 2e-5
137 block_capacity = 9
138 prefix = _randn_with_row_padding(num_tokens, hidden_size, padding=7)
139 delta = (
140 _randn_with_row_padding(num_tokens, hidden_size, padding=11)
141 if has_delta
142 else None
143 )
144 blocks = _randn_with_row_padding(
145 num_tokens, block_capacity, hidden_size, padding=13
146 )
147 norm_weight = 1 + 0.1 * torch.randn(
148 hidden_size, device="cuda", dtype=torch.bfloat16
149 )
150 qk_weight = (
151 torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16) / hidden_size**0.5
152 )
153 output_norm_weight = (
154 1 + 0.1 * torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16)
155 if apply_output_norm
156 else None
157 )
158 expected_prefix = prefix.clone()
159 if delta is not None:
160 expected_prefix = expected_prefix + delta
161 values = torch.cat(
162 (blocks[:, :num_blocks].clone(), expected_prefix.unsqueeze(1)), dim=1
163 )
164 keys = F.rms_norm(values.float(), (hidden_size,), norm_weight.float(), eps)
165 probs = (keys @ qk_weight.float()).softmax(dim=-1)
166 expected = torch.matmul(probs.unsqueeze(1), values.float()).squeeze(1)
167 if output_norm_weight is not None:
168 expected = F.rms_norm(
169 expected, (hidden_size,), output_norm_weight.float(), output_eps
170 )
171 expected = expected.to(prefix.dtype)
172 original_blocks = blocks.clone()
173 block_write_idx = num_blocks if write_block else -1
174
175 actual = attn_res(
176 prefix,
177 delta,
178 blocks,
179 norm_weight,
180 qk_weight,
181 output_norm_weight,
182 num_blocks,
183 block_write_idx,
184 eps,
185 output_eps,
186 )
187
188 torch.testing.assert_close(actual, expected, atol=8e-2, rtol=3e-2)
189 torch.testing.assert_close(prefix, expected_prefix, atol=0, rtol=0)
190 if write_block:
191 original_blocks[:, block_write_idx].copy_(expected_prefix)
192 torch.testing.assert_close(blocks, original_blocks, atol=0, rtol=0)
193 assert actual.is_contiguous()