test_mla_backends.py 31.9 KB
Newer Older
Matthew Bonanni's avatar
Matthew Bonanni committed
1
2
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
6
7
8
9
"""Tests for v1 MLA backends without GPUModelRunner dependency.

Known Issues:
- FLASH_ATTN_MLA backend occasionally produces NaN values in
  test_backend_correctness[mixed_small] when run after
  test_backend_correctness[small_prefill], but passes when run alone.
"""
10

Matthew Bonanni's avatar
Matthew Bonanni committed
11
12
13
import pytest
import torch

14
15
16
17
from tests.v1.attention.utils import (
    BatchSpec,
    create_common_attn_metadata,
    create_vllm_config,
18
    try_get_attention_backend,
19
)
20
from vllm import _custom_ops as ops
21
from vllm.attention.backends.registry import AttentionBackendEnum
22
from vllm.attention.ops.flashmla import is_flashmla_dense_supported
23
from vllm.attention.utils.fa_utils import flash_attn_supports_mla
24
from vllm.config.vllm import set_current_vllm_config
25
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
26
from vllm.utils.math_utils import cdiv
27
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
28
from vllm.v1.attention.backends.mla.common import QueryLenSupport
Matthew Bonanni's avatar
Matthew Bonanni committed
29
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
30
from vllm.v1.kv_cache_interface import MLAAttentionSpec
Matthew Bonanni's avatar
Matthew Bonanni committed
31
32

BACKENDS_TO_TEST = [
33
34
35
36
37
    AttentionBackendEnum.CUTLASS_MLA,
    AttentionBackendEnum.FLASHMLA,
    AttentionBackendEnum.FLASH_ATTN_MLA,
    AttentionBackendEnum.FLASHINFER_MLA,
    AttentionBackendEnum.TRITON_MLA,
Matthew Bonanni's avatar
Matthew Bonanni committed
38
39
]

40
# Remove sm100 backends from the list if not using sm100
41
if not torch.cuda.is_available() or torch.cuda.get_device_properties(0).major < 10:
42
43
    BACKENDS_TO_TEST.remove(AttentionBackendEnum.CUTLASS_MLA)
    BACKENDS_TO_TEST.remove(AttentionBackendEnum.FLASHINFER_MLA)
44
45
46

# Remove FLASH_ATTN_MLA from the list if not supported
if not flash_attn_supports_mla():
47
    BACKENDS_TO_TEST.remove(AttentionBackendEnum.FLASH_ATTN_MLA)
Matthew Bonanni's avatar
Matthew Bonanni committed
48

49
50
# Remove FLASHMLA from the list if not supported
if not is_flashmla_dense_supported()[0]:
51
    BACKENDS_TO_TEST.remove(AttentionBackendEnum.FLASHMLA)
52

53
54
55
56
57
58
59
60
61
62
63
SPEC_DECODE_BACKENDS = []
for backend in BACKENDS_TO_TEST:
    builder_cls, _ = try_get_attention_backend(backend)
    query_len_support = getattr(
        builder_cls, "query_len_support", QueryLenSupport.SINGLE_ONLY
    )
    if query_len_support != QueryLenSupport.SINGLE_ONLY:
        SPEC_DECODE_BACKENDS.append(backend)

BACKEND_BLOCK_SIZES = {}
for backend in BACKENDS_TO_TEST:
64
    supported_sizes = backend.get_class().get_supported_kernel_block_sizes()
65
66
67
68
69
70
71
72
73
    if supported_sizes:
        default_size = supported_sizes[0]
        block_size = (
            default_size if isinstance(default_size, int) else default_size.base
        )
    else:
        block_size = 16
    BACKEND_BLOCK_SIZES[backend] = block_size

Matthew Bonanni's avatar
Matthew Bonanni committed
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
torch.manual_seed(42)


def _convert_dtype_to_torch(dtype):
    """Convert ModelDType to torch.dtype."""
    if isinstance(dtype, str):
        if dtype == "auto":
            return torch.float16  # Default dtype for testing
        elif dtype in STR_DTYPE_TO_TORCH_DTYPE:
            return STR_DTYPE_TO_TORCH_DTYPE[dtype]
        else:
            raise ValueError(f"Unknown dtype: {dtype}")
    elif isinstance(dtype, torch.dtype):
        return dtype
    else:
        raise ValueError(f"Unknown dtype: {dtype}")


