[ROCm][CI] Adding nixl multiconn (#47481)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jul 6, 2026 at 02:04 UTC
8f0e75e16b96e07ae58283e440128b491f729972
4 files changed
+88
-7
.buildkite/test-amd.yaml
+53
@@ -2439,6 +2439,59 @@ steps:
2439
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
2440
- HYBRID_SSM=1 ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
2441
2442
+- label: Hybrid SSM NixlConnector PD prefix cache test (2 GPUs) # TBD
2443
+ timeout_in_minutes: 180
2444
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
2445
+ agent_pool: mi300_2
2446
+ num_gpus: 2
2447
+ optional: true
2448
+ working_dir: "/vllm-workspace/tests"
2449
+ source_file_dependencies:
2450
+ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
2451
+ - vllm/v1/core/sched/
2452
+ - vllm/v1/core/kv_cache_coordinator.py
2453
+ - tests/v1/kv_connector/nixl_integration/
2454
+ - vllm/platforms/rocm.py
2455
+ commands:
2456
+ - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
2457
+ - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh
2458
+
2459
+- label: MultiConnector (Nixl+Offloading) PD accuracy (2 GPUs) # TBD
2460
+ timeout_in_minutes: 180
2461
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
2462
+ agent_pool: mi300_2
2463
+ num_gpus: 2
2464
+ optional: true
2465
+ working_dir: "/vllm-workspace/tests"
2466
+ source_file_dependencies:
2467
+ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
2468
+ - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
2469
+ - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
2470
+ - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
2471
+ - tests/v1/kv_connector/nixl_integration/
2472
+ - vllm/platforms/rocm.py
2473
+ commands:
2474
+ - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
2475
+ - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
2476
+
2477
+- label: MultiConnector (Nixl+Offloading) PD edge cases (2 GPUs) # TBD
2478
+ timeout_in_minutes: 180
2479
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
2480
+ agent_pool: mi300_2
2481
+ num_gpus: 2
2482
+ optional: true
2483
+ working_dir: "/vllm-workspace/tests"
2484
+ source_file_dependencies:
2485
+ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
2486
+ - vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py
2487
+ - vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
2488
+ - vllm/distributed/kv_transfer/kv_connector/v1/offloading/
2489
+ - tests/v1/kv_connector/nixl_integration/
2490
+ - vllm/platforms/rocm.py
2491
+ commands:
2492
+ - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
2493
+ - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
2494
+
2495
- label: V1 e2e (4 GPUs) # TBD
2496
timeout_in_minutes: 180
2497
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh
+15
-5
@@ -9,8 +9,10 @@ PREFILL_GPU_ID=${PREFILL_GPU_ID:-0}
9
DECODE_GPU_ID=${DECODE_GPU_ID:-1}
10
MODEL=${MODEL:-"ibm-granite/granite-4.0-h-tiny"}
11
GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.8}
12
+VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-}
13
+ATTENTION_BACKEND=${ATTENTION_BACKEND:-FLASHINFER}
14
13
-echo "Running Mamba prefix cache test (GPUs: P=$PREFILL_GPU_ID, D=$DECODE_GPU_ID, model=$MODEL)"
15
+echo "Running Mamba prefix cache test (GPUs: P=$PREFILL_GPU_ID, D=$DECODE_GPU_ID, model=$MODEL, backend=$ATTENTION_BACKEND)"
16
17
KV_CONFIG='{"kv_connector":"NixlConnector","kv_role":"kv_both"}'
18
@@ -36,6 +38,14 @@ cleanup_instances() {
38
39
cleanup_instances
40
41
+EXTRA_ARGS=()
42
+if [[ -n "$VLLM_SERVE_EXTRA_ARGS" ]]; then
43
+ IFS=',' read -r -a EXTRA_ARGS <<< "$VLLM_SERVE_EXTRA_ARGS"
44
+fi
45
+if [[ -n "$ATTENTION_BACKEND" ]]; then
46
+ EXTRA_ARGS+=(--attention-backend "$ATTENTION_BACKEND")
47
+fi
48
+
49
# Start prefill instance
50
PREFILL_PORT=8001
51
CUDA_VISIBLE_DEVICES=$PREFILL_GPU_ID \
@@ -51,8 +61,8 @@ vllm serve $MODEL \
61
--trust-remote-code \
62
--enable-prefix-caching \
63
--mamba-cache-mode all \
54
- --attention-backend FLASHINFER \
55
- --kv-transfer-config "$KV_CONFIG" &
64
+ --kv-transfer-config "$KV_CONFIG" \
65
+ "${EXTRA_ARGS[@]}" &
66
67
# Start decode instance
68
DECODE_PORT=8002
@@ -69,8 +79,8 @@ vllm serve $MODEL \
79
--trust-remote-code \
80
--enable-prefix-caching \
81
--mamba-cache-mode all \
72
- --attention-backend FLASHINFER \
73
- --kv-transfer-config "$KV_CONFIG" &
82
+ --kv-transfer-config "$KV_CONFIG" \
83
+ "${EXTRA_ARGS[@]}" &
84
85
echo "Waiting for prefill instance on port $PREFILL_PORT..."
86
wait_for_server "$PREFILL_PORT"
tests/v1/kv_connector/nixl_integration/run_multi_connector_accuracy_test.sh
+10
-1
@@ -18,6 +18,7 @@
18
# Environment variables:
19
# MODEL_NAMES - model to test (default: Qwen/Qwen3-0.6B)
20
# GPU_MEMORY_UTILIZATION - GPU memory fraction (default: 0.6)
21
+# ATTENTION_BACKEND - optional attention backend for vllm serve
22
# VLLM_SERVE_EXTRA_ARGS - comma-separated extra args for vllm serve
23
# SKIP_CROSS_LAYERS - set to 1 to skip the cross-layer layout test
24
# SKIP_NORMAL_LAYOUT - set to 1 to skip the normal layout test
@@ -34,9 +35,11 @@ fi
35
36
GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.6}
37
BLOCK_SIZE=${BLOCK_SIZE:-128}
38
+ATTENTION_BACKEND=${ATTENTION_BACKEND:-}
39
VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-}
40
39
-GIT_ROOT=$(git rev-parse --show-toplevel)
41
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
42
+GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}"
43
SMI_BIN=$(which nvidia-smi || which rocm-smi || echo "")
44
45
# ── KV transfer configs ─────────────────────────────────────────────────
@@ -139,6 +142,9 @@ run_tests_for_model() {
142
BASE_CMD="${BASE_CMD} $arg"
143
done
144
fi
145
+ if [[ -n "$ATTENTION_BACKEND" ]]; then
146
+ BASE_CMD="${BASE_CMD} --attention-backend $ATTENTION_BACKEND"
147
+ fi
148
eval "$BASE_CMD &"
149
150
# ── Start decode instance ──
@@ -161,6 +167,9 @@ run_tests_for_model() {
167
BASE_CMD="${BASE_CMD} $arg"
168
done
169
fi
170
+ if [[ -n "$ATTENTION_BACKEND" ]]; then
171
+ BASE_CMD="${BASE_CMD} --attention-backend $ATTENTION_BACKEND"
172
+ fi
173
eval "$BASE_CMD &"
174
175
# ── Wait for servers ──
tests/v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh
+10
-1
@@ -19,6 +19,7 @@
19
# MODEL_NAMES - model to test (default: Qwen/Qwen3-0.6B)
20
# KV_CACHE_MEMORY_BYTES - GPU KV cache size in bytes (default: 268435456 = 256 MiB)
21
# BLOCK_SIZE - KV cache block size (default: 128)
22
+# ATTENTION_BACKEND - optional attention backend for vllm serve
23
# VLLM_SERVE_EXTRA_ARGS - comma-separated extra args for vllm serve
24
set -xe
25
@@ -34,9 +35,11 @@ fi
35
KV_CACHE_MEMORY_BYTES=${KV_CACHE_MEMORY_BYTES:-268435456} # 256 MiB
36
MAX_MODEL_LEN=${MAX_MODEL_LEN:-2048}
37
BLOCK_SIZE=${BLOCK_SIZE:-128}
38
+ATTENTION_BACKEND=${ATTENTION_BACKEND:-}
39
VLLM_SERVE_EXTRA_ARGS=${VLLM_SERVE_EXTRA_ARGS:-}
40
39
-GIT_ROOT=$(git rev-parse --show-toplevel)
41
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
42
+GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}"
43
44
# ── KV transfer config ──────────────────────────────────────────────────
45
@@ -110,6 +113,9 @@ run_tests_for_model() {
113
BASE_CMD="${BASE_CMD} $arg"
114
done
115
fi
116
+ if [[ -n "$ATTENTION_BACKEND" ]]; then
117
+ BASE_CMD="${BASE_CMD} --attention-backend $ATTENTION_BACKEND"
118
+ fi
119
eval "$BASE_CMD &"
120
121
# ── Start decode instance ──
@@ -133,6 +139,9 @@ run_tests_for_model() {
139
BASE_CMD="${BASE_CMD} $arg"
140
done
141
fi
142
+ if [[ -n "$ATTENTION_BACKEND" ]]; then
143
+ BASE_CMD="${BASE_CMD} --attention-backend $ATTENTION_BACKEND"
144
+ fi
145
eval "$BASE_CMD &"
146
147
# ── Wait for servers ──