[ROCm] Rank TRITON_ATTN above unified attention on gfx950 until the aiter pin has the #62606 fix

tests/model_executor/layers/test_fused_shared_expert.py

@@ -875,6 +875,7 @@ def test_models_fse_init(
         quantization_config=None,
         runner_type="generate",
         is_moe=True,
+        is_hybrid=False,
         logits_processors=None,
         rswa_window=None,
     )

tests/v1/attention/test_rocm_attention_backends_selection.py

@@ -731,7 +731,9 @@ def test_aiter_fa_requires_mi3xx(mock_vllm_config):
 def turboquant_run_config():
     """Current config of a run whose KV cache dtype is a turboquant_* preset."""
     config = SimpleNamespace(
-        cache_config=SimpleNamespace(cache_dtype="turboquant_k8v4")
+        model_config=None,
+        speculative_config=None,
+        cache_config=SimpleNamespace(cache_dtype="turboquant_k8v4"),
     )
     with patch(
         "vllm.config.get_current_vllm_config_or_none",
@@ -993,3 +995,81 @@ def test_unified_attn_drops_lhbnc_with_kv_connector():
 
     assert layouts
     assert all(layout.is_block_compact for layout in layouts)
+
+
+@pytest.mark.parametrize(
+    "speculative, hybrid, aiter_found, expected_backend",
+    [
+        (False, False, True, AttentionBackendEnum.ROCM_ATTN),
+        (True, False, True, AttentionBackendEnum.TRITON_ATTN),
+        (True, False, False, AttentionBackendEnum.TRITON_ATTN),
+        (False, True, True, AttentionBackendEnum.TRITON_ATTN),
+        (False, True, False, AttentionBackendEnum.TRITON_ATTN),
+    ],
+)
+def test_auto_selection_avoids_rocm_attn_triton_fallbacks(
+    speculative, hybrid, aiter_found, expected_backend, mock_get_cdna_version
+):
+    """Speculative decoding and hybrid models would run ROCM_ATTN's Triton
+    fallbacks, so TRITON_ATTN ranks above it."""
+    from vllm.platforms.rocm import RocmPlatform
+
+    config = SimpleNamespace(
+        speculative_config=SimpleNamespace() if speculative else None,
+        model_config=SimpleNamespace(is_hybrid=hybrid),
+        cache_config=SimpleNamespace(cache_dtype="auto"),
+    )
+    attn_selector_config = AttentionSelectorConfig(
+        head_size=128,
+        dtype=torch.bfloat16,
+        kv_cache_dtype="auto",
+        block_size=16,
+        use_mla=False,
+        has_sink=False,
+        use_sparse=False,
+    )
+    with (
+        patch("vllm.config.get_current_vllm_config_or_none", return_value=config),
+        patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=aiter_found),
+    ):
+        backend_path = RocmPlatform.get_attn_backend_cls(
+            selected_backend=None, attn_selector_config=attn_selector_config
+        )
+
+    assert backend_path == expected_backend.get_path()
+
+
+@pytest.mark.parametrize(
+    "use_spec_decode, is_hybrid, use_kv_connector",
+    [
+        (False, False, False),
+        (False, False, True),
+        (True, False, False),
+        (False, True, False),
+        (True, True, False),
+        (True, False, True),
+    ],
+)
+def test_backend_priorities_order(use_spec_decode, is_hybrid, use_kv_connector):
+    from vllm.platforms.rocm import _get_backend_priorities
+
+    with patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=True):
+        priorities = _get_backend_priorities(
+            use_mla=False,
+            use_sparse=False,
+            use_kv_connector=use_kv_connector,
+            use_spec_decode=use_spec_decode,
+            is_hybrid=is_hybrid,
+        )
+
+    rocm_attn = AttentionBackendEnum.ROCM_ATTN
+    unified = AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN
+    triton_attn = AttentionBackendEnum.TRITON_ATTN
+    order = [b for b in priorities if b in (rocm_attn, unified, triton_attn)]
+    if use_spec_decode or is_hybrid:
+        expected = [triton_attn, unified, rocm_attn]
+    else:
+        expected = [rocm_attn, unified, triton_attn]
+    if use_kv_connector:
+        expected.remove(rocm_attn)
+    assert order == expected

vllm/platforms/rocm.py

@@ -490,6 +490,8 @@ def _get_backend_priorities(
     use_mla: bool,
     use_sparse: bool,
     use_kv_connector: bool = False,
+    use_spec_decode: bool = False,
+    is_hybrid: bool = False,
 ) -> list[AttentionBackendEnum]:
     from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
 
@@ -511,15 +513,26 @@ def _get_backend_priorities(
     backends = []
     # Keep ROCM_ATTN disabled for KV connectors until connector transfer
     # semantics are validated for its asymmetric native K/V cache views.
-    if not use_kv_connector:
+    use_rocm_attn = not use_kv_connector
+    # Spec decode and hybrid models put ROCM_ATTN on its slower Triton path.
+    spec_or_hybrid = use_spec_decode or is_hybrid
+    if use_rocm_attn and not spec_or_hybrid:
         backends.append(AttentionBackendEnum.ROCM_ATTN)
     if rocm_aiter_ops.is_mha_enabled():
         backends.append(AttentionBackendEnum.ROCM_AITER_FA)
+    # ROCM_AITER_UNIFIED_ATTN gives wrong output on gfx950 for some KV cache
+    # sizes until the aiter pin includes the fix (#62606), so spec decode and
+    # hybrid models rank TRITON_ATTN above it for now.
+    if spec_or_hybrid:
+        backends.append(AttentionBackendEnum.TRITON_ATTN)
     if is_aiter_found_and_supported():
         backends.append(AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN)
     elif rocm_aiter_ops.is_rdna_aiter_enabled():
         backends.insert(0, AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN)
-    backends.append(AttentionBackendEnum.TRITON_ATTN)
+    if not spec_or_hybrid:
+        backends.append(AttentionBackendEnum.TRITON_ATTN)
+    if use_rocm_attn and spec_or_hybrid:
+        backends.append(AttentionBackendEnum.ROCM_ATTN)
     backends.append(AttentionBackendEnum.TURBOQUANT)
     backends.append(AttentionBackendEnum.ULTRAQUANT)
 
@@ -691,17 +704,23 @@ def get_valid_backends(
         valid_backends_priorities = []
         invalid_reasons = {}
 
+        from vllm.config import get_current_vllm_config_or_none
+
+        vllm_config = get_current_vllm_config_or_none()
+        model_config = vllm_config.model_config if vllm_config is not None else None
+        use_spec_decode = (
+            vllm_config is not None and vllm_config.speculative_config is not None
+        )
+        is_hybrid = model_config is not None and model_config.is_hybrid
         backend_priorities = _get_backend_priorities(
             attn_selector_config.use_mla,
             attn_selector_config.use_sparse,
             attn_selector_config.use_kv_connector,
+            use_spec_decode=use_spec_decode,
+            is_hybrid=is_hybrid,
         )
-        from vllm.config import get_current_vllm_config_or_none
-
-        vllm_config = get_current_vllm_config_or_none()
         is_encoder_decoder = (
-            getattr(getattr(vllm_config, "model_config", None), "attn_type", None)
-            == "encoder_decoder"
+            getattr(model_config, "attn_type", None) == "encoder_decoder"
         )
         # ROCM_ATTN still uses a legacy attention layout (KV is the outer
         # dimension) that is incompatible with the encoder backend layouts. The