# Define common batch configurations
BATCH_SPECS = {
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
    "small_decode": BatchSpec(seq_lens=[32, 40], query_lens=[1, 1]),
    "small_prefill": BatchSpec(seq_lens=[32, 40], query_lens=[8, 8]),
    "mixed_small": BatchSpec(seq_lens=[32, 40, 48, 56], query_lens=[1, 1, 5, 5]),
    "medium_decode": BatchSpec(
        seq_lens=[128, 256, 512, 1024, 128, 256, 512, 1024],
        query_lens=[1, 1, 1, 1, 1, 1, 1, 1],
    ),
    "medium_prefill": BatchSpec(
        seq_lens=[256, 512, 1024, 2048], query_lens=[16, 16, 16, 16]
    ),
    "mixed_medium": BatchSpec(
        seq_lens=[512, 1024, 2048, 512, 1024, 2048], query_lens=[1, 1, 1, 7, 7, 7]
    ),
    "large_decode": BatchSpec(seq_lens=[2048] * 32, query_lens=[1] * 32),
    "large_prefill": BatchSpec(seq_lens=[4096] * 8, query_lens=[32] * 8),
    "single_decode": BatchSpec(seq_lens=[1024], query_lens=[1]),
    "single_prefill": BatchSpec(seq_lens=[1024], query_lens=[64]),
111
112
113
114
115
116
    "spec_decode_small": BatchSpec(
        seq_lens=[128, 256, 512, 1024], query_lens=[4, 4, 4, 4]
    ),
    "spec_decode_medium": BatchSpec(
        seq_lens=[512, 1024, 2048, 512, 1024, 2048], query_lens=[8, 8, 8, 8, 8, 8]
    ),
Matthew Bonanni's avatar
Matthew Bonanni committed
117
118
119
120
}


def create_and_prepopulate_kv_cache(
121
122
123
124
125
126
127
128
129
    kv_c_contexts: list[torch.Tensor],
    k_pe_contexts: list[torch.Tensor],
    block_size: int,
    head_size: int,
    dtype: torch.dtype,
    device: torch.device,
    num_blocks: int,
    common_attn_metadata: CommonAttentionMetadata,
    randomize_blocks: bool = True,
130
131
    kv_cache_dtype: str | None = None,
    scale: float | torch.Tensor = 1.0,
132
) -> torch.Tensor:
Matthew Bonanni's avatar
Matthew Bonanni committed
133
    """Create and prepopulate an MLA KV cache with context data.
134

Matthew Bonanni's avatar
Matthew Bonanni committed
135
136
137
138
139
140
141
142
143
144
    Args:
        kv_c_contexts: List of latent KV context tensors for each sequence
        k_pe_contexts: List of key positional embedding context tensors
                       for each sequence
        block_size: Size of each block
        head_size: Size of each head (latent dimension)
        dtype: Data type for the cache
        device: Device to create the cache on
        num_blocks: Total number of blocks in the cache
        common_attn_metadata: Common attention metadata
145
        randomize_blocks: Whether to randomly permute blocks
Matthew Bonanni's avatar
Matthew Bonanni committed
146
                          or use sequential order
147
148
149
150
151
        kv_cache_dtype: Optional kv cache dtype string. When set to
                        "fp8_ds_mla" the cache is populated using the
                        fp8 DeepSeek MLA layout via concat_and_cache_mla.
        scale: Scaling factor forwarded to concat_and_cache_mla when the
               fp8 cache layout is requested.
152

Matthew Bonanni's avatar
Matthew Bonanni committed
153
154
155
156
157
    Returns:
        MLA KV cache tensor
    """
    batch_size = len(kv_c_contexts)
    seq_lens = common_attn_metadata.seq_lens_cpu
158
159
160
161
    query_lens = (
        common_attn_metadata.query_start_loc_cpu[1:]
        - common_attn_metadata.query_start_loc_cpu[:-1]
    )
Matthew Bonanni's avatar
Matthew Bonanni committed
162
163
164
165
    context_lens = common_attn_metadata.num_computed_tokens_cpu
    block_table = common_attn_metadata.block_table_tensor
    slot_mapping = common_attn_metadata.slot_mapping

166
167
168
169
    use_fp8_ds_mla = kv_cache_dtype == "fp8_ds_mla"

    if use_fp8_ds_mla:
        if not kv_c_contexts:
170
171
172
            raise ValueError(
                "kv_c_contexts cannot be empty when using fp8_ds_mla cache dtype"
            )
173
174
175
        kv_lora_rank = kv_c_contexts[0].shape[-1]
        rope_dim = k_pe_contexts[0].shape[-1]
        entry_size = kv_lora_rank + 4 * 4 + 2 * rope_dim
176
177
178
179
180
181
182
183
        kv_cache = torch.zeros(
            num_blocks, block_size, entry_size, dtype=torch.uint8, device=device
        )
        scale_tensor = (
            scale
            if isinstance(scale, torch.Tensor)
            else torch.tensor(scale, dtype=torch.float32, device=device)
        )
184
185
186
        scale_tensor = scale_tensor.to(device=device, dtype=torch.float32)
    else:
        # Create MLA KV cache: (num_blocks, block_size, head_size)
187
        kv_cache = torch.zeros(
188
189
            num_blocks, block_size, head_size, dtype=dtype, device=device
        )
190
        kv_cache_flat = kv_cache.view(-1, head_size)
Matthew Bonanni's avatar
Matthew Bonanni committed
191
192
193
194
195
196

    # Populate the cache with the context tokens
    # Start from block_id=1 since block_id=0 is considered the null block
    start_block_idx = 1
    for i in range(batch_size):
        kv_c_context, k_pe_context = kv_c_contexts[i], k_pe_contexts[i]
