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 14, 05:30 UTC
[CI/Build] Fix AMD CI: test_cpu_gpu.py (#27388)
Signed-off-by: zhewenli <zhewenli@meta.com>
Zhewen Li committed
Oct 23, 2025 at 00:55 UTC
50b788a17a8a059dad354640bebcaaee7dd72f3f
1 file changed
+13
-4
tests/v1/kv_offload/test_cpu_gpu.py
+13
-4
@@ -8,11 +8,20 @@ import torch
8
9
from vllm.platforms import current_platform
10
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
11
-from vllm.v1.attention.backends.flashinfer import FlashInferBackend
12
-from vllm.v1.attention.backends.mla.flashattn_mla import FlashAttnMLABackend
11
from vllm.v1.kv_offload.mediums import CPULoadStoreSpec, GPULoadStoreSpec
12
from vllm.v1.kv_offload.worker.cpu_gpu import CpuGpuOffloadingHandler
13
14
+BACKENDS_TO_TEST = [FlashAttentionBackend]
15
+
16
+if not current_platform.is_rocm():
17
+ from vllm.v1.attention.backends.flashinfer import FlashInferBackend
18
+
19
+ BACKENDS_TO_TEST.append(FlashInferBackend)
20
+
21
+ from vllm.v1.attention.backends.mla.flashattn_mla import FlashAttnMLABackend
22
+
23
+ BACKENDS_TO_TEST.append(FlashAttnMLABackend)
24
+
25
NUM_GPU_BLOCKS = [64]
26
NUM_CPU_BLOCKS = [256]
27
GPU_BLOCK_SIZES = [16]
@@ -55,8 +64,8 @@ def test_transfer(
64
) -> None:
65
current_platform.seed_everything(seed)
66
58
- # create per-layer GPU KV caches
59
- attn_backends_list = [FlashAttentionBackend, FlashInferBackend, FlashAttnMLABackend]
67
+ # create per-layer GPU KV caches based on available attn_backends
68
+ attn_backends_list = BACKENDS_TO_TEST
69
70
gpu_caches = {}
71
attn_backends = {}