@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e78e367c6e

[ROCm] Give turboquant boundary layers a layout-compatible backend (#54988)

Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: Codex <noreply@openai.com>

stefankoncarevic committed Sep 24, 2026 at 05:00 UTC e78e367c6e22740fac4b027fc6431756bb014ec6
3 files changed +232 -16
tests/v1/attention/test_rocm_attention_backends_selection.py
+143
@@ -2,6 +2,7 @@
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
@@ -34,6 +35,15 @@ def mock_get_cdna_version():
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,
@@ -434,6 +444,139 @@ def test_aiter_fa_requires_mi3xx(mock_vllm_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
vllm/platforms/rocm.py
+82 -16
@@ -23,10 +23,14 @@ if TYPE_CHECKING:
23 from vllm.config import VllmConfig
24 from vllm.config.kernel import IrOpPriorityConfig
25 from vllm.utils.argparse_utils import FlexibleArgumentParser
26 + from vllm.v1.attention.backend import AttentionBackend
27 from vllm.v1.attention.selector import AttentionSelectorConfig
28
29 logger = init_logger(__name__)
30
31 +_KV_CACHE_DTYPE_REASON = "kv_cache_dtype not supported"
32 +_TURBOQUANT_LAYOUT_REASON = "no KV cache layout in common with TURBOQUANT"
33 +
34 try:
35 from amdsmi import (
36 AmdSmiException,
@@ -505,6 +509,53 @@ def _get_backend_priorities(
509 return backends
510
511
512 +def _uses_turboquant(vllm_config: "VllmConfig | None") -> bool:
513 + """Whether the run's KV cache dtype is one of the turboquant_* presets."""
514 + cache_config = getattr(vllm_config, "cache_config", None)
515 + return cache_config is not None and str(cache_config.cache_dtype).startswith(
516 + "turboquant_"
517 + )
518 +
519 +
520 +def _shares_layout_with_turboquant(backend_class: type["AttentionBackend"]) -> bool:
521 + """Whether this backend reads a KV cache layout TURBOQUANT reads too.
522 +
523 + A turboquant_* run keeps its boundary layers at the native dtype (see
524 + TurboQuantConfig.get_boundary_skip_layers), so those layers pick a backend of
525 + their own while every other layer picks TURBOQUANT. One layout has to serve
526 + the whole worker, and get_supported_kv_cache_layouts() turns an empty
527 + intersection into a hard error at engine startup.
528 + """
529 + layouts = backend_class.supported_kv_cache_layouts()
530 + turboquant_layouts = (
531 + AttentionBackendEnum.TURBOQUANT.get_class().supported_kv_cache_layouts()
532 + )
533 + if layouts is None or turboquant_layouts is None:
534 + return True
535 + return not set(layouts).isdisjoint(turboquant_layouts)
536 +
537 +
538 +def _get_invalid_reasons(
539 + backend_class: type["AttentionBackend"],
540 + device_capability: DeviceCapability,
541 + attn_selector_config: "AttentionSelectorConfig",
542 + *,
543 + is_turboquant_run: bool,
544 +) -> list[str]:
545 + """Why this backend cannot serve the layer, empty when it can."""
546 + invalid_reasons = backend_class.validate_configuration(
547 + device_capability=device_capability,
548 + **attn_selector_config._asdict(),
549 + )
550 + if (
551 + not invalid_reasons
552 + and is_turboquant_run
553 + and not _shares_layout_with_turboquant(backend_class)
554 + ):
555 + invalid_reasons = [_TURBOQUANT_LAYOUT_REASON]
556 + return invalid_reasons
557 +
558 +
559 class RocmPlatform(Platform):
560 _enum = PlatformEnum.ROCM
561 device_name: str = "rocm"
@@ -611,12 +662,14 @@ class RocmPlatform(Platform):
662 # TODO: Make this explicit in the selector in a future PR.
663 if is_encoder_decoder and AttentionBackendEnum.ROCM_ATTN in backend_priorities:
664 backend_priorities.remove(AttentionBackendEnum.ROCM_ATTN)
665 + is_turboquant_run = _uses_turboquant(vllm_config)
666 for priority, backend in enumerate(backend_priorities):
667 try:
616 - backend_class = backend.get_class()
617 - invalid_reasons_i = backend_class.validate_configuration(
618 - device_capability=device_capability,
619 - **attn_selector_config._asdict(),
668 + invalid_reasons_i = _get_invalid_reasons(
669 + backend.get_class(),
670 + device_capability,
671 + attn_selector_config,
672 + is_turboquant_run=is_turboquant_run,
673 )
674 except ImportError:
675 invalid_reasons_i = ["ImportError"]
@@ -639,11 +692,16 @@ class RocmPlatform(Platform):
692
693 # First try checking just the selected backend, if there is one.
694 if selected_backend is not None:
695 + # Keep lazy: vllm.config imports current_platform during initialization.
696 + from vllm.config import get_current_vllm_config_or_none
697 +
698 + is_turboquant_run = _uses_turboquant(get_current_vllm_config_or_none())
699 try:
643 - backend_class = selected_backend.get_class()
644 - sel_invalid_reasons = backend_class.validate_configuration(
645 - device_capability=device_capability,
646 - **attn_selector_config._asdict(),
700 + sel_invalid_reasons = _get_invalid_reasons(
701 + selected_backend.get_class(),
702 + device_capability,
703 + attn_selector_config,
704 + is_turboquant_run=is_turboquant_run,
705 )
706 except ImportError:
707 sel_invalid_reasons = ["ImportError"]
@@ -653,21 +711,29 @@ class RocmPlatform(Platform):
711 selected_backend.name,
712 )
713 return selected_backend.get_path()
656 - # Only tolerate the mismatch for turboquant_* KV-cache layers:
657 - # boundary layers keep the native dtype (served by the selected
658 - # backend) while turboquant_* layers need TURBOQUANT, so no single
659 - # --attention-backend can serve every layer. For any other dtype
660 - # the explicit selection is genuinely invalid -> fail loud.
714 + # Only tolerate the mismatch when turboquant is in play: boundary
715 + # layers keep the native dtype while every other layer needs
716 + # TURBOQUANT, so no single --attention-backend can serve every layer.
717 + # For any other dtype the selection is genuinely invalid -> fail loud.
718 kv_dtype = attn_selector_config.kv_cache_dtype
662 - if not (kv_dtype is not None and str(kv_dtype).startswith("turboquant")):
719 + layer_is_turboquant = kv_dtype is not None and str(kv_dtype).startswith(
720 + "turboquant"
721 + )
722 + is_turboquant_fallback = (
723 + is_turboquant_run or layer_is_turboquant
724 + ) and sel_invalid_reasons in (
725 + [_KV_CACHE_DTYPE_REASON],
726 + [_TURBOQUANT_LAYOUT_REASON],
727 + )
728 + if not is_turboquant_fallback:
729 raise ValueError(
730 f"Selected backend {selected_backend} is not valid for "
731 f"this configuration. Reason: {sel_invalid_reasons}"
732 )
733 # NOTE: pass a str (not the list) -- info_once hashes its args.
734 logger.info_once(
669 - "Selected backend %s is incompatible with this turboquant "
670 - "layer (%s); using the auto-selected per-layer backend. "
735 + "Selected backend %s is incompatible with this layer (%s) of "
736 + "the turboquant run; using the auto-selected per-layer backend. "
737 "Reason: %s",
738 selected_backend.name,
739 attn_selector_config.attn_type,
vllm/v1/attention/selector.py
+7
@@ -194,6 +194,9 @@ def get_attn_backend(
194 backend=backend,
195 attn_selector_config=attn_selector_config,
196 num_heads=num_heads,
197 + _run_kv_cache_dtype=(
198 + cache_config.cache_dtype if cache_config is not None else None
199 + ),
200 )
201
202
@@ -202,7 +205,11 @@ def _cached_get_attn_backend(
205 backend,
206 attn_selector_config: AttentionSelectorConfig,
207 num_heads: int | None = None,
208 + *,
209 + _run_kv_cache_dtype: CacheDType | None = None,
210 ) -> type[AttentionBackend]:
211 + # Some platform selectors inspect the run-wide cache dtype. Keep it in the
212 + # cache key even though they read the value from the current config.
213 from vllm.platforms import current_platform
214
215 attention_cls = current_platform.get_attn_backend_cls(