197
198
199
200
201
        context_len = kv_c_context.shape[0]
        if context_len == 0:
            start_block_idx += cdiv(int(seq_lens[i]), block_size)
            continue

Matthew Bonanni's avatar
Matthew Bonanni committed
202
        start = start_block_idx * block_size
203
204

        if use_fp8_ds_mla:
205
            slots = torch.arange(context_len, device=device, dtype=torch.long) + start
206
207
208
209
210
211
212
213
214
            ops.concat_and_cache_mla(
                kv_c_context,
                k_pe_context.squeeze(1),
                kv_cache,
                slots,
                kv_cache_dtype="fp8_ds_mla",
                scale=scale_tensor,
            )
        else:
215
            kv_context = torch.cat([kv_c_context, k_pe_context.squeeze(1)], dim=-1)
216
217
            end = start + kv_context.shape[0]
            kv_cache_flat[start:end, ...] = kv_context
Matthew Bonanni's avatar
Matthew Bonanni committed
218
219
220
221
222
223
224
225

        # Stay block aligned and allocate enough blocks for the new tokens
        start_block_idx += cdiv(int(seq_lens[i]), block_size)

    blocks_end = start_block_idx

    # Permute the context blocks (excluding block 0 which is null)
    if randomize_blocks:
226
227
228
        perm = (
            torch.randperm(blocks_end - 1) + 1
        )  # Random permutation starting from block 1
Matthew Bonanni's avatar
Matthew Bonanni committed
229
    else:
230
        perm = torch.arange(1, blocks_end)  # Sequential order starting from block 1
Matthew Bonanni's avatar
Matthew Bonanni committed
231
232

    inv_perm = torch.zeros(blocks_end, dtype=torch.long, device=device)
233
    inv_perm[1:] = torch.argsort(perm) + 1  # Add 1 to account for starting from block 1
Matthew Bonanni's avatar
Matthew Bonanni committed
234
235
236
237
238
239
240
241
242
243
    kv_cache[1:blocks_end, ...] = kv_cache[perm, ...]

    # Construct the right block table
    # Start from block_id=1 since block_id=0 is considered the null block
    start_block_idx = 1
    for i in range(batch_size):
        num_blocks_for_seq = cdiv(int(seq_lens[i]), block_size)
        start = start_block_idx
        end = start + num_blocks_for_seq
        block_table[i, :num_blocks_for_seq] = inv_perm[start:end]
244
        block_table[i, num_blocks_for_seq:] = 0
Matthew Bonanni's avatar
Matthew Bonanni committed
245
246
247
248
249
250
251
252
253
254
        start_block_idx += num_blocks_for_seq

        # Create a realistic slot mapping that corresponds to the block table
    for i in range(batch_size):
        token_offsets = torch.arange(int(query_lens[i])) + int(context_lens[i])
        block_indices = token_offsets // block_size
        token_inter_block_offsets = token_offsets % block_size
        start = common_attn_metadata.query_start_loc_cpu[i]
        end = common_attn_metadata.query_start_loc_cpu[i + 1]
        slot_mapping[start:end] = block_table[
255
256
            i, block_indices
        ] * block_size + token_inter_block_offsets.to(device)
Matthew Bonanni's avatar
Matthew Bonanni committed
257
258
259
260
261
262
263
264
265
266
267

    return kv_cache


class MockAttentionLayer:
    """A mock attention layer for testing."""

    def __init__(self, device: torch.device):
        self._q_scale = torch.tensor(1.0, device=device)
        self._k_scale = torch.tensor(1.0, device=device)
        self._v_scale = torch.tensor(1.0, device=device)
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
        self._prob_scale = torch.tensor(1.0, device=device)
        self._q_scale_float = 1.0
        self._k_scale_float = 1.0
        self._v_scale_float = 1.0

    def forward(self, *_args, **_kwargs):
        raise NotImplementedError


class MockMLAAttentionLayer(AttentionLayerBase):
    """A mock MLA attention layer for populating static_forward_context."""

    def __init__(self, impl):
        self.impl = impl

    def get_attn_backend(self):
        raise NotImplementedError

    def get_kv_cache_spec(self, vllm_config):
        raise NotImplementedError
Matthew Bonanni's avatar
Matthew Bonanni committed
288
289


290
def run_attention_backend(
291
    backend: AttentionBackendEnum,
292
    kv_cache_spec: MLAAttentionSpec,
293
294
295
296
297
298
299
300
301
302
303
304
305
306
    layer_names: list[str],
    vllm_config,
    device: torch.device,
    common_attn_metadata: CommonAttentionMetadata,
    query: torch.Tensor,
    kv_c: torch.Tensor,
    k_pe: torch.Tensor,
    kv_cache: torch.Tensor,
    kv_lora_rank: int,
    qk_nope_head_dim: int,
    qk_rope_head_dim: int,
    v_head_dim: int,
    mock_kv_b_proj,
) -> torch.Tensor:
Matthew Bonanni's avatar
Matthew Bonanni committed
307
308
    """Run attention computation using the specified backend's AttentionImpl."""

309
    builder_cls, impl_cls = try_get_attention_backend(backend)
