@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 316 lines 11 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Smoke tests for the CPU MLA backend.
4
5 Two levels of coverage are provided:
6
7 * :func:`test_kv_cache_cpu_write` checks that the kv-cache write logic
8 inlined in the CPU MLA backend writes the latent KV cache in the same
9 layout the CPU decode kernel expects. This runs without downloading
10 any model weights and is safe for CI.
11 * :func:`test_cpu_mla_backend_smoke` exercises the full end-to-end path
12 (LLM engine, model construction, prefill + decode) against a shrunk
13 DeepSeek-V2-Lite with `hf_overrides`. It uses `dummy` weights so the
14 outputs are not meaningful, only shape / plumbing correctness is
15 asserted. Marked `cpu_model` so it is skipped on non-CPU CI.
16 """
17
18 import pytest
19 import torch
20
21 from vllm.platforms import current_platform
22 from vllm.v1.attention.backends.mla.prefill.cpu_sdpa import (
23 CPUSDPAMLAPrefillBackend,
24 )
25 from vllm.v1.attention.backends.mla.prefill.selector import (
26 get_mla_prefill_backend,
27 )
28 from vllm.v1.attention.backends.mla.prefill.zen_cpu_sdpa import (
29 ZenCPUSDPAMLAPrefillBackend,
30 )
31
32
33 def _reference_ragged_attention(
34 q: torch.Tensor,
35 k: torch.Tensor,
36 v: torch.Tensor,
37 q_cu_seq_lens: torch.Tensor,
38 kv_cu_seq_lens: torch.Tensor,
39 *,
40 scale: float,
41 causal: bool,
42 ) -> tuple[torch.Tensor, torch.Tensor]:
43 outputs: list[torch.Tensor] = []
44 lses: list[torch.Tensor] = []
45 num_reqs = q_cu_seq_lens.numel() - 1
46 for req_idx in range(num_reqs):
47 q_start = int(q_cu_seq_lens[req_idx])
48 q_end = int(q_cu_seq_lens[req_idx + 1])
49 kv_start = int(kv_cu_seq_lens[req_idx])
50 kv_end = int(kv_cu_seq_lens[req_idx + 1])
51
52 q_r = q[q_start:q_end].transpose(0, 1).float()
53 k_r = k[kv_start:kv_end].transpose(0, 1).float()
54 v_r = v[kv_start:kv_end].transpose(0, 1).float()
55
56 scores = torch.matmul(q_r, k_r.transpose(-2, -1)) * scale
57 if causal:
58 causal_mask = torch.ones(
59 (q_end - q_start, kv_end - kv_start),
60 dtype=torch.bool,
61 device=scores.device,
62 ).triu(1)
63 scores.masked_fill_(causal_mask.unsqueeze(0), float("-inf"))
64
65 probs = torch.softmax(scores, dim=-1)
66 outputs.append(torch.matmul(probs, v_r).transpose(0, 1).to(v.dtype))
67 lses.append(torch.logsumexp(scores, dim=-1))
68
69 return torch.cat(outputs, dim=0), torch.cat(lses, dim=1)
70
71
72 @pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
73 def test_kv_cache_cpu_write() -> None:
74 kv_lora_rank = 512
75 pe_dim = 64
76 block_size = 16
77 num_blocks = 4
78 num_tokens = 8
79 dtype = torch.float32
80
81 kv_c = torch.randn(num_tokens, kv_lora_rank, dtype=dtype)
82 k_pe = torch.randn(num_tokens, pe_dim, dtype=dtype)
83 kv_cache = torch.zeros(num_blocks, block_size, kv_lora_rank + pe_dim, dtype=dtype)
84 slot_mapping = torch.arange(num_tokens, dtype=torch.long)
85 # Mark the last slot as padding to make sure we honour negative slots.
86 slot_mapping[-1] = -1
87
88 flat = kv_cache.view(-1, kv_lora_rank + pe_dim)
89 slots = slot_mapping.to(torch.long)
90 valid_mask = slots >= 0
91 if not valid_mask.all().item():
92 slots = slots[valid_mask]
93 kv_c = kv_c[valid_mask]
94 k_pe = k_pe[valid_mask]
95 target_dtype = flat.dtype
96 flat[slots, :kv_lora_rank] = kv_c.to(target_dtype)
97 flat[slots, kv_lora_rank:] = k_pe.to(target_dtype)
98
99 flat = kv_cache.view(-1, kv_lora_rank + pe_dim)
100 for i in range(num_tokens - 1):
101 assert torch.equal(flat[i, :kv_lora_rank], kv_c[i])
102 assert torch.equal(flat[i, kv_lora_rank:], k_pe[i])
103 # The padding slot must remain untouched.
104 assert torch.all(flat[num_tokens - 1] == 0)
105
106
107 @pytest.mark.cpu_model
108 @pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
109 def test_cpu_mla_backend_smoke(tmp_path) -> None:
110 # Shrink the architecture aggressively so the test is CI-friendly.
111 from vllm import LLM, SamplingParams
112
113 hf_overrides = {
114 "num_hidden_layers": 2,
115 "n_routed_experts": 4,
116 "first_k_dense_replace": 0,
117 "num_experts_per_tok": 2,
118 }
119 llm = LLM(
120 model="deepseek-ai/DeepSeek-V2-Lite",
121 skip_tokenizer_init=True,
122 trust_remote_code=True,
123 load_format="dummy",
124 enforce_eager=True,
125 max_model_len=128,
126 max_num_seqs=2,
127 block_size=16,
128 gpu_memory_utilization=0.25,
129 hf_overrides=hf_overrides,
130 kv_transfer_config={
131 "kv_connector": "ExampleConnector",
132 "kv_role": "kv_both",
133 "kv_connector_extra_config": {
134 "shared_storage_path": str(tmp_path),
135 },
136 },
137 )
138
139 sampling_params = SamplingParams(max_tokens=4, temperature=0.0)
140 outputs = llm.generate([[1, 2, 3, 4], [5, 6, 7]], sampling_params, use_tqdm=False)
141 assert len(outputs) == 2
142 for output in outputs:
143 # `dummy` weights do not produce meaningful text, but the number
144 # of generated tokens must match what we asked for.
145 assert len(output.outputs[0].token_ids) == 4
146
147 prefix_hit_prompt = [1] * 20
148 outputs = llm.generate([prefix_hit_prompt], sampling_params, use_tqdm=False)
149 assert len(outputs[0].outputs[0].token_ids) == 4
150 assert llm.reset_prefix_cache(reset_connector=False)
151
152 # ExampleConnector persists KV externally at block granularity. With a
153 # 20-token prompt and block_size=16, the second run reloads 16 cached
154 # tokens and forces the 4-token suffix through MLA prefill with context.
155 outputs = llm.generate([prefix_hit_prompt], sampling_params, use_tqdm=False)
156 assert len(outputs[0].outputs[0].token_ids) == 4
157
158
159 @pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
160 def test_cpu_mla_prefill_backend_selected() -> None:
161 # On CPU (no device capability) the MLA prefill backend must be the
162 # SDPA-based CPU backend, not flash-attn which is unavailable on CPU.
163 # Accelerator-specific subclasses (e.g. the zentorch-backed Zen CPU
164 # backend) are also valid selections.
165 backend_cls = get_mla_prefill_backend(None)
166 assert issubclass(backend_cls, CPUSDPAMLAPrefillBackend)
167
168
169 @pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
170 def test_cpu_mla_prefill_new_tokens() -> None:
171 # The SDPA prefill must exactly match a reference causal attention over
172 # the ragged (varlen) q/k/v layout, with per-request padding handled.
173 backend = CPUSDPAMLAPrefillBackend(
174 num_heads=4,
175 scale=0.25,
176 kv_lora_rank=16,
177 qk_nope_head_dim=16,
178 qk_rope_head_dim=0,
179 v_head_dim=16,
180 vllm_config=None,
181 )
182
183 class _PrefillMeta:
184 query_start_loc = torch.tensor([0, 3, 6], dtype=torch.int32)
185
186 backend.prepare_metadata(_PrefillMeta())
187 q = torch.randn(6, 4, 16, dtype=torch.bfloat16)
188 k = torch.randn(6, 4, 16, dtype=torch.bfloat16)
189 v = torch.randn(6, 4, 16, dtype=torch.bfloat16)
190 out = backend.run_prefill_new_tokens(q, k, v, return_softmax_lse=False)
191 ref, _ = _reference_ragged_attention(
192 q,
193 k,
194 v,
195 _PrefillMeta.query_start_loc,
196 _PrefillMeta.query_start_loc,
197 scale=0.25,
198 causal=True,
199 )
200 assert out.float().allclose(ref.float(), atol=2e-2)
201 assert out.shape == (6, 4, 16)
202
203 out_lse, lse = backend.run_prefill_new_tokens(q, k, v, return_softmax_lse=True)
204 ref_out, ref_lse = _reference_ragged_attention(
205 q,
206 k,
207 v,
208 _PrefillMeta.query_start_loc,
209 _PrefillMeta.query_start_loc,
210 scale=0.25,
211 causal=True,
212 )
213 assert out_lse.float().allclose(ref_out.float(), atol=2e-2)
214 assert lse.float().allclose(ref_lse.float(), atol=2e-2)
215 assert lse.shape == (4, 6)
216
217
218 @pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
219 def test_cpu_mla_prefill_context_chunk() -> None:
220 backend = CPUSDPAMLAPrefillBackend(
221 num_heads=4,
222 scale=0.25,
223 kv_lora_rank=16,
224 qk_nope_head_dim=16,
225 qk_rope_head_dim=0,
226 v_head_dim=16,
227 vllm_config=None,
228 )
229
230 class _PrefillMeta:
231 query_start_loc = torch.tensor([0, 5], dtype=torch.int32)
232
233 class _Chunk:
234 query_start_loc = torch.tensor([0, 2, 5], dtype=torch.int32)
235 cu_seq_lens = torch.tensor([0, 3, 7], dtype=torch.int32)
236
237 backend.prepare_metadata(_PrefillMeta())
238 q = torch.randn(5, 4, 16, dtype=torch.bfloat16)
239 k = torch.randn(7, 4, 16, dtype=torch.bfloat16)
240 v = torch.randn(7, 4, 16, dtype=torch.bfloat16)
241
242 out, lse = backend.run_prefill_context_chunk(_Chunk(), q, k, v)
243 ref_out, ref_lse = _reference_ragged_attention(
244 q,
245 k,
246 v,
247 _Chunk.query_start_loc,
248 _Chunk.cu_seq_lens,
249 scale=0.25,
250 causal=False,
251 )
252 assert out.float().allclose(ref_out.float(), atol=2e-2)
253 assert lse.float().allclose(ref_lse.float(), atol=2e-2)
254 assert out.shape == (5, 4, 16)
255 assert lse.shape == (4, 5)
256
257
258 @pytest.mark.skipif(not current_platform.is_cpu(), reason="CPU only")
259 @pytest.mark.skipif(
260 not ZenCPUSDPAMLAPrefillBackend.is_available(),
261 reason="requires a Zen CPU with the zentorch_sdpa op registered",
262 )
263 def test_zen_cpu_mla_prefill_new_tokens() -> None:
264 # Runs the zentorch kernel itself. The head dims are deliberately
265 # asymmetric (v_head_dim < qk_head_dim, as in real MLA) so that the V
266 # padding and output slicing are exercised; the generic-CPU cases above
267 # use equal dims and never reach that code.
268 qk_head_dim = 24
269 v_head_dim = 16
270 backend = ZenCPUSDPAMLAPrefillBackend(
271 num_heads=4,
272 scale=0.25,
273 kv_lora_rank=16,
274 qk_nope_head_dim=16,
275 qk_rope_head_dim=8,
276 v_head_dim=v_head_dim,
277 vllm_config=None,
278 )
279
280 class _PrefillMeta:
281 query_start_loc = torch.tensor([0, 3, 6], dtype=torch.int32)
282
283 backend.prepare_metadata(_PrefillMeta())
284 q = torch.randn(6, 4, qk_head_dim, dtype=torch.bfloat16)
285 k = torch.randn(6, 4, qk_head_dim, dtype=torch.bfloat16)
286 v = torch.randn(6, 4, v_head_dim, dtype=torch.bfloat16)
287
288 out = backend.run_prefill_new_tokens(q, k, v, return_softmax_lse=False)
289 ref, _ = _reference_ragged_attention(
290 q,
291 k,
292 v,
293 _PrefillMeta.query_start_loc,
294 _PrefillMeta.query_start_loc,
295 scale=0.25,
296 causal=True,
297 )
298 assert out.shape == (6, 4, v_head_dim)
299 assert out.float().allclose(ref.float(), atol=2e-2)
300
301 # zentorch_sdpa does not return a usable LSE, so that path must fall back
302 # to the fp32 reference and match the generic backend exactly.
303 generic = CPUSDPAMLAPrefillBackend(
304 num_heads=4,
305 scale=0.25,
306 kv_lora_rank=16,
307 qk_nope_head_dim=16,
308 qk_rope_head_dim=8,
309 v_head_dim=v_head_dim,
310 vllm_config=None,
311 )
312 generic.prepare_metadata(_PrefillMeta())
313 zen_out, zen_lse = backend.run_prefill_new_tokens(q, k, v, return_softmax_lse=True)
314 gen_out, gen_lse = generic.run_prefill_new_tokens(q, k, v, return_softmax_lse=True)
315 assert torch.equal(zen_out, gen_out)
316 assert torch.equal(zen_lse, gen_lse)