#!/usr/bin/env python3 """MB2: measure KV transfer time between two vLLM instances over Mooncake. Pattern (adapted from microbench/connector_tax/cache_sweep/smoke_test_migrate_cache.py): 1. Prefill on A: do_remote_decode with max_tokens=1 (A computes & caches KV) 2. Pull to B: do_remote_prefill on B with kv_transfer_params from step 1 (this is the operation that performs the KV transfer) 3. Verify: send a follow-up to B; cached_tokens should equal the prompt length (confirms the KV landed on B) We time step 2 — that gives us E2E "transfer + B's prefill check" latency. By sweeping input_length we trace T_transfer(KV_size). The follow-up step gives us a sanity check (correctness) but isn't timed. """ from __future__ import annotations import argparse import asyncio import json import statistics import time import uuid from pathlib import Path import httpx MODEL_PATH = "/home/admin/cpfs/wjh/models/Qwen/Qwen3-Coder-30B-A3B-Instruct" async def get_engine_id(client: httpx.AsyncClient, host: str, bootstrap_port: int) -> str: """The /query endpoint lives on the mooncake bootstrap port, not the vLLM HTTP serving port.""" r = await client.get(f"http://{host}:{bootstrap_port}/query") r.raise_for_status() data = r.json() return data["0"]["engine_id"] async def completion( client: httpx.AsyncClient, host: str, port: int, prompt_token_ids: list[int], max_tokens: int, kv_transfer_params: dict | None = None, ) -> tuple[float, dict]: payload = { "model": MODEL_PATH, "prompt": prompt_token_ids, "max_tokens": max_tokens, "min_tokens": max_tokens if max_tokens == 1 else 1, "temperature": 0.0, "stream": False, } if kv_transfer_params: payload["kv_transfer_params"] = kv_transfer_params t0 = time.perf_counter() r = await client.post( f"http://{host}:{port}/v1/completions", json=payload, timeout=600.0, ) elapsed_s = time.perf_counter() - t0 r.raise_for_status() return elapsed_s, r.json() def synth_prompt(rng_seed: int, n_tokens: int) -> list[int]: """Deterministic token-id sequence, far enough from special tokens.""" import random rng = random.Random(rng_seed) return [rng.randint(100, 150000) for _ in range(n_tokens)] async def measure_one( client: httpx.AsyncClient, src_host: str, src_port: int, dst_host: str, dst_port: int, src_eid: str, dst_eid: str, src_bootstrap_addr: str, input_tokens: int, rng_seed: int, skip_verify: bool = False, ) -> dict: """Three-step measurement (step 3 is sanity, optional for strict PD-disagg where the dst is a kv_consumer-only instance that cannot serve a plain request). vLLM-shipped MooncakeConnector kv_transfer_params schema (vllm/distributed/.../v1/mooncake/mooncake_connector.py:385): do_remote_decode : {transfer_id} do_remote_prefill : {transfer_id, remote_engine_id, remote_bootstrap_addr} """ prompt = synth_prompt(rng_seed, input_tokens) session = uuid.uuid4().hex transfer_id = uuid.uuid4().hex t_step1_client = time.time() t_prefill_s, prefill_resp = await completion( client, src_host, src_port, prompt, max_tokens=1, kv_transfer_params={ "do_remote_decode": True, "transfer_id": transfer_id, }, ) t_step2_client = time.time() t_transfer_s, pull_resp = await completion( client, dst_host, dst_port, prompt, max_tokens=1, kv_transfer_params={ "do_remote_prefill": True, "transfer_id": transfer_id, "remote_engine_id": src_eid, "remote_bootstrap_addr": src_bootstrap_addr, }, ) t_step2_end_client = time.time() cached_followup = None t_followup_s = None if not skip_verify: t_followup_s, follow_resp = await completion( client, dst_host, dst_port, prompt, max_tokens=1, ) usage = (follow_resp.get("usage") or {}) details = usage.get("prompt_tokens_details") or {} cached_followup = details.get("cached_tokens", 0) or usage.get("cached_tokens", 0) pull_usage = (pull_resp.get("usage") or {}) pull_completion_tokens = pull_usage.get("completion_tokens", 0) ok = pull_completion_tokens >= 1 if not skip_verify and cached_followup is not None: ok = ok and (cached_followup >= input_tokens * 0.9) return { "input_tokens": input_tokens, "session": session, "t_step1_client_unix": t_step1_client, "t_step2_client_unix": t_step2_client, "t_step2_end_unix": t_step2_end_client, "t_prefill_s": t_prefill_s, "t_transfer_s": t_transfer_s, "t_followup_s": t_followup_s, "cached_followup": cached_followup, "pull_completion_tokens": pull_completion_tokens, "ok": ok, } async def main_async(args: argparse.Namespace) -> None: sizes_str = args.sizes sizes = [int(s) for s in sizes_str.split(",")] repeats = args.repeats src_host, dst_host = args.src_host, args.dst_host src_port, dst_port = args.src_port, args.dst_port src_bp, dst_bp = args.src_bp, args.dst_bp limits = httpx.Limits(max_connections=10, max_keepalive_connections=10) async with httpx.AsyncClient(limits=limits, trust_env=False) as client: src_eid = await get_engine_id(client, src_host, src_bp) try: dst_eid = await get_engine_id(client, dst_host, dst_bp) except Exception as exc: print(f"[mb2] dst bootstrap ({dst_host}:{dst_bp}) unreachable ({exc.__class__.__name__}); " f"running in strict-PD mode (dst is kv_consumer, no bootstrap).") dst_eid = None print(f"[mb2] src={src_host}:{src_port} bp={src_bp} eid={src_eid[:16]}...") print(f"[mb2] dst={dst_host}:{dst_port} bp={dst_bp} eid=" f"{(dst_eid[:16] + '...') if dst_eid else 'N/A (kv_consumer)'}") src_bootstrap_addr = f"http://{src_host}:{src_bp}" results = [] for sz in sizes: for r in range(repeats): row = await measure_one( client, src_host, src_port, dst_host, dst_port, src_eid, dst_eid, src_bootstrap_addr=src_bootstrap_addr, input_tokens=sz, rng_seed=sz * 1000 + r, skip_verify=args.skip_verify, ) cached = row.get("cached_followup") cached_str = f"{cached}/{sz}" if cached is not None else "skip" print(f" size={sz:>6} rep={r} " f"transfer={row['t_transfer_s']*1000:7.1f}ms " f"followup_cached={cached_str} " f"ok={row['ok']}") results.append(row) # Summarise per-size summary = [] for sz in sizes: ts = [r["t_transfer_s"] for r in results if r["input_tokens"] == sz and r["ok"]] if not ts: continue summary.append({ "input_tokens": sz, "n_ok": len(ts), "transfer_s_mean": statistics.mean(ts), "transfer_s_p50": statistics.median(ts), "transfer_s_p90": statistics.quantiles(ts, n=10)[-1] if len(ts) >= 10 else max(ts), "transfer_s_min": min(ts), "transfer_s_max": max(ts), }) out = { "model": MODEL_PATH, "kv_bytes_per_token": 98304, "src_host": src_host, "src_port": src_port, "dst_host": dst_host, "dst_port": dst_port, "config_label": args.label, "raw": results, "summary": summary, } Path(args.out).write_text(json.dumps(out, indent=2)) print(f"[mb2] wrote {args.out}") for s in summary: sz = s["input_tokens"] kv_mib = sz * 98304 / 1024 / 1024 print(f" {sz:>6} tok ({kv_mib:>7.1f} MiB KV): " f"mean {s['transfer_s_mean']*1000:7.1f} ms · " f"p50 {s['transfer_s_p50']*1000:7.1f} · " f"p90 {s['transfer_s_p90']*1000:7.1f} " f"(n_ok={s['n_ok']})") def main() -> None: p = argparse.ArgumentParser() p.add_argument("--src-host", default="127.0.0.1", help="hostname/IP of the producer (A); use a routable " "address for inter-node tests") p.add_argument("--dst-host", default="127.0.0.1", help="hostname/IP of the consumer (B)") p.add_argument("--src-port", type=int, default=8000, help="vLLM HTTP port on the producer side (A)") p.add_argument("--dst-port", type=int, default=8001, help="vLLM HTTP port on the consumer side (B)") p.add_argument("--src-bp", type=int, default=8998, help="Mooncake bootstrap port on A (serves /query)") p.add_argument("--dst-bp", type=int, default=8999, help="Mooncake bootstrap port on B (serves /query)") p.add_argument( "--sizes", default="512,1024,2048,4096,8192,16384,32768,65536", help="Comma-separated input_token sizes to sweep", ) p.add_argument("--repeats", type=int, default=5) p.add_argument("--label", default="intra-node", help="Label written into the output (e.g. intra-node / inter-node)") p.add_argument("--out", default="mb2_result.json") p.add_argument("--skip-verify", action="store_true", help="Skip the step-3 verify completion (required for " "strict PD-disagg where dst is kv_consumer-only).") args = p.parse_args() asyncio.run(main_async(args)) if __name__ == "__main__": main()