[ROCm][CI] Prep Tests For Change To ROCM_ATTN As New Default Backend On ROCm (#36025)

Signed-off-by: Micah Williamson <micah.williamson@amd.com>

tests/entrypoints/openai/test_tensorizer_entrypoint.py

@@ -15,6 +15,7 @@ from vllm.model_executor.model_loader.tensorizer import (
     tensorize_lora_adapter,
     tensorize_vllm_model,
 )
+from vllm.platforms import current_platform
 
 from ...utils import RemoteOpenAIServer
 
@@ -74,6 +75,8 @@ def server(model_uri, tensorize_model_and_lora):
         MODEL_NAME,
         "--enable-lora",
     ]
+    if current_platform.is_rocm():
+        args += ["--attention-backend", "TRITON_ATTN"]
 
     model_dir = os.path.dirname(model_uri)
     with RemoteOpenAIServer(model_dir, args) as remote_server:

tests/models/language/pooling_mteb_test/test_gte.py

@@ -8,6 +8,7 @@ from tests.models.utils import (
     EmbedModelInfo,
     RerankModelInfo,
 )
+from vllm.platforms import current_platform
 
 from .mteb_embed_utils import mteb_test_embed_models
 from .mteb_score_utils import mteb_test_rerank_models
@@ -142,4 +143,9 @@ def test_embed_models_correctness(
 
 @pytest.mark.parametrize("model_info", RERANK_MODELS)
 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
-    mteb_test_rerank_models(vllm_runner, model_info)
+    vllm_extra_kwargs = {}
+    if current_platform.is_rocm():
+        vllm_extra_kwargs["attention_backend"] = "TRITON_ATTN"
+    mteb_test_rerank_models(
+        vllm_runner, model_info, vllm_extra_kwargs=vllm_extra_kwargs
+    )

tests/models/multimodal/generation/test_common.py

@@ -173,6 +173,9 @@ VLM_TEST_SETTINGS = {
         marks=[
             pytest.mark.core_model,
         ],
+        vllm_runner_kwargs={"attention_backend": "TRITON_ATTN"}
+        if current_platform.is_rocm()
+        else {},
     ),
     "ultravox": VLMTestInfo(
         models=["fixie-ai/ultravox-v0_5-llama-3_2-1b"],

tests/test_regression.py

@@ -13,6 +13,7 @@ import pytest
 import torch
 
 from vllm import LLM, SamplingParams
+from vllm.platforms import current_platform
 
 
 @pytest.mark.skip(reason="In V1, we reject tokens > max_seq_len")
@@ -65,7 +66,8 @@ def test_model_from_modelscope(monkeypatch: pytest.MonkeyPatch):
         # Don't use HF_TOKEN for ModelScope repos, otherwise it will fail
         # with 400 Client Error: Bad Request.
         m.setenv("HF_TOKEN", "")
-        llm = LLM(model="qwen/Qwen1.5-0.5B-Chat")
+        attn_backend = "TRITON_ATTN" if current_platform.is_rocm() else "auto"
+        llm = LLM(model="qwen/Qwen1.5-0.5B-Chat", attention_backend=attn_backend)
 
         prompts = [
             "Hello, my name is",

tests/v1/e2e/test_kv_sharing_fast_prefill.py

@@ -91,6 +91,7 @@ def test_kv_sharing_fast_prefill(
             compilation_config=compilation_config,
             seed=SEED,
             kv_sharing_fast_prefill=kv_sharing_fast_prefill,
+            attention_backend="TRITON_ATTN",
         )
         responses = llm.generate(prompts, sampling_params)
         check_answers(

tests/v1/e2e/test_spec_decode.py

@@ -732,11 +732,13 @@ def test_mtp_correctness(
         method, model_name, tp_size = model_setup
         _skip_if_insufficient_gpus_for_tp(tp_size)
 
+        attn_backend = "TRITON_ATTN" if current_platform.is_rocm() else "auto"
         ref_llm = LLM(
             model=model_name,
             max_model_len=2048,
             tensor_parallel_size=tp_size,
             trust_remote_code=True,
+            attention_backend=attn_backend,
         )
         ref_outputs = ref_llm.chat(test_prompts, sampling_config)
         evaluate_llm_for_gsm8k(
@@ -756,6 +758,7 @@ def test_mtp_correctness(
                 "max_model_len": 2048,
             },
             max_model_len=2048,
+            attention_backend=attn_backend,
         )
         evaluate_llm_for_gsm8k(
             spec_llm, expected_accuracy_threshold=expected_accuracy_threshold

tests/v1/sample/test_logprobs.py

@@ -42,9 +42,7 @@ SAMPLE_PROMPT = BatchLogprobsComposition.SAMPLE_PROMPT
 # Force LLM instances into an identical, deterministic execution
 # mode so the test isolates spec-decode correctness only:
 ROCM_DETERMINISM_KWARGS: dict = (
-    dict(
-        max_num_seqs=1,
-    )
+    dict(max_num_seqs=1, attention_backend="TRITON_ATTN")
     if current_platform.is_rocm()
     else {}
 )