Matthew Bonanni's avatar
Matthew Bonanni committed
310

311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
    # Set the current vllm config so that get_current_vllm_config() works
    # in the backend implementations
    with set_current_vllm_config(vllm_config):
        # Instantiate MLA implementation
        num_heads = vllm_config.model_config.get_num_attention_heads(
            vllm_config.parallel_config
        )
        num_kv_heads = vllm_config.model_config.get_num_kv_heads(
            vllm_config.parallel_config
        )
        head_size = vllm_config.model_config.get_head_size()
        scale = 1.0 / (head_size**0.5)
        impl = impl_cls(
            num_heads=num_heads,
            head_size=head_size,
            scale=scale,
            num_kv_heads=num_kv_heads,
            alibi_slopes=None,
            sliding_window=None,
            kv_cache_dtype="auto",
            logits_soft_cap=None,
            attn_type="decoder",
            kv_sharing_target_layer_name=None,
            q_lora_rank=None,
            kv_lora_rank=kv_lora_rank,
            qk_nope_head_dim=qk_nope_head_dim,
            qk_rope_head_dim=qk_rope_head_dim,
            qk_head_dim=qk_nope_head_dim + qk_rope_head_dim,
            v_head_dim=v_head_dim,
            kv_b_proj=mock_kv_b_proj,
        )
Matthew Bonanni's avatar
Matthew Bonanni committed
342

343
344
345
        # Process weights to create W_UK_T and W_UV attributes needed by MLA
        act_dtype = _convert_dtype_to_torch(vllm_config.model_config.dtype)
        impl.process_weights_after_loading(act_dtype)
Matthew Bonanni's avatar
Matthew Bonanni committed
346

347
348
349
350
351
352
353
354
355
356
357
358
359
        # Populate static_forward_context with mock attention layers
        for layer_name in layer_names:
            vllm_config.compilation_config.static_forward_context[layer_name] = (
                MockMLAAttentionLayer(impl)
            )

        # Build metadata
        builder = builder_cls(kv_cache_spec, layer_names, vllm_config, device)
        attn_metadata = builder.build(
            common_prefix_len=0,
            common_attn_metadata=common_attn_metadata,
        )

360
361
362
363
364
365
        # Create mock layer and output buffer
        mock_layer = MockAttentionLayer(device)
        num_tokens = query.shape[0]
        output = torch.empty(
            num_tokens, num_heads * v_head_dim, dtype=query.dtype, device=query.device
        )
Matthew Bonanni's avatar
Matthew Bonanni committed
366

367
368
369
370
371
372
        # Run forward pass
        # NOTE: The query, key, and value are already shaped correctly
        # in the calling test function.
        output = impl.forward(
            mock_layer, query, kv_c, k_pe, kv_cache, attn_metadata, output=output
        )
Matthew Bonanni's avatar
Matthew Bonanni committed
373

374
        return output
Matthew Bonanni's avatar
Matthew Bonanni committed
375
376


377
378
379
380
381
382
383
384
385
386
387
388
389
@pytest.mark.parametrize(
    "batch_spec_name",
    [
        "small_decode",
        "small_prefill",
        "mixed_small",
        "medium_decode",
        "medium_prefill",
        "mixed_medium",
        "large_decode",
        "large_prefill",
        "single_decode",
        "single_prefill",
390
391
        "spec_decode_small",
        "spec_decode_medium",
392
393
    ],
)
394
395
396
397
398
@pytest.mark.parametrize("model", ["deepseek-ai/DeepSeek-R1"])
@pytest.mark.parametrize("tensor_parallel_size", [1, 4, 8, 16])
def test_backend_correctness(
    dist_init, batch_spec_name: str, model: str, tensor_parallel_size: int
):
Matthew Bonanni's avatar
Matthew Bonanni committed
399
400
401
402
403
404
405
406
407
408
409
410
411
412
    """
    Test that all backends produce similar outputs to a reference implementation
    using torch.nn.functional.scaled_dot_product_attention.

    This test works by:
    1. Generating a batch of sequences with specified context and query lengths.
    2. Computing a ground-truth attention output using torch.sdpa on
       contiguous Q, K, and V tensors.
    3. Simulating vLLM's paged KV cache: It takes the context portion of the
       K/V tensors and manually places them into a paged buffer according to
       the test's (randomly generated) block table.
    4. Running each vLLM attention backend with the new queries and the
       simulated paged KV cache.
    5. Comparing the vLLM backend's output to the ground-truth SDPA output.
413
414
415
416
417

    Note: When tensor_parallel_size > 1, we simulate the head partitioning
    by overriding the model config to use fewer heads, without requiring
    multiple GPUs. This tests that backends work correctly with different
    head counts.
Matthew Bonanni's avatar
Matthew Bonanni committed
418
    """
419

Matthew Bonanni's avatar
Matthew Bonanni committed
420
    batch_spec = BATCH_SPECS[batch_spec_name]
421
    is_spec_decode_test = batch_spec_name.startswith("spec_decode")
