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 12, 21:00 UTC
[CI/Build] Make test_attention_selector.py run tests on correct platform (#29064)
Signed-off-by: Randall Smith <ransmith@amd.com> Signed-off-by: rasmith <Randall.Smith@amd.com> Co-authored-by: Randall Smith <ransmith@amd.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
rasmith committed
Nov 20, 2025 at 14:45 UTC
5e5a7eb16f121f05e19c8bdf88247744ab9d1b83
1 file changed
+4
-1
tests/kernels/attention/test_attention_selector.py
+4
-1
@@ -7,6 +7,7 @@ import pytest
7
import torch
8
9
from vllm.attention.selector import _cached_get_attn_backend, get_attn_backend
10
+from vllm.platforms import current_platform
11
from vllm.platforms.cpu import CpuPlatform
12
from vllm.platforms.cuda import CudaPlatform
13
from vllm.platforms.rocm import RocmPlatform
@@ -47,9 +48,11 @@ DEVICE_MLA_BLOCK_SIZES = {
48
49
50
def generate_params():
51
+ is_rocm = current_platform.is_rocm()
52
params = []
53
+ device_list = ["cuda", "cpu"] if not is_rocm else ["hip", "cpu"]
54
for use_mla in [True, False]:
52
- for device in ["cuda", "hip", "cpu"]:
55
+ for device in device_list:
56
backends = (
57
DEVICE_MLA_BACKENDS[device]
58
if use_mla