| 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 | with contextlib.ExitStack() as stack: |
| 162 | bootstrap_port = 0 |
| 163 | if connector == "mooncake": |
| 164 | from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_utils import ( # noqa: E501 |
| 165 | MooncakeBootstrapServer, |
| 166 | ) |
| 167 | |
| 168 | bootstrap = MooncakeBootstrapServer("127.0.0.1", 0) |
| 169 | stack.callback(bootstrap.shutdown) |
| 170 | bootstrap.start() |
| 171 | bootstrap_port = bootstrap.port |
| 172 | extra["bootstrap_server_address"] = f"127.0.0.1:{bootstrap_port}" |
| 173 | servers: list[RemoteOpenAIServer] = [] |
| 174 | stack.callback(RemoteOpenAIServer.shutdown_many, servers) |
| 175 | for role, tp, assigned in ( |
| 176 | ("kv_producer", prefill_tp, devices[:prefill_tp]), |
| 177 | ("kv_consumer", decode_tp, devices[prefill_tp : prefill_tp + decode_tp]), |
| 178 | ): |
| 179 | config = { |
| 180 | "kv_connector": connector_name, |
| 181 | "kv_role": role, |
| 182 | "kv_connector_extra_config": extra, |
| 183 | } |
| 184 | env = { |
| 185 | **common_env, |
| 186 | selector: ",".join(assigned), |
| 187 | "VLLM_PORT": str(get_open_port()), |
| 188 | "VLLM_NIXL_SIDE_CHANNEL_PORT": str(get_open_port()), |
| 189 | } |
| 190 | servers.append( |
| 191 | RemoteOpenAIServer( |
| 192 | model, |
| 193 | [ |
| 194 | *common_args, |
| 195 | "--tensor-parallel-size", |
| 196 | str(tp), |
| 197 | "--kv-transfer-config", |
| 198 | json.dumps(config), |
| 199 | ], |
| 200 | env_dict=env, |
| 201 | max_wait_seconds=1200, |
| 202 | ) |
| 203 | ) |
| 204 | prefill, decode = servers |
| 205 | proxy_url = stack.enter_context( |
| 206 | _proxy(connector, prefill, decode, bootstrap_port) |
| 207 | ) |
| 208 | |
| 209 | # Reuse the full upstream evaluation and its model-specific assertions. |
| 210 | accuracy = importlib.import_module( |
| 211 | f"tests.v1.kv_connector.{connector}_integration.test_accuracy" |
| 212 | ) |
| 213 | monkeypatch.setattr(accuracy, "MODEL_NAME", model) |
| 214 | monkeypatch.setattr(accuracy, "BASE_URL", proxy_url) |
| 215 | assert model in accuracy.EXPECTED_VALUES |
| 216 | |
| 217 | def transfer_totals(): |
| 218 | totals = { |
| 219 | "external_tokens": _metric_total( |
| 220 | decode, |
| 221 | "vllm:prompt_tokens_by_source_total", |
| 222 | source="external_kv_transfer", |
| 223 | ) |
| 224 | } |
| 225 | if connector == "nixl": |
| 226 | totals["nixl_bytes"] = _metric_total( |
| 227 | decode, "vllm:nixl_bytes_transferred_sum" |
| 228 | ) |
| 229 | return totals |
| 230 | |
| 231 | simple_evaluate = accuracy.lm_eval.simple_evaluate |
| 232 | |
| 233 | def evaluate_with_transfer_checks(*args, **kwargs): |
| 234 | # The upstream helper runs a smoke prompt before GSM8K. Wait until |
| 235 | # its counters appear, then require new transfers during evaluation. |
| 236 | deadline = time.monotonic() + 15 |
| 237 | before = transfer_totals() |
| 238 | while not all(value > 0 for value in before.values()): |
| 239 | assert time.monotonic() < deadline, ( |
| 240 | f"Smoke prompt did not transfer KV: {before}" |
| 241 | ) |
| 242 | time.sleep(1) |
| 243 | before = transfer_totals() |
| 244 | |
| 245 | # The API adapter loads a local tokenizer because these proxies do |
| 246 | # not expose /tokenize. Match it to the server's pinned checkpoint. |
| 247 | kwargs["model_args"] += f",revision={_REVISIONS[model]}" |
| 248 | results = simple_evaluate(*args, **kwargs) |
| 249 | deadline = time.monotonic() + 15 |
| 250 | after = transfer_totals() |
| 251 | while not all(after[key] > value for key, value in before.items()): |
| 252 | assert time.monotonic() < deadline, ( |
| 253 | f"GSM8K did not transfer KV: before={before}, after={after}" |
| 254 | ) |
| 255 | time.sleep(1) |
| 256 | after = transfer_totals() |
| 257 | print(f"GSM8K transfer counters: before={before}, after={after}") |
| 258 | return results |
| 259 | |
| 260 | monkeypatch.setattr( |
| 261 | accuracy.lm_eval, "simple_evaluate", evaluate_with_transfer_checks |
| 262 | ) |
| 263 | accuracy.test_accuracy() |