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 6, 17:00 UTC
[ROCm] Fix KV copy methods and auto-select attention backend for ROCm (#36845)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Mar 16, 2026 at 03:07 UTC
911355e216d34d08ae6fec11be118d5817c4e5fd
2 files changed
+75
-17
tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh
+51
-17
@@ -21,6 +21,11 @@
21
# MODEL_NAME - target model (default: meta-llama/Llama-3.1-8B-Instruct)
22
# NUM_SPEC_TOKENS - number of speculative tokens (default: 3)
23
# GPU_MEMORY_UTILIZATION - (default: 0.7)
24
+# ATTENTION_BACKEND - attention backend to use
25
+# Default: TRITON_ATTN on ROCm, FLASH_ATTN on NVIDIA
26
+# ROCm options: TRITON_ATTN, ROCM_ATTN, ROCM_AITER_FA,
27
+# ROCM_AITER_UNIFIED_ATTN
28
+# NVIDIA options: FLASH_ATTN, FLASHINFER
29
set -x
30
31
# ── Model & spec decode config ──────────────────────────────────────────
@@ -51,6 +56,28 @@ GIT_ROOT=$(git rev-parse --show-toplevel)
56
57
SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "")
58
59
+# ── Detect platform (NVIDIA vs ROCm) ────────────────────────────────────
60
+
61
+if [[ "$SMI_BIN" == *"rocm"* ]]; then
62
+ GPU_PLATFORM="rocm"
63
+ GPU_DEVICE_VAR="HIP_VISIBLE_DEVICES"
64
+else
65
+ GPU_PLATFORM="nvidia"
66
+ GPU_DEVICE_VAR="CUDA_VISIBLE_DEVICES"
67
+fi
68
+echo "Detected GPU platform: ${GPU_PLATFORM} (using ${GPU_DEVICE_VAR})"
69
+
70
+# ── Attention backend config ─────────────────────────────────────────────
71
+
72
+if [[ -z "${ATTENTION_BACKEND:-}" ]]; then
73
+ if [[ "$GPU_PLATFORM" == "rocm" ]]; then
74
+ ATTENTION_BACKEND="TRITON_ATTN"
75
+ else
76
+ ATTENTION_BACKEND="FLASH_ATTN"
77
+ fi
78
+fi
79
+echo "Using attention backend: ${ATTENTION_BACKEND}"
80
+
81
cleanup_instances() {
82
echo ""
83
echo "Cleaning up..."
@@ -84,13 +111,16 @@ wait_for_server() {
111
112
# ── Resolve GPU list ─────────────────────────────────────────────────────
113
87
-if [[ -n "${CUDA_VISIBLE_DEVICES:-}" ]]; then
88
- IFS=',' read -ra ALL_GPUS <<< "$CUDA_VISIBLE_DEVICES"
114
+# Accept either CUDA_VISIBLE_DEVICES or HIP_VISIBLE_DEVICES
115
+VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-${HIP_VISIBLE_DEVICES:-}}"
116
+
117
+if [[ -n "${VISIBLE_DEVICES}" ]]; then
118
+ IFS=',' read -ra ALL_GPUS <<< "$VISIBLE_DEVICES"
119
else
120
ALL_GPUS=()
91
- if [[ "$SMI_BIN" == *"nvidia"* ]]; then
121
+ if [[ "$GPU_PLATFORM" == "nvidia" ]]; then
122
num=$($SMI_BIN --query-gpu=name --format=csv,noheader | wc -l)
93
- elif [[ "$SMI_BIN" == *"rocm"* ]]; then
123
+ elif [[ "$GPU_PLATFORM" == "rocm" ]]; then
124
num=$($SMI_BIN -l | grep -c GPU)
125
else
126
num=1
@@ -100,7 +130,7 @@ fi
130
131
TOTAL_GPUS_NEEDED=$(( (NUM_PREFILL_INSTANCES * PREFILLER_TP_SIZE) + (NUM_DECODE_INSTANCES * DECODER_TP_SIZE) ))
132
if [[ ${#ALL_GPUS[@]} -lt $TOTAL_GPUS_NEEDED ]]; then
103
- echo "FAIL: Need $TOTAL_GPUS_NEEDED GPUs but only have ${#ALL_GPUS[@]} (CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-not set})"
133
+ echo "FAIL: Need $TOTAL_GPUS_NEEDED GPUs but only have ${#ALL_GPUS[@]} (visible devices=${VISIBLE_DEVICES:-not set})"
134
exit 1
135
fi
136
@@ -119,12 +149,14 @@ run_test_for_device() {
149
echo "================================================================"
150
echo "NixlConnector PD + Spec Decode Acceptance Test (kv_buffer_device=${kv_device})"
151
echo "================================================================"
122
- echo "Model: ${MODEL_NAME}"
123
- echo "SD method: ${SD_METHOD}"
124
- echo "SD model: ${SD_MODEL}"
125
- echo "Spec tokens: ${NUM_SPEC_TOKENS}"
126
- echo "KV buffer device: ${kv_device}"
127
- echo "GPUs available: ${ALL_GPUS[*]}"
152
+ echo "Model: ${MODEL_NAME}"
153
+ echo "SD method: ${SD_METHOD}"
154
+ echo "SD model: ${SD_MODEL}"
155
+ echo "Spec tokens: ${NUM_SPEC_TOKENS}"
156
+ echo "KV buffer device: ${kv_device}"
157
+ echo "Attention backend: ${ATTENTION_BACKEND}"
158
+ echo "GPU platform: ${GPU_PLATFORM}"
159
+ echo "GPUs available: ${ALL_GPUS[*]}"
160
echo "================================================================"
161
162
local PREFILL_HOSTS=()
@@ -146,7 +178,8 @@ run_test_for_device() {
178
local SIDE_CHANNEL_PORT=$((5559 + i))
179
180
echo "Starting prefill instance $i on GPU $GPU_ID, port $PORT"
149
- CUDA_VISIBLE_DEVICES=$GPU_ID \
181
+ env \
182
+ ${GPU_DEVICE_VAR}=$GPU_ID \
183
VLLM_KV_CACHE_LAYOUT='HND' \
184
UCX_NET_DEVICES=all \
185
VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
@@ -159,7 +192,7 @@ run_test_for_device() {
192
--tensor-parallel-size $PREFILLER_TP_SIZE \
193
--kv-transfer-config "$kv_config" \
194
--speculative-config "$PREFILL_SPEC_CONFIG" \
162
- --attention-backend FLASH_ATTN &
195
+ --attention-backend $ATTENTION_BACKEND &
196
197
PREFILL_HOSTS+=("localhost")
198
PREFILL_PORTS+=("$PORT")
@@ -178,7 +211,8 @@ run_test_for_device() {
211
local SIDE_CHANNEL_PORT=$((5659 + i * $DECODER_TP_SIZE))
212
213
echo "Starting decode instance $i on GPU $GPU_ID, port $PORT"
181
- CUDA_VISIBLE_DEVICES=$GPU_ID \
214
+ env \
215
+ ${GPU_DEVICE_VAR}=$GPU_ID \
216
VLLM_KV_CACHE_LAYOUT='HND' \
217
UCX_NET_DEVICES=all \
218
VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \
@@ -191,7 +225,7 @@ run_test_for_device() {
225
--tensor-parallel-size $DECODER_TP_SIZE \
226
--kv-transfer-config "$kv_config" \
227
--speculative-config "$DECODE_SPEC_CONFIG" \
194
- --attention-backend FLASH_ATTN &
228
+ --attention-backend $ATTENTION_BACKEND &
229
230
DECODE_HOSTS+=("localhost")
231
DECODE_PORTS+=("$PORT")
@@ -218,7 +252,7 @@ run_test_for_device() {
252
sleep 5
253
254
# Run test
221
- echo "Running spec decode acceptance test (kv_buffer_device=${kv_device})..."
255
+ echo "Running spec decode acceptance test (kv_buffer_device=${kv_device}, backend=${ATTENTION_BACKEND})..."
256
DECODE_PORT=${DECODE_PORTS[0]} \
257
TEST_MODEL=$MODEL_NAME \
258
python3 -m pytest -s -x "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/test_spec_decode_acceptance.py"
@@ -234,4 +268,4 @@ for device in $KV_BUFFER_DEVICES; do
268
run_test_for_device "$device"
269
done
270
237
-echo "=== All spec decode acceptance tests passed ==="
271
+echo "=== All spec decode acceptance tests passed (backend=${ATTENTION_BACKEND}) ==="
vllm/platforms/rocm.py
+24
@@ -851,6 +851,30 @@ class RocmPlatform(Platform):
851
"`dtype` flag in CLI, for example: --dtype=half."
852
)
853
854
+ @classmethod
855
+ def insert_blocks_to_device(
856
+ cls,
857
+ src_cache: torch.Tensor,
858
+ dst_cache: torch.Tensor,
859
+ src_block_indices: torch.Tensor,
860
+ dst_block_indices: torch.Tensor,
861
+ ) -> None:
862
+ """Copy blocks from src_cache to dst_cache on GPU."""
863
+ _src_cache = src_cache[:, src_block_indices]
864
+ dst_cache[:, dst_block_indices] = _src_cache.to(dst_cache.device)
865
+
866
+ @classmethod
867
+ def swap_out_blocks_to_host(
868
+ cls,
869
+ src_cache: torch.Tensor,
870
+ dst_cache: torch.Tensor,
871
+ src_block_indices: torch.Tensor,
872
+ dst_block_indices: torch.Tensor,
873
+ ) -> None:
874
+ """Copy blocks from GPU to host (CPU)."""
875
+ _src_cache = src_cache[:, src_block_indices]
876
+ dst_cache[:, dst_block_indices] = _src_cache.cpu()
877
+
878
@classmethod
879
def support_hybrid_kv_cache(cls) -> bool:
880
return True