Gahow Wang
91673f1fb8
This commit closes the loop on the fresh-venv MB2 path. Three corrections
on top of the previous scaffold made the bench fire successfully on
dash1 GPU 0+1 with kv_both connector roles:
1. Re-target instrumentation patch to vLLM's shipped MooncakeConnector
(vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py).
The mooncake-package's own mooncake_connector_v1.py turned out not to
be the implementation vLLM 0.18.1 loads — the
'{"kv_connector": "MooncakeConnector"}' config picks up the vLLM-shipped
one. Patches go at _send_blocks (P-side) and receive_kv_from_single_worker
(D-side, async, both entry and FINISH branch).
2. /query lives on the mooncake bootstrap port, not the vLLM HTTP port.
Add --src-bp / --dst-bp args; default 8998 / 8999.
3. kv_transfer_params schema for the vanilla connector:
do_remote_decode → {transfer_id}
do_remote_prefill → {transfer_id, remote_engine_id, remote_bootstrap_addr}
where remote_bootstrap_addr must include the http:// scheme. The dash0
smoke_test_migrate_cache.py was written for the patched build, which
used a different field-name set (remote_host, remote_port,
remote_block_ids); those are rejected here.
Also discovered (and worked around): vLLM 0.18.1 with kv_role=kv_consumer
raises AttributeError on `self.bootstrap_server` because that attribute
is only assigned conditionally inside `if not self.is_kv_consumer`. We
sidestep by running kv_both for the microbench — transfer mechanics are
identical (same batch_transfer_sync_write call); the role gate only
affects which request types each instance accepts. For §5 strict PD-disagg
baseline we'll need either to fix this bug or front the pair with a
role-aware proxy.
Sanity smoke (3 sizes × 2 repeats, dash1 GPU 0+1, kv_both intra-node):
input KV-MiB send_blocks_ms (P) receive_kv_ms (D) client_step2_ms
512 48 5–23 7–33 18–91
2048 192 21 23 37
8192 768 85 88 110
=> intra-node bandwidth ~9 GB/s on the actual transfer for 768 MiB,
which is well below NVLink p2p; likely PCIe-staged. Worth verifying.
Next step (in flight): full sweep 512..128k tokens × 5 repeats with
the per-stage analyzer.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
250 lines
9.0 KiB
Python
Executable File
250 lines
9.0 KiB
Python
Executable File
#!/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, 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://127.0.0.1:{bootstrap_port}/query")
|
|
r.raise_for_status()
|
|
data = r.json()
|
|
return data["0"]["engine_id"]
|
|
|
|
|
|
async def completion(
|
|
client: httpx.AsyncClient,
|
|
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://127.0.0.1:{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_port: int, 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_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_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_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_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_bp)
|
|
# In strict PD-disagg only the producer binds /query; consumers don't.
|
|
try:
|
|
dst_eid = await get_engine_id(client, dst_bp)
|
|
except Exception as exc:
|
|
print(f"[mb2] dst bootstrap ({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_eid={src_eid[:16]}... dst_eid="
|
|
f"{(dst_eid[:16] + '...') if dst_eid else 'N/A (kv_consumer)'}")
|
|
|
|
src_bootstrap_addr = f"http://127.0.0.1:{src_bp}"
|
|
results = []
|
|
for sz in sizes:
|
|
for r in range(repeats):
|
|
row = await measure_one(
|
|
client, src_port, 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_port": src_port,
|
|
"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-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()
|