A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 00:15 UTC
[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>
Micah Williamson committed
Mar 9, 2026 at 13:27 UTC
4ff9b045fe7a9da9b5a7737407ed4e7ef203ffad
10 files changed
+32
-10
.buildkite/lm-eval-harness/test_lm_eval_correctness.py
+8
-2
@@ -13,9 +13,10 @@ import os
13
from contextlib import contextmanager
14
15
import lm_eval
16
-import numpy as np
16
import yaml
17
18
+from vllm.platforms import current_platform
19
+
20
DEFAULT_RTOL = 0.08
21
22
@@ -63,6 +64,9 @@ def launch_lm_eval(eval_config, tp_size):
64
"allow_deprecated_quantization=True,"
65
)
66
67
+ if current_platform.is_rocm() and "Nemotron-3" in eval_config["model_name"]:
68
+ model_args += "attention_backend=TRITON_ATTN"
69
+
70
env_vars = eval_config.get("env_vars", None)
71
with scoped_env_vars(env_vars):
72
results = lm_eval.simple_evaluate(
@@ -102,6 +106,8 @@ def test_lm_eval_correctness_param(config_filename, tp_size):
106
f"ground_truth={ground_truth:.3f} | "
107
f"measured={measured_value:.3f} | rtol={rtol}"
108
)
105
- success = success and np.isclose(ground_truth, measured_value, rtol=rtol)
109
+
110
+ min_acceptable = ground_truth * (1 - rtol)
111
+ success = success and measured_value >= min_acceptable
112
113
assert success
.buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh
+1
-1
@@ -24,7 +24,7 @@ if command -v rocm-smi &> /dev/null || [[ -d /opt/rocm ]] || [[ -n "${ROCM_PATH:
24
BACKENDS=("allgather_reducescatter")
25
# Disable MOE padding for ROCm since it is causing eplb to fail
26
export VLLM_ROCM_MOE_PADDING=0
27
- PLATFORM_ARGS=("--no-async-scheduling")
27
+ PLATFORM_ARGS=("--no-async-scheduling" "--attention-backend=TRITON_ATTN")
28
echo "Disabled async scheduling for ROCm platform due to issues with spec decode."
29
else
30
# Non-ROCm platform (CUDA/other)
.buildkite/test-amd.yaml
+2
-2
@@ -529,7 +529,7 @@ steps:
529
commands:
530
- pip install tensorizer # for tensorizer test
531
# for basic
532
- - python3 basic/offline_inference/chat.py
532
+ - python3 basic/offline_inference/chat.py --attention-backend TRITON_ATTN
533
- python3 basic/offline_inference/generate.py --model facebook/opt-125m
534
- python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10
535
- python3 basic/offline_inference/classify.py
@@ -2208,7 +2208,7 @@ steps:
2208
commands:
2209
- pip install tensorizer # for tensorizer test
2210
# for basic
2211
- - python3 basic/offline_inference/chat.py
2211
+ - python3 basic/offline_inference/chat.py --attention-backend TRITON_ATTN
2212
- python3 basic/offline_inference/generate.py --model facebook/opt-125m
2213
- python3 basic/offline_inference/generate.py --model meta-llama/Llama-2-13b-chat-hf --cpu-offload-gb 10
2214
- python3 basic/offline_inference/classify.py
tests/entrypoints/openai/test_tensorizer_entrypoint.py
+3
@@ -15,6 +15,7 @@ from vllm.model_executor.model_loader.tensorizer import (
15
tensorize_lora_adapter,
16
tensorize_vllm_model,
17
)
18
+from vllm.platforms import current_platform
19
20
from ...utils import RemoteOpenAIServer
21
@@ -74,6 +75,8 @@ def server(model_uri, tensorize_model_and_lora):
75
MODEL_NAME,
76
"--enable-lora",
77
]
78
+ if current_platform.is_rocm():
79
+ args += ["--attention-backend", "TRITON_ATTN"]
80
81
model_dir = os.path.dirname(model_uri)
82
with RemoteOpenAIServer(model_dir, args) as remote_server:
tests/models/language/pooling_mteb_test/test_gte.py
+7
-1
@@ -8,6 +8,7 @@ from tests.models.utils import (
8
EmbedModelInfo,
9
RerankModelInfo,
10
)
11
+from vllm.platforms import current_platform
12
13
from .mteb_embed_utils import mteb_test_embed_models
14
from .mteb_score_utils import mteb_test_rerank_models
@@ -142,4 +143,9 @@ def test_embed_models_correctness(
143
144
@pytest.mark.parametrize("model_info", RERANK_MODELS)
145
def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
145
- mteb_test_rerank_models(vllm_runner, model_info)
146
+ vllm_extra_kwargs = {}
147
+ if current_platform.is_rocm():
148
+ vllm_extra_kwargs["attention_backend"] = "TRITON_ATTN"
149
+ mteb_test_rerank_models(
150
+ vllm_runner, model_info, vllm_extra_kwargs=vllm_extra_kwargs
151
+ )
tests/models/multimodal/generation/test_common.py
+3
@@ -173,6 +173,9 @@ VLM_TEST_SETTINGS = {
173
marks=[
174
pytest.mark.core_model,
175
],
176
+ vllm_runner_kwargs={"attention_backend": "TRITON_ATTN"}
177
+ if current_platform.is_rocm()
178
+ else {},
179
),
180
"ultravox": VLMTestInfo(
181
models=["fixie-ai/ultravox-v0_5-llama-3_2-1b"],
tests/test_regression.py
+3
-1
@@ -13,6 +13,7 @@ import pytest
13
import torch
14
15
from vllm import LLM, SamplingParams
16
+from vllm.platforms import current_platform
17
18
19
@pytest.mark.skip(reason="In V1, we reject tokens > max_seq_len")
@@ -65,7 +66,8 @@ def test_model_from_modelscope(monkeypatch: pytest.MonkeyPatch):
66
# Don't use HF_TOKEN for ModelScope repos, otherwise it will fail
67
# with 400 Client Error: Bad Request.
68
m.setenv("HF_TOKEN", "")
68
- llm = LLM(model="qwen/Qwen1.5-0.5B-Chat")
69
+ attn_backend = "TRITON_ATTN" if current_platform.is_rocm() else "auto"
70
+ llm = LLM(model="qwen/Qwen1.5-0.5B-Chat", attention_backend=attn_backend)
71
72
prompts = [
73
"Hello, my name is",
tests/v1/e2e/test_kv_sharing_fast_prefill.py
+1
@@ -91,6 +91,7 @@ def test_kv_sharing_fast_prefill(
91
compilation_config=compilation_config,
92
seed=SEED,
93
kv_sharing_fast_prefill=kv_sharing_fast_prefill,
94
+ attention_backend="TRITON_ATTN",
95
)
96
responses = llm.generate(prompts, sampling_params)
97
check_answers(
tests/v1/e2e/test_spec_decode.py
+3
@@ -732,11 +732,13 @@ def test_mtp_correctness(
732
method, model_name, tp_size = model_setup
733
_skip_if_insufficient_gpus_for_tp(tp_size)
734
735
+ attn_backend = "TRITON_ATTN" if current_platform.is_rocm() else "auto"
736
ref_llm = LLM(
737
model=model_name,
738
max_model_len=2048,
739
tensor_parallel_size=tp_size,
740
trust_remote_code=True,
741
+ attention_backend=attn_backend,
742
)
743
ref_outputs = ref_llm.chat(test_prompts, sampling_config)
744
evaluate_llm_for_gsm8k(
@@ -756,6 +758,7 @@ def test_mtp_correctness(
758
"max_model_len": 2048,
759
},
760
max_model_len=2048,
761
+ attention_backend=attn_backend,
762
)
763
evaluate_llm_for_gsm8k(
764
spec_llm, expected_accuracy_threshold=expected_accuracy_threshold
tests/v1/sample/test_logprobs.py
+1
-3
@@ -42,9 +42,7 @@ SAMPLE_PROMPT = BatchLogprobsComposition.SAMPLE_PROMPT
42
# Force LLM instances into an identical, deterministic execution
43
# mode so the test isolates spec-decode correctness only:
44
ROCM_DETERMINISM_KWARGS: dict = (
45
- dict(
46
- max_num_seqs=1,
47
- )
45
+ dict(max_num_seqs=1, attention_backend="TRITON_ATTN")
46
if current_platform.is_rocm()
47
else {}
48
)