422
423
    unique_block_sizes = sorted(set(BACKEND_BLOCK_SIZES.values()))
    default_block_size = unique_block_sizes[0]
424
    required_blocks = sum(
425
426
        (seq_len + default_block_size - 1) // default_block_size
        for seq_len in batch_spec.seq_lens
427
428
429
430
    )
    # Add 1 for null block at index 0, and some buffer
    num_gpu_blocks = required_blocks + 1 + 100

431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
    hf_config_override = None
    if tensor_parallel_size > 1:
        from vllm.config import ModelConfig

        temp_config = ModelConfig(model=model, max_model_len=1)
        original_num_heads = temp_config.hf_text_config.num_attention_heads
        original_num_kv_heads = getattr(
            temp_config.hf_text_config, "num_key_value_heads", None
        )
        hf_config_override = {
            "num_attention_heads": original_num_heads // tensor_parallel_size,
        }
        if original_num_kv_heads is not None:
            hf_config_override["num_key_value_heads"] = max(
                1, original_num_kv_heads // tensor_parallel_size
            )

448
    vllm_config = create_vllm_config(
449
        model_name=model,
450
        tensor_parallel_size=1,  # Always use TP=1 to avoid multi-GPU requirements
451
452
        max_model_len=max(batch_spec.seq_lens),
        num_gpu_blocks=num_gpu_blocks,
453
        block_size=default_block_size,
454
        hf_config_override=hf_config_override,
455
    )
456
457
458
459
460
461
462
463
464
465
466
467
468
469

    # For spec decode tests, add a speculative_config to set the reorder_batch_threshold
    if is_spec_decode_test:
        from vllm.config import SpeculativeConfig

        # Get the query length from the batch spec (they should all be uniform)
        query_len = batch_spec.query_lens[0]
        # Set num_speculative_tokens to query_len - 1
        # (since threshold is 1 + num_spec_tokens)
        # Use ngram method which doesn't require a draft model
        vllm_config.speculative_config = SpeculativeConfig(
            method="ngram", num_speculative_tokens=query_len - 1
        )

Matthew Bonanni's avatar
Matthew Bonanni committed
470
471
472
473
474
475
476
    device = torch.device("cuda:0")

    # 1. Setup
    batch_size = batch_spec.batch_size
    seq_lens = batch_spec.seq_lens
    query_lens = batch_spec.query_lens
    num_q_heads = vllm_config.model_config.get_num_attention_heads(
477
478
        vllm_config.parallel_config
    )
Matthew Bonanni's avatar
Matthew Bonanni committed
479
480
481
482
483
484
485
    head_size = vllm_config.model_config.get_head_size()
    dtype = _convert_dtype_to_torch(vllm_config.model_config.dtype)
    kv_lora_rank = 512
    qk_rope_head_dim = 64
    qk_nope_head_dim = 128
    v_head_dim = 128
    total_head_size = kv_lora_rank + qk_rope_head_dim
486
    assert kv_lora_rank + qk_rope_head_dim == head_size, (
Matthew Bonanni's avatar
Matthew Bonanni committed
487
        f"MLA dimensions don't match: {total_head_size} != {head_size}"
488
    )
Matthew Bonanni's avatar
Matthew Bonanni committed
489
490
491
492
    scale = 1.0 / (total_head_size**0.5)

    # 2. Generate data and compute SDPA reference output for MLA
    all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], []
493
    all_sdpa_outputs: list[list[torch.Tensor]] = []
Matthew Bonanni's avatar
Matthew Bonanni committed
494
495
496
    kv_c_contexts, k_pe_contexts = [], []

    # Create shared MLA weight matrices for consistency across all sequences
497
498
499
500
501
502
    W_UK = torch.randn(
        kv_lora_rank, num_q_heads, qk_nope_head_dim, dtype=dtype, device=device
    )
    W_UV = torch.randn(
        kv_lora_rank, num_q_heads, v_head_dim, dtype=dtype, device=device
    )
Matthew Bonanni's avatar
Matthew Bonanni committed
503
504
    kv_b_proj_weight = torch.cat([W_UK, W_UV], dim=-1)

505
506
507
    for i, backend in enumerate(BACKENDS_TO_TEST):
        all_sdpa_outputs.append([])

Matthew Bonanni's avatar
Matthew Bonanni committed
508
509
510
511
512
513
514
515
    for i in range(batch_size):
        s_len = seq_lens[i]
        q_len = query_lens[i]
        context_len = s_len - q_len

        # Generate MLA tensors
        # Q has both nope and rope components:
        # [q_len, num_heads, qk_nope_head_dim + qk_rope_head_dim]
516
517
518
519
520
521
522
        q_c = torch.randn(
            q_len,
            num_q_heads,
            qk_nope_head_dim + qk_rope_head_dim,
            dtype=dtype,
            device=device,
        )
Matthew Bonanni's avatar
Matthew Bonanni committed
523
524

        # KV_C (latent K/V): [s_len, kv_lora_rank]
525
        kv_c_full = torch.randn(s_len, kv_lora_rank, dtype=dtype, device=device)
