| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests for attention backend selectors.""" |
| 4 | |
| 5 | from types import SimpleNamespace |
| 6 | from unittest.mock import MagicMock, patch |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | |
| 11 | from vllm.platforms import current_platform |
| 12 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 13 | from vllm.v1.attention.selector import AttentionSelectorConfig |
| 14 | |
| 15 | # ROCm-specific attention backend selection tests |
| 16 | pytestmark = pytest.mark.skipif( |
| 17 | not current_platform.is_rocm(), reason="ROCm-specific tests" |
| 18 | ) |
| 19 | |
| 20 | |
| 21 | @pytest.fixture |
| 22 | def mock_vllm_config(): |
| 23 | """Create a mock VllmConfig for testing.""" |
| 24 | config = MagicMock() |
| 25 | config.model_config.dtype = torch.float16 |
| 26 | config.model_config.hf_config.architectures = ["LlamaForCausalLM"] |
| 27 | config.cache_config.block_size = 16 |
| 28 | return config |
| 29 | |
| 30 | |
| 31 | @pytest.fixture |
| 32 | def mock_get_cdna_version(): |
| 33 | """Mock cdna version arch detection to return True.""" |
| 34 | with patch("vllm.platforms.rocm.get_cdna_version", return_value=3): |
| 35 | yield |
| 36 | |
| 37 | |
| 38 | @pytest.fixture |
| 39 | def cleared_attention_selector_cache(): |
| 40 | from vllm.v1.attention.selector import _cached_get_attn_backend |
| 41 | |
| 42 | _cached_get_attn_backend.cache_clear() |
| 43 | yield |
| 44 | _cached_get_attn_backend.cache_clear() |
| 45 | |
| 46 | |
| 47 | def test_aiter_unified_attention_uses_dedicated_metadata_builder(): |
| 48 | from vllm.v1.attention.backends.rocm_aiter_unified_attn import ( |
| 49 | RocmAiterUnifiedAttentionBackend, |
| 50 | RocmAiterUnifiedAttentionMetadataBuilder, |
| 51 | ) |
| 52 | from vllm.v1.attention.backends.rocm_attn import ( |
| 53 | RocmAttentionBackend, |
| 54 | RocmAttentionMetadataBuilder, |
| 55 | ) |
| 56 | |
| 57 | assert RocmAttentionBackend.get_builder_cls() is RocmAttentionMetadataBuilder |
| 58 | assert ( |
| 59 | RocmAiterUnifiedAttentionBackend.get_builder_cls() |
| 60 | is RocmAiterUnifiedAttentionMetadataBuilder |
| 61 | ) |
| 62 | |
| 63 | |
| 64 | def test_aiter_unified_attention_capture_preserves_query_start_locations(): |
| 65 | from vllm.v1.attention.backends.rocm_aiter_unified_attn import ( |
| 66 | RocmAiterUnifiedAttentionMetadataBuilder, |
| 67 | ) |
| 68 | |
| 69 | builder = object.__new__(RocmAiterUnifiedAttentionMetadataBuilder) |
| 70 | metadata = MagicMock() |
| 71 | metadata.seq_lens = torch.tensor([1048576, 524288], dtype=torch.int32) |
| 72 | builder.build = MagicMock(return_value=metadata) |
| 73 | common = MagicMock() |
| 74 | expected_query_start_loc = torch.tensor([0, 2, 5], dtype=torch.int32) |
| 75 | common.query_start_loc = expected_query_start_loc.clone() |
| 76 | metadata.query_start_loc = common.query_start_loc |
| 77 | |
| 78 | actual = builder.build_for_cudagraph_capture(common) |
| 79 | |
| 80 | builder.build.assert_called_once_with(0, common) |
| 81 | assert actual is metadata |
| 82 | assert torch.equal(actual.seq_lens, torch.ones_like(actual.seq_lens)) |
| 83 | assert actual.query_start_loc is common.query_start_loc |
| 84 | assert torch.equal(actual.query_start_loc, expected_query_start_loc) |
| 85 | |
| 86 | |
| 87 | @pytest.mark.parametrize("use_dcp", [False, True]) |
| 88 | @pytest.mark.parametrize( |
| 89 | "env_vars, selected_backend, expected_backend_path", |
| 90 | [ |
| 91 | # Test Case: Explicit FLEX_ATTENTION backend |
| 92 | ( |
| 93 | {}, |
| 94 | "FLEX_ATTENTION", |
| 95 | AttentionBackendEnum.FLEX_ATTENTION.get_path(), |
| 96 | ), |
| 97 | # Test Case 1: Default (no env vars, no explicit backend) |
| 98 | ( |
| 99 | {}, |
| 100 | None, |
| 101 | AttentionBackendEnum.ROCM_ATTN.get_path(), |
| 102 | ), |
| 103 | # Test Case 2: Explicit TRITON_ATTN backend |
| 104 | ( |
| 105 | {}, |
| 106 | "TRITON_ATTN", |
| 107 | AttentionBackendEnum.TRITON_ATTN.get_path(), |
| 108 | ), |
| 109 | # Test Case 3: Explicit ROCM_ATTN backend |
| 110 | ( |
| 111 | {}, |
| 112 | "ROCM_ATTN", |
| 113 | AttentionBackendEnum.ROCM_ATTN.get_path(), |
| 114 | ), |
| 115 | # Test Case 4: Explicit ROCM_AITER_FA backend |
| 116 | ( |
| 117 | {}, |
| 118 | "ROCM_AITER_FA", |
| 119 | AttentionBackendEnum.ROCM_AITER_FA.get_path(), |
| 120 | ), |
| 121 | # Test Case 5: Explicit ROCM_AITER_UNIFIED_ATTN backend |
| 122 | ( |
| 123 | {}, |
| 124 | "ROCM_AITER_UNIFIED_ATTN", |
| 125 | AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path(), |
| 126 | ), |
| 127 | # Test Case 6: VLLM_ROCM_USE_AITER=1 |
| 128 | ( |
| 129 | {"VLLM_ROCM_USE_AITER": "1"}, |
| 130 | None, |
| 131 | AttentionBackendEnum.ROCM_ATTN.get_path(), |
| 132 | ), |
| 133 | # Test Case 7: VLLM_ROCM_USE_AITER=1 + explicit TRITON_ATTN |
| 134 | ( |
| 135 | {"VLLM_ROCM_USE_AITER": "1"}, |
| 136 | "TRITON_ATTN", |
| 137 | AttentionBackendEnum.TRITON_ATTN.get_path(), |
| 138 | ), |
| 139 | # Test Case 8: VLLM_ROCM_USE_AITER=1 + VLLM_ROCM_USE_AITER_MHA=0 |
| 140 | ( |
| 141 | {"VLLM_ROCM_USE_AITER": "1", "VLLM_ROCM_USE_AITER_MHA": "0"}, |
| 142 | None, |
| 143 | AttentionBackendEnum.ROCM_ATTN.get_path(), |
| 144 | ), |
| 145 | # Test Case 9: VLLM_ROCM_USE_AITER=1 + explicit ROCM_ATTN |
| 146 | ( |
| 147 | {"VLLM_ROCM_USE_AITER": "1"}, |
| 148 | "ROCM_ATTN", |
| 149 | AttentionBackendEnum.ROCM_ATTN.get_path(), |
| 150 | ), |
| 151 | ], |
| 152 | ) |
| 153 | def test_standard_attention_backend_selection( |
| 154 | env_vars, |
| 155 | selected_backend, |
| 156 | expected_backend_path, |
| 157 | use_dcp, |
| 158 | mock_vllm_config, |
| 159 | mock_get_cdna_version, |
| 160 | monkeypatch, |
| 161 | ): |
| 162 | """Standard ROCm backends remain selectable without DCP and reject DCP.""" |
| 163 | # Set environment variables |
| 164 | for key, value in env_vars.items(): |
| 165 | monkeypatch.setenv(key, value) |
| 166 | |
| 167 | # Import after setting env vars to ensure they're picked up |
| 168 | # Reload envs to pick up new environment variables |
| 169 | import importlib |
| 170 | |
| 171 | import vllm.envs as envs |
| 172 | |
| 173 | importlib.reload(envs) |
| 174 | |
| 175 | # Convert string backend to enum if provided |
| 176 | backend_enum = None |
| 177 | if selected_backend: |
| 178 | backend_enum = getattr(AttentionBackendEnum, selected_backend) |
| 179 | |
| 180 | # Get the backend class path |
| 181 | from vllm.platforms.rocm import RocmPlatform |
| 182 | |
| 183 | attn_selector_config = AttentionSelectorConfig( |
| 184 | head_size=128, |
| 185 | dtype=torch.float16, |
| 186 | kv_cache_dtype="auto", |
| 187 | block_size=16, |
| 188 | use_mla=False, |
| 189 | has_sink=False, |
| 190 | use_sparse=False, |
| 191 | use_dcp=use_dcp, |
| 192 | ) |
| 193 | |
| 194 | if use_dcp: |
| 195 | with pytest.raises(ValueError, match="DCP not supported"): |
| 196 | RocmPlatform.get_attn_backend_cls(backend_enum, attn_selector_config) |
| 197 | return |
| 198 | |
| 199 | backend_path = RocmPlatform.get_attn_backend_cls( |
| 200 | selected_backend=backend_enum, attn_selector_config=attn_selector_config |
| 201 | ) |
| 202 | |
| 203 | assert backend_path == expected_backend_path |
| 204 | |
| 205 | |
| 206 | @pytest.mark.parametrize("use_dcp", [False, True]) |
| 207 | @pytest.mark.parametrize( |
| 208 | "env_vars, selected_backend, block_size, expected_backend_path, should_raise", |
| 209 | [ |
| 210 | # Test Case 1: TRITON_MLA with block_size != 1 |
| 211 | ( |
| 212 | {}, |
| 213 | "TRITON_MLA", |
| 214 | 16, |
| 215 | AttentionBackendEnum.TRITON_MLA.get_path(), |
| 216 | False, |
| 217 | ), |
| 218 | # Test Case 2: TRITON_MLA with block_size == 1 (should raise) |
| 219 | ( |
| 220 | {}, |
| 221 | "TRITON_MLA", |
| 222 | 1, |
| 223 | None, |
| 224 | True, |
| 225 | ), |
| 226 | # Test Case 3: ROCM_AITER_MLA with block_size == 1 |
| 227 | ( |
| 228 | {}, |
| 229 | "ROCM_AITER_MLA", |
| 230 | 1, |
| 231 | AttentionBackendEnum.ROCM_AITER_MLA.get_path(), |
| 232 | False, |
| 233 | ), |
| 234 | # Test Case 4: ROCM_AITER_MLA with block_size != 1 (should raise) |
| 235 | ( |
| 236 | {}, |
| 237 | "ROCM_AITER_MLA", |
| 238 | 16, |
| 239 | AttentionBackendEnum.ROCM_AITER_MLA.get_path(), |
| 240 | False, |
| 241 | ), |
| 242 | # Test Case 5: VLLM_ROCM_USE_AITER=1 with block_size == 1 |
| 243 | ( |
| 244 | {"VLLM_ROCM_USE_AITER": "1"}, |
| 245 | None, |
| 246 | 1, |
| 247 | AttentionBackendEnum.ROCM_AITER_MLA.get_path(), |
| 248 | False, |
| 249 | ), |
| 250 | # Test Case 6: VLLM_ROCM_USE_AITER=1 with block_size == 16 |
| 251 | # (should use ROCM_AITER_MLA now, as it supports block_size 16) |
| 252 | ( |
| 253 | {"VLLM_ROCM_USE_AITER": "1"}, |
| 254 | None, |
| 255 | 16, |
| 256 | AttentionBackendEnum.ROCM_AITER_MLA.get_path(), |
| 257 | False, |
| 258 | ), |
| 259 | # Test Case 7: VLLM_ROCM_USE_AITER=1 + explicit TRITON_MLA |
| 260 | ( |
| 261 | {"VLLM_ROCM_USE_AITER": "1"}, |
| 262 | "TRITON_MLA", |
| 263 | 16, |
| 264 | AttentionBackendEnum.TRITON_MLA.get_path(), |
| 265 | False, |
| 266 | ), |
| 267 | # Test Case 8: Explicit ROCM_AITER_TRITON_MLA |
| 268 | ( |
| 269 | {}, |
| 270 | "ROCM_AITER_TRITON_MLA", |
| 271 | 16, |
| 272 | AttentionBackendEnum.ROCM_AITER_TRITON_MLA.get_path(), |
| 273 | False, |
| 274 | ), |
| 275 | ], |
| 276 | ) |
| 277 | def test_mla_backend_selection( |
| 278 | env_vars, |
| 279 | selected_backend, |
| 280 | block_size, |
| 281 | expected_backend_path, |
| 282 | should_raise, |
| 283 | use_dcp, |
| 284 | mock_vllm_config, |
| 285 | monkeypatch, |
| 286 | ): |
| 287 | """Dense MLA remains selectable with DCP for valid block sizes.""" |
| 288 | # Set environment variables |
| 289 | for key, value in env_vars.items(): |
| 290 | monkeypatch.setenv(key, value) |
| 291 | |
| 292 | # Import after setting env vars |
| 293 | # Reload envs |
| 294 | import importlib |
| 295 | |
| 296 | import vllm.envs as envs |
| 297 | |
| 298 | importlib.reload(envs) |
| 299 | |
| 300 | # Mock is_aiter_mla_enabled based on env vars and block_size |
| 301 | aiter_enabled = env_vars.get("VLLM_ROCM_USE_AITER") == "1" |
| 302 | |
| 303 | mock_rocm_ops = MagicMock() |
| 304 | mock_rocm_ops.is_mla_enabled.return_value = aiter_enabled |
| 305 | mock_aiter_module = MagicMock() |
| 306 | mock_aiter_module.rocm_aiter_ops = mock_rocm_ops |
| 307 | |
| 308 | with patch.dict("sys.modules", {"vllm._aiter_ops": mock_aiter_module}): |
| 309 | # Convert string backend to enum if provided |
| 310 | backend_enum = None |
| 311 | if selected_backend: |
| 312 | backend_enum = getattr(AttentionBackendEnum, selected_backend) |
| 313 | |
| 314 | from vllm.platforms.rocm import RocmPlatform |
| 315 | |
| 316 | if should_raise: |
| 317 | with pytest.raises(ValueError): |
| 318 | attn_selector_config = AttentionSelectorConfig( |
| 319 | head_size=128, |
| 320 | dtype=torch.float16, |
| 321 | kv_cache_dtype="auto", |
| 322 | block_size=block_size, |
| 323 | use_mla=True, |
| 324 | has_sink=False, |
| 325 | use_sparse=False, |
| 326 | use_dcp=use_dcp, |
| 327 | ) |
| 328 | attn_selector_config = AttentionSelectorConfig( |
| 329 | head_size=128, |
| 330 | dtype=torch.float16, |
| 331 | kv_cache_dtype="auto", |
| 332 | block_size=block_size, |
| 333 | use_mla=True, |
| 334 | has_sink=False, |
| 335 | use_sparse=False, |
| 336 | use_dcp=use_dcp, |
| 337 | ) |
| 338 | backend_path = RocmPlatform.get_attn_backend_cls( |
| 339 | selected_backend=backend_enum, |
| 340 | attn_selector_config=attn_selector_config, |
| 341 | ) |
| 342 | |
| 343 | else: |
| 344 | attn_selector_config = AttentionSelectorConfig( |
| 345 | head_size=128, |
| 346 | dtype=torch.float16, |
| 347 | kv_cache_dtype="auto", |
| 348 | block_size=block_size, |
| 349 | use_mla=True, |
| 350 | has_sink=False, |
| 351 | use_sparse=False, |
| 352 | use_dcp=use_dcp, |
| 353 | ) |
| 354 | |
| 355 | backend_path = RocmPlatform.get_attn_backend_cls( |
| 356 | selected_backend=backend_enum, attn_selector_config=attn_selector_config |
| 357 | ) |
| 358 | |
| 359 | assert backend_path == expected_backend_path |
| 360 | |
| 361 | |
| 362 | @pytest.mark.parametrize("use_dcp", [False, True]) |
| 363 | @pytest.mark.parametrize( |
| 364 | "selected_backend", [None, AttentionBackendEnum.ROCM_AITER_MLA_SPARSE] |
| 365 | ) |
| 366 | def test_sparse_mla_backend_rejects_dcp(selected_backend, use_dcp): |
| 367 | """Sparse MLA remains selectable without DCP and fails early with DCP.""" |
| 368 | from vllm.platforms.rocm import RocmPlatform |
| 369 | |
| 370 | selector_config = AttentionSelectorConfig( |
| 371 | head_size=576, |
| 372 | dtype=torch.bfloat16, |
| 373 | kv_cache_dtype="auto", |
| 374 | block_size=16, |
| 375 | use_mla=True, |
| 376 | use_sparse=True, |
| 377 | use_dcp=use_dcp, |
| 378 | ) |
| 379 | if use_dcp: |
| 380 | with pytest.raises(ValueError, match="DCP not supported"): |
| 381 | RocmPlatform.get_attn_backend_cls(selected_backend, selector_config) |
| 382 | else: |
| 383 | assert RocmPlatform.get_attn_backend_cls(selected_backend, selector_config) == ( |
| 384 | AttentionBackendEnum.ROCM_AITER_MLA_SPARSE.get_path() |
| 385 | ) |
| 386 | |
| 387 | |
| 388 | @pytest.mark.parametrize("use_dcp", [False, True]) |
| 389 | @pytest.mark.parametrize( |
| 390 | "selected_backend, head_size, kv_cache_dtype", |
| 391 | [ |
| 392 | (AttentionBackendEnum.TRITON_ATTN_DIFFKV, 192, "bfloat16"), |
| 393 | # A 128-dimensional head uses 118 packed bytes, or 59 fp16 elements. |
| 394 | (AttentionBackendEnum.TURBOQUANT, 59, "turboquant_k3v4_nc"), |
| 395 | ], |
| 396 | ) |
| 397 | def test_specialized_attention_backends_reject_dcp( |
| 398 | selected_backend, head_size, kv_cache_dtype, use_dcp |
| 399 | ): |
| 400 | """Valid DiffKV and compressed-cache configurations must reject DCP.""" |
| 401 | from vllm.platforms.rocm import RocmPlatform |
| 402 | |
| 403 | selector_config = AttentionSelectorConfig( |
| 404 | head_size=head_size, |
| 405 | dtype=torch.bfloat16, |
| 406 | kv_cache_dtype=kv_cache_dtype, |
| 407 | block_size=16, |
| 408 | use_dcp=use_dcp, |
| 409 | ) |
| 410 | if use_dcp: |
| 411 | with pytest.raises(ValueError, match="DCP not supported"): |
| 412 | RocmPlatform.get_attn_backend_cls(selected_backend, selector_config) |
| 413 | else: |
| 414 | assert RocmPlatform.get_attn_backend_cls(selected_backend, selector_config) == ( |
| 415 | selected_backend.get_path() |
| 416 | ) |
| 417 | |
| 418 | |
| 419 | def test_aiter_fa_requires_mi3xx(mock_vllm_config): |
| 420 | """Test that ROCM_AITER_FA requires CDNA3+ architecture.""" |
| 421 | from vllm.platforms.rocm import RocmPlatform |
| 422 | |
| 423 | # Mock cdna version to return 1 (used by supports_compute_capability) |
| 424 | with ( |
| 425 | patch("vllm.platforms.rocm.get_cdna_version", return_value=1), |
| 426 | pytest.raises( |
| 427 | ValueError, |
| 428 | match="compute capability not supported", |
| 429 | ), |
| 430 | ): |
| 431 | attn_selector_config = AttentionSelectorConfig( |
| 432 | head_size=128, |
| 433 | dtype=torch.float16, |
| 434 | kv_cache_dtype="auto", |
| 435 | block_size=16, |
| 436 | use_mla=False, |
| 437 | has_sink=False, |
| 438 | use_sparse=False, |
| 439 | ) |
| 440 | |
| 441 | RocmPlatform.get_attn_backend_cls( |
| 442 | selected_backend=AttentionBackendEnum.ROCM_AITER_FA, |
| 443 | attn_selector_config=attn_selector_config, |
| 444 | ) |
| 445 | |
| 446 | |
| 447 | @pytest.fixture |
| 448 | def turboquant_run_config(): |
| 449 | """Current config of a run whose KV cache dtype is a turboquant_* preset.""" |
| 450 | config = SimpleNamespace( |
| 451 | cache_config=SimpleNamespace(cache_dtype="turboquant_k8v4") |
| 452 | ) |
| 453 | with patch( |
| 454 | "vllm.config.get_current_vllm_config_or_none", |
| 455 | return_value=config, |
| 456 | ): |
| 457 | yield config |
| 458 | |
| 459 | |
| 460 | def test_turboquant_boundary_selection_is_not_cached_from_ordinary_run( |
| 461 | cleared_attention_selector_cache, |
| 462 | ): |
| 463 | from vllm.config import CacheConfig, VllmConfig, set_current_vllm_config |
| 464 | from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts |
| 465 | from vllm.v1.attention.selector import get_attn_backend |
| 466 | from vllm.v1.kv_cache_layout import KVCacheLayout |
| 467 | |
| 468 | ordinary_config = VllmConfig(cache_config=CacheConfig(cache_dtype="auto")) |
| 469 | turboquant_config = VllmConfig( |
| 470 | cache_config=CacheConfig(cache_dtype="turboquant_k8v4") |
| 471 | ) |
| 472 | |
| 473 | backend_priorities = [ |
| 474 | AttentionBackendEnum.ROCM_ATTN, |
| 475 | AttentionBackendEnum.TRITON_ATTN, |
| 476 | AttentionBackendEnum.TURBOQUANT, |
| 477 | ] |
| 478 | with patch( |
| 479 | "vllm.platforms.rocm._get_backend_priorities", |
| 480 | return_value=backend_priorities, |
| 481 | ): |
| 482 | with set_current_vllm_config(ordinary_config): |
| 483 | ordinary_backend = get_attn_backend(128, torch.float16, "auto") |
| 484 | |
| 485 | with set_current_vllm_config(turboquant_config): |
| 486 | boundary_backend = get_attn_backend(128, torch.float16, "auto") |
| 487 | quantized_backend = get_attn_backend(128, torch.float16, "turboquant_k8v4") |
| 488 | |
| 489 | assert ordinary_backend is AttentionBackendEnum.ROCM_ATTN.get_class() |
| 490 | assert quantized_backend is AttentionBackendEnum.TURBOQUANT.get_class() |
| 491 | assert get_supported_kv_cache_layouts([boundary_backend, quantized_backend]) == [ |
| 492 | KVCacheLayout.LBNHC |
| 493 | ] |
| 494 | |
| 495 | |
| 496 | def test_turboquant_run_does_not_mask_unsupported_selected_backend( |
| 497 | turboquant_run_config, |
| 498 | ): |
| 499 | from vllm.platforms.rocm import RocmPlatform |
| 500 | |
| 501 | attn_selector_config = AttentionSelectorConfig( |
| 502 | head_size=128, |
| 503 | dtype=torch.float16, |
| 504 | kv_cache_dtype="auto", |
| 505 | block_size=16, |
| 506 | ) |
| 507 | |
| 508 | with ( |
| 509 | patch("vllm.platforms.rocm.get_cdna_version", return_value=1), |
| 510 | pytest.raises(ValueError, match="compute capability not supported"), |
| 511 | ): |
| 512 | RocmPlatform.get_attn_backend_cls( |
| 513 | selected_backend=AttentionBackendEnum.ROCM_AITER_FA, |
| 514 | attn_selector_config=attn_selector_config, |
| 515 | ) |
| 516 | |
| 517 | |
| 518 | def test_turboquant_layout_check_respects_backend_override(): |
| 519 | from vllm.platforms.rocm import _shares_layout_with_turboquant |
| 520 | from vllm.v1.kv_cache_layout import KVCacheLayout |
| 521 | |
| 522 | boundary_backend = MagicMock() |
| 523 | boundary_backend.supported_kv_cache_layouts.return_value = (KVCacheLayout.LBHNC,) |
| 524 | turboquant_override = MagicMock() |
| 525 | turboquant_override.supported_kv_cache_layouts.return_value = (KVCacheLayout.LBHNC,) |
| 526 | |
| 527 | with patch.object( |
| 528 | AttentionBackendEnum.TURBOQUANT, |
| 529 | "get_class", |
| 530 | return_value=turboquant_override, |
| 531 | ) as get_turboquant_class: |
| 532 | assert _shares_layout_with_turboquant(boundary_backend) |
| 533 | get_turboquant_class.assert_called_once_with() |
| 534 | |
| 535 | |
| 536 | @pytest.mark.parametrize( |
| 537 | "selected_backend", |
| 538 | [None, "ROCM_ATTN", "ROCM_AITER_FA", "TURBOQUANT"], |
| 539 | ) |
| 540 | def test_turboquant_boundary_layers_share_a_layout( |
| 541 | selected_backend, |
| 542 | turboquant_run_config, |
| 543 | mock_get_cdna_version, |
| 544 | ): |
| 545 | """A turboquant_* run must resolve to backends with a common KV layout. |
| 546 | |
| 547 | The boundary layers keep the native dtype and pick their own backend while |
| 548 | every other layer picks TURBOQUANT. Engine startup hard-errors when the two |
| 549 | share no layout, so the boundary layers must not land on a native ROCm or |
| 550 | AITER backend, whichever backend the run asked for. |
| 551 | """ |
| 552 | from vllm.platforms.rocm import RocmPlatform |
| 553 | from vllm.utils.import_utils import resolve_obj_by_qualname |
| 554 | from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts |
| 555 | |
| 556 | backend_enum = ( |
| 557 | getattr(AttentionBackendEnum, selected_backend) if selected_backend else None |
| 558 | ) |
| 559 | |
| 560 | def resolve(kv_cache_dtype): |
| 561 | path = RocmPlatform.get_attn_backend_cls( |
| 562 | selected_backend=backend_enum, |
| 563 | attn_selector_config=AttentionSelectorConfig( |
| 564 | head_size=128, |
| 565 | dtype=torch.float16, |
| 566 | kv_cache_dtype=kv_cache_dtype, |
| 567 | block_size=16, |
| 568 | ), |
| 569 | ) |
| 570 | return resolve_obj_by_qualname(path) |
| 571 | |
| 572 | boundary_backend = resolve("auto") |
| 573 | quantized_backend = resolve("turboquant_k8v4") |
| 574 | |
| 575 | assert quantized_backend is AttentionBackendEnum.TURBOQUANT.get_class() |
| 576 | # Raises when the intersection is empty, which is the startup failure. |
| 577 | assert get_supported_kv_cache_layouts([boundary_backend, quantized_backend]) |
| 578 | |
| 579 | |
| 580 | def test_sparse_not_supported(mock_vllm_config): |
| 581 | """Test that sparse MLA without use_mla flag raises an error.""" |
| 582 | from vllm.platforms.rocm import RocmPlatform |
| 583 | |
| 584 | with pytest.raises( |
| 585 | ValueError, |
| 586 | match="No valid attention backend found", |
| 587 | ): |
| 588 | attn_selector_config = AttentionSelectorConfig( |
| 589 | head_size=128, |
| 590 | dtype=torch.float16, |
| 591 | kv_cache_dtype="auto", |
| 592 | block_size=16, |
| 593 | use_mla=False, |
| 594 | has_sink=False, |
| 595 | use_sparse=True, |
| 596 | ) |
| 597 | |
| 598 | RocmPlatform.get_attn_backend_cls( |
| 599 | selected_backend=None, attn_selector_config=attn_selector_config |
| 600 | ) |
| 601 | |
| 602 | |
| 603 | def _kv_connector_selector_config() -> AttentionSelectorConfig: |
| 604 | return AttentionSelectorConfig( |
| 605 | head_size=128, |
| 606 | dtype=torch.float16, |
| 607 | kv_cache_dtype="auto", |
| 608 | block_size=16, |
| 609 | use_mla=False, |
| 610 | has_sink=False, |
| 611 | use_sparse=False, |
| 612 | use_kv_connector=True, |
| 613 | ) |
| 614 | |
| 615 | |
| 616 | def test_unified_attn_declares_kv_connector_support(): |
| 617 | """ROCM_AITER_UNIFIED_ATTN opts into KV connectors and ROCM_ATTN does not.""" |
| 618 | from vllm.v1.attention.backends.rocm_aiter_unified_attn import ( |
| 619 | RocmAiterUnifiedAttentionBackend, |
| 620 | ) |
| 621 | from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend |
| 622 | |
| 623 | assert RocmAiterUnifiedAttentionBackend.supports_kv_connector() is True |
| 624 | assert RocmAttentionBackend.supports_kv_connector() is False |
| 625 | |
| 626 | |
| 627 | def test_unified_attn_supports_kv_connector(mock_vllm_config, mock_get_cdna_version): |
| 628 | """ROCM_AITER_UNIFIED_ATTN can be selected with KV connectors.""" |
| 629 | from vllm.platforms.rocm import RocmPlatform |
| 630 | |
| 631 | backend_path = RocmPlatform.get_attn_backend_cls( |
| 632 | selected_backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN, |
| 633 | attn_selector_config=_kv_connector_selector_config(), |
| 634 | ) |
| 635 | |
| 636 | assert backend_path == AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path() |
| 637 | |
| 638 | |
| 639 | def test_rocm_attn_rejects_kv_connector(mock_vllm_config, mock_get_cdna_version): |
| 640 | """Selecting ROCM_ATTN with a KV connector is illegal.""" |
| 641 | from vllm.platforms.rocm import RocmPlatform |
| 642 | |
| 643 | attn_selector_config = _kv_connector_selector_config() |
| 644 | |
| 645 | with pytest.raises(ValueError, match="KV connector not supported"): |
| 646 | RocmPlatform.get_attn_backend_cls( |
| 647 | selected_backend=AttentionBackendEnum.ROCM_ATTN, |
| 648 | attn_selector_config=attn_selector_config, |
| 649 | ) |
| 650 | |
| 651 | |
| 652 | @pytest.mark.parametrize( |
| 653 | "aiter_found, expected_backend", |
| 654 | [ |
| 655 | (True, AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN), |
| 656 | (False, AttentionBackendEnum.TRITON_ATTN), |
| 657 | ], |
| 658 | ) |
| 659 | def test_auto_selection_for_kv_connector( |
| 660 | aiter_found, expected_backend, mock_vllm_config, mock_get_cdna_version |
| 661 | ): |
| 662 | """Auto-selection with a KV connector and AITER enabled resolves to unified attn, |
| 663 | and to triton attn if AITER not enabled.""" |
| 664 | from vllm.platforms.rocm import RocmPlatform |
| 665 | |
| 666 | with patch( |
| 667 | "vllm._aiter_ops.is_aiter_found_and_supported", return_value=aiter_found |
| 668 | ): |
| 669 | backend_path = RocmPlatform.get_attn_backend_cls( |
| 670 | selected_backend=None, |
| 671 | attn_selector_config=_kv_connector_selector_config(), |
| 672 | ) |
| 673 | |
| 674 | assert backend_path == expected_backend.get_path() |
| 675 | |
| 676 | |
| 677 | def test_unified_attn_prefers_block_contiguous_layout(): |
| 678 | """Unified attn prefers a block-first KV layout, hence ok with kv connectors.""" |
| 679 | from vllm.v1.attention.backends.rocm_aiter_unified_attn import ( |
| 680 | RocmAiterUnifiedAttentionBackend, |
| 681 | ) |
| 682 | from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend |
| 683 | |
| 684 | unified_preferred = RocmAiterUnifiedAttentionBackend.supported_kv_cache_layouts()[0] |
| 685 | rocm_attn_preferred = RocmAttentionBackend.supported_kv_cache_layouts()[0] |
| 686 | |
| 687 | assert unified_preferred.is_block_contiguous is True |
| 688 | assert rocm_attn_preferred.is_block_contiguous is False |
| 689 | |
| 690 | |
| 691 | def test_unified_attn_drops_lhbnc_with_kv_connector(): |
| 692 | """Connectors move a block as one contiguous byte range, which LHBNC breaks.""" |
| 693 | from vllm.config import KVTransferConfig, VllmConfig, set_current_vllm_config |
| 694 | from vllm.v1.attention.backends.rocm_aiter_unified_attn import ( |
| 695 | RocmAiterUnifiedAttentionBackend, |
| 696 | ) |
| 697 | from vllm.v1.kv_cache_interface import KVCacheLayout |
| 698 | |
| 699 | assert KVCacheLayout.LHBNC in ( |
| 700 | RocmAiterUnifiedAttentionBackend.supported_kv_cache_layouts() |
| 701 | ) |
| 702 | |
| 703 | config = VllmConfig( |
| 704 | kv_transfer_config=KVTransferConfig( |
| 705 | kv_connector="ExampleConnector", kv_role="kv_both" |
| 706 | ) |
| 707 | ) |
| 708 | with set_current_vllm_config(config): |
| 709 | layouts = RocmAiterUnifiedAttentionBackend.supported_kv_cache_layouts() |
| 710 | |
| 711 | assert layouts |
| 712 | assert all(layout.is_block_compact for layout in layouts) |