@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 310 lines 11.3 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """RoPE equivalence tests for the fused Kimi-K3 MLA epilogues."""
4
5 import pytest
6 import torch
7
8 from vllm.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import (
9 fused_mla_decode_q_concat_kv_cache_insert,
10 fused_mla_key_concat_ds_mla_insert,
11 fused_mla_key_concat_kv_cache_insert,
12 fused_mla_kv_concat,
13 fused_mla_kv_concat_quant_fp8,
14 fused_mla_qkv_quant_kv_cache_fp8_insert,
15 )
16 from vllm.platforms import current_platform
17
18 pytestmark = pytest.mark.skipif(
19 not current_platform.is_cuda(), reason="Kimi-K3 fused MLA requires CUDA"
20 )
21
22 _DTYPE = torch.bfloat16
23 _NUM_TOKENS = 3
24 _NUM_HEADS = 4
25 _BLOCK_SIZE = 8
26 _POSITIONS = (1, 7, 13)
27 _SLOTS = (0, 3, 9)
28
29
30 def _randn(*shape: int, dtype: torch.dtype = _DTYPE) -> torch.Tensor:
31 return torch.randn(*shape, device="cuda", dtype=dtype) * 0.2
32
33
34 def _rope_cache(max_position: int = 32) -> torch.Tensor:
35 inv_freq = 1.0 / (
36 50000 ** (torch.arange(0, 64, 2, dtype=torch.float32, device="cuda") / 64)
37 )
38 positions = torch.arange(max_position, dtype=torch.float32, device="cuda")
39 freqs = torch.outer(positions, inv_freq)
40 # The fused epilogue reads the cos/sin table in fp32 (RoPE math runs in fp32).
41 return torch.cat((freqs.cos(), freqs.sin()), dim=-1)
42
43
44 def _apply_gptj_rope(
45 x: torch.Tensor, positions: torch.Tensor, cos_sin_cache: torch.Tensor
46 ) -> torch.Tensor:
47 cos, sin = cos_sin_cache.index_select(0, positions).chunk(2, dim=-1)
48 for _ in range(x.ndim - 2):
49 cos = cos.unsqueeze(1)
50 sin = sin.unsqueeze(1)
51 x1 = x[..., ::2].float()
52 x2 = x[..., 1::2].float()
53 out1 = x1 * cos.float() - x2 * sin.float()
54 out2 = x2 * cos.float() + x1 * sin.float()
55 return torch.stack((out1, out2), dim=-1).flatten(-2).to(x.dtype)
56
57
58 def _cache_rows(cache: torch.Tensor, slots: torch.Tensor) -> torch.Tensor:
59 return cache.reshape(-1, cache.shape[-1]).index_select(0, slots)
60
61
62 def _assert_fp8_close(actual: torch.Tensor, expected: torch.Tensor) -> None:
63 torch.testing.assert_close(
64 actual.float(),
65 expected.to(torch.float8_e4m3fn).float(),
66 atol=0.03125,
67 rtol=0.15,
68 )
69
70
71 def _strided_context_inputs(
72 num_tokens: int, num_heads: int, dtype: torch.dtype, k_pe_fp8: bool
73 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
74 """The real prefill-context layouts: ``k_nope``/``v`` as the two strided
75 halves of one kv_b_proj output and ``k_pe`` as a column slice of the gather
76 workspace (left in the fp8 cache layout for a plain fp8 cache)."""
77 kv_nope = _randn(num_tokens, num_heads, 256, dtype=dtype)
78 k_nope, v = kv_nope.split((128, 128), dim=-1)
79 workspace = _randn(num_tokens, 576, dtype=dtype)
80 if k_pe_fp8:
81 workspace = workspace.to(torch.float8_e4m3fn)
82 return k_nope, workspace[:, 512:].unsqueeze(1), v
83
84
85 @pytest.mark.parametrize("num_tokens", [0, _NUM_TOKENS])
86 @pytest.mark.parametrize("num_heads", [3, _NUM_HEADS])
87 @pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float16])
88 @torch.inference_mode()
89 def test_context_kv_concat_accepts_strided_inputs(
90 num_tokens: int, num_heads: int, input_dtype: torch.dtype
91 ) -> None:
92 k_nope, k_pe, _ = _strided_context_inputs(
93 num_tokens, num_heads, input_dtype, k_pe_fp8=False
94 )
95
96 k = fused_mla_kv_concat(k_nope, k_pe)
97
98 assert k.is_contiguous()
99 assert k.dtype == input_dtype
100 torch.testing.assert_close(k[..., :128], k_nope, atol=0, rtol=0)
101 torch.testing.assert_close(
102 k[..., 128:], k_pe.expand(-1, num_heads, -1), atol=0, rtol=0
103 )
104
105
106 @pytest.mark.parametrize("k_pe_dtype", ["input", "fp8"])
107 @pytest.mark.parametrize("num_tokens", [0, _NUM_TOKENS])
108 @pytest.mark.parametrize("num_heads", [3, _NUM_HEADS])
109 @pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float16])
110 @torch.inference_mode()
111 def test_context_kv_pack_quantizes_strided_inputs(
112 k_pe_dtype: str,
113 num_tokens: int,
114 num_heads: int,
115 input_dtype: torch.dtype,
116 ) -> None:
117 k_nope, k_pe, v = _strided_context_inputs(
118 num_tokens, num_heads, input_dtype, k_pe_fp8=k_pe_dtype == "fp8"
119 )
120
121 k_actual, v_actual = fused_mla_kv_concat_quant_fp8(k_nope, k_pe, v)
122
123 fp8 = torch.float8_e4m3fn
124 k_expected = torch.empty_like(k_actual)
125 k_expected[..., :128] = k_nope.to(fp8)
126 k_expected[..., 128:] = k_pe.to(fp8)
127 assert k_actual.is_contiguous()
128 assert v_actual.is_contiguous()
129 # The pack casts with the native pairwise converters, so it must be
130 # bit-identical to torch's `.to(fp8)` rather than merely close.
131 torch.testing.assert_close(k_actual.float(), k_expected.float(), atol=0, rtol=0)
132 torch.testing.assert_close(v_actual.float(), v.to(fp8).float(), atol=0, rtol=0)
133
134
135 @pytest.mark.parametrize("cache_kind", ["bf16", "fp8", "fp8_ds_mla"])
136 @torch.inference_mode()
137 def test_prefill_epilogue_fuses_gptj_rope(cache_kind: str) -> None:
138 torch.manual_seed(0)
139 positions = torch.tensor(_POSITIONS, device="cuda", dtype=torch.int64)
140 slots = torch.tensor(_SLOTS, device="cuda", dtype=torch.int64)
141 cos_sin_cache = _rope_cache()
142 q = _randn(_NUM_TOKENS, _NUM_HEADS, 192)
143 k_nope = _randn(_NUM_TOKENS, _NUM_HEADS, 128)
144 k_pe = _randn(_NUM_TOKENS, 64)
145 kv_c = _randn(_NUM_TOKENS, 512)
146 v = _randn(_NUM_TOKENS, _NUM_HEADS, 128)
147
148 q_expected = q.clone()
149 q_expected[..., 128:] = _apply_gptj_rope(
150 q_expected[..., 128:], positions, cos_sin_cache
151 )
152 k_pe_expected = _apply_gptj_rope(k_pe, positions, cos_sin_cache)
153 k_expected = torch.cat(
154 (k_nope, k_pe_expected[:, None, :].expand(-1, _NUM_HEADS, -1)), dim=-1
155 )
156 cache_expected = torch.cat((kv_c, k_pe_expected), dim=-1)
157
158 if cache_kind == "bf16":
159 cache = torch.zeros(2, _BLOCK_SIZE, 576, device="cuda", dtype=_DTYPE)
160 q_actual = q.clone()
161 k_actual = fused_mla_key_concat_kv_cache_insert(
162 q_actual,
163 k_nope,
164 k_pe,
165 kv_c,
166 cache,
167 slots,
168 positions,
169 cos_sin_cache,
170 )
171 torch.testing.assert_close(q_actual, q_expected)
172 torch.testing.assert_close(k_actual, k_expected)
173 torch.testing.assert_close(_cache_rows(cache, slots), cache_expected)
174 elif cache_kind == "fp8":
175 cache = torch.zeros(
176 2, _BLOCK_SIZE, 576, device="cuda", dtype=torch.float8_e4m3fn
177 )
178 one = torch.ones(1, device="cuda", dtype=torch.float32)
179 q_actual, k_actual, v_actual = fused_mla_qkv_quant_kv_cache_fp8_insert(
180 q,
181 k_nope,
182 k_pe,
183 kv_c,
184 v,
185 cache,
186 slots,
187 one,
188 one,
189 one,
190 one,
191 positions,
192 cos_sin_cache,
193 )
194 _assert_fp8_close(q_actual, q_expected)
195 _assert_fp8_close(k_actual, k_expected)
196 _assert_fp8_close(v_actual, v)
197 _assert_fp8_close(_cache_rows(cache, slots), cache_expected)
198 else:
199 cache = torch.zeros(2, _BLOCK_SIZE, 656, device="cuda", dtype=torch.uint8)
200 q_actual = q.clone()
201 k_actual = fused_mla_key_concat_ds_mla_insert(
202 q_actual,
203 k_nope,
204 k_pe,
205 kv_c,
206 cache,
207 slots,
208 positions,
209 cos_sin_cache,
210 )
211 rope_cache = _cache_rows(cache, slots)[:, 528:656].view(_DTYPE)
212 torch.testing.assert_close(q_actual, q_expected)
213 torch.testing.assert_close(k_actual, k_expected)
214 torch.testing.assert_close(rope_cache, k_pe_expected)
215
216
217 @pytest.mark.parametrize("cache_kind", ["bf16", "fp8", "fp8_ds_mla"])
218 @torch.inference_mode()
219 def test_decode_epilogue_fuses_gptj_rope(cache_kind: str) -> None:
220 torch.manual_seed(1)
221 positions = torch.tensor(_POSITIONS, device="cuda", dtype=torch.int64)
222 slots = torch.tensor(_SLOTS, device="cuda", dtype=torch.int64)
223 cos_sin_cache = _rope_cache()
224 ql_nope = _randn(_NUM_TOKENS, _NUM_HEADS, 512)
225 q_pe = _randn(_NUM_TOKENS, _NUM_HEADS, 64)
226 kv_c = _randn(_NUM_TOKENS, 512)
227 k_pe = _randn(_NUM_TOKENS, 64)
228
229 q_pe_expected = _apply_gptj_rope(q_pe, positions, cos_sin_cache)
230 k_pe_expected = _apply_gptj_rope(k_pe, positions, cos_sin_cache)
231 q_expected = torch.cat((ql_nope, q_pe_expected), dim=-1)
232 cache_expected = torch.cat((kv_c, k_pe_expected), dim=-1)
233
234 kwargs = {"positions": positions, "cos_sin_cache": cos_sin_cache}
235 if cache_kind == "bf16":
236 cache = torch.zeros(2, _BLOCK_SIZE, 576, device="cuda", dtype=_DTYPE)
237 q_actual = fused_mla_decode_q_concat_kv_cache_insert(
238 ql_nope, q_pe, kv_c, k_pe, cache, slots, **kwargs
239 )
240 torch.testing.assert_close(q_actual, q_expected)
241 torch.testing.assert_close(_cache_rows(cache, slots), cache_expected)
242 elif cache_kind == "fp8":
243 cache = torch.zeros(
244 2, _BLOCK_SIZE, 576, device="cuda", dtype=torch.float8_e4m3fn
245 )
246 one = torch.ones(1, device="cuda", dtype=torch.float32)
247 q_actual = fused_mla_decode_q_concat_kv_cache_insert(
248 ql_nope,
249 q_pe,
250 kv_c,
251 k_pe,
252 cache,
253 slots,
254 q_scale_inv=one,
255 cache_scale_inv=one,
256 **kwargs,
257 )
258 _assert_fp8_close(q_actual, q_expected)
259 _assert_fp8_close(_cache_rows(cache, slots), cache_expected)
260 else:
261 cache = torch.zeros(2, _BLOCK_SIZE, 656, device="cuda", dtype=torch.uint8)
262 q_actual = fused_mla_decode_q_concat_kv_cache_insert(
263 ql_nope, q_pe, kv_c, k_pe, cache, slots, ds_mla=True, **kwargs
264 )
265 rope_cache = _cache_rows(cache, slots)[:, 528:656].view(_DTYPE)
266 torch.testing.assert_close(q_actual, q_expected)
267 torch.testing.assert_close(rope_cache, k_pe_expected)
268
269
270 @torch.inference_mode()
271 def test_decode_epilogue_preserves_nope_path() -> None:
272 torch.manual_seed(2)
273 slots = torch.tensor(_SLOTS, device="cuda", dtype=torch.int64)
274 ql_nope = _randn(_NUM_TOKENS, _NUM_HEADS, 512)
275 q_pe = _randn(_NUM_TOKENS, _NUM_HEADS, 64)
276 kv_c = _randn(_NUM_TOKENS, 512)
277 k_pe = _randn(_NUM_TOKENS, 64)
278 cache = torch.zeros(2, _BLOCK_SIZE, 576, device="cuda", dtype=_DTYPE)
279
280 q_actual = fused_mla_decode_q_concat_kv_cache_insert(
281 ql_nope, q_pe, kv_c, k_pe, cache, slots
282 )
283
284 torch.testing.assert_close(q_actual, torch.cat((ql_nope, q_pe), dim=-1))
285 torch.testing.assert_close(
286 _cache_rows(cache, slots), torch.cat((kv_c, k_pe), dim=-1)
287 )
288
289
290 @torch.inference_mode()
291 @pytest.mark.parametrize("num_tokens", [64, 65])
292 def test_decode_epilogue_row_split_boundary(num_tokens: int) -> None:
293 """Both sides of the M-based warps-per-row dispatch produce the same rows."""
294 torch.manual_seed(3)
295 num_blocks = (num_tokens + _BLOCK_SIZE - 1) // _BLOCK_SIZE
296 slots = torch.randperm(num_blocks * _BLOCK_SIZE, device="cuda")[:num_tokens].long()
297 ql_nope = _randn(num_tokens, _NUM_HEADS, 512)
298 q_pe = _randn(num_tokens, _NUM_HEADS, 64)
299 kv_c = _randn(num_tokens, 512)
300 k_pe = _randn(num_tokens, 64)
301 cache = torch.zeros(num_blocks, _BLOCK_SIZE, 576, device="cuda", dtype=_DTYPE)
302
303 q_actual = fused_mla_decode_q_concat_kv_cache_insert(
304 ql_nope, q_pe, kv_c, k_pe, cache, slots
305 )
306
307 torch.testing.assert_close(q_actual, torch.cat((ql_nope, q_pe), dim=-1))
308 torch.testing.assert_close(
309 _cache_rows(cache, slots), torch.cat((kv_c, k_pe), dim=-1)
310 )