@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f2840120f6

Fork 0

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 4, 00:00 UTC

[ROCm][CI] Fix NIXL spec-decode acceptance startup and diagnostics (#41313)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

Andreas Karatzas committed May 10, 2026 at 01:50 UTC f2840120f632c55a972dc2ec482241db89b6fc0f

5 files changed +137 -27

.buildkite/test-amd.yaml

+2 -2

@@ -880,7 +880,7 @@ steps:

880 - vllm/platforms/rocm.py

881 commands:

882 - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt

883 - - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh

883 + - ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh

884

885 - label: V1 e2e (2 GPUs) # TBD

886 timeout_in_minutes: 180

@@ -3318,7 +3318,7 @@ steps:

3318 - vllm/platforms/rocm.py

3319 commands:

3320 - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt

3321 - - ROCM_ATTN=1 bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh

3321 + - ATTENTION_BACKEND=ROCM_ATTN bash v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh

3322

3323 - label: Distributed NixlConnector PD accuracy (4 GPUs) # TBD

3324 timeout_in_minutes: 180

tests/v1/kv_connector/nixl_integration/nixl_side_channel_probe.py new

+48

@@ -0,0 +1,48 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Probe a NIXL side-channel socket for handshake metadata readiness."""

4 +

5 +import argparse

6 +import ipaddress

7 +

8 +import msgspec

9 +import zmq

10 +

11 +GET_META_MSG = b"get_meta_msg"

12 +

13 +

14 +def parse_args() -> argparse.Namespace:

15 + parser = argparse.ArgumentParser()

16 + parser.add_argument("--host", required=True)

17 + parser.add_argument("--port", required=True, type=int)

18 + parser.add_argument("--rank", default=0, type=int)

19 + parser.add_argument("--timeout-ms", default=1000, type=int)

20 + return parser.parse_args()

21 +

22 +

23 +def make_zmq_path(host: str, port: int) -> str:

24 + try:

25 + if isinstance(ipaddress.ip_address(host), ipaddress.IPv6Address):

26 + return f"tcp://[{host}]:{port}"

27 + except ValueError:

28 + pass

29 + return f"tcp://{host}:{port}"

30 +

31 +

32 +def main() -> None:

33 + args = parse_args()

34 + ctx = zmq.Context()

35 + sock = ctx.socket(zmq.REQ)

36 + sock.setsockopt(zmq.LINGER, 0)

37 + sock.setsockopt(zmq.RCVTIMEO, args.timeout_ms)

38 + try:

39 + sock.connect(make_zmq_path(args.host, args.port))

40 + sock.send(msgspec.msgpack.encode((GET_META_MSG, args.rank)))

41 + sock.recv()

42 + finally:

43 + sock.close()

44 + ctx.term()

45 +

46 +

47 +if __name__ == "__main__":

48 + main()

tests/v1/kv_connector/nixl_integration/spec_decode_acceptance_test.sh

+64 -18

@@ -51,6 +51,8 @@ PREFILLER_TP_SIZE=${PREFILLER_TP_SIZE:-1}

51 DECODER_TP_SIZE=${DECODER_TP_SIZE:-1}

52 GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.7}

53 BLOCK_SIZE=${BLOCK_SIZE:-16}

54 +SERVER_HOST="${SERVER_HOST:-127.0.0.1}"

55 +NIXL_SIDE_CHANNEL_HOST="${NIXL_SIDE_CHANNEL_HOST:-$SERVER_HOST}"

57 # Resolve the repository root from the script location instead of `.git`.

58 SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"

@@ -96,18 +98,57 @@ trap 'echo " Interrupted."; exit 130' INT TERM

99 wait_for_server() {

100 local port=$1

99 - local deadline=600

101 + local server_pid=$2

102 + local server_name=$3

103 + local endpoint=${4:-/v1/completions}

104 + local deadline=${5:-600}

105 local elapsed=0

101 - echo "Waiting for server on port ${port}..."

106 + echo "Waiting for ${server_name} on port ${port}..."

107 while [ $elapsed -lt $deadline ]; do

103 - if curl -s "localhost:${port}/v1/completions" > /dev/null 2>&1; then

104 - echo "Server on port ${port} ready"

108 + if ! ps -p "$server_pid" > /dev/null 2>&1; then

109 + local status=0

110 + wait "$server_pid" || status=$?

111 + echo "FAIL: ${server_name} process ${server_pid} exited with status ${status} before port ${port} became ready"

112 + exit 1

113 + fi

114 + if curl -s "http://${SERVER_HOST}:${port}${endpoint}" > /dev/null 2>&1; then

115 + echo "${server_name} on port ${port} ready"

116 return 0

117 fi

118 sleep 2

119 elapsed=$((elapsed + 2))

120 done

110 - echo "FAIL: Server on port ${port} did not start within ${deadline}s"

121 + echo "FAIL: ${server_name} on port ${port} did not start within ${deadline}s"

122 + exit 1

123 +}

124 +

125 +wait_for_nixl_side_channel() {

126 + local host=$1

127 + local port=$2

128 + local server_pid=$3

129 + local server_name=$4

130 + local deadline=120

131 + local elapsed=0

132 + echo "Waiting for ${server_name} NIXL side channel on ${host}:${port}..."

133 + while [ $elapsed -lt $deadline ]; do

134 + if ! ps -p "$server_pid" > /dev/null 2>&1; then

135 + local status=0

136 + wait "$server_pid" || status=$?

137 + echo "FAIL: ${server_name} server process ${server_pid} exited with status ${status} before NIXL side channel ${host}:${port} became ready"

138 + exit 1

139 + fi

140 + if python3 "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/nixl_side_channel_probe.py" \

141 + --host "$host" \

142 + --port "$port" \

143 + --timeout-ms 1000 > /dev/null 2>&1

144 + then

145 + echo "${server_name} NIXL side channel on ${host}:${port} ready"

146 + return 0

147 + fi

148 + sleep 2

149 + elapsed=$((elapsed + 2))

150 + done

151 + echo "FAIL: ${server_name} NIXL side channel ${host}:${port} did not start within ${deadline}s"

152 exit 1

153 }

154

@@ -158,6 +199,8 @@ run_test_for_device() {

199 echo "KV buffer device: ${kv_device}"

200 echo "Attention backend: ${ATTENTION_BACKEND}"

201 echo "GPU platform: ${GPU_PLATFORM}"

202 + echo "Server host: ${SERVER_HOST}"

203 + echo "NIXL side channel: ${NIXL_SIDE_CHANNEL_HOST}"

204 echo "GPUs available: ${ALL_GPUS[*]}"

205 echo "================================================================"

206

@@ -167,7 +210,8 @@ run_test_for_device() {

210 local DECODE_PORTS=()

211 local GPU_IDX=0

212

170 - # Start prefill instances

213 + # Start prefill instances and wait for each one before allocating the next

214 + # server. This keeps failures from leaving extra model servers spinning.

215 for i in $(seq 0 $((NUM_PREFILL_INSTANCES-1))); do

216 local GPU_ID="${ALL_GPUS[$GPU_IDX]}"

217 GPU_IDX=$((GPU_IDX + 1))

@@ -184,6 +228,7 @@ run_test_for_device() {

228 ${GPU_DEVICE_VAR}=$GPU_ID \

229 VLLM_KV_CACHE_LAYOUT='HND' \

230 UCX_NET_DEVICES=all \

231 + VLLM_NIXL_SIDE_CHANNEL_HOST=$NIXL_SIDE_CHANNEL_HOST \

232 VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \

233 vllm serve $MODEL_NAME \

234 --port $PORT \

@@ -195,12 +240,15 @@ run_test_for_device() {

240 --kv-transfer-config "$kv_config" \

241 --speculative-config "$PREFILL_SPEC_CONFIG" \

242 --attention-backend $ATTENTION_BACKEND &

243 + local SERVER_PID=$!

244

199 - PREFILL_HOSTS+=("localhost")

245 + PREFILL_HOSTS+=("$SERVER_HOST")

246 PREFILL_PORTS+=("$PORT")

247 + wait_for_server "$PORT" "$SERVER_PID" "prefill"

248 + wait_for_nixl_side_channel "$NIXL_SIDE_CHANNEL_HOST" "$SIDE_CHANNEL_PORT" "$SERVER_PID" "prefill"

249 done

250

203 - # Start decode instances

251 + # Start decode instances after prefill is ready.

252 for i in $(seq 0 $((NUM_DECODE_INSTANCES-1))); do

253 local GPU_ID="${ALL_GPUS[$GPU_IDX]}"

254 GPU_IDX=$((GPU_IDX + 1))

@@ -217,6 +265,7 @@ run_test_for_device() {

265 ${GPU_DEVICE_VAR}=$GPU_ID \

266 VLLM_KV_CACHE_LAYOUT='HND' \

267 UCX_NET_DEVICES=all \

268 + VLLM_NIXL_SIDE_CHANNEL_HOST=$NIXL_SIDE_CHANNEL_HOST \

269 VLLM_NIXL_SIDE_CHANNEL_PORT=$SIDE_CHANNEL_PORT \

270 vllm serve $MODEL_NAME \

271 --port $PORT \

@@ -228,17 +277,12 @@ run_test_for_device() {

277 --kv-transfer-config "$kv_config" \

278 --speculative-config "$DECODE_SPEC_CONFIG" \

279 --attention-backend $ATTENTION_BACKEND &

280 + local SERVER_PID=$!

281

232 - DECODE_HOSTS+=("localhost")

282 + DECODE_HOSTS+=("$SERVER_HOST")

283 DECODE_PORTS+=("$PORT")

234 - done

235 -

236 - # Wait for servers

237 - for PORT in "${PREFILL_PORTS[@]}"; do

238 - wait_for_server "$PORT"

239 - done

240 - for PORT in "${DECODE_PORTS[@]}"; do

241 - wait_for_server "$PORT"

284 + wait_for_server "$PORT" "$SERVER_PID" "decode"

285 + wait_for_nixl_side_channel "$NIXL_SIDE_CHANNEL_HOST" "$SIDE_CHANNEL_PORT" "$SERVER_PID" "decode"

286 done

287

288 # Start proxy

@@ -250,12 +294,14 @@ run_test_for_device() {

294 --prefiller-ports ${PREFILL_PORTS[*]} \

295 --decoder-hosts ${DECODE_HOSTS[*]} \

296 --decoder-ports ${DECODE_PORTS[*]} &

297 + local PROXY_PID=$!

298

254 - sleep 5

299 + wait_for_server "$PROXY_PORT" "$PROXY_PID" "proxy" "/healthcheck" 60

300

301 # Run test

302 echo "Running spec decode acceptance test (kv_buffer_device=${kv_device}, backend=${ATTENTION_BACKEND})..."

303 DECODE_PORT=${DECODE_PORTS[0]} \

304 + SERVER_HOST=$SERVER_HOST \

305 TEST_MODEL=$MODEL_NAME \

306 python3 -m pytest -s -x "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/test_spec_decode_acceptance.py"

307

tests/v1/kv_connector/nixl_integration/test_spec_decode_acceptance.py

+4 -3

@@ -27,7 +27,8 @@ from transformers import AutoTokenizer

28 from vllm.benchmarks.datasets import get_samples

30 -PROXY_BASE_URL = "http://localhost:8192/v1"

30 +SERVER_HOST = os.environ.get("SERVER_HOST", "127.0.0.1")

31 +PROXY_BASE_URL = f"http://{SERVER_HOST}:8192/v1"

32 DECODE_PORT = os.environ.get("DECODE_PORT", "8200")

33 MODEL_NAME = os.environ.get("TEST_MODEL", "meta-llama/Llama-3.1-8B-Instruct")

@@ -101,7 +102,7 @@ def _get_mt_bench_prompts() -> list[str]:

102

103 def _fetch_metric(metric_name: str) -> float:

104 """Fetch a single counter metric from the decode server's /metrics."""

104 - url = f"http://localhost:{DECODE_PORT}/metrics"

105 + url = f"http://{SERVER_HOST}:{DECODE_PORT}/metrics"

106 body = urlopen(url).read().decode()

107 for line in body.split("\n"):

108 if line.startswith(metric_name + "{") or line.startswith(metric_name + " "):

@@ -111,7 +112,7 @@ def _fetch_metric(metric_name: str) -> float:

112

113 def _fetch_per_position_acceptance() -> dict[int, float]:

114 """Fetch per-position acceptance counts from decode /metrics."""

114 - url = f"http://localhost:{DECODE_PORT}/metrics"

115 + url = f"http://{SERVER_HOST}:{DECODE_PORT}/metrics"

116 body = urlopen(url).read().decode()

117 counts: dict[int, float] = {}

118 for line in body.split("\n"):

vllm/distributed/kv_transfer/kv_connector/v1/nixl/worker.py

+19 -4

@@ -413,7 +413,10 @@ class NixlConnectorWorker:

413

414 self.kv_cache_layout = get_kv_cache_layout()

415 self.host_buffer_kv_cache_layout = self.kv_cache_layout

416 - logger.info("Detected attention backend %s", self.backend_name)

416 + logger.info(

417 + "Detected attention backend(s) %s",

418 + [backend.get_name() for backend in self.attn_backends],

419 + )

420 logger.info("Detected kv cache layout %s", self.kv_cache_layout)

421

422 # lazy initialized in register_kv_caches

@@ -870,9 +873,21 @@ class NixlConnectorWorker:

873 else:

874 self.block_len_per_layer.append(physical_page_size)

875

873 - assert cache.shape[0] == num_blocks, (

874 - "All kv cache tensors must have the same number of blocks"

875 - )

876 + if cache.shape[0] != num_blocks:

877 + raise AssertionError(

878 + "All kv cache tensors must have the same number of "

879 + f"blocks; layer={layer_name}, "

880 + f"expected_num_blocks={num_blocks}, "

881 + f"cache_shape={tuple(cache.shape)}, "

882 + f"cache_stride={tuple(cache.stride())}, "

883 + f"layer_spec={type(layer_spec).__name__}, "

884 + f"backend={self.backend_name}, "

885 + "all_backends="

886 + f"{[backend.get_name() for backend in self.attn_backends]}, "

887 + f"kv_cache_layout={self.kv_cache_layout}, "

888 + "blocks_first="

889 + f"{self.transfer_topo.is_kv_layout_blocks_first}"

890 + )

891

892 if not self.use_mla:

893 # Different kv cache shape is not supported by HeteroTP.