| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | |
| 6 | import pytest |
| 7 | import torch |
| 8 | |
| 9 | from vllm.model_executor.layers.attention.mla_attention import ( |
| 10 | MLACommonMetadata, |
| 11 | MLACommonMetadataBuilder, |
| 12 | QueryLenSupport, |
| 13 | ) |
| 14 | from vllm.v1.attention.backend import CommonAttentionMetadata |
| 15 | from vllm.v1.kv_cache_interface import MLAAttentionSpec |
| 16 | |
| 17 | |
| 18 | class _NonCausalMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): |
| 19 | supports_non_causal_multi_token_decode = True |
| 20 | |
| 21 | |
| 22 | def _metadata( |
| 23 | query_start_loc: list[int], |
| 24 | num_tokens: int | None = None, |
| 25 | causal: bool = False, |
| 26 | ) -> CommonAttentionMetadata: |
| 27 | num_reqs = len(query_start_loc) - 1 |
| 28 | num_tokens = query_start_loc[-1] if num_tokens is None else num_tokens |
| 29 | return CommonAttentionMetadata( |
| 30 | query_start_loc=torch.tensor(query_start_loc, dtype=torch.int32), |
| 31 | query_start_loc_cpu=torch.tensor(query_start_loc, dtype=torch.int32), |
| 32 | seq_lens=torch.arange(1, num_reqs + 1, dtype=torch.int32) * 100 + 8, |
| 33 | num_reqs=num_reqs, |
| 34 | num_actual_tokens=num_tokens, |
| 35 | max_query_len=max( |
| 36 | end - start for start, end in zip(query_start_loc, query_start_loc[1:]) |
| 37 | ), |
| 38 | max_seq_len=num_reqs * 100 + 8, |
| 39 | block_table_tensor=torch.arange(num_reqs * 3, dtype=torch.int32).view( |
| 40 | num_reqs, 3 |
| 41 | ), |
| 42 | slot_mapping=torch.arange(num_tokens), |
| 43 | causal=causal, |
| 44 | seq_lens_cpu_upper_bound=None, |
| 45 | ) |
| 46 | |
| 47 | |
| 48 | def _builder(marked: bool = True) -> _NonCausalMLAMetadataBuilder: |
| 49 | builder = object.__new__(_NonCausalMLAMetadataBuilder) |
| 50 | builder.device = torch.device("cpu") |
| 51 | builder.reorder_batch_threshold = 1 |
| 52 | builder.query_len_support = QueryLenSupport.SINGLE_ONLY |
| 53 | builder.non_causal_multi_token_decode = marked |
| 54 | builder.dcp_world_size = 1 |
| 55 | builder.use_pcp = False |
| 56 | builder.metadata_cls = MLACommonMetadata |
| 57 | builder.model_config = SimpleNamespace( |
| 58 | dtype=torch.bfloat16, get_head_size=lambda: 576 |
| 59 | ) |
| 60 | return builder |
| 61 | |
| 62 | |
| 63 | def test_group_capability_keeps_runtime_causality_per_step(): |
| 64 | builder = _builder() |
| 65 | |
| 66 | target_metadata = builder.build(0, _metadata([0, 1, 2], causal=True)) |
| 67 | assert ( |
| 68 | target_metadata.num_decodes, |
| 69 | target_metadata.num_decode_tokens, |
| 70 | target_metadata.num_prefills, |
| 71 | target_metadata.causal, |
| 72 | ) == (2, 2, 0, True) |
| 73 | |
| 74 | common_metadata = _metadata([0, 8, 16]) |
| 75 | metadata = builder.build(0, common_metadata) |
| 76 | |
| 77 | assert metadata.num_decodes == 2 |
| 78 | assert metadata.num_decode_tokens == 16 |
| 79 | assert metadata.num_prefills == 0 |
| 80 | assert metadata.prefill is None |
| 81 | assert metadata.decode is not None |
| 82 | assert metadata.decode.block_table.shape == (2, 3) |
| 83 | assert metadata.decode.seq_lens.shape == (2,) |
| 84 | assert torch.equal(metadata.decode.block_table, common_metadata.block_table_tensor) |
| 85 | assert torch.equal(metadata.decode.seq_lens, common_metadata.seq_lens) |
| 86 | assert not metadata.causal |
| 87 | |
| 88 | |
| 89 | def test_noncausal_support_is_explicit_and_uniform(): |
| 90 | with pytest.raises(ValueError, match="explicitly supported"): |
| 91 | _builder(marked=False).build(0, _metadata([0, 8, 16])) |
| 92 | with pytest.raises(ValueError, match="uniform query block"): |
| 93 | _builder().build(0, _metadata([0, 3, 8])) |
| 94 | |
| 95 | |
| 96 | def test_noncausal_block_allows_trailing_cudagraph_padding(): |
| 97 | common_metadata = _metadata([0, 8, 16, 16], num_tokens=24) |
| 98 | common_metadata.seq_lens[-1] = 0 |
| 99 | |
| 100 | metadata = _builder().build(0, common_metadata) |
| 101 | |
| 102 | assert metadata.num_decodes == 3 |
| 103 | assert metadata.num_decode_tokens == 24 |
| 104 | assert metadata.decode is not None |
| 105 | assert metadata.decode.seq_lens.tolist() == [108, 208, 0] |
| 106 | |
| 107 | |
| 108 | def test_noncausal_block_rejects_non_trailing_padding(): |
| 109 | with pytest.raises(ValueError, match="uniform query block"): |
| 110 | _builder().build(0, _metadata([0, 8, 8, 16], num_tokens=24)) |
| 111 | |
| 112 | |
| 113 | def test_noncausal_decode_metadata_keeps_live_request_buffers(): |
| 114 | common_metadata = _metadata([0, 8, 16]) |
| 115 | metadata = _builder().build(0, common_metadata) |
| 116 | |
| 117 | assert metadata.decode is not None |
| 118 | assert metadata.decode.seq_lens.data_ptr() == common_metadata.seq_lens.data_ptr() |
| 119 | assert ( |
| 120 | metadata.decode.block_table.data_ptr() |
| 121 | == common_metadata.block_table_tensor.data_ptr() |
| 122 | ) |
| 123 | |
| 124 | |
| 125 | def test_mla_cache_marker_is_promoted_to_group_capability(): |
| 126 | kwargs = { |
| 127 | "block_size": 64, |
| 128 | "num_kv_heads": 1, |
| 129 | "head_size": 576, |
| 130 | "dtype": torch.bfloat16, |
| 131 | } |
| 132 | marked = MLAAttentionSpec(**kwargs, non_causal_multi_token_decode=True) |
| 133 | unmarked = MLAAttentionSpec(**kwargs) |
| 134 | |
| 135 | assert MLAAttentionSpec.merge([marked, marked]).non_causal_multi_token_decode |
| 136 | assert not MLAAttentionSpec.merge( |
| 137 | [unmarked, unmarked] |
| 138 | ).non_causal_multi_token_decode |
| 139 | assert MLAAttentionSpec.merge([unmarked, marked]).non_causal_multi_token_decode |