diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index 9ecf63b028d1..5409169ccd78 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -933,21 +933,24 @@ def test_rocm_attn_rejects_kv_connector(mock_vllm_config, mock_get_cdna_version) @pytest.mark.parametrize( - "aiter_found, expected_backend", + "aiter_found, gfx950, expected_backend", [ - (True, AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN), - (False, AttentionBackendEnum.TRITON_ATTN), + (True, False, AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN), + (True, True, AttentionBackendEnum.TRITON_ATTN), + (False, False, AttentionBackendEnum.TRITON_ATTN), + (False, True, AttentionBackendEnum.TRITON_ATTN), ], ) def test_auto_selection_for_kv_connector( - aiter_found, expected_backend, mock_vllm_config, mock_get_cdna_version + aiter_found, gfx950, expected_backend, mock_vllm_config, mock_get_cdna_version ): """Auto-selection with a KV connector and AITER enabled resolves to unified attn, - and to triton attn if AITER not enabled.""" + and to triton attn if AITER not enabled or on gfx950.""" from vllm.platforms.rocm import RocmPlatform - with patch( - "vllm._aiter_ops.is_aiter_found_and_supported", return_value=aiter_found + with ( + patch("vllm._aiter_ops.is_aiter_found_and_supported", return_value=aiter_found), + patch("vllm.platforms.rocm.on_gfx950", return_value=gfx950), ): backend_path = RocmPlatform.get_attn_backend_cls( selected_backend=None, diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index 0ce608bd729a..1c98dd0ac1f7 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -520,6 +520,14 @@ def _get_backend_priorities( elif rocm_aiter_ops.is_rdna_aiter_enabled(): backends.insert(0, AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN) backends.append(AttentionBackendEnum.TRITON_ATTN) + + # TODO(#62606): remove once the aiter pin includes ROCm/aiter#6661. + # ROCM_AITER_UNIFIED_ATTN gives wrong output on gfx950 for some KV cache + # sizes, so rank TRITON_ATTN right above it there. + unified = AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN + if on_gfx950() and unified in backends: + backends.remove(AttentionBackendEnum.TRITON_ATTN) + backends.insert(backends.index(unified), AttentionBackendEnum.TRITON_ATTN) backends.append(AttentionBackendEnum.TURBOQUANT) backends.append(AttentionBackendEnum.ULTRAQUANT)