| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import MethodType |
| 5 | from unittest.mock import MagicMock |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from vllm.platforms import current_platform |
| 11 | |
| 12 | requires_flashinfer_mla = pytest.mark.skipif( |
| 13 | not current_platform.has_device_capability(100), |
| 14 | reason="FlashInfer MLA requires compute capability 10 or above.", |
| 15 | ) |
| 16 | |
| 17 | |
| 18 | def test_mla_dcp_gathered_query_reserves_backend_head_storage(): |
| 19 | from vllm.v1.attention.ops.dcp import reserve_query_head_storage |
| 20 | |
| 21 | query = torch.randn(3, 24, 576, dtype=torch.bfloat16) |
| 22 | |
| 23 | padded = reserve_query_head_storage(query, 128) |
| 24 | |
| 25 | assert padded.shape == query.shape |
| 26 | assert padded.stride() == query.stride() |
| 27 | assert padded.untyped_storage().nbytes() >= 3 * 128 * 576 * 2 |
| 28 | assert torch.equal(padded, query) |
| 29 | |
| 30 | |
| 31 | @requires_flashinfer_mla |
| 32 | def test_flashinfer_mla_selects_backend_for_gathered_heads(): |
| 33 | from vllm.v1.attention.backends.mla.flashinfer_mla import ( |
| 34 | _select_mla_decode_backend, |
| 35 | ) |
| 36 | |
| 37 | assert _select_mla_decode_backend(6) is None |
| 38 | assert _select_mla_decode_backend(24) == "cute-dsl" |
| 39 | |
| 40 | |
| 41 | @requires_flashinfer_mla |
| 42 | def test_flashinfer_mla_forward_uses_gathered_head_count(monkeypatch): |
| 43 | import vllm.v1.attention.backends.mla.flashinfer_mla as flashinfer_mla |
| 44 | |
| 45 | impl = MagicMock() |
| 46 | impl.bmm1_scale = 1.0 |
| 47 | impl.bmm2_scale = 1.0 |
| 48 | impl.need_to_return_lse_for_decode = True |
| 49 | impl.dcp_world_size = 1 |
| 50 | impl.dcp_rank = 0 |
| 51 | impl.cp_kv_cache_interleave_size = 1 |
| 52 | impl.num_heads = 6 |
| 53 | impl.qk_nope_head_dim = 128 |
| 54 | impl.kv_lora_rank = 512 |
| 55 | impl.qk_rope_head_dim = 64 |
| 56 | impl.kv_cache_dtype = "auto" |
| 57 | |
| 58 | attn_metadata = MagicMock() |
| 59 | attn_metadata.causal = True |
| 60 | attn_metadata.num_decode_tokens = 2 |
| 61 | attn_metadata.num_decodes = 2 |
| 62 | attn_metadata.max_seq_len = 1 |
| 63 | attn_metadata.decode.block_table = torch.zeros(2, 1, dtype=torch.int32) |
| 64 | attn_metadata.decode.seq_lens = torch.ones(2, dtype=torch.int32) |
| 65 | attn_metadata.decode.max_query_len = 1 |
| 66 | |
| 67 | query = torch.ones(2, 24, 576, dtype=torch.bfloat16) |
| 68 | kv_cache = torch.ones(1, 128, 576, dtype=torch.bfloat16) |
| 69 | kernel = MagicMock( |
| 70 | return_value=( |
| 71 | torch.ones(2, 1, 24, 512, dtype=torch.bfloat16), |
| 72 | torch.ones(2, 24, dtype=torch.float32), |
| 73 | ) |
| 74 | ) |
| 75 | monkeypatch.setattr(flashinfer_mla, "_get_workspace_buffer", MagicMock()) |
| 76 | monkeypatch.setattr(flashinfer_mla, "trtllm_batch_decode_with_kv_cache_mla", kernel) |
| 77 | layer = MagicMock() |
| 78 | |
| 79 | output, lse = flashinfer_mla.FlashInferMLAImpl.forward_mqa( |
| 80 | impl, query, kv_cache, attn_metadata, layer |
| 81 | ) |
| 82 | |
| 83 | assert kernel.call_args.kwargs["backend"] == "cute-dsl" |
| 84 | assert output.shape == (2, 24, 512) |
| 85 | assert lse is not None and lse.shape == (2, 24) |
| 86 | |
| 87 | |
| 88 | @requires_flashinfer_mla |
| 89 | @pytest.mark.parametrize("causal", [True, False], ids=["causal", "noncausal"]) |
| 90 | def test_flashinfer_mla_forward_uses_native_dcp_api(monkeypatch, causal): |
| 91 | import vllm.v1.attention.backends.mla.flashinfer_mla as flashinfer_mla |
| 92 | |
| 93 | impl = MagicMock() |
| 94 | impl.bmm1_scale = 1.0 |
| 95 | impl.bmm2_scale = 1.0 |
| 96 | impl.need_to_return_lse_for_decode = True |
| 97 | impl.num_heads = 6 |
| 98 | impl.qk_nope_head_dim = 128 |
| 99 | impl.kv_lora_rank = 512 |
| 100 | impl.qk_rope_head_dim = 64 |
| 101 | impl.kv_cache_dtype = "auto" |
| 102 | impl.dcp_world_size = 8 |
| 103 | impl.dcp_rank = 3 |
| 104 | impl.cp_kv_cache_interleave_size = 1 |
| 105 | impl._flattened_decode_metadata = MethodType( |
| 106 | flashinfer_mla.FlashInferMLAImpl._flattened_decode_metadata, impl |
| 107 | ) |
| 108 | |
| 109 | num_reqs, query_len = 2, 3 |
| 110 | num_tokens = num_reqs * query_len |
| 111 | block_table = torch.tensor([[10], [20]], dtype=torch.int32) |
| 112 | seq_lens = torch.tensor([4, 5], dtype=torch.int32) |
| 113 | global_causal_seq_lens = torch.tensor([31, 35], dtype=torch.int32) |
| 114 | attn_metadata = MagicMock() |
| 115 | attn_metadata.causal = causal |
| 116 | attn_metadata.num_decode_tokens = num_tokens |
| 117 | attn_metadata.num_decodes = num_reqs |
| 118 | attn_metadata.max_seq_len = 9 |
| 119 | attn_metadata.decode.block_table = block_table |
| 120 | attn_metadata.decode.seq_lens = seq_lens |
| 121 | attn_metadata.decode.dcp_tot_seq_lens = global_causal_seq_lens |
| 122 | attn_metadata.decode.max_query_len = query_len |
| 123 | attn_metadata.decode.query_start_loc = torch.arange( |
| 124 | 0, num_tokens + 1, query_len, dtype=torch.int32 |
| 125 | ) |
| 126 | attn_metadata.decode.query_len = 0 |
| 127 | |
| 128 | query = torch.ones(num_tokens, 24, 576, dtype=torch.bfloat16) |
| 129 | kv_cache = torch.ones(2, 128, 576, dtype=torch.bfloat16) |
| 130 | kernel_batch = num_reqs if causal else num_tokens |
| 131 | kernel_query_len = query_len if causal else 1 |
| 132 | kernel = MagicMock( |
| 133 | return_value=( |
| 134 | torch.ones( |
| 135 | kernel_batch, |
| 136 | kernel_query_len, |
| 137 | 24, |
| 138 | 512, |
| 139 | dtype=torch.bfloat16, |
| 140 | ), |
| 141 | torch.ones( |
| 142 | kernel_batch, |
| 143 | kernel_query_len, |
| 144 | 24, |
| 145 | dtype=torch.float32, |
| 146 | ), |
| 147 | ) |
| 148 | ) |
| 149 | monkeypatch.setattr(flashinfer_mla, "_get_workspace_buffer", MagicMock()) |
| 150 | monkeypatch.setattr(flashinfer_mla, "trtllm_batch_decode_with_kv_cache_mla", kernel) |
| 151 | |
| 152 | output, lse = flashinfer_mla.FlashInferMLAImpl.forward_mqa( |
| 153 | impl, query, kv_cache, attn_metadata, MagicMock() |
| 154 | ) |
| 155 | |
| 156 | call = kernel.call_args.kwargs |
| 157 | assert call["query"].shape == (kernel_batch, kernel_query_len, 24, 576) |
| 158 | expected_block_table = ( |
| 159 | block_table if causal else block_table.repeat_interleave(query_len, dim=0) |
| 160 | ) |
| 161 | expected_seq_lens = seq_lens if causal else seq_lens.repeat_interleave(query_len) |
| 162 | expected_global_seq_lens = ( |
| 163 | global_causal_seq_lens |
| 164 | if causal |
| 165 | else global_causal_seq_lens.repeat_interleave(query_len) |
| 166 | ) |
| 167 | torch.testing.assert_close(call["block_tables"], expected_block_table) |
| 168 | torch.testing.assert_close(call["seq_lens"], expected_seq_lens) |
| 169 | assert call["backend"] == "cute-dsl" |
| 170 | assert call["enable_dcp"] is True |
| 171 | assert call["cp_world"] == 8 |
| 172 | assert call["cp_rank"] == 3 |
| 173 | torch.testing.assert_close( |
| 174 | call["causal_seqlens_kv_global"], expected_global_seq_lens |
| 175 | ) |
| 176 | assert "multi_ctas_kv_counter_buffer" not in call |
| 177 | assert flashinfer_mla.FlashInferMLAImpl.lse_base_on_e |
| 178 | assert output.shape == (num_tokens, 24, 512) |
| 179 | assert lse is not None and lse.shape == (num_tokens, 24) |