@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 10cc12ba66
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

Feature/mla tests (#23195)

Signed-off-by: Matthew Bonanni <mbonanni001@gmail.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Matthew Bonanni committed Aug 20, 2025 at 17:46 UTC 10cc12ba66834e33659f1ce3a00235506db20dd5
4 files changed +551 -24
tests/v1/attention/test_attention_backends.py
+11 -15
@@ -150,15 +150,15 @@ def create_and_prepopulate_kv_cache(
150
151 # Permute the context blocks (excluding block 0 which is null)
152 if randomize_blocks:
153 - perm = torch.randperm(
154 - blocks_end - 1) + 1 # Random permutation starting from block 1
153 + # Random permutation starting from block 1
154 + perm = torch.randperm(blocks_end - 1) + 1
155 else:
156 - perm = torch.arange(
157 - 1, blocks_end) # Sequential order starting from block 1
156 + # Sequential order starting from block 1
157 + perm = torch.arange(1, blocks_end)
158
159 inv_perm = torch.zeros(blocks_end, dtype=torch.long, device=device)
160 - inv_perm[1:] = torch.argsort(
161 - perm) + 1 # Add 1 to account for starting from block 1
160 + # Add 1 to account for starting from block 1
161 + inv_perm[1:] = torch.argsort(perm) + 1
162 kv_cache[:, 1:blocks_end, ...] = kv_cache[:, perm, ...]
163
164 # Construct the right block table
@@ -281,7 +281,8 @@ def run_attention_backend(backend: _Backend, kv_cache_spec: FullAttentionSpec,
281
282 @pytest.mark.parametrize("batch_spec_name", [
283 "small_decode", "small_prefill", "mixed_small", "medium_decode",
284 - "medium_prefill", "mixed_medium"
284 + "medium_prefill", "mixed_medium", "large_decode", "large_prefill",
285 + "single_decode", "single_prefill"
286 ])
287 @pytest.mark.parametrize("model", ["meta-llama/Meta-Llama-3-8B"])
288 def test_backend_correctness(batch_spec_name: str, model: str):
@@ -302,7 +303,8 @@ def test_backend_correctness(batch_spec_name: str, model: str):
303 """
304 batch_spec = BATCH_SPECS[batch_spec_name]
305 vllm_config = create_vllm_config(model_name=model,
305 - max_model_len=max(batch_spec.seq_lens))
306 + max_model_len=max(batch_spec.seq_lens),
307 + num_gpu_blocks=8192)
308 device = torch.device("cuda:0")
309
310 kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
@@ -465,12 +467,6 @@ def test_backend_correctness(batch_spec_name: str, model: str):
467 rtol=rtol,
468 atol=atol)
469
468 - if not all_close:
469 - print(f"[{backend_name}] output differs from SDPA baseline. "
470 - f"Max diff: {max_diff:.6f} (rel: {max_rel_diff:.6f})")
471 - print(f"[{backend_name}] output: {backend_output}")
472 - print(f"[{backend_name}] SDPA baseline: {sdpa_output}")
473 -
470 assert all_close, (
471 f"[{backend_name}] output differs from SDPA baseline. "
476 - f"Max diff: {max_diff:.6f} (rel: {max_rel_diff:.6f})")
472 + f"Max diff: {max_diff:.6f}, max rel diff: {max_rel_diff:.6f})")
\ No newline at end of file
tests/v1/attention/test_mla_backends.py new
+522
@@ -0,0 +1,522 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for v1 MLA backends without GPUModelRunner dependency."""
4 +
5 +import pytest
6 +import torch
7 +
8 +from tests.v1.attention.utils import (BatchSpec, _Backend,
9 + create_common_attn_metadata,
10 + create_standard_kv_cache_spec,
11 + create_vllm_config,
12 + get_attention_backend)
13 +from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv
14 +from vllm.v1.attention.backends.utils import CommonAttentionMetadata
15 +from vllm.v1.kv_cache_interface import FullAttentionSpec
16 +
17 +BACKENDS_TO_TEST = [
18 + _Backend.CUTLASS_MLA, _Backend.FLASHMLA_VLLM_V1,
19 + _Backend.TRITON_MLA_VLLM_V1
20 +]
21 +
22 +# Remove CUTLASS_MLA from the list if not using sm100
23 +if not torch.cuda.is_available() or torch.cuda.get_device_properties(
24 + 0).major < 10:
25 + BACKENDS_TO_TEST.remove(_Backend.CUTLASS_MLA)
26 +
27 +torch.manual_seed(42)
28 +
29 +
30 +def _convert_dtype_to_torch(dtype):
31 + """Convert ModelDType to torch.dtype."""
32 + if isinstance(dtype, str):
33 + if dtype == "auto":
34 + return torch.float16 # Default dtype for testing
35 + elif dtype in STR_DTYPE_TO_TORCH_DTYPE:
36 + return STR_DTYPE_TO_TORCH_DTYPE[dtype]
37 + else:
38 + raise ValueError(f"Unknown dtype: {dtype}")
39 + elif isinstance(dtype, torch.dtype):
40 + return dtype
41 + else:
42 + raise ValueError(f"Unknown dtype: {dtype}")
43 +
44 +
45 +# Define common batch configurations
46 +BATCH_SPECS = {
47 + "small_decode":
48 + BatchSpec(seq_lens=[32, 40], query_lens=[1, 1]),
49 + "small_prefill":
50 + BatchSpec(seq_lens=[32, 40], query_lens=[8, 8]),
51 + "mixed_small":
52 + BatchSpec(seq_lens=[32, 40, 48, 56], query_lens=[1, 1, 5, 5]),
53 + "medium_decode":
54 + BatchSpec(seq_lens=[128, 256, 512, 1024, 128, 256, 512, 1024],
55 + query_lens=[1, 1, 1, 1, 1, 1, 1, 1]),
56 + "medium_prefill":
57 + BatchSpec(seq_lens=[256, 512, 1024, 2048], query_lens=[16, 16, 16, 16]),
58 + "mixed_medium":
59 + BatchSpec(seq_lens=[512, 1024, 2048, 512, 1024, 2048],
60 + query_lens=[1, 1, 1, 7, 7, 7]),
61 + "large_decode":
62 + BatchSpec(seq_lens=[2048] * 32, query_lens=[1] * 32),
63 + "large_prefill":
64 + BatchSpec(seq_lens=[4096] * 8, query_lens=[32] * 8),
65 + "single_decode":
66 + BatchSpec(seq_lens=[1024], query_lens=[1]),
67 + "single_prefill":
68 + BatchSpec(seq_lens=[1024], query_lens=[64]),
69 +}
70 +
71 +
72 +def create_dummy_kv_cache(kv_cache_spec: FullAttentionSpec,
73 + device: torch.device,
74 + num_blocks: int = 100) -> torch.Tensor:
75 + """Create a dummy KV cache tensor for testing."""
76 + kv_cache = torch.randn(
77 + num_blocks,
78 + kv_cache_spec.block_size,
79 + kv_cache_spec.head_size, # latent dimension
80 + dtype=_convert_dtype_to_torch(kv_cache_spec.dtype),
81 + device=device,
82 + )
83 + return kv_cache
84 +
85 +
86 +def create_and_prepopulate_kv_cache(
87 + kv_c_contexts: list[torch.Tensor],
88 + k_pe_contexts: list[torch.Tensor],
89 + block_size: int,
90 + num_kv_heads: int,
91 + head_size: int,
92 + dtype: torch.dtype,
93 + device: torch.device,
94 + num_blocks: int,
95 + common_attn_metadata: CommonAttentionMetadata,
96 + randomize_blocks: bool = True) -> torch.Tensor:
97 + """Create and prepopulate an MLA KV cache with context data.
98 +
99 + Args:
100 + kv_c_contexts: List of latent KV context tensors for each sequence
101 + k_pe_contexts: List of key positional embedding context tensors
102 + for each sequence
103 + block_size: Size of each block
104 + num_kv_heads: Number of KV heads (should be 1 for MLA)
105 + head_size: Size of each head (latent dimension)
106 + dtype: Data type for the cache
107 + device: Device to create the cache on
108 + num_blocks: Total number of blocks in the cache
109 + common_attn_metadata: Common attention metadata
110 + randomize_blocks: Whether to randomly permute blocks
111 + or use sequential order
112 +
113 + Returns:
114 + MLA KV cache tensor
115 + """
116 + batch_size = len(kv_c_contexts)
117 + seq_lens = common_attn_metadata.seq_lens_cpu
118 + query_lens = common_attn_metadata.query_start_loc_cpu[
119 + 1:] - common_attn_metadata.query_start_loc_cpu[:-1]
120 + context_lens = common_attn_metadata.num_computed_tokens_cpu
121 + block_table = common_attn_metadata.block_table_tensor
122 + slot_mapping = common_attn_metadata.slot_mapping
123 +
124 + # Create MLA KV cache: (num_blocks, block_size, head_size)
125 + kv_cache = torch.empty(num_blocks,
126 + block_size,
127 + head_size,
128 + dtype=dtype,
129 + device=device)
130 + kv_cache_flat = kv_cache.view(-1, head_size)
131 +
132 + # Populate the cache with the context tokens
133 + # Start from block_id=1 since block_id=0 is considered the null block
134 + start_block_idx = 1
135 + for i in range(batch_size):
136 + kv_c_context, k_pe_context = kv_c_contexts[i], k_pe_contexts[i]
137 + kv_context = torch.cat([kv_c_context, k_pe_context.squeeze(1)], dim=-1)
138 + start = start_block_idx * block_size
139 + end = start + kv_context.shape[0]
140 + kv_cache_flat[start:end, ...] = kv_context
141 +
142 + # Stay block aligned and allocate enough blocks for the new tokens
143 + start_block_idx += cdiv(int(seq_lens[i]), block_size)
144 +
145 + blocks_end = start_block_idx
146 +
147 + # Permute the context blocks (excluding block 0 which is null)
148 + if randomize_blocks:
149 + perm = torch.randperm(
150 + blocks_end - 1) + 1 # Random permutation starting from block 1
151 + else:
152 + perm = torch.arange(
153 + 1, blocks_end) # Sequential order starting from block 1
154 +
155 + inv_perm = torch.zeros(blocks_end, dtype=torch.long, device=device)
156 + inv_perm[1:] = torch.argsort(
157 + perm) + 1 # Add 1 to account for starting from block 1
158 + kv_cache[1:blocks_end, ...] = kv_cache[perm, ...]
159 +
160 + # Construct the right block table
161 + # Start from block_id=1 since block_id=0 is considered the null block
162 + start_block_idx = 1
163 + for i in range(batch_size):
164 + num_blocks_for_seq = cdiv(int(seq_lens[i]), block_size)
165 + start = start_block_idx
166 + end = start + num_blocks_for_seq
167 + block_table[i, :num_blocks_for_seq] = inv_perm[start:end]
168 + start_block_idx += num_blocks_for_seq
169 +
170 + # Create a realistic slot mapping that corresponds to the block table
171 + for i in range(batch_size):
172 + token_offsets = torch.arange(int(query_lens[i])) + int(context_lens[i])
173 + block_indices = token_offsets // block_size
174 + token_inter_block_offsets = token_offsets % block_size
175 + start = common_attn_metadata.query_start_loc_cpu[i]
176 + end = common_attn_metadata.query_start_loc_cpu[i + 1]
177 + slot_mapping[start:end] = block_table[
178 + i,
179 + block_indices] * block_size + token_inter_block_offsets.to(device)
180 +
181 + return kv_cache
182 +
183 +
184 +class MockAttentionLayer:
185 + """A mock attention layer for testing."""
186 +
187 + def __init__(self, device: torch.device):
188 + self._q_scale = torch.tensor(1.0, device=device)
189 + self._k_scale = torch.tensor(1.0, device=device)
190 + self._v_scale = torch.tensor(1.0, device=device)
191 +
192 +
193 +def run_attention_backend(backend: _Backend, kv_cache_spec: FullAttentionSpec,
194 + layer_names: list[str], vllm_config,
195 + device: torch.device,
196 + common_attn_metadata: CommonAttentionMetadata,
197 + query: torch.Tensor, kv_c: torch.Tensor,
198 + k_pe: torch.Tensor, kv_cache: torch.Tensor,
199 + kv_lora_rank: int, qk_nope_head_dim: int,
200 + qk_rope_head_dim: int, v_head_dim: int,
201 + mock_kv_b_proj) -> torch.Tensor:
202 + """Run attention computation using the specified backend's AttentionImpl."""
203 +
204 + builder_cls, impl_cls = get_attention_backend(backend)
205 +
206 + # Build metadata
207 + builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device)
208 + attn_metadata = builder.build(
209 + common_prefix_len=0,
210 + common_attn_metadata=common_attn_metadata,
211 + )
212 +
213 + # Instantiate MLA implementation
214 + num_heads = vllm_config.model_config.get_num_attention_heads(
215 + vllm_config.parallel_config)
216 + num_kv_heads = vllm_config.model_config.get_num_kv_heads(
217 + vllm_config.parallel_config)
218 + head_size = vllm_config.model_config.get_head_size()
219 + scale = 1.0 / (head_size**0.5)
220 + impl = impl_cls(
221 + num_heads=num_heads,
222 + head_size=head_size,
223 + scale=scale,
224 + num_kv_heads=num_kv_heads,
225 + alibi_slopes=None,
226 + sliding_window=None,
227 + kv_cache_dtype="auto",
228 + logits_soft_cap=None,
229 + attn_type="decoder",
230 + kv_sharing_target_layer_name=None,
231 + q_lora_rank=None,
232 + kv_lora_rank=kv_lora_rank,
233 + qk_nope_head_dim=qk_nope_head_dim,
234 + qk_rope_head_dim=qk_rope_head_dim,
235 + qk_head_dim=qk_nope_head_dim + qk_rope_head_dim,
236 + v_head_dim=v_head_dim,
237 + kv_b_proj=mock_kv_b_proj,
238 + )
239 +
240 + # Process weights to create W_UK_T and W_UV attributes needed by MLA
241 + act_dtype = _convert_dtype_to_torch(vllm_config.model_config.dtype)
242 + impl.process_weights_after_loading(act_dtype)
243 +
244 + # Create mock layer and output buffer
245 + mock_layer = MockAttentionLayer(device)
246 + num_tokens = query.shape[0]
247 + output = torch.empty(num_tokens,
248 + num_heads * v_head_dim,
249 + dtype=query.dtype,
250 + device=query.device)
251 +
252 + # Run forward pass
253 + # NOTE: The query, key, and value are already shaped correctly
254 + # in the calling test function.
255 + output = impl.forward(mock_layer,
256 + query,
257 + kv_c,
258 + k_pe,
259 + kv_cache,
260 + attn_metadata,
261 + output=output)
262 +
263 + return output
264 +
265 +
266 +@pytest.mark.parametrize("batch_spec_name", [
267 + "small_decode", "small_prefill", "mixed_small", "medium_decode",
268 + "medium_prefill", "mixed_medium", "large_decode", "large_prefill",
269 + "single_decode", "single_prefill"
270 +])
271 +@pytest.mark.parametrize("model", ["deepseek-ai/DeepSeek-V2-Lite-Chat"])
272 +def test_backend_correctness(dist_init, batch_spec_name: str, model: str):
273 + """
274 + Test that all backends produce similar outputs to a reference implementation
275 + using torch.nn.functional.scaled_dot_product_attention.
276 +
277 + This test works by:
278 + 1. Generating a batch of sequences with specified context and query lengths.
279 + 2. Computing a ground-truth attention output using torch.sdpa on
280 + contiguous Q, K, and V tensors.
281 + 3. Simulating vLLM's paged KV cache: It takes the context portion of the
282 + K/V tensors and manually places them into a paged buffer according to
283 + the test's (randomly generated) block table.
284 + 4. Running each vLLM attention backend with the new queries and the
285 + simulated paged KV cache.
286 + 5. Comparing the vLLM backend's output to the ground-truth SDPA output.
287 + """
288 + batch_spec = BATCH_SPECS[batch_spec_name]
289 + vllm_config = create_vllm_config(model_name=model,
290 + max_model_len=max(batch_spec.seq_lens),
291 + num_gpu_blocks=2048)
292 + device = torch.device("cuda:0")
293 +
294 + kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
295 +
296 + # 1. Setup
297 + batch_size = batch_spec.batch_size
298 + seq_lens = batch_spec.seq_lens
299 + query_lens = batch_spec.query_lens
300 + num_q_heads = vllm_config.model_config.get_num_attention_heads(
301 + vllm_config.parallel_config)
302 + num_kv_heads = vllm_config.model_config.get_num_kv_heads(
303 + vllm_config.parallel_config)
304 + head_size = vllm_config.model_config.get_head_size()
305 + dtype = _convert_dtype_to_torch(vllm_config.model_config.dtype)
306 + block_size = vllm_config.cache_config.block_size
307 + kv_lora_rank = 512
308 + qk_rope_head_dim = 64
309 + qk_nope_head_dim = 128
310 + v_head_dim = 128
311 + total_head_size = kv_lora_rank + qk_rope_head_dim
312 + assert kv_lora_rank + qk_rope_head_dim == head_size, \
313 + f"MLA dimensions don't match: {total_head_size} != {head_size}"
314 + scale = 1.0 / (total_head_size**0.5)
315 +
316 + # 2. Generate data and compute SDPA reference output for MLA
317 + all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], []
318 + all_sdpa_outputs = []
319 + kv_c_contexts, k_pe_contexts = [], []
320 +
321 + # Create shared MLA weight matrices for consistency across all sequences
322 + W_UK = torch.randn(kv_lora_rank,
323 + num_q_heads,
324 + qk_nope_head_dim,
325 + dtype=dtype,
326 + device=device)
327 + W_UV = torch.randn(kv_lora_rank,
328 + num_q_heads,
329 + v_head_dim,
330 + dtype=dtype,
331 + device=device)
332 + kv_b_proj_weight = torch.cat([W_UK, W_UV], dim=-1)
333 +
334 + for i in range(batch_size):
335 + s_len = seq_lens[i]
336 + q_len = query_lens[i]
337 + context_len = s_len - q_len
338 +
339 + # Generate MLA tensors
340 + # Q has both nope and rope components:
341 + # [q_len, num_heads, qk_nope_head_dim + qk_rope_head_dim]
342 + q_c = torch.randn(q_len,
343 + num_q_heads,
344 + qk_nope_head_dim + qk_rope_head_dim,
345 + dtype=dtype,
346 + device=device)
347 +
348 + # KV_C (latent K/V): [s_len, kv_lora_rank]
349 + kv_c_full = torch.randn(s_len,
350 + kv_lora_rank,
351 + dtype=dtype,
352 + device=device)
353 +
354 + # K_PE (rope component): [s_len, 1, qk_rope_head_dim]
355 + k_pe_full = torch.randn(s_len,
356 + 1,
357 + qk_rope_head_dim,
358 + dtype=dtype,
359 + device=device)
360 +
361 + # Determine if this is decode (single token)
362 + # or prefill (multiple tokens)
363 + is_decode = q_len == 1
364 +
365 + # Split q into nope and rope components
366 + q_nope, q_pe = q_c.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1)
367 +
368 + if is_decode:
369 + # Decode path: MQA-style attention in latent space
370 + # Transform q_nope to latent space: q_nope @ W_UK
371 + # q_nope: [1, num_heads, qk_nope_head_dim]
372 + # W_UK: [kv_lora_rank, num_heads, qk_nope_head_dim]
373 + ql_nope = torch.einsum("qnh,lnh->qnl", q_nope,
374 + W_UK) # [1, num_heads, kv_lora_rank]
375 +
376 + # Build MQA attention inputs
377 + # Q: [1, num_heads, kv_lora_rank + qk_rope_head_dim]
378 + q_mqa = torch.cat([ql_nope, q_pe], dim=-1)
379 + # K: [s_len, kv_lora_rank + qk_rope_head_dim]
380 + # (broadcasted to all heads)
381 + k_mqa = torch.cat([kv_c_full, k_pe_full.squeeze(1)], dim=-1)
382 + k_mqa = k_mqa.unsqueeze(1).expand(-1, num_q_heads, -1)
383 + # V: [s_len, kv_lora_rank] (broadcasted to all heads)
384 + v_mqa = kv_c_full.unsqueeze(1).expand(-1, num_q_heads, -1)
385 +
386 + # SDPA expects (N, H, L, D)
387 + q_sdpa_in = q_mqa.unsqueeze(0).transpose(1, 2)
388 + k_sdpa_in = k_mqa.unsqueeze(0).transpose(1, 2)
389 + v_sdpa_in = v_mqa.unsqueeze(0).transpose(1, 2)
390 +
391 + sdpa_out_i = torch.nn.functional.scaled_dot_product_attention(
392 + q_sdpa_in, k_sdpa_in, v_sdpa_in, is_causal=False, scale=scale)
393 + sdpa_out_i = sdpa_out_i.transpose(1, 2).squeeze(
394 + 0) # [1, num_heads, kv_lora_rank]
395 +
396 + # Project back to output space: sdpa_out @ W_UV
397 + sdpa_out_i = torch.einsum("qnl,lnv->qnv", sdpa_out_i, W_UV)
398 + sdpa_out_i = sdpa_out_i.flatten(start_dim=-2)
399 + else:
400 + # Prefill path: MHA-style attention with full sequence
401 + # Apply kv_b_proj to the full kv_c tensor
402 + kv_nope_full = torch.einsum("sl,lnh->snh", kv_c_full,
403 + kv_b_proj_weight)
404 + k_nope_full, v_full = kv_nope_full.split(
405 + [qk_nope_head_dim, v_head_dim], dim=-1)
406 +
407 + # Build attention inputs for full sequence
408 + q_mha = torch.cat([q_nope, q_pe],
409 + dim=-1) # [q_len, num_heads, total_dim]
410 + k_pe_full_expanded = k_pe_full.expand(-1, num_q_heads, -1)
411 + k_full = torch.cat([k_nope_full, k_pe_full_expanded], dim=-1)
412 +
413 + # Create custom attention mask:
414 + # - Query tokens can attend to all context tokens
415 + # - Query tokens can only attend to query tokens up to their pos
416 + attn_mask = torch.ones(q_len,
417 + s_len,
418 + dtype=torch.bool,
419 + device=device)
420 + # Apply causal mask only to the query portion (context_len onwards)
421 + causal_mask = torch.tril(torch.ones(q_len, q_len, device=device))
422 + attn_mask[:, context_len:] = causal_mask
423 +
424 + # SDPA expects (N, H, L, D)
425 + q_sdpa_in = q_mha.unsqueeze(0).transpose(1, 2)
426 + k_sdpa_in = k_full.unsqueeze(0).transpose(1, 2)
427 + v_sdpa_in = v_full.unsqueeze(0).transpose(1, 2)
428 +
429 + # Single attention call with custom mask
430 + sdpa_out_i = torch.nn.functional.scaled_dot_product_attention(
431 + q_sdpa_in,
432 + k_sdpa_in,
433 + v_sdpa_in,
434 + attn_mask=attn_mask,
435 + scale=scale)
436 + sdpa_out_i = sdpa_out_i.transpose(1, 2).squeeze(0)
437 + sdpa_out_i = sdpa_out_i.flatten(start_dim=-2)
438 +
439 + all_sdpa_outputs.append(sdpa_out_i)
440 +
441 + # Inputs for vLLM MLA backends are just the new tokens
442 + all_q_vllm.append(q_c)
443 + all_kv_c_vllm.append(kv_c_full[context_len:]) # New kv_c tokens
444 + all_k_pe_vllm.append(k_pe_full[context_len:]) # New k_pe tokens
445 +
446 + # Contextual K/V data used to populate the paged cache (MLA format)
447 + kv_c_contexts.append(kv_c_full[:context_len])
448 + k_pe_contexts.append(k_pe_full[:context_len])
449 +
450 + # Concatenate all sequences (no reordering needed)
451 + query_vllm = torch.cat(all_q_vllm, dim=0)
452 + kv_c_vllm = torch.cat(all_kv_c_vllm, dim=0)
453 + k_pe_vllm = torch.cat(all_k_pe_vllm, dim=0)
454 + sdpa_output = torch.cat(all_sdpa_outputs, dim=0)
455 +
456 + # Create mock kv_b_proj using the same weights as reference implementation
457 + from vllm.model_executor.layers.linear import ColumnParallelLinear
458 + mock_kv_b_proj = ColumnParallelLinear(input_size=kv_lora_rank,
459 + output_size=num_q_heads *
460 + (qk_nope_head_dim + v_head_dim),
461 + bias=False).to(device=device,
462 + dtype=dtype)
463 +
464 + # Set the mock weights to match our reference implementation
465 + # Reshape W_UK and W_UV to match the expected kv_b_proj format
466 + # [kv_lora_rank, num_heads, qk_nope_head_dim + v_head_dim]
467 + kv_b_proj_weight = kv_b_proj_weight.view(
468 + kv_lora_rank, num_q_heads * (qk_nope_head_dim + v_head_dim))
469 + mock_kv_b_proj.weight = torch.nn.Parameter(kv_b_proj_weight.T)
470 +
471 + # Create metadata using original batch spec
472 + common_attn_metadata = create_common_attn_metadata(
473 + batch_spec, vllm_config.cache_config.block_size, device)
474 +
475 + # 3. Simulate Paged KV Cache and a realistic slot_mapping
476 + kv_cache = create_and_prepopulate_kv_cache(
477 + kv_c_contexts=kv_c_contexts,
478 + k_pe_contexts=k_pe_contexts,
479 + block_size=block_size,
480 + num_kv_heads=num_kv_heads,
481 + head_size=head_size,
482 + dtype=dtype,
483 + device=device,
484 + num_blocks=vllm_config.cache_config.num_gpu_blocks,
485 + common_attn_metadata=common_attn_metadata,
486 + randomize_blocks=True)
487 +
488 + # 4. Run vLLM backends and compare
489 + for backend_name in BACKENDS_TO_TEST:
490 + backend_output = run_attention_backend(
491 + backend_name, kv_cache_spec, ["placeholder"], vllm_config, device,
492 + common_attn_metadata, query_vllm, kv_c_vllm, k_pe_vllm, kv_cache,
493 + kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim,
494 + mock_kv_b_proj)
495 +
496 + # Check shape and dtype consistency
497 + assert backend_output.shape == sdpa_output.shape, (
498 + f"[{backend_name}] shape {backend_output.shape} != "
499 + f"SDPA shape {sdpa_output.shape}")
500 + assert backend_output.dtype == sdpa_output.dtype, (
501 + f"[{backend_name}] dtype {backend_output.dtype} != "
502 + f"SDPA dtype {sdpa_output.dtype}")
503 +
504 + assert torch.isfinite(backend_output).all(), (
505 + f"[{backend_name}] produced non-finite values")
506 +
507 + # Check numerical similarity
508 + rtol = 1e-2
509 + atol = 5e-1
510 +
511 + max_diff = torch.max(torch.abs(backend_output - sdpa_output)).item()
512 + max_rel_diff = torch.max(
513 + torch.abs(backend_output - sdpa_output) /
514 + torch.abs(sdpa_output)).item()
515 + all_close = torch.allclose(backend_output,
516 + sdpa_output,
517 + rtol=rtol,
518 + atol=atol)
519 +
520 + assert all_close, (
521 + f"[{backend_name}] output differs from SDPA baseline. "
522 + f"Max diff: {max_diff:.6f}, max rel diff: {max_rel_diff:.6f})")
tests/v1/attention/utils.py
+10 -1
@@ -135,6 +135,12 @@ def get_attention_backend(backend_name: _Backend):
135 "vllm.v1.attention.backends.tree_attn.TreeAttentionBackend",
136 _Backend.XFORMERS_VLLM_V1:
137 "vllm.v1.attention.backends.xformers.XFormersAttentionBackend",
138 + _Backend.CUTLASS_MLA:
139 + "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend",
140 + _Backend.FLASHMLA_VLLM_V1:
141 + "vllm.v1.attention.backends.mla.flashmla.FlashMLABackend",
142 + _Backend.TRITON_MLA_VLLM_V1:
143 + "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend",
144 }
145
146 if backend_name not in backend_map:
@@ -167,9 +173,11 @@ def create_vllm_config(model_name: str = "meta-llama/Meta-Llama-3-8B",
173 tensor_parallel_size: int = 1,
174 max_model_len: int = 1024,
175 dtype: Union[ModelDType, torch.dtype] = "auto",
176 + num_gpu_blocks: int = 1000,
177 block_size: int = 16,
178 max_num_seqs: int = 256,
179 max_num_batched_tokens: int = 8192,
180 + enable_chunked_prefill: bool = True,
181 add_mock_model_methods: bool = True) -> VllmConfig:
182 """Create a VllmConfig for testing with reasonable defaults."""
183
@@ -189,7 +197,7 @@ def create_vllm_config(model_name: str = "meta-llama/Meta-Llama-3-8B",
197 )
198 # Set cache blocks for testing
199 # (these may be set during initialization normally)
192 - cache_config.num_gpu_blocks = 1000
200 + cache_config.num_gpu_blocks = num_gpu_blocks
201 cache_config.num_cpu_blocks = 0
202
203 parallel_config = ParallelConfig(
@@ -198,6 +206,7 @@ def create_vllm_config(model_name: str = "meta-llama/Meta-Llama-3-8B",
206 scheduler_config = SchedulerConfig(
207 max_num_seqs=max_num_seqs,
208 max_num_batched_tokens=max_num_batched_tokens,
209 + enable_chunked_prefill=enable_chunked_prefill,
210 )
211
212 device_config = DeviceConfig()
vllm/v1/attention/backends/mla/common.py
+8 -8
@@ -24,7 +24,7 @@ Main reference: DeepseekV2 paper, and FlashInfer Implementation
24 (https://arxiv.org/abs/2405.04434 and https://github.com/flashinfer-ai/flashinfer/pull/551).
25
26 Deepseek's MLA attention works the following way:
27 -* Use a single latent vector to represent the per-token entry of the KV cache.
27 +* Use a single latent vector to represent the per-token entry of the KV cache.
28 * For decode (i.e. the memory friendly approach) the attention "simulates" a
29 multi-head attention, while the compute is similar to multi-query attention.
30
@@ -82,7 +82,7 @@ spda_o = scaled_dot_product_attention(
82 torch.cat([q_nope, q_pe], dim=-1),
83 torch.cat([k_nope, k_pe.unsqueeze(1).expand(-1, N, -1)], dim=-1),
84 v
85 -)
85 +)
86 return spda_o @ W_O
87
88 NOTE: in the actual code,
@@ -120,20 +120,20 @@ return o.view(-1, N * V) @ self.num_heads @ W_O
120
121 ## Chunked Prefill
122
123 -For chunked prefill we want to use the compute friendly algorithm. We are
124 -assuming sufficiently large Sq / Skv ratio, in the future may want to switch to
123 +For chunked prefill we want to use the compute friendly algorithm. We are
124 +assuming sufficiently large Sq / Skv ratio, in the future may want to switch to
125 the data-movement friendly approach if the chunk (i.e. `Sq`) is small.
126
127 However, the compute-friendly approach can potentially run out of memory if Skv
128 is large due to: `k_nope = (kv_c @ W_UK).view(Skv, N, P)`
129
130 -To mitigate this, we chunk the computation of attention with respect to the
131 -current context (i.e. `cache_kv_c` and `cache_k_pe`) so that we can used a
130 +To mitigate this, we chunk the computation of attention with respect to the
131 +current context (i.e. `cache_kv_c` and `cache_k_pe`) so that we can used a
132 fixed workspace size.
133
134 The chunked prefill approach is as follows:
135
136 -MCC Max chunk of context to process per iter, computed dynamically,
136 +MCC Max chunk of context to process per iter, computed dynamically,
137 used to bound the memory usage
138
139 q_c = h_t @ W_DQ
@@ -155,7 +155,7 @@ curr_o, curr_lse = scaled_dot_product_attention(
155 new_v,
156 casual=True,
157 return_softmax_lse=True
158 -)
158 +)
159
160 // Compute attention with the already existing context
161 for chunk_idx in range(cdiv(C, MCC)):