@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 471 lines 17.8 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Correctness tests for the AMX-only CPU MLA backend: the vendored
4 decode/extend/bmm kernels, the KV cache write, and the ``AMXMLAImpl``
5 backend built on top of them.
6 """
7
8 import pytest
9 import torch
10
11 from vllm.platforms import current_platform
12
13 if not current_platform.is_cpu():
14 pytest.skip("skipping CPU-only tests", allow_module_level=True)
15
16 if not torch.cpu._is_amx_tile_supported():
17 pytest.skip("AMX MLA requires an AMX-capable host", allow_module_level=True)
18
19 torch.cpu._init_amx()
20
21 from vllm import _custom_ops as ops # noqa: E402
22 from vllm.utils.torch_utils import set_random_seed # noqa: E402
23 from vllm.v1.attention.backends.mla.amx_mla import ( # noqa: E402
24 AMXMLAImpl,
25 _compute_num_kv_splits,
26 _expand_block_table,
27 )
28
29 KV_LORA_RANK = 64
30 QK_ROPE_HEAD_DIM = 32
31 HEAD_SIZE = KV_LORA_RANK + QK_ROPE_HEAD_DIM # 96, kv-cache row width
32 # The extend/decode kernels transpose by 512-bit lanes: both the cache row
33 # width and the (kv_lora_rank-wide) value width must be 32-element aligned.
34 DTYPE = torch.bfloat16
35 ATOL = 2e-2
36 RTOL = 2e-2
37 # Backend-level tests chain absorb -> attend -> de-absorb through bf16,
38 # accumulating more rounding error than a single kernel op.
39 _IMPL_ATOL = 1.5e-1
40 _IMPL_RTOL = 1.5e-1
41
42
43 def _flatten_cache(kv_cache: torch.Tensor) -> torch.Tensor:
44 """(num_blocks, block_size, head_size) -> (num_slots, head_size)."""
45 return kv_cache.view(-1, kv_cache.size(-1))
46
47
48 def _random_paged_cache(num_slots: int, head_size: int) -> torch.Tensor:
49 return torch.randn(num_slots, head_size, dtype=DTYPE)
50
51
52 def _ref_latent_attn(
53 q: torch.Tensor,
54 keys: torch.Tensor,
55 scale: float,
56 kv_lora_rank: int,
57 ) -> torch.Tensor:
58 """Plain-PyTorch causal-free latent-space MQA attention for one request.
59
60 q: (num_heads, head_size); keys: (ctx_len, head_size). Values are the
61 first ``kv_lora_rank`` columns of the same latent rows (MLA's K/V
62 aliasing), matching what decode_attention_cpu/extend_attention_cpu
63 compute internally for the MLA-shaped case.
64 """
65 values = keys[:, :kv_lora_rank]
66 scores = (q.float() @ keys.float().T) * scale
67 probs = torch.softmax(scores, dim=-1)
68 return (probs @ values.float()).to(q.dtype)
69
70
71 def test_bmm_cpu_matches_torch_bmm():
72 set_random_seed(0)
73 # bmm_cpu requires the output's last dim (mat2's out-features) to be a
74 # multiple of 32 (tinygemm's tile width); the contraction dim is free.
75 n, b, p, l = 8, 32, 32, 32 # noqa: E741
76 mat1 = torch.randn(n, b, p, dtype=DTYPE)
77 mat2 = torch.randn(n, l, p, dtype=DTYPE) # (N, OUT, IN) Linear convention
78
79 ref = torch.bmm(mat1.float(), mat2.float().transpose(1, 2)).to(DTYPE)
80
81 out = torch.empty(n, b, l, dtype=DTYPE)
82 ops.bmm_cpu(out, mat1, mat2, False, None)
83 torch.testing.assert_close(out, ref, atol=ATOL, rtol=RTOL)
84
85 packed = torch.ops._C.convert_weight_packed(mat2.contiguous())
86 out_vnni = torch.empty(n, b, l, dtype=DTYPE)
87 ops.bmm_cpu(out_vnni, mat1, packed, True, None)
88 torch.testing.assert_close(out_vnni, ref, atol=ATOL, rtol=RTOL)
89
90
91 def test_concat_and_cache_mla_round_trip():
92 set_random_seed(0)
93 num_tokens = 37
94 num_blocks, block_size = 8, 16
95 kv_cache = torch.zeros(num_blocks, block_size, HEAD_SIZE, dtype=DTYPE)
96
97 kv_c_normed = torch.randn(num_tokens, KV_LORA_RANK, dtype=DTYPE)
98 k_pe = torch.randn(num_tokens, QK_ROPE_HEAD_DIM, dtype=DTYPE)
99 slot_mapping = torch.randperm(num_blocks * block_size)[:num_tokens].to(torch.int64)
100
101 ops.amx_mla_concat_and_cache(kv_c_normed, k_pe, kv_cache, slot_mapping)
102
103 flat = _flatten_cache(kv_cache)
104 read_back = flat[slot_mapping]
105 expected = torch.cat([kv_c_normed, k_pe], dim=-1)
106 torch.testing.assert_close(read_back, expected)
107
108 untouched = torch.ones(num_blocks * block_size, dtype=torch.bool)
109 untouched[slot_mapping] = False
110 assert (flat[untouched] == 0).all()
111
112
113 @pytest.mark.parametrize("seq_lens", [[5], [1, 300, 33], [513, 1, 129, 7]])
114 def test_decode_attention_cpu_matches_reference(seq_lens):
115 set_random_seed(0)
116 num_seqs = len(seq_lens)
117 num_heads = 4
118 block_size = 32
119 max_len = max(seq_lens)
120 max_blocks = (max_len + block_size - 1) // block_size
121 num_blocks = num_seqs * max_blocks
122
123 kv_cache = _random_paged_cache(num_blocks * block_size, HEAD_SIZE)
124 block_table = torch.randperm(num_blocks).view(num_seqs, max_blocks).to(torch.int64)
125 req_to_token = _expand_block_table(block_table, block_size)
126 req_pool_indices = torch.arange(num_seqs, dtype=torch.int64)
127 seq_lens_t = torch.tensor(seq_lens, dtype=torch.int64)
128 scale = HEAD_SIZE**-0.5
129
130 q = torch.randn(num_seqs, num_heads, HEAD_SIZE, dtype=DTYPE)
131 kv_cache_flat = kv_cache.view(-1, 1, HEAD_SIZE)
132 v_buffer = kv_cache_flat[..., :KV_LORA_RANK]
133
134 o = torch.zeros(num_seqs, num_heads, KV_LORA_RANK, dtype=DTYPE)
135 attn_logits = torch.zeros(
136 num_seqs, num_heads, 4, KV_LORA_RANK + 1, dtype=torch.float32
137 )
138 ops.cpu_mla_decode(
139 q,
140 kv_cache_flat,
141 v_buffer,
142 o,
143 None,
144 None,
145 None,
146 attn_logits,
147 req_to_token,
148 req_pool_indices,
149 seq_lens_t,
150 scale,
151 0.0,
152 False,
153 0,
154 None,
155 None,
156 )
157
158 flat_cache = _flatten_cache(kv_cache)
159 for i, seq_len in enumerate(seq_lens):
160 token_ids = req_to_token[i, :seq_len]
161 keys = flat_cache[token_ids]
162 ref = _ref_latent_attn(q[i], keys, scale, KV_LORA_RANK)
163 torch.testing.assert_close(o[i], ref, atol=ATOL, rtol=RTOL)
164
165
166 @pytest.mark.parametrize(
167 "prefix_lens,extend_lens",
168 [
169 ([0], [9]), # fresh prefill, no cached prefix
170 ([64], [17]), # continuation of a cached prefix
171 ([0, 40, 128], [23, 5, 31]), # mixed batch
172 ],
173 )
174 def test_extend_attention_cpu_matches_reference(prefix_lens, extend_lens):
175 set_random_seed(0)
176 num_seqs = len(prefix_lens)
177 num_heads = 4
178 block_size = 32
179 seq_lens = [p + e for p, e in zip(prefix_lens, extend_lens)]
180 max_len = max(seq_lens)
181 max_blocks = (max_len + block_size - 1) // block_size
182 num_blocks = num_seqs * max_blocks
183
184 kv_cache = _random_paged_cache(num_blocks * block_size, HEAD_SIZE)
185 block_table = torch.randperm(num_blocks).view(num_seqs, max_blocks).to(torch.int64)
186 req_to_token = _expand_block_table(block_table, block_size)
187 req_pool_indices = torch.arange(num_seqs, dtype=torch.int64)
188 seq_lens_t = torch.tensor(seq_lens, dtype=torch.int64)
189 extend_seq_lens_t = torch.tensor(extend_lens, dtype=torch.int64)
190 extend_start_loc = torch.cumsum(
191 torch.tensor([0, *extend_lens[:-1]], dtype=torch.int64), dim=0
192 )
193 scale = HEAD_SIZE**-0.5
194 total_new_tokens = sum(extend_lens)
195
196 # The new tokens' latent K/V, already written into the cache (mirrors
197 # do_kv_cache_update running before forward_mha in real usage).
198 k_extend_flat = torch.randn(total_new_tokens, HEAD_SIZE, dtype=DTYPE)
199 flat_cache = _flatten_cache(kv_cache)
200 for i in range(num_seqs):
201 prefix_len = prefix_lens[i]
202 ext_len = extend_lens[i]
203 new_token_ids = req_to_token[i, prefix_len : prefix_len + ext_len]
204 start = extend_start_loc[i].item()
205 flat_cache[new_token_ids] = k_extend_flat[start : start + ext_len]
206
207 q_extend = torch.randn(total_new_tokens, num_heads, HEAD_SIZE, dtype=DTYPE)
208 k_extend = k_extend_flat.unsqueeze(1)
209 v_extend = k_extend[..., :KV_LORA_RANK]
210 kv_cache_flat = kv_cache.view(-1, 1, HEAD_SIZE)
211 v_buffer = kv_cache_flat[..., :KV_LORA_RANK]
212
213 o_extend = torch.empty(total_new_tokens, num_heads, KV_LORA_RANK, dtype=DTYPE)
214 ops.cpu_mla_extend(
215 q_extend,
216 k_extend,
217 v_extend,
218 o_extend,
219 kv_cache_flat,
220 v_buffer,
221 req_to_token,
222 req_pool_indices,
223 seq_lens_t,
224 extend_seq_lens_t,
225 extend_start_loc,
226 max(extend_lens),
227 scale,
228 0.0,
229 False,
230 0,
231 None,
232 None,
233 None,
234 )
235
236 for i in range(num_seqs):
237 prefix_len = prefix_lens[i]
238 ext_len = extend_lens[i]
239 start = extend_start_loc[i].item()
240 for t in range(ext_len):
241 visible = prefix_len + t + 1
242 token_ids = req_to_token[i, :visible]
243 keys = flat_cache[token_ids]
244 ref = _ref_latent_attn(q_extend[start + t], keys, scale, KV_LORA_RANK)
245 torch.testing.assert_close(o_extend[start + t], ref, atol=ATOL, rtol=RTOL)
246
247
248 class _FakeLinear:
249 """Minimal duck-typed stand-in for kv_b_proj: only .weight and
250 .quant_method are read by get_and_maybe_dequant_weights()."""
251
252 def __init__(self, weight: torch.Tensor):
253 self.weight = weight
254 self.quant_method = None
255
256
257 def _make_amx_mla_impl(num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank):
258 weight = torch.randn(
259 num_heads * (qk_nope_head_dim + v_head_dim), kv_lora_rank, dtype=DTYPE
260 )
261 kv_b_proj = _FakeLinear(weight)
262 impl = AMXMLAImpl(
263 num_heads=num_heads,
264 head_size=kv_lora_rank + QK_ROPE_HEAD_DIM,
265 scale=(qk_nope_head_dim + QK_ROPE_HEAD_DIM) ** -0.5,
266 num_kv_heads=1,
267 alibi_slopes=None,
268 sliding_window=None,
269 kv_cache_dtype="auto",
270 logits_soft_cap=None,
271 attn_type="decoder",
272 kv_sharing_target_layer_name=None,
273 q_lora_rank=None,
274 kv_lora_rank=kv_lora_rank,
275 qk_nope_head_dim=qk_nope_head_dim,
276 qk_rope_head_dim=QK_ROPE_HEAD_DIM,
277 qk_head_dim=qk_nope_head_dim + QK_ROPE_HEAD_DIM,
278 v_head_dim=v_head_dim,
279 kv_b_proj=kv_b_proj,
280 )
281 impl.process_weights_after_loading(DTYPE)
282 return impl, weight
283
284
285 class _FakePrefillMetadata:
286 def __init__(
287 self, block_table, cpu_seq_lens, query_start_loc, req_to_token, req_pool_indices
288 ):
289 self.block_table = block_table
290 self.cpu_seq_lens = cpu_seq_lens
291 query_start_loc_i64 = query_start_loc.to(torch.int64)
292 extend_seq_lens = query_start_loc_i64[1:] - query_start_loc_i64[:-1]
293 self.extend_seq_lens = extend_seq_lens
294 self.extend_start_loc = query_start_loc_i64[:-1]
295 self.max_len_extend = int(extend_seq_lens.max().item())
296 self.req_to_token = req_to_token
297 self.req_pool_indices = req_pool_indices
298
299
300 class _FakeDecodeMetadata:
301 def __init__(self, block_table, seq_lens, req_to_token, req_pool_indices):
302 self.block_table = block_table
303 self.seq_lens_i64 = seq_lens.to(torch.int64)
304 self.req_to_token = req_to_token
305 self.req_pool_indices = req_pool_indices
306 self.num_kv_splits = _compute_num_kv_splits(
307 int(seq_lens.max().item()), current_platform.num_compute_units()
308 )
309
310
311 class _FakeAttnMetadata:
312 def __init__(self, decode=None, prefill=None, max_seq_len=0):
313 self.decode = decode
314 self.prefill = prefill
315 self.max_seq_len = max_seq_len
316
317
318 def test_amx_mla_impl_forward_mqa_matches_reference(default_vllm_config):
319 """forward_mqa receives already-absorbed Q (as MLAAttention.forward_impl
320 would produce via layer.W_UK_T) and must reproduce plain per-head MLA
321 decode attention through W_UK/W_UV."""
322 set_random_seed(1)
323 num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank = 4, 32, 32, 64
324 impl, kv_b_weight = _make_amx_mla_impl(
325 num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank
326 )
327 w_uk, w_uv = kv_b_weight.T.view(
328 kv_lora_rank, num_heads, qk_nope_head_dim + v_head_dim
329 ).split([qk_nope_head_dim, v_head_dim], dim=-1)
330
331 seq_lens = [1, 40, 5]
332 num_seqs = len(seq_lens)
333 block_size = 32
334 max_blocks = (max(seq_lens) + block_size - 1) // block_size
335 num_blocks = num_seqs * max_blocks
336 head_size = kv_lora_rank + QK_ROPE_HEAD_DIM
337
338 kv_cache = torch.randn(num_blocks, block_size, head_size, dtype=DTYPE)
339 block_table = torch.randperm(num_blocks).view(num_seqs, max_blocks).to(torch.int64)
340 req_to_token = _expand_block_table(block_table, block_size)
341
342 q_nope = torch.randn(num_seqs, num_heads, qk_nope_head_dim, dtype=DTYPE)
343 q_pe = torch.randn(num_seqs, num_heads, QK_ROPE_HEAD_DIM, dtype=DTYPE)
344 # Pre-absorb, mirroring MLAAttention.forward_impl's own bmm against W_UK_T.
345 ql_nope = torch.einsum("bnp,lnp->bnl", q_nope.float(), w_uk.float()).to(DTYPE)
346
347 attn_metadata = _FakeAttnMetadata(
348 decode=_FakeDecodeMetadata(
349 block_table=block_table,
350 seq_lens=torch.tensor(seq_lens, dtype=torch.int64),
351 req_to_token=req_to_token,
352 req_pool_indices=torch.arange(num_seqs, dtype=torch.int64),
353 ),
354 max_seq_len=max(seq_lens),
355 )
356 o, lse = impl.forward_mqa((ql_nope, q_pe), kv_cache, attn_metadata, layer=None)
357 assert lse is None
358 # De-absorb, mirroring MLAAttention.forward_impl's own bmm against
359 # layer._v_up_proj, to get back into real per-head V space for
360 # comparison against a from-scratch reference.
361 out_real = torch.einsum("bnl,lnv->bnv", o.float(), w_uv.float())
362
363 flat_cache = _flatten_cache(kv_cache)
364 scale = impl.scale
365 for i, seq_len in enumerate(seq_lens):
366 token_ids = req_to_token[i, :seq_len]
367 keys_latent = flat_cache[token_ids] # (L, head_size)
368 k_nope_latent, k_pe = keys_latent.split(
369 [kv_lora_rank, QK_ROPE_HEAD_DIM], dim=-1
370 )
371 k_real = torch.einsum("lp,pnd->nld", k_nope_latent.float(), w_uk.float())
372 k_pe_b = k_pe.float().unsqueeze(0).expand(num_heads, -1, -1)
373 k_full = torch.cat([k_real, k_pe_b], dim=-1) # (num_heads, L, qk_head_dim)
374 q_full = torch.cat(
375 [q_nope[i], q_pe[i]], dim=-1
376 ).float() # (num_heads, qk_head_dim)
377 scores = torch.einsum("nd,nld->nl", q_full, k_full) * scale
378 probs = torch.softmax(scores, dim=-1)
379 v_real = torch.einsum("lp,pnv->nlv", k_nope_latent.float(), w_uv.float())
380 ref = torch.einsum("nl,nlv->nv", probs, v_real)
381 torch.testing.assert_close(out_real[i], ref, atol=_IMPL_ATOL, rtol=_IMPL_RTOL)
382
383
384 def test_amx_mla_impl_forward_mha_matches_reference(default_vllm_config):
385 """forward_mha receives raw unabsorbed Q and must attend correctly for a
386 prefill batch with a mix of fresh and continued (cached-prefix)
387 sequences."""
388 set_random_seed(2)
389 num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank = 4, 32, 32, 64
390 impl, kv_b_weight = _make_amx_mla_impl(
391 num_heads, qk_nope_head_dim, v_head_dim, kv_lora_rank
392 )
393 w_uk, w_uv = kv_b_weight.T.view(
394 kv_lora_rank, num_heads, qk_nope_head_dim + v_head_dim
395 ).split([qk_nope_head_dim, v_head_dim], dim=-1)
396
397 prefix_lens = [0, 20]
398 extend_lens = [13, 7]
399 seq_lens = [p + e for p, e in zip(prefix_lens, extend_lens)]
400 num_seqs = len(prefix_lens)
401 block_size = 32
402 max_blocks = (max(seq_lens) + block_size - 1) // block_size
403 num_blocks = num_seqs * max_blocks
404 head_size = kv_lora_rank + QK_ROPE_HEAD_DIM
405
406 kv_cache = torch.randn(num_blocks, block_size, head_size, dtype=DTYPE)
407 block_table = torch.randperm(num_blocks).view(num_seqs, max_blocks).to(torch.int64)
408 req_to_token = _expand_block_table(block_table, block_size)
409 flat_cache = _flatten_cache(kv_cache)
410
411 total_new_tokens = sum(extend_lens)
412 query_start_loc = torch.tensor(
413 [0, *torch.cumsum(torch.tensor(extend_lens), dim=0).tolist()],
414 dtype=torch.int64,
415 )
416
417 q_nope = torch.randn(total_new_tokens, num_heads, qk_nope_head_dim, dtype=DTYPE)
418 q_pe = torch.randn(total_new_tokens, num_heads, QK_ROPE_HEAD_DIM, dtype=DTYPE)
419 q = torch.cat([q_nope, q_pe], dim=-1)
420
421 kv_c_normed = torch.randn(total_new_tokens, kv_lora_rank, dtype=DTYPE)
422 k_pe = torch.randn(total_new_tokens, 1, QK_ROPE_HEAD_DIM, dtype=DTYPE)
423
424 # Write the new tokens' latent K/V into the cache ahead of time, mirroring
425 # do_kv_cache_update running before forward_mha in real usage.
426 for i in range(num_seqs):
427 prefix_len = prefix_lens[i]
428 ext_len = extend_lens[i]
429 new_token_ids = req_to_token[i, prefix_len : prefix_len + ext_len]
430 start = query_start_loc[i].item()
431 flat_cache[new_token_ids, :kv_lora_rank] = kv_c_normed[start : start + ext_len]
432 flat_cache[new_token_ids, kv_lora_rank:] = k_pe[start : start + ext_len, 0]
433
434 attn_metadata = _FakeAttnMetadata(
435 prefill=_FakePrefillMetadata(
436 block_table=block_table,
437 cpu_seq_lens=torch.tensor(seq_lens, dtype=torch.int64),
438 query_start_loc=query_start_loc,
439 req_to_token=req_to_token,
440 req_pool_indices=torch.arange(num_seqs, dtype=torch.int64),
441 )
442 )
443 output = torch.empty(total_new_tokens, num_heads * v_head_dim, dtype=DTYPE)
444 impl.forward_mha(
445 q, kv_c_normed, k_pe, kv_cache, attn_metadata, k_scale=None, output=output
446 )
447 out_view = output.view(total_new_tokens, num_heads, v_head_dim).float()
448
449 scale = impl.scale
450 for i in range(num_seqs):
451 prefix_len = prefix_lens[i]
452 ext_len = extend_lens[i]
453 start = query_start_loc[i].item()
454 for t in range(ext_len):
455 visible = prefix_len + t + 1
456 token_ids = req_to_token[i, :visible]
457 keys_latent = flat_cache[token_ids]
458 k_nope_latent, k_pe_ref = keys_latent.split(
459 [kv_lora_rank, QK_ROPE_HEAD_DIM], dim=-1
460 )
461 k_real = torch.einsum("lp,pnd->nld", k_nope_latent.float(), w_uk.float())
462 k_pe_b = k_pe_ref.float().unsqueeze(0).expand(num_heads, -1, -1)
463 k_full = torch.cat([k_real, k_pe_b], dim=-1)
464 q_full = torch.cat([q_nope[start + t], q_pe[start + t]], dim=-1).float()
465 scores = torch.einsum("nd,nld->nl", q_full, k_full) * scale
466 probs = torch.softmax(scores, dim=-1)
467 v_real = torch.einsum("lp,pnv->nlv", k_nope_latent.float(), w_uv.float())
468 ref = torch.einsum("nl,nlv->nv", probs, v_real)
469 torch.testing.assert_close(
470 out_view[start + t], ref, atol=_IMPL_ATOL, rtol=_IMPL_RTOL
471 )