Matthew Bonanni's avatar
Matthew Bonanni committed
526
527

        # K_PE (rope component): [s_len, 1, qk_rope_head_dim]
528
        k_pe_full = torch.randn(s_len, 1, qk_rope_head_dim, dtype=dtype, device=device)
Matthew Bonanni's avatar
Matthew Bonanni committed
529

530
531
532
533
534
535
536
537
        # Determine if this sequence uses the decode pipeline or prefill
        # pipeline for each backend
        # NOTE: For spec decode tests with uniform query_len > 1, backends that
        # support spec decode (FLASH_ATTN_MLA with varlen support, FLASHMLA with
        # uniform support) will use the decode pipeline (MQA-style), while
        # backends that only support single-token queries will use the prefill
        # pipeline (MHA-style). This ensures the reference implementation
        # matches each backend's actual decode/prefill pipeline path.
538
        is_decode = []
539
        for backend_idx, backend in enumerate(BACKENDS_TO_TEST):
540
            builder_cls, _ = try_get_attention_backend(backend)
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
            if is_spec_decode_test:
                query_len_support = getattr(
                    builder_cls, "query_len_support", QueryLenSupport.SINGLE_ONLY
                )
                supports_spec = query_len_support != QueryLenSupport.SINGLE_ONLY
                is_decode.append(supports_spec)
            else:
                threshold = getattr(builder_cls, "reorder_batch_threshold", None)
                query_len_support = getattr(
                    builder_cls, "query_len_support", QueryLenSupport.SINGLE_ONLY
                )
                within_threshold = q_len <= threshold if threshold else False
                if (
                    within_threshold
                    and query_len_support == QueryLenSupport.UNIFORM
                    and i > 0
                ):
                    first_q_len = query_lens[0]
                    within_threshold = q_len == first_q_len
                is_decode.append(within_threshold)
Matthew Bonanni's avatar
Matthew Bonanni committed
561
562
563
564

        # Split q into nope and rope components
        q_nope, q_pe = q_c.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1)

565
566
567
568
569
        #######################################################
        # Decode path: MQA-style attention in latent space
        # Transform q_nope to latent space: q_nope @ W_UK
        # q_nope: [1, num_heads, qk_nope_head_dim]
        # W_UK: [kv_lora_rank, num_heads, qk_nope_head_dim]
570
571
572
        ql_nope = torch.einsum(
            "qnh,lnh->qnl", q_nope, W_UK
        )  # [1, num_heads, kv_lora_rank]
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597

        # Build MQA attention inputs
        # Q: [1, num_heads, kv_lora_rank + qk_rope_head_dim]
        q_mqa = torch.cat([ql_nope, q_pe], dim=-1)
        # K: [s_len, kv_lora_rank + qk_rope_head_dim]
        # (broadcasted to all heads)
        k_mqa = torch.cat([kv_c_full, k_pe_full.squeeze(1)], dim=-1)
        k_mqa = k_mqa.unsqueeze(1).expand(-1, num_q_heads, -1)
        # V: [s_len, kv_lora_rank] (broadcasted to all heads)
        v_mqa = kv_c_full.unsqueeze(1).expand(-1, num_q_heads, -1)

        # Create custom attention mask for decode path:
        # - Query tokens can attend to all context tokens
        # - Query tokens can only attend to query tokens up to their position
        attn_mask = torch.ones(q_len, s_len, dtype=torch.bool, device=device)
        # Apply causal mask only to the query portion (context_len onwards)
        causal_mask = torch.tril(torch.ones(q_len, q_len, device=device))
        attn_mask[:, context_len:] = causal_mask

        # SDPA expects (N, H, L, D)
        q_sdpa_in = q_mqa.unsqueeze(0).transpose(1, 2)
        k_sdpa_in = k_mqa.unsqueeze(0).transpose(1, 2)
        v_sdpa_in = v_mqa.unsqueeze(0).transpose(1, 2)

        sdpa_out_i_decode = torch.nn.functional.scaled_dot_product_attention(
598
599
            q_sdpa_in, k_sdpa_in, v_sdpa_in, attn_mask=attn_mask, scale=scale
        )
600
        sdpa_out_i_decode = sdpa_out_i_decode.transpose(1, 2).squeeze(
601
602
            0
        )  # [1, num_heads, kv_lora_rank]
603
604

        # Project back to output space: sdpa_out @ W_UV
605
        sdpa_out_i_decode = torch.einsum("qnl,lnv->qnv", sdpa_out_i_decode, W_UV)
606
607
608
609
610
611
        sdpa_out_i_decode = sdpa_out_i_decode.flatten(start_dim=-2)

        #######################################################
        # Prefill path: MHA-style attention with full sequence
        # Apply kv_b_proj to the full kv_c tensor
        kv_nope_full = torch.einsum("sl,lnh->snh", kv_c_full, kv_b_proj_weight)
612
        k_nope_full, v_full = kv_nope_full.split([qk_nope_head_dim, v_head_dim], dim=-1)
613
614

        # Build attention inputs for full sequence
