[ROCm][CI] Validate Mooncake and NIXL prefill/decode accuracy (#58095)
Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Andreas Karatzas committed
Sep 23, 2026 at 11:16 UTC
8033bd08c1d3626d499de040279d7476b6cd2a4a
7 files changed
+409
.buildkite/test-amd.yaml
+43
@@ -3562,6 +3562,49 @@ steps:
3562
- uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
3563
- ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
3564
3565
+- label: ":amd: (MI355) Distributed AITER NixlConnector PD accuracy"
3566
+ timeout_in_minutes: 60
3567
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
3568
+ dind: false
3569
+ agent_pool: mi355_4
3570
+ num_gpus: 4
3571
+ optional: true
3572
+ working_dir: "/vllm-workspace/tests"
3573
+ source_file_dependencies:
3574
+ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
3575
+ - tests/v1/kv_connector/nixl_integration/
3576
+ - tests/v1/kv_connector/rocm_pd_accuracy_utils.py
3577
+ - tests/utils.py
3578
+ - vllm/v1/attention/backends/rocm_aiter_unified_attn.py
3579
+ - vllm/platforms/rocm.py
3580
+ - requirements/common.txt
3581
+ - requirements/kv_connectors_rocm.txt
3582
+ commands:
3583
+ - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
3584
+ - pytest -v -s v1/kv_connector/nixl_integration/test_accuracy_rocm.py
3585
+
3586
+- label: ":amd: (MI355) Distributed MooncakeConnector HIP PD accuracy"
3587
+ timeout_in_minutes: 30
3588
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
3589
+ dind: false
3590
+ agent_pool: mi355_4
3591
+ num_gpus: 4
3592
+ optional: true
3593
+ working_dir: "/vllm-workspace/tests"
3594
+ source_file_dependencies:
3595
+ - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/
3596
+ - tests/v1/kv_connector/mooncake_integration/
3597
+ - tests/v1/kv_connector/rocm_pd_accuracy_utils.py
3598
+ - tests/utils.py
3599
+ - examples/disaggregated/mooncake_connector/
3600
+ - vllm/v1/attention/backends/rocm_aiter_unified_attn.py
3601
+ - vllm/platforms/rocm.py
3602
+ - requirements/common.txt
3603
+ - requirements/kv_connectors_rocm.txt
3604
+ commands:
3605
+ - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
3606
+ - pytest -v -s v1/kv_connector/mooncake_integration/test_accuracy_rocm.py
3607
+
3608
- label: ":amd: (MI355) DP EP Distributed NixlConnector PD accuracy" # TBD
3609
timeout_in_minutes: 40
3610
mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
.buildkite/test_areas/disaggregated.yaml
+23
@@ -42,6 +42,29 @@ steps:
42
commands:
43
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
44
- FLASHINFER=1 bash v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh
45
+ mirror:
46
+ amd:
47
+ label: ":amd: (MI355) Distributed AITER NixlConnector PD accuracy"
48
+ dind: false
49
+ device: mi355_4
50
+ num_devices: 4
51
+ timeout_in_minutes: 60
52
+ working_dir: "/vllm-workspace/tests"
53
+ depends_on:
54
+ - image-build-amd
55
+ source_file_dependencies:
56
+ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
57
+ - tests/v1/kv_connector/nixl_integration/
58
+ - tests/v1/kv_connector/rocm_pd_accuracy_utils.py
59
+ - tests/utils.py
60
+ - vllm/v1/attention/backends/rocm_aiter_unified_attn.py
61
+ - vllm/platforms/rocm.py
62
+ - requirements/common.txt
63
+ - requirements/kv_connectors_rocm.txt
64
+ commands:
65
+ - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
66
+ # Preserve all six model/TP pairs with native AITER attention.
67
+ - pytest -v -s v1/kv_connector/nixl_integration/test_accuracy_rocm.py
68
69
- label: ":nvidia: (L4) Push NixlConnector PP prefill PD accuracy"
70
key: push-nixlconnector-pp-prefill-pd-accuracy-4-gpus
.buildkite/test_areas/disaggregated_mooncake.yaml
+24
@@ -19,6 +19,30 @@ steps:
19
commands:
20
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
21
- bash v1/kv_connector/mooncake_integration/config_sweep_accuracy_test.sh
22
+ mirror:
23
+ amd:
24
+ label: ":amd: (MI355) Distributed MooncakeConnector HIP PD accuracy"
25
+ dind: false
26
+ device: mi355_4
27
+ num_devices: 4
28
+ timeout_in_minutes: 30
29
+ working_dir: "/vllm-workspace/tests"
30
+ depends_on:
31
+ - image-build-amd
32
+ source_file_dependencies:
33
+ - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/
34
+ - tests/v1/kv_connector/mooncake_integration/
35
+ - tests/v1/kv_connector/rocm_pd_accuracy_utils.py
36
+ - tests/utils.py
37
+ - examples/disaggregated/mooncake_connector/
38
+ - vllm/v1/attention/backends/rocm_aiter_unified_attn.py
39
+ - vllm/platforms/rocm.py
40
+ - requirements/common.txt
41
+ - requirements/kv_connectors_rocm.txt
42
+ commands:
43
+ - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
44
+ # Same full accuracy/TP sweep, using same-host HIP GPU transport.
45
+ - pytest -v -s v1/kv_connector/mooncake_integration/test_accuracy_rocm.py
46
47
- label: ":nvidia: (L4) Mooncake EC TCP E2E"
48
key: mooncake-ec-tcp-e2e-2-gpus
tests/v1/kv_connector/mooncake_integration/test_accuracy_rocm.py
new
+29
@@ -0,0 +1,29 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Mooncake P/D accuracy with same-host HIP transport on ROCm."""
4
+
5
+import pytest
6
+
7
+from tests.utils import spawn_new_process_for_each_test
8
+from tests.v1.kv_connector.rocm_pd_accuracy_utils import (
9
+ QWEN_MODEL,
10
+ run_rocm_pd_accuracy,
11
+)
12
+from vllm.platforms import current_platform
13
+
14
+pytestmark = [
15
+ pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm P/D equivalents"),
16
+ pytest.mark.distributed(num_gpus=4),
17
+]
18
+
19
+
20
+@pytest.mark.parametrize(
21
+ "model,prefill_tp,decode_tp",
22
+ [
23
+ pytest.param(QWEN_MODEL, 1, 1, id="mooncake-hip-qwen-p1-d1"),
24
+ pytest.param(QWEN_MODEL, 2, 2, id="mooncake-hip-qwen-p2-d2"),
25
+ ],
26
+)
27
+@spawn_new_process_for_each_test
28
+def test_rocm_pd_accuracy(model, prefill_tp, decode_tp, monkeypatch):
29
+ run_rocm_pd_accuracy("mooncake", model, prefill_tp, decode_tp, monkeypatch)
tests/v1/kv_connector/nixl_integration/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/v1/kv_connector/nixl_integration/test_accuracy_rocm.py
new
+34
@@ -0,0 +1,34 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""NIXL P/D accuracy with AITER unified attention on ROCm."""
4
+
5
+import pytest
6
+
7
+from tests.utils import spawn_new_process_for_each_test
8
+from tests.v1.kv_connector.rocm_pd_accuracy_utils import (
9
+ DEEPSEEK_MODEL,
10
+ QWEN_MODEL,
11
+ run_rocm_pd_accuracy,
12
+)
13
+from vllm.platforms import current_platform
14
+
15
+pytestmark = [
16
+ pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm P/D equivalents"),
17
+ pytest.mark.distributed(num_gpus=4),
18
+]
19
+
20
+
21
+@pytest.mark.parametrize(
22
+ "model,prefill_tp,decode_tp",
23
+ [
24
+ pytest.param(QWEN_MODEL, 2, 2, id="nixl-aiter-qwen-p2-d2"),
25
+ pytest.param(QWEN_MODEL, 1, 2, id="nixl-aiter-qwen-p1-d2"),
26
+ pytest.param(QWEN_MODEL, 2, 1, id="nixl-aiter-qwen-p2-d1"),
27
+ pytest.param(DEEPSEEK_MODEL, 1, 1, id="nixl-aiter-deepseek-p1-d1"),
28
+ pytest.param(DEEPSEEK_MODEL, 1, 2, id="nixl-aiter-deepseek-p1-d2"),
29
+ pytest.param(DEEPSEEK_MODEL, 2, 1, id="nixl-aiter-deepseek-p2-d1"),
30
+ ],
31
+)
32
+@spawn_new_process_for_each_test
33
+def test_rocm_pd_accuracy(model, prefill_tp, decode_tp, monkeypatch):
34
+ run_rocm_pd_accuracy("nixl", model, prefill_tp, decode_tp, monkeypatch)
tests/v1/kv_connector/rocm_pd_accuracy_utils.py
new
+254
@@ -0,0 +1,254 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Shared ROCm P/D accuracy harness for Mooncake and NIXL integration tests.
4
+
5
+Keep the models, TP pairs, full GSM8K workload, and acceptance checks of the
6
+Mooncake and FlashInfer/NIXL sweeps. Mooncake uses same-host HIP IPC; NIXL uses
7
+AITER unified attention. These cases do not exercise RDMA or FlashInfer.
8
+"""
9
+
10
+import contextlib
11
+import importlib
12
+import json
13
+import os
14
+import signal
15
+import subprocess
16
+import sys
17
+import time
18
+from pathlib import Path
19
+
20
+import requests
21
+from prometheus_client.parser import text_string_to_metric_families
22
+
23
+from tests.utils import RemoteOpenAIServer
24
+from vllm.platforms import current_platform
25
+from vllm.utils.network_utils import get_open_port
26
+
27
+_ROOT = Path(__file__).resolve().parents[3]
28
+QWEN_MODEL = "Qwen/Qwen3-0.6B"
29
+DEEPSEEK_MODEL = "deepseek-ai/deepseek-vl2-tiny"
30
+_REVISIONS = {
31
+ QWEN_MODEL: "c1899de289a04d12100db370d81485cdf75e47ca",
32
+ DEEPSEEK_MODEL: "66c54660eae7e90c9ba259bfdf92d07d6e3ce8aa",
33
+}
34
+
35
+
36
+def _metric_total(server, name: str, **labels: str) -> float:
37
+ response = requests.get(server.url_for("metrics"), timeout=10)
38
+ response.raise_for_status()
39
+ return sum(
40
+ sample.value
41
+ for family in text_string_to_metric_families(response.text)
42
+ for sample in family.samples
43
+ if sample.name == name
44
+ and all(sample.labels.get(key) == value for key, value in labels.items())
45
+ )
46
+
47
+
48
+@contextlib.contextmanager
49
+def _proxy(connector: str, prefill, decode, bootstrap_port: int):
50
+ port = get_open_port()
51
+ if connector == "mooncake":
52
+ script = (
53
+ _ROOT / "examples/disaggregated/mooncake_connector/"
54
+ "mooncake_connector_proxy.py"
55
+ )
56
+ args = [
57
+ "--prefill",
58
+ prefill.url_root,
59
+ str(bootstrap_port),
60
+ "--decode",
61
+ decode.url_root,
62
+ ]
63
+ else:
64
+ script = Path(__file__).parent / "nixl_integration/toy_proxy_server.py"
65
+ args = [
66
+ "--prefiller-hosts",
67
+ "127.0.0.1",
68
+ "--prefiller-ports",
69
+ str(prefill.port),
70
+ "--decoder-hosts",
71
+ "127.0.0.1",
72
+ "--decoder-ports",
73
+ str(decode.port),
74
+ ]
75
+ process = subprocess.Popen(
76
+ [sys.executable, str(script), "--port", str(port), *args],
77
+ start_new_session=True,
78
+ )
79
+ try:
80
+ deadline = time.monotonic() + 60
81
+ while time.monotonic() < deadline:
82
+ assert process.poll() is None, "P/D proxy exited during startup"
83
+ try:
84
+ requests.get(f"http://127.0.0.1:{port}/", timeout=1)
85
+ break
86
+ except requests.ConnectionError:
87
+ time.sleep(0.2)
88
+ else:
89
+ raise TimeoutError("P/D proxy did not start")
90
+ yield f"http://127.0.0.1:{port}/v1"
91
+ finally:
92
+ if process.poll() is None:
93
+ os.killpg(process.pid, signal.SIGTERM)
94
+ try:
95
+ process.wait(timeout=10)
96
+ except subprocess.TimeoutExpired:
97
+ os.killpg(process.pid, signal.SIGKILL)
98
+ process.wait(timeout=10)
99
+
100
+
101
+def run_rocm_pd_accuracy(connector, model, prefill_tp, decode_tp, monkeypatch):
102
+ """Transferred KV must feed decoding while full upstream accuracy holds."""
103
+ # Preserve the caller's allocation instead of querying host-wide SMI indices.
104
+ selectors = ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES")
105
+ selector = next((key for key in selectors if key in os.environ), selectors[0])
106
+ visible = os.getenv(selector)
107
+ devices = (
108
+ visible.split(",")
109
+ if visible is not None
110
+ else [str(i) for i in range(current_platform.device_count())]
111
+ )
112
+ assert len(devices) >= prefill_tp + decode_tp, "Insufficient allocated GPUs"
113
+ for key in selectors:
114
+ if key != selector:
115
+ monkeypatch.delenv(key, raising=False)
116
+
117
+ if connector == "mooncake":
118
+ import mooncake.engine
119
+
120
+ assert mooncake.engine.SUPPORT_HIP, "Mooncake requires its ROCm/HIP build"
121
+ connector_name = "MooncakeConnector"
122
+ extra = {"mooncake_protocol": "hip", "device_name": ""}
123
+ else:
124
+ from vllm.distributed.nixl_utils import NixlWrapper
125
+
126
+ assert NixlWrapper is not None, "NIXL ROCm bindings must import"
127
+ connector_name = "NixlConnector"
128
+ extra = {}
129
+
130
+ common_args = [
131
+ "--revision",
132
+ _REVISIONS[model],
133
+ "--dtype",
134
+ "bfloat16",
135
+ "--enforce-eager",
136
+ "--distributed-executor-backend",
137
+ "mp",
138
+ "--block-size",
139
+ "128",
140
+ "--max-num-seqs",
141
+ "128",
142
+ "--max-num-batched-tokens",
143
+ "4096",
144
+ "--gpu-memory-utilization",
145
+ "0.2",
146
+ "--kv-cache-memory-bytes",
147
+ str(8 * 1024**3),
148
+ "--moe-backend",
149
+ "triton",
150
+ "--attention-backend",
151
+ "ROCM_AITER_UNIFIED_ATTN",
152
+ ]
153
+ common_env = {
154
+ "PATH": str(Path(sys.executable).parent) + os.pathsep + os.environ["PATH"],
155
+ "VLLM_HOST_IP": "127.0.0.1",
156
+ "VLLM_WORKER_MULTIPROC_METHOD": "spawn",
157
+ "VLLM_USE_V2_MODEL_RUNNER": "1",
158
+ "VLLM_ROCM_USE_AITER": "1",
159
+ "VLLM_KV_CACHE_LAYOUT": "HND",
160
+ }
161
+ bootstrap_port = get_open_port()
162
+ with contextlib.ExitStack() as stack:
163
+ servers: list[RemoteOpenAIServer] = []
164
+ stack.callback(RemoteOpenAIServer.shutdown_many, servers)
165
+ for role, tp, assigned in (
166
+ ("kv_producer", prefill_tp, devices[:prefill_tp]),
167
+ ("kv_consumer", decode_tp, devices[prefill_tp : prefill_tp + decode_tp]),
168
+ ):
169
+ config = {
170
+ "kv_connector": connector_name,
171
+ "kv_role": role,
172
+ "kv_connector_extra_config": extra,
173
+ }
174
+ env = {
175
+ **common_env,
176
+ selector: ",".join(assigned),
177
+ "VLLM_PORT": str(get_open_port()),
178
+ "VLLM_NIXL_SIDE_CHANNEL_PORT": str(get_open_port()),
179
+ "VLLM_MOONCAKE_BOOTSTRAP_PORT": str(bootstrap_port),
180
+ }
181
+ servers.append(
182
+ RemoteOpenAIServer(
183
+ model,
184
+ [
185
+ *common_args,
186
+ "--tensor-parallel-size",
187
+ str(tp),
188
+ "--kv-transfer-config",
189
+ json.dumps(config),
190
+ ],
191
+ env_dict=env,
192
+ max_wait_seconds=1200,
193
+ )
194
+ )
195
+ prefill, decode = servers
196
+ proxy_url = stack.enter_context(
197
+ _proxy(connector, prefill, decode, bootstrap_port)
198
+ )
199
+
200
+ # Reuse the full upstream evaluation and its model-specific assertions.
201
+ accuracy = importlib.import_module(
202
+ f"tests.v1.kv_connector.{connector}_integration.test_accuracy"
203
+ )
204
+ monkeypatch.setattr(accuracy, "MODEL_NAME", model)
205
+ monkeypatch.setattr(accuracy, "BASE_URL", proxy_url)
206
+ assert model in accuracy.EXPECTED_VALUES
207
+
208
+ def transfer_totals():
209
+ totals = {
210
+ "external_tokens": _metric_total(
211
+ decode,
212
+ "vllm:prompt_tokens_by_source_total",
213
+ source="external_kv_transfer",
214
+ )
215
+ }
216
+ if connector == "nixl":
217
+ totals["nixl_bytes"] = _metric_total(
218
+ decode, "vllm:nixl_bytes_transferred_sum"
219
+ )
220
+ return totals
221
+
222
+ simple_evaluate = accuracy.lm_eval.simple_evaluate
223
+
224
+ def evaluate_with_transfer_checks(*args, **kwargs):
225
+ # The upstream helper runs a smoke prompt before GSM8K. Wait until
226
+ # its counters appear, then require new transfers during evaluation.
227
+ deadline = time.monotonic() + 15
228
+ before = transfer_totals()
229
+ while not all(value > 0 for value in before.values()):
230
+ assert time.monotonic() < deadline, (
231
+ f"Smoke prompt did not transfer KV: {before}"
232
+ )
233
+ time.sleep(1)
234
+ before = transfer_totals()
235
+
236
+ # The API adapter loads a local tokenizer because these proxies do
237
+ # not expose /tokenize. Match it to the server's pinned checkpoint.
238
+ kwargs["model_args"] += f",revision={_REVISIONS[model]}"
239
+ results = simple_evaluate(*args, **kwargs)
240
+ deadline = time.monotonic() + 15
241
+ after = transfer_totals()
242
+ while not all(after[key] > value for key, value in before.items()):
243
+ assert time.monotonic() < deadline, (
244
+ f"GSM8K did not transfer KV: before={before}, after={after}"
245
+ )
246
+ time.sleep(1)
247
+ after = transfer_totals()
248
+ print(f"GSM8K transfer counters: before={before}, after={after}")
249
+ return results
250
+
251
+ monkeypatch.setattr(
252
+ accuracy.lm_eval, "simple_evaluate", evaluate_with_transfer_checks
253
+ )
254
+ accuracy.test_accuracy()