| 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.models.deepseek_v4.amd.rocm import ( |
| 10 | DeepseekV4ROCMAiterMLASparseMetadataBuilder, |
| 11 | DeepseekV4ROCMAiterSparseSWAMetadataBuilder, |
| 12 | ) |
| 13 | from vllm.v1.attention.backend import AttentionCGSupport |
| 14 | from vllm.v1.attention.backends.mla import indexer |
| 15 | from vllm.v1.attention.backends.mla.indexer import ( |
| 16 | DeepseekV4IndexerBackend, |
| 17 | DeepseekV32IndexerMetadataBuilder, |
| 18 | DeepseekV41IndexerBackend, |
| 19 | ) |
| 20 | |
| 21 | |
| 22 | def _make_indexer_config(*, architecture: str, adaptive: bool): |
| 23 | return SimpleNamespace( |
| 24 | model_config=SimpleNamespace(architectures=[architecture]), |
| 25 | num_speculative_tokens=1, |
| 26 | speculative_config=SimpleNamespace(enable_adaptive_verification=adaptive), |
| 27 | ) |
| 28 | |
| 29 | |
| 30 | def _mock_rocm_platform(monkeypatch: pytest.MonkeyPatch) -> None: |
| 31 | # Platform predicates are mutually exclusive in production. Override both |
| 32 | # so these ROCm policy tests do not inherit CUDA capabilities from the CI |
| 33 | # host running them. |
| 34 | monkeypatch.setattr(indexer.current_platform, "is_cuda", lambda: False) |
| 35 | monkeypatch.setattr(indexer.current_platform, "is_rocm", lambda: True) |
| 36 | |
| 37 | |
| 38 | def _make_indexer_builder(*, adaptive: bool, capacity: int = 12): |
| 39 | builder = DeepseekV32IndexerMetadataBuilder.__new__( |
| 40 | DeepseekV32IndexerMetadataBuilder |
| 41 | ) |
| 42 | builder.vllm_config = SimpleNamespace( |
| 43 | speculative_config=SimpleNamespace(enable_adaptive_verification=adaptive) |
| 44 | ) |
| 45 | builder.supports_varlen = False |
| 46 | builder.decode_seq_lens_buffer = torch.zeros(capacity, dtype=torch.int32) |
| 47 | builder.expanded_block_table_buffer = torch.zeros((capacity, 2), dtype=torch.int32) |
| 48 | builder.decode_lens_buffer = torch.zeros(capacity, dtype=torch.int32) |
| 49 | builder.arange_buffer = torch.arange(capacity, dtype=torch.int32) |
| 50 | return builder |
| 51 | |
| 52 | |
| 53 | @pytest.mark.cpu_test |
| 54 | def test_deepseek_v4_rocm_adaptive_builders_support_varlen_full_graphs(): |
| 55 | adaptive_config = SimpleNamespace( |
| 56 | speculative_config=SimpleNamespace(enable_adaptive_verification=True) |
| 57 | ) |
| 58 | fixed_config = SimpleNamespace( |
| 59 | speculative_config=SimpleNamespace(enable_adaptive_verification=False) |
| 60 | ) |
| 61 | |
| 62 | for builder_cls in ( |
| 63 | DeepseekV4ROCMAiterMLASparseMetadataBuilder, |
| 64 | DeepseekV4ROCMAiterSparseSWAMetadataBuilder, |
| 65 | ): |
| 66 | assert ( |
| 67 | builder_cls.get_cudagraph_support(adaptive_config, SimpleNamespace()) |
| 68 | == AttentionCGSupport.ALWAYS |
| 69 | ) |
| 70 | assert ( |
| 71 | builder_cls.get_cudagraph_support(fixed_config, SimpleNamespace()) |
| 72 | == AttentionCGSupport.UNIFORM_BATCH |
| 73 | ) |
| 74 | |
| 75 | |
| 76 | @pytest.mark.cpu_test |
| 77 | def test_deepseek_v4_rocm_adaptive_indexer_support(monkeypatch: pytest.MonkeyPatch): |
| 78 | _mock_rocm_platform(monkeypatch) |
| 79 | adaptive_config = _make_indexer_config( |
| 80 | architecture="DeepseekV4ForCausalLM", adaptive=True |
| 81 | ) |
| 82 | fixed_config = _make_indexer_config( |
| 83 | architecture="DeepseekV4ForCausalLM", adaptive=False |
| 84 | ) |
| 85 | |
| 86 | assert DeepseekV4IndexerBackend.supports_device_cpu_query_lens_mismatch() |
| 87 | assert indexer._rocm_supports_flattened_device_query_lens(adaptive_config) |
| 88 | assert ( |
| 89 | DeepseekV4IndexerBackend.get_builder_cls() is DeepseekV32IndexerMetadataBuilder |
| 90 | ) |
| 91 | assert indexer._use_flattening(adaptive_config) |
| 92 | assert ( |
| 93 | DeepseekV32IndexerMetadataBuilder.get_cudagraph_support( |
| 94 | adaptive_config, SimpleNamespace() |
| 95 | ) |
| 96 | == AttentionCGSupport.ALWAYS |
| 97 | ) |
| 98 | assert not indexer._use_flattening(fixed_config) |
| 99 | assert ( |
| 100 | DeepseekV32IndexerMetadataBuilder.get_cudagraph_support( |
| 101 | fixed_config, SimpleNamespace() |
| 102 | ) |
| 103 | == AttentionCGSupport.UNIFORM_BATCH |
| 104 | ) |
| 105 | |
| 106 | |
| 107 | @pytest.mark.cpu_test |
| 108 | @pytest.mark.parametrize( |
| 109 | "architecture", |
| 110 | ["DeepseekV32ForCausalLM", "DeepseekV41ForCausalLM", "GlmMoeDsaForCausalLM"], |
| 111 | ) |
| 112 | def test_rocm_adaptive_indexer_flattening_is_scoped_to_deepseek_v4( |
| 113 | monkeypatch: pytest.MonkeyPatch, |
| 114 | architecture: str, |
| 115 | ): |
| 116 | _mock_rocm_platform(monkeypatch) |
| 117 | adaptive_config = _make_indexer_config(architecture=architecture, adaptive=True) |
| 118 | |
| 119 | assert not indexer._rocm_supports_flattened_device_query_lens(adaptive_config) |
| 120 | assert not indexer._use_flattening(adaptive_config) |
| 121 | assert ( |
| 122 | DeepseekV32IndexerMetadataBuilder.get_cudagraph_support( |
| 123 | adaptive_config, SimpleNamespace() |
| 124 | ) |
| 125 | == AttentionCGSupport.UNIFORM_BATCH |
| 126 | ) |
| 127 | |
| 128 | if architecture == "DeepseekV41ForCausalLM": |
| 129 | assert not DeepseekV41IndexerBackend.supports_device_cpu_query_lens_mismatch() |
| 130 | |
| 131 | |
| 132 | @pytest.mark.cpu_test |
| 133 | def test_rocm_adaptive_indexer_preserves_single_request_uniform_path( |
| 134 | monkeypatch: pytest.MonkeyPatch, |
| 135 | ): |
| 136 | builder = _make_indexer_builder(adaptive=True, capacity=8) |
| 137 | |
| 138 | class FakeUniformKernel: |
| 139 | called = False |
| 140 | |
| 141 | def __call__( |
| 142 | self, |
| 143 | seq_lens, |
| 144 | decode_seq_lens, |
| 145 | block_table, |
| 146 | expanded_block_table, |
| 147 | decode_lens, |
| 148 | num_decode_tokens, |
| 149 | max_decode_len, |
| 150 | ): |
| 151 | self.called = True |
| 152 | assert num_decode_tokens == 4 |
| 153 | assert block_table.stride(0) == expanded_block_table.stride(0) == 2 |
| 154 | decode_seq_lens[:max_decode_len] = torch.arange( |
| 155 | seq_lens[0] - max_decode_len + 1, |
| 156 | seq_lens[0] + 1, |
| 157 | dtype=torch.int32, |
| 158 | ) |
| 159 | expanded_block_table[:max_decode_len] = block_table[0] |
| 160 | decode_lens[:max_decode_len] = 1 |
| 161 | |
| 162 | fake_kernel = FakeUniformKernel() |
| 163 | monkeypatch.setattr(indexer, "_PREPARE_UNIFORM_DECODE_KERNEL", fake_kernel) |
| 164 | |
| 165 | seq_lens, block_table, decode_lens, batch_size, requires_padding = ( |
| 166 | builder._prepare_decode_tensors( |
| 167 | seq_lens=torch.tensor([10], dtype=torch.int32), |
| 168 | block_table=torch.tensor([[1, 2]], dtype=torch.int32), |
| 169 | decode_lens=torch.tensor([4], dtype=torch.int32), |
| 170 | decode_lens_cpu=torch.tensor([4], dtype=torch.int32), |
| 171 | query_start_loc=torch.tensor([0], dtype=torch.int32), |
| 172 | num_decodes=1, |
| 173 | num_decode_tokens=4, |
| 174 | use_native=False, |
| 175 | next_n=8, |
| 176 | max_decode_len=4, |
| 177 | ) |
| 178 | ) |
| 179 | |
| 180 | assert fake_kernel.called |
| 181 | torch.testing.assert_close(seq_lens, torch.tensor([7, 8, 9, 10], dtype=torch.int32)) |
| 182 | torch.testing.assert_close( |
| 183 | block_table, |
| 184 | torch.tensor([[1, 2], [1, 2], [1, 2], [1, 2]], dtype=torch.int32), |
| 185 | ) |
| 186 | torch.testing.assert_close(decode_lens, torch.ones(4, dtype=torch.int32)) |
| 187 | assert batch_size == 4 |
| 188 | assert not requires_padding |
| 189 | |
| 190 | |
| 191 | @pytest.mark.cpu_test |
| 192 | def test_rocm_adaptive_indexer_replays_changed_allocations_with_stable_buffers(): |
| 193 | builder = _make_indexer_builder(adaptive=True) |
| 194 | buffer_ptrs = ( |
| 195 | builder.decode_seq_lens_buffer.data_ptr(), |
| 196 | builder.expanded_block_table_buffer.data_ptr(), |
| 197 | builder.decode_lens_buffer.data_ptr(), |
| 198 | ) |
| 199 | block_table = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int32) |
| 200 | decode_lens_cpu = torch.tensor([4, 4, 0], dtype=torch.int32) |
| 201 | |
| 202 | first = builder._prepare_decode_tensors( |
| 203 | seq_lens=torch.tensor([20, 20, 0], dtype=torch.int32), |
| 204 | block_table=block_table, |
| 205 | decode_lens=torch.tensor([7, 1, 0], dtype=torch.int32), |
| 206 | decode_lens_cpu=decode_lens_cpu, |
| 207 | query_start_loc=torch.tensor([0, 7, 8], dtype=torch.int32), |
| 208 | num_decodes=3, |
| 209 | num_decode_tokens=10, |
| 210 | use_native=False, |
| 211 | next_n=8, |
| 212 | max_decode_len=4, |
| 213 | ) |
| 214 | torch.testing.assert_close( |
| 215 | first[0], |
| 216 | torch.tensor([14, 15, 16, 17, 18, 19, 20, 20, 0, 0], dtype=torch.int32), |
| 217 | ) |
| 218 | torch.testing.assert_close( |
| 219 | first[1][:, 0], |
| 220 | torch.tensor([1, 1, 1, 1, 1, 1, 1, 3, 0, 0], dtype=torch.int32), |
| 221 | ) |
| 222 | |
| 223 | second = builder._prepare_decode_tensors( |
| 224 | seq_lens=torch.tensor([20, 20, 0], dtype=torch.int32), |
| 225 | block_table=block_table, |
| 226 | decode_lens=torch.tensor([1, 7, 0], dtype=torch.int32), |
| 227 | decode_lens_cpu=decode_lens_cpu, |
| 228 | query_start_loc=torch.tensor([0, 1, 8], dtype=torch.int32), |
| 229 | num_decodes=3, |
| 230 | num_decode_tokens=10, |
| 231 | use_native=False, |
| 232 | next_n=8, |
| 233 | max_decode_len=4, |
| 234 | ) |
| 235 | torch.testing.assert_close( |
| 236 | second[0], |
| 237 | torch.tensor([20, 14, 15, 16, 17, 18, 19, 20, 0, 0], dtype=torch.int32), |
| 238 | ) |
| 239 | torch.testing.assert_close( |
| 240 | second[1][:, 0], |
| 241 | torch.tensor([1, 3, 3, 3, 3, 3, 3, 3, 0, 0], dtype=torch.int32), |
| 242 | ) |
| 243 | torch.testing.assert_close(second[2], torch.ones(10, dtype=torch.int32)) |
| 244 | assert second[3:] == (10, False) |
| 245 | assert buffer_ptrs == ( |
| 246 | builder.decode_seq_lens_buffer.data_ptr(), |
| 247 | builder.expanded_block_table_buffer.data_ptr(), |
| 248 | builder.decode_lens_buffer.data_ptr(), |
| 249 | ) |