615
        q_mha = torch.cat([q_nope, q_pe], dim=-1)  # [q_len, num_heads, total_dim]
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
        k_pe_full_expanded = k_pe_full.expand(-1, num_q_heads, -1)
        k_full = torch.cat([k_nope_full, k_pe_full_expanded], dim=-1)

        # Create custom attention mask:
        # - Query tokens can attend to all context tokens
        # - Query tokens can only attend to query tokens up to their pos
        attn_mask = torch.ones(q_len, s_len, dtype=torch.bool, device=device)
        # Apply causal mask only to the query portion (context_len onwards)
        causal_mask = torch.tril(torch.ones(q_len, q_len, device=device))
        attn_mask[:, context_len:] = causal_mask

        # SDPA expects (N, H, L, D)
        q_sdpa_in = q_mha.unsqueeze(0).transpose(1, 2)
        k_sdpa_in = k_full.unsqueeze(0).transpose(1, 2)
        v_sdpa_in = v_full.unsqueeze(0).transpose(1, 2)

        # Single attention call with custom mask
        sdpa_out_i_prefill = torch.nn.functional.scaled_dot_product_attention(
634
635
            q_sdpa_in, k_sdpa_in, v_sdpa_in, attn_mask=attn_mask, scale=scale
        )
636
637
638
        sdpa_out_i_prefill = sdpa_out_i_prefill.transpose(1, 2).squeeze(0)
        sdpa_out_i_prefill = sdpa_out_i_prefill.flatten(start_dim=-2)

639
640
641
        for backend_idx, backend in enumerate(BACKENDS_TO_TEST):
            if is_decode[backend_idx]:
                all_sdpa_outputs[backend_idx].append(sdpa_out_i_decode)
642
            else:
643
                all_sdpa_outputs[backend_idx].append(sdpa_out_i_prefill)
Matthew Bonanni's avatar
Matthew Bonanni committed
644
645
646
647
648
649
650
651
652
653
654
655
656
657

        # Inputs for vLLM MLA backends are just the new tokens
        all_q_vllm.append(q_c)
        all_kv_c_vllm.append(kv_c_full[context_len:])  # New kv_c tokens
        all_k_pe_vllm.append(k_pe_full[context_len:])  # New k_pe tokens

        # Contextual K/V data used to populate the paged cache (MLA format)
        kv_c_contexts.append(kv_c_full[:context_len])
        k_pe_contexts.append(k_pe_full[:context_len])

    # Concatenate all sequences (no reordering needed)
    query_vllm = torch.cat(all_q_vllm, dim=0)
    kv_c_vllm = torch.cat(all_kv_c_vllm, dim=0)
    k_pe_vllm = torch.cat(all_k_pe_vllm, dim=0)
658
659
660
    sdpa_outputs = {}
    for backend_idx, backend in enumerate(BACKENDS_TO_TEST):
        sdpa_outputs[backend] = torch.cat(all_sdpa_outputs[backend_idx], dim=0)
Matthew Bonanni's avatar
Matthew Bonanni committed
661
662
663

    # Create mock kv_b_proj using the same weights as reference implementation
    from vllm.model_executor.layers.linear import ColumnParallelLinear
664
665
666
667
668
669

    mock_kv_b_proj = ColumnParallelLinear(
        input_size=kv_lora_rank,
        output_size=num_q_heads * (qk_nope_head_dim + v_head_dim),
        bias=False,
    ).to(device=device, dtype=dtype)
Matthew Bonanni's avatar
Matthew Bonanni committed
670
671
672
673
674

    # Set the mock weights to match our reference implementation
    # Reshape W_UK and W_UV to match the expected kv_b_proj format
    # [kv_lora_rank, num_heads, qk_nope_head_dim + v_head_dim]
    kv_b_proj_weight = kv_b_proj_weight.view(
675
676
        kv_lora_rank, num_q_heads * (qk_nope_head_dim + v_head_dim)
    )
677
    mock_kv_b_proj.weight = torch.nn.Parameter(kv_b_proj_weight.T, requires_grad=False)
Matthew Bonanni's avatar
Matthew Bonanni committed
678

679
680
681
682
    # 3. Create metadata and KV caches for each block size
    # Group backends by block size and test each group
    metadata_per_block_size = {}
    kv_cache_per_block_size = {}
Matthew Bonanni's avatar
Matthew Bonanni committed
683

684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
    for block_size in unique_block_sizes:
        # Create metadata for this block size
        common_attn_metadata = create_common_attn_metadata(
            batch_spec, block_size, device
        )

        # Pad block table to meet requirement:
        # block_num % (128 / block_size) == 0
        required_divisor = int(128 / block_size)
        current_block_num = common_attn_metadata.block_table_tensor.shape[1]
        if current_block_num % required_divisor != 0:
            # Pad to next multiple of required_divisor
            padded_block_num = (
                (current_block_num + required_divisor - 1) // required_divisor
            ) * required_divisor
            padding_cols = padded_block_num - current_block_num
            padding = torch.zeros(
                (common_attn_metadata.block_table_tensor.shape[0], padding_cols),
                dtype=torch.int32,
                device=device,
            )
            common_attn_metadata.block_table_tensor = torch.cat(
                [common_attn_metadata.block_table_tensor, padding], dim=1
            )

        metadata_per_block_size[block_size] = common_attn_metadata

        # Create KV cache for this block size
        required_blocks_for_size = sum(
            (seq_len + block_size - 1) // block_size for seq_len in batch_spec.seq_lens
        )
        num_blocks_for_size = required_blocks_for_size + 1 + 100

        kv_cache = create_and_prepopulate_kv_cache(
            kv_c_contexts=kv_c_contexts,
            k_pe_contexts=k_pe_contexts,
            block_size=block_size,
            head_size=head_size,
            dtype=dtype,
            device=device,
            num_blocks=num_blocks_for_size,
            common_attn_metadata=common_attn_metadata,
            randomize_blocks=True,
        )
        kv_cache_per_block_size[block_size] = kv_cache
Matthew Bonanni's avatar
Matthew Bonanni committed
729
730

    # 4. Run vLLM backends and compare
731
    failures = []
732
733
    for backend_idx, backend_name in enumerate(BACKENDS_TO_TEST):
        # Skip backends that don't support spec decode for spec decode tests
734
        if is_spec_decode_test and backend_name not in SPEC_DECODE_BACKENDS:
735
736
            continue

737
738
739
740
741
742
        # Get the appropriate block_size, metadata, and cache for this backend
        block_size = BACKEND_BLOCK_SIZES[backend_name]
        common_attn_metadata = metadata_per_block_size[block_size]
        kv_cache = kv_cache_per_block_size[block_size]

        # Create kv_cache_spec with the correct block_size for this backend
743
        backend_kv_cache_spec = MLAAttentionSpec(
744
745
746
747
748
749
750
            block_size=block_size,
            num_kv_heads=vllm_config.model_config.get_num_kv_heads(
                vllm_config.parallel_config
            ),
            head_size=vllm_config.model_config.get_head_size(),
            dtype=vllm_config.model_config.dtype,
            sliding_window=vllm_config.model_config.get_sliding_window(),
751
            cache_dtype_str=vllm_config.cache_config.cache_dtype,
752
753
        )

Matthew Bonanni's avatar
Matthew Bonanni committed
754
        backend_output = run_attention_backend(
755
            backend_name,
756
            backend_kv_cache_spec,
757
758
759
760
761
762
763
764
765
766
767
768
769
770
            ["placeholder"],
            vllm_config,
            device,
            common_attn_metadata,
            query_vllm,
            kv_c_vllm,
            k_pe_vllm,
            kv_cache,
            kv_lora_rank,
            qk_nope_head_dim,
            qk_rope_head_dim,
            v_head_dim,
            mock_kv_b_proj,
        )
Matthew Bonanni's avatar
Matthew Bonanni committed
771

772
773
774
        # Use backend_idx to get the correct SDPA output for this backend
        expected_output = sdpa_outputs[backend_name]

Matthew Bonanni's avatar
Matthew Bonanni committed
775
        # Check shape and dtype consistency
776
777
778
779
780
781
782
783
784
        try:
            assert backend_output.shape == expected_output.shape, (
                f"[{backend_name}] shape {backend_output.shape} != "
                f"SDPA shape {expected_output.shape}"
            )
            assert backend_output.dtype == expected_output.dtype, (
                f"[{backend_name}] dtype {backend_output.dtype} != "
                f"SDPA dtype {expected_output.dtype}"
            )
Matthew Bonanni's avatar
Matthew Bonanni committed
785

786
787
788
            assert torch.isfinite(backend_output).all(), (
                f"[{backend_name}] produced non-finite values"
            )
Matthew Bonanni's avatar
Matthew Bonanni committed
789

790
791
792
            # Check numerical similarity
            rtol = 1e-2
            atol = 5e-1
Matthew Bonanni's avatar
Matthew Bonanni committed
793

794
795
796
797
798
799
800
            max_diff = torch.max(torch.abs(backend_output - expected_output)).item()
            max_rel_diff = torch.max(
                torch.abs(backend_output - expected_output) / torch.abs(expected_output)
            ).item()
            all_close = torch.allclose(
                backend_output, expected_output, rtol=rtol, atol=atol
            )
Matthew Bonanni's avatar
Matthew Bonanni committed
801

802
803
804
805
806
807
808
809
810
811
812
813
            assert all_close, (
                f"[{backend_name}] output differs from SDPA baseline. "
                f"Max diff: {max_diff:.6f}, max rel diff: {max_rel_diff:.6f})"
            )
        except AssertionError as e:
            failures.append(str(e))

    # Report all failures at once
    if failures:
        # Create a summary for the single-line failure message
        backend_names = []
        for f in failures:
814
            if "[AttentionBackendEnum." in f:
815
816
817
818
819
820
                backend_name = f.split("[")[1].split("]")[0]
                backend_names.append(backend_name)

        summary = f"{len(failures)} backend(s) failed: {', '.join(backend_names)}"
        detailed_msg = "\n".join(failures)
        pytest.fail(f"{summary}\n{detailed_msg}")