Migration transfer-cost study: KV transfer is slow on busy GPUs

MIGRATION_TRANSFER_COST.md: under real load, migration KV transfer runs at
~3 GB/s vs ~10 GB/s idle. Decomposed (instruments + MB6 microbench) into
~55% RDMA-actual (HBM/PCIe contention with running kernels: 7.6->4.0 GB/s)
+ ~45% control-plane GIL starvation during long prefills. Reproduced on a
fresh upstream venv (byte-identical transfer path) -> upstream/hardware
inherent, not our patch. Layerwise is the wrong lever; the tax is structural
on a loaded agentic cluster. Includes mb6_transfer_under_load + run_mb6,
instrument_dst_migration/mooncake, and the dst/transfer decomposition analyzers.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-05-29 11:53:01 +08:00
parent 67fcec7933
commit 1262c9c22e
11 changed files with 2055 additions and 0 deletions

View File

@@ -0,0 +1,198 @@
"""SLO-goodput analyzer + PD_advantage for the PD-disagg crossover study.
Reads per-arm replayer output (replay_metrics.jsonl) and computes, per arm:
- completion rate (error-free fraction)
- raw TTFT / TPOT / E2E percentiles (over successes — reported for context
only; NEVER the verdict metric, since failing arms have a small success set)
- SLO-goodput: fraction of OFFERED requests that are error-free AND meet a
(TTFT, TPOT) SLO. This is the verdict metric.
The two arms must replay the IDENTICAL trace (same seed), so they are paired
request-for-request. PD_advantage = goodput(arm) / goodput(baseline); y=1 is
the crossover line — PD_advantage >= 1 means PD-disagg wins.
Goodput is computed over a grid of SLO thresholds so the conclusion does not
hinge on one arbitrary cutoff.
Usage:
python analyze_goodput.py \
--arm 8C-proxy .../8C-proxy/replay_metrics.jsonl \
--arm 4P+4D .../4P+4D/replay_metrics.jsonl \
--baseline 8C-proxy \
--ttft-slo 0.5 1 2 5 --tpot-slo 0.05 0.1 0.2
"""
from __future__ import annotations
import argparse
import json
import statistics
from pathlib import Path
def load_metrics(path: Path) -> list[dict]:
rows = []
with path.open("r", encoding="utf-8") as fh:
for line in fh:
line = line.strip()
if line:
rows.append(json.loads(line))
return rows
def percentile(sorted_vals: list[float], pct: float) -> float:
n = len(sorted_vals)
if n == 0:
return float("nan")
if n == 1:
return sorted_vals[0]
rank = pct * (n - 1)
lo = int(rank)
hi = min(lo + 1, n - 1)
frac = rank - lo
return sorted_vals[lo] * (1 - frac) + sorted_vals[hi] * frac
def pstats(vals: list[float]) -> dict:
clean = sorted(v for v in vals if v is not None)
if not clean:
return {"count": 0}
return {
"count": len(clean),
"mean": statistics.fmean(clean),
"p50": percentile(clean, 0.50),
"p90": percentile(clean, 0.90),
"p99": percentile(clean, 0.99),
}
def offered_window_s(rows: list[dict]) -> float:
ts = [r.get("trace_timestamp_s") for r in rows if r.get("trace_timestamp_s") is not None]
if len(ts) < 2:
return 0.0
return max(ts) - min(ts)
def meets_slo(r: dict, ttft_slo: float, tpot_slo: float) -> bool:
if r.get("error") is not None:
return False
ttft = r.get("ttft_s")
tpot = r.get("tpot_s")
if ttft is None:
return False
if ttft > ttft_slo:
return False
# tpot=0 happens only for single-token outputs; treat as meeting any SLO.
if tpot is not None and tpot > tpot_slo:
return False
return True
def load_summary(jsonl_path: Path) -> dict:
"""Read the sibling replay_metrics.summary.json (wall-clock, amplification)."""
sp = jsonl_path.with_suffix(".summary.json")
if sp.exists():
try:
return json.loads(sp.read_text())
except Exception:
return {}
return {}
def summarize_arm(name: str, jsonl_path: Path, rows: list[dict]) -> dict:
n = len(rows)
ok = [r for r in rows if r.get("error") is None]
window = offered_window_s(rows)
summ = load_summary(jsonl_path)
return {
"name": name,
"n_offered": n,
"n_success": len(ok),
"completion_rate": len(ok) / n if n else 0.0,
"offered_window_s": window,
"offered_qps": n / window if window > 0 else 0.0,
# Throughput: how much longer than the offered window it took to drain.
# ~1.0 = keeps up; >1 = falling behind (the cleanest PD-collapse signal).
"wall_clock_s": summ.get("wall_clock_s"),
"amplification": summ.get("amplification"),
"ttft": pstats([r.get("ttft_s") for r in ok]),
"tpot": pstats([r.get("tpot_s") for r in ok]),
"e2e": pstats([r.get("latency_s") for r in ok]),
"_rows": rows,
}
def main() -> None:
p = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--arm", nargs=2, action="append", metavar=("NAME", "PATH"),
required=True, help="arm name + replay_metrics.jsonl path (repeatable)")
p.add_argument("--baseline", required=True, help="arm name to use as PD_advantage denominator")
p.add_argument("--ttft-slo", nargs="+", type=float, default=[0.5, 1.0, 2.0, 5.0])
p.add_argument("--tpot-slo", nargs="+", type=float, default=[0.05, 0.1, 0.2])
p.add_argument("--out-json", type=Path, default=None)
args = p.parse_args()
arms = {}
for name, path in args.arm:
arms[name] = summarize_arm(name, Path(path), load_metrics(Path(path)))
if args.baseline not in arms:
raise SystemExit(f"baseline {args.baseline!r} not among arms {list(arms)}")
# ---- per-arm overview ------------------------------------------------
print("=" * 78)
print("PER-ARM OVERVIEW (latency stats over successes only — context, not verdict)")
print("=" * 78)
hdr = f"{'arm':<12}{'offered':>8}{'compl%':>8}{'ampl':>6}{'oQPS':>7}" \
f"{'TTFTp50':>9}{'TTFTp90':>9}{'TPOTp50':>9}{'TPOTp99':>9}{'E2Ep90':>9}"
print(hdr)
for name, a in arms.items():
t, tp, e = a["ttft"], a["tpot"], a["e2e"]
ampl = a.get("amplification")
ampl_s = f"{ampl:>6.2f}" if isinstance(ampl, (int, float)) else f"{'--':>6}"
print(f"{name:<12}{a['n_offered']:>8}{100*a['completion_rate']:>7.1f}%"
f"{ampl_s}{a['offered_qps']:>7.2f}"
f"{t.get('p50', float('nan')):>9.2f}{t.get('p90', float('nan')):>9.2f}"
f"{1000*tp.get('p50', float('nan')):>8.0f}m{1000*tp.get('p99', float('nan')):>8.0f}m"
f"{e.get('p90', float('nan')):>9.2f}")
# ---- SLO-goodput grid + PD_advantage --------------------------------
base = arms[args.baseline]
grid = []
print()
print("=" * 78)
print(f"SLO-GOODPUT (attainment = error-free AND TTFT<=slo AND TPOT<=slo)")
print(f"PD_advantage = attainment(arm) / attainment(baseline={args.baseline}); "
f">=1 means arm wins")
print("=" * 78)
for ttft_slo in args.ttft_slo:
for tpot_slo in args.tpot_slo:
row = {"ttft_slo_s": ttft_slo, "tpot_slo_s": tpot_slo, "arms": {}}
base_n = sum(1 for r in base["_rows"] if meets_slo(r, ttft_slo, tpot_slo))
base_att = base_n / base["n_offered"] if base["n_offered"] else 0.0
line = f"TTFT<={ttft_slo:>4}s TPOT<={int(1000*tpot_slo):>4}ms | "
cells = []
for name, a in arms.items():
n_slo = sum(1 for r in a["_rows"] if meets_slo(r, ttft_slo, tpot_slo))
att = n_slo / a["n_offered"] if a["n_offered"] else 0.0
adv = (att / base_att) if base_att > 0 else float("nan")
row["arms"][name] = {"attainment": att, "pd_advantage": adv, "n_slo": n_slo}
tag = "" if name == args.baseline else f" adv={adv:.2f}"
cells.append(f"{name}={100*att:>5.1f}%{tag}")
print(line + " ".join(cells))
grid.append(row)
if args.out_json:
out = {
"baseline": args.baseline,
"arms": {n: {k: v for k, v in a.items() if k != "_rows"}
for n, a in arms.items()},
"slo_grid": grid,
}
args.out_json.write_text(json.dumps(out, indent=2))
print(f"\nwrote {args.out_json}")
if __name__ == "__main__":
main()

View File

@@ -0,0 +1,360 @@
#!/usr/bin/env python3
"""Instrument vLLM V1 scheduler to dump per-request DST-side migration timeline.
For each request that arrives at the engine with `kv_transfer_params`
containing `do_remote_prefill=True` (i.e., the decode-target of an
EAR v3 migration), record:
t_arrival_unix — Scheduler.add_request() entry
t_wait_for_kvs_unix — status set to WAITING_FOR_REMOTE_KVS (KV pull start)
t_kv_recv_done_unix — req_id added to finished_recving_kv_req_ids
t_first_scheduled_unix — first time req appears in self.running after KV done
t_first_token_unix — first new_token_ids appended in update_from_output
arrival_state — {n_running, n_waiting, pending_prefill_tok,
n_waiting_for_kvs}
We complement the proxy `breakdown.json` (t_decode_sent_unix /
t_first_token_unix) to attribute the migration's dst-side wait into:
HTTP relay + admission_pre_kv + KV pull + admission_post_kv + first_iter
One JSONL per EngineCore PID at $DM_LOG_DIR/dm_mig_pid<pid>.jsonl
(default DM_LOG_DIR=/tmp). Records are flushed when t_first_token is
reached or when the request is aborted/finished.
Co-exists with MB5 KV snapshot patches (different START/END markers).
Usage:
python instrument_dst_migration.py --apply [--venv PATH]
python instrument_dst_migration.py --revert [--venv PATH]
python instrument_dst_migration.py --check [--venv PATH]
"""
from __future__ import annotations
import argparse
import re
from pathlib import Path
DEFAULT_VENV = Path("/home/admin/cpfs/wjh/agentic-kv/.venv")
TARGET_REL = "lib/python3.12/site-packages/vllm/v1/core/sched/scheduler.py"
START_MARK = "# DM_INSTRUMENT_START"
END_MARK = "# DM_INSTRUMENT_END"
# ---------- Patch 1: module-level header (helpers + globals) -----------------
# Anchor: the very first `class Scheduler(SchedulerInterface):` line. We insert
# the entire helper block immediately before that, so MB5's prior block (if
# present) is preserved and our block lives in module scope. The anchor must
# stay outside our own START/END markers so revert() can re-find it.
HEADER_ANCHOR = "class Scheduler(SchedulerInterface):"
HEADER_INSERT = f"""{START_MARK}
import json as _dm_json
import os as _dm_os
import threading as _dm_threading
import time as _dm_time
_DM_LOG_DIR = _dm_os.environ.get("DM_LOG_DIR", "/tmp")
try:
_dm_os.makedirs(_DM_LOG_DIR, exist_ok=True)
except Exception:
pass
_DM_LOG_PATH = _dm_os.path.join(
_DM_LOG_DIR, f"dm_mig_pid{{_dm_os.getpid()}}.jsonl"
)
_DM_LOG_FILE = None
_DM_LOG_LOCK = _dm_threading.Lock()
# req_id -> in-flight record. We pop and flush when t_first_token lands or on
# finish/abort.
_DM_DATA: dict = {{}}
def _dm_write_event(d: dict) -> None:
global _DM_LOG_FILE
if _DM_LOG_FILE is None:
_DM_LOG_FILE = open(_DM_LOG_PATH, "a", buffering=1)
with _DM_LOG_LOCK:
_DM_LOG_FILE.write(_dm_json.dumps(d) + "\\n")
def _dm_is_migrated(request) -> bool:
ktp = getattr(request, "kv_transfer_params", None)
if not isinstance(ktp, dict):
return False
return bool(ktp.get("do_remote_prefill"))
def _dm_snapshot_arrival(scheduler) -> dict:
try:
n_running = len(scheduler.running)
except Exception:
n_running = -1
try:
n_waiting_main = len(scheduler.waiting)
except Exception:
n_waiting_main = -1
try:
n_skipped = len(scheduler.skipped_waiting)
except Exception:
n_skipped = 0
pending_tok = 0
n_kv = 0
try:
from vllm.v1.request import RequestStatus as _RS
for r in list(scheduler.waiting):
try:
if getattr(r, "status", None) == _RS.WAITING_FOR_REMOTE_KVS:
n_kv += 1
npr = int(getattr(r, "num_prompt_tokens", 0))
nct = int(getattr(r, "num_computed_tokens", 0))
pending_tok += max(0, npr - nct)
except Exception:
pass
for r in list(scheduler.skipped_waiting):
try:
if getattr(r, "status", None) == _RS.WAITING_FOR_REMOTE_KVS:
n_kv += 1
except Exception:
pass
except Exception:
pass
return {{
"n_running": int(n_running),
"n_waiting": int(n_waiting_main + n_skipped),
"pending_prefill_tok": int(pending_tok),
"n_waiting_for_kvs": int(n_kv),
}}
def _dm_emit_and_drop(req_id: str, reason: str = "first_token") -> None:
rec = _DM_DATA.pop(req_id, None)
if rec is None:
return
rec["flush_reason"] = reason
rec["t_flush_unix"] = _dm_time.time()
_dm_write_event(rec)
{END_MARK}
"""
# ---------- Patch 2: add_request() hook --------------------------------------
# Right after self.requests[request.request_id] = request (line ~1927) and the
# if self.log_stats: block. Anchor includes the QUEUED record_event line so it
# is uniquely matchable.
ADD_REQUEST_ANCHOR = """ self._enqueue_waiting_request(request)
self.requests[request.request_id] = request
if self.log_stats:
request.record_event(EngineCoreEventType.QUEUED)
"""
ADD_REQUEST_REPLACE = f""" self._enqueue_waiting_request(request)
self.requests[request.request_id] = request
if self.log_stats:
request.record_event(EngineCoreEventType.QUEUED)
{START_MARK}
try:
if _dm_is_migrated(request):
_DM_DATA[request.request_id] = {{
"req_id": str(request.request_id),
"is_migrated": True,
"n_prompt_tokens": int(getattr(request, "num_prompt_tokens", 0)),
"t_arrival_unix": _dm_time.time(),
"t_wait_for_kvs_unix": None,
"t_kv_recv_done_unix": None,
"t_first_scheduled_unix": None,
"t_first_token_unix": None,
"arrival_state": _dm_snapshot_arrival(self),
}}
except Exception:
pass
{END_MARK}
"""
# ---------- Patch 3: WAITING_FOR_REMOTE_KVS transition -----------------------
WAIT_KV_ANCHOR = """ request.status = RequestStatus.WAITING_FOR_REMOTE_KVS
step_skipped_waiting.prepend_request(request)
"""
WAIT_KV_REPLACE = f""" request.status = RequestStatus.WAITING_FOR_REMOTE_KVS
{START_MARK}
try:
_rec = _DM_DATA.get(request.request_id)
if _rec is not None and _rec["t_wait_for_kvs_unix"] is None:
_rec["t_wait_for_kvs_unix"] = _dm_time.time()
except Exception:
pass
{END_MARK}
step_skipped_waiting.prepend_request(request)
"""
# ---------- Patch 4: finished_recving signal ---------------------------------
FINISHED_RECV_ANCHOR = """ if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
self.finished_recving_kv_req_ids.add(req_id)
elif RequestStatus.is_finished(req.status):
"""
FINISHED_RECV_REPLACE = f""" if req.status == RequestStatus.WAITING_FOR_REMOTE_KVS:
self.finished_recving_kv_req_ids.add(req_id)
{START_MARK}
try:
_rec = _DM_DATA.get(req_id)
if _rec is not None and _rec["t_kv_recv_done_unix"] is None:
_rec["t_kv_recv_done_unix"] = _dm_time.time()
except Exception:
pass
{END_MARK}
elif RequestStatus.is_finished(req.status):
"""
# ---------- Patch 5: first scheduled (sweep at end of schedule()) ------------
# Co-exists with the MB5 snapshot inserted at the same location.
SCHED_END_ANCHOR = """ # MB5_INSTRUMENT_START
_mb5_snapshot(self)
# MB5_INSTRUMENT_END
return scheduler_output
"""
SCHED_END_REPLACE = f""" # MB5_INSTRUMENT_START
_mb5_snapshot(self)
# MB5_INSTRUMENT_END
{START_MARK}
try:
if _DM_DATA:
_now_dm = _dm_time.time()
for _r in self.running:
_rec = _DM_DATA.get(_r.request_id)
if _rec is not None and _rec["t_first_scheduled_unix"] is None:
_rec["t_first_scheduled_unix"] = _now_dm
except Exception:
pass
{END_MARK}
return scheduler_output
"""
# ---------- Patch 6: first new token in update_from_output -------------------
FIRST_TOK_ANCHOR = """ # Check for stop and update request status.
if new_token_ids:
new_token_ids, stopped = self._update_request_with_output(
request, new_token_ids
)
"""
FIRST_TOK_REPLACE = f""" # Check for stop and update request status.
if new_token_ids:
{START_MARK}
try:
_rec = _DM_DATA.get(request.request_id)
if _rec is not None and _rec["t_first_token_unix"] is None:
_rec["t_first_token_unix"] = _dm_time.time()
_dm_emit_and_drop(request.request_id, reason="first_token")
except Exception:
pass
{END_MARK}
new_token_ids, stopped = self._update_request_with_output(
request, new_token_ids
)
"""
# ---------- Patch 7: abort/finish — flush partial record ---------------------
FINISH_ANCHOR = """ request.status = finished_status
self._free_request(request, delay_free_blocks=delay_free_blocks)
"""
FINISH_REPLACE = f""" request.status = finished_status
{START_MARK}
try:
if request.request_id in _DM_DATA:
_dm_emit_and_drop(request.request_id, reason="finish_or_abort")
except Exception:
pass
{END_MARK}
self._free_request(request, delay_free_blocks=delay_free_blocks)
"""
PATCHES = [
("header", HEADER_ANCHOR, HEADER_INSERT + HEADER_ANCHOR),
("add_request", ADD_REQUEST_ANCHOR, ADD_REQUEST_REPLACE),
("wait_for_kvs", WAIT_KV_ANCHOR, WAIT_KV_REPLACE),
("finished_recving", FINISHED_RECV_ANCHOR, FINISHED_RECV_REPLACE),
("first_scheduled", SCHED_END_ANCHOR, SCHED_END_REPLACE),
("first_token", FIRST_TOK_ANCHOR, FIRST_TOK_REPLACE),
("finish_flush", FINISH_ANCHOR, FINISH_REPLACE),
]
def find_target(venv_or_path: Path) -> Path:
candidates = [venv_or_path / TARGET_REL, DEFAULT_VENV / TARGET_REL]
for c in candidates:
if c.is_file():
return c
raise FileNotFoundError(f"cannot find {TARGET_REL} under {venv_or_path}")
def is_patched(text: str) -> bool:
return START_MARK in text
def apply(target: Path) -> None:
text = target.read_text()
if is_patched(text):
print(f"[dm-instr] already patched: {target}")
return
new = text
for name, src, dst in PATCHES:
if src not in new:
raise RuntimeError(
f"patch {name!r}: anchor not found in {target}. "
f"Anchor head: {src.splitlines()[0]!r}"
)
new = new.replace(src, dst, 1)
target.write_text(new)
print(f"[dm-instr] applied {len(PATCHES)} patches -> {target}")
def revert(target: Path) -> None:
text = target.read_text()
if not is_patched(text):
print(f"[dm-instr] not patched (nothing to revert): {target}")
return
# Strip our DM_* block, including the trailing newline that
# terminated the END_MARK line. We do NOT collapse other blank-line
# runs (MB5_* whitespace and original spacing between methods are
# preserved).
pat = re.compile(
r"[ \t]*" + re.escape(START_MARK) + r".*?" + re.escape(END_MARK) + r"\n",
flags=re.DOTALL,
)
new = pat.sub("", text)
# The header insert added a leading "# DM_INSTRUMENT_START\n" with
# two trailing blank lines and the anchor; revert removed the block
# plus its trailing newline, leaving one extra blank line before the
# class — harmless. We additionally collapse the very narrow case of
# "\n\n\nclass Scheduler" -> "\n\nclass Scheduler" so revert is
# byte-identical for that anchor.
new = re.sub(r"\n{3,}class Scheduler\(", "\n\nclass Scheduler(", new)
target.write_text(new)
print(f"[dm-instr] reverted: {target}")
def main() -> None:
p = argparse.ArgumentParser()
p.add_argument("--apply", action="store_true")
p.add_argument("--revert", action="store_true")
p.add_argument("--check", action="store_true")
p.add_argument("--venv", type=Path, default=DEFAULT_VENV)
args = p.parse_args()
target = find_target(args.venv)
if args.apply:
apply(target)
elif args.revert:
revert(target)
elif args.check:
state = "PATCHED" if is_patched(target.read_text()) else "CLEAN"
print(f"[dm-instr] {state}: {target}")
else:
p.error("specify --apply / --revert / --check")
if __name__ == "__main__":
main()

View File

@@ -151,11 +151,65 @@ RECV_FINISH_REPLACE = f""" if response.status == MooncakeXfer
{END_MARK}
break"""
# ---- Patch 5: send_kv_to_decode entry (P-side, producer receives pull req) ----
SEND_ENTRY_TARGET = """ async def send_kv_to_decode(
self, identity: bytes, sock: zmq.asyncio.Socket, meta: MooncakeXferMetadata
):
pending_reqs: dict[ReqId, SendBlockMeta] = {}"""
SEND_ENTRY_REPLACE = f""" async def send_kv_to_decode(
self, identity: bytes, sock: zmq.asyncio.Socket, meta: MooncakeXferMetadata
):
pending_reqs: dict[ReqId, SendBlockMeta] = {{}}
{START_MARK}
try:
_mb2_log_event({{"event": "send_kv_to_decode_enter",
"d_req_ids": [str(r) for r in meta.req_blocks],
"t_start_unix": _mb2_time.time(),
"tp_rank": getattr(self, "tp_rank", -1)}})
except Exception:
pass
{END_MARK}"""
# ---- Patch 6: wait_and_ret ready-wait timing (P-side, src KV commit wait) ----
READY_WAIT_TARGET = """ async def wait_and_ret(
d_req_id: ReqId, send_meta: SendBlockMeta
) -> tuple[ReqId, SendBlockMeta]:
await send_meta.ready.wait()
return d_req_id, send_meta"""
READY_WAIT_REPLACE = f""" async def wait_and_ret(
d_req_id: ReqId, send_meta: SendBlockMeta
) -> tuple[ReqId, SendBlockMeta]:
{START_MARK}
_mb2_rw_start = _mb2_time.perf_counter()
_mb2_rw_start_unix = _mb2_time.time()
_mb2_rw_already = send_meta.ready.is_set()
{END_MARK}
await send_meta.ready.wait()
{START_MARK}
try:
_mb2_log_event({{"event": "ready_wait",
"d_req_id": str(d_req_id),
"transfer_id": str(getattr(send_meta, "transfer_id", "")),
"ready_already_set": bool(_mb2_rw_already),
"ready_wait_s": _mb2_time.perf_counter() - _mb2_rw_start,
"t_start_unix": _mb2_rw_start_unix,
"tp_rank": getattr(self, "tp_rank", -1)}})
except Exception:
pass
{END_MARK}
return d_req_id, send_meta"""
PATCHES = [
("header", HEADER_ANCHOR, HEADER_ANCHOR + HEADER_INSERT),
("_send_blocks", SEND_TARGET, SEND_REPLACE),
("receive_kv (entry)", RECV_ENTRY_TARGET, RECV_ENTRY_REPLACE),
("receive_kv (FINISH)", RECV_FINISH_TARGET, RECV_FINISH_REPLACE),
("send_kv (entry)", SEND_ENTRY_TARGET, SEND_ENTRY_REPLACE),
("ready_wait", READY_WAIT_TARGET, READY_WAIT_REPLACE),
]

View File

@@ -0,0 +1,261 @@
#!/usr/bin/env python3
"""MB6: KV-transfer bandwidth vs instance busy-ness.
Confirms the causal hypothesis from the v3 breakdown: the migration
transfer runs far below wire speed because it happens between instances
that are concurrently busy with compute (GIL-starved control plane +
HBM/NIC contention), NOT because of a wire/NIC limit.
Method (reuses the MB2 transfer primitive):
prefill on A (do_remote_decode, max_tokens=1) -> migrate to B
(do_remote_prefill). Time step 2 = the KV transfer.
For each background-load level B in --bg-loads, we hold B concurrent
long-decode streams on BOTH instances to keep them busy, then run
--repeats measured transfers per size. With the MB2 mooncake instrument
applied (MB2_LOG_DIR set), the analyzer can split the e2e transfer into
RDMA-actual (`send_blocks`) vs control-plane.
Expected: bg=0 reproduces MB2 (~10 GB/s); higher bg degrades toward the
~2-3 GB/s seen in the v3 trace.
Usage:
python mb6_transfer_under_load.py \
--src-port 8000 --dst-port 8001 --src-bp 8998 --dst-bp 8999 \
--sizes 16384,65536 --bg-loads 0,8,24 --repeats 4 \
--out mb6_result.json
"""
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"
KV_PER_TOK = 98304 # Qwen3-30B-A3B est bytes/token
def synth_prompt(seed: int, n: int) -> list[int]:
import random
rng = random.Random(seed)
return [rng.randint(100, 150000) for _ in range(n)]
async def get_engine_id(client, host, bp):
r = await client.get(f"http://{host}:{bp}/query")
r.raise_for_status()
return r.json()["0"]["engine_id"]
async def completion(client, host, port, prompt, max_tokens, ktp=None, stream=False):
payload = {
"model": MODEL_PATH, "prompt": prompt, "max_tokens": max_tokens,
"min_tokens": max_tokens if max_tokens == 1 else 1,
"temperature": 0.0, "stream": stream,
}
if ktp:
payload["kv_transfer_params"] = ktp
t0 = time.perf_counter()
if stream:
# consume the stream to keep the instance decoding
async with client.stream("POST", f"http://{host}:{port}/v1/completions",
json=payload, timeout=600.0) as r:
r.raise_for_status()
async for _ in r.aiter_bytes():
pass
return time.perf_counter() - t0, {}
r = await client.post(f"http://{host}:{port}/v1/completions",
json=payload, timeout=600.0)
elapsed = time.perf_counter() - t0
r.raise_for_status()
return elapsed, r.json()
async def num_running(client, host, port) -> int:
"""Read vLLM running-request gauge from /metrics."""
try:
r = await client.get(f"http://{host}:{port}/metrics", timeout=5.0)
for line in r.text.splitlines():
if line.startswith("vllm:num_requests_running"):
return int(float(line.split()[-1]))
except Exception:
pass
return -1
class BackgroundLoad:
"""Maintain N concurrent long-decode streams on a set of (host,port)."""
def __init__(self, client, endpoints, concurrency, prompt_tokens=2000,
out_tokens=6000):
self.client = client
self.endpoints = endpoints
self.concurrency = concurrency
self.prompt_tokens = prompt_tokens
self.out_tokens = out_tokens
self._stop = asyncio.Event()
self._tasks: list[asyncio.Task] = []
async def _worker(self, idx):
host, port = self.endpoints[idx % len(self.endpoints)]
seed = 900000 + idx
while not self._stop.is_set():
prompt = synth_prompt(seed, self.prompt_tokens)
seed += 1
try:
await completion(self.client, host, port, prompt,
max_tokens=self.out_tokens, stream=True)
except Exception:
await asyncio.sleep(0.5)
def start(self):
self._tasks = [asyncio.create_task(self._worker(i))
for i in range(self.concurrency)]
async def stop(self):
self._stop.set()
for t in self._tasks:
t.cancel()
await asyncio.gather(*self._tasks, return_exceptions=True)
self._tasks = []
async def measure_transfer(client, src_host, src_port, dst_host, dst_port,
src_eid, src_bootstrap_addr, input_tokens, seed):
prompt = synth_prompt(seed, input_tokens)
transfer_id = uuid.uuid4().hex
# step 1: prefill on A
await completion(client, src_host, src_port, prompt, max_tokens=1,
ktp={"do_remote_decode": True, "transfer_id": transfer_id})
# step 2: migrate to B (this is the timed transfer)
t_start_unix = time.time()
t_xfer, _ = await completion(
client, dst_host, dst_port, prompt, max_tokens=1,
ktp={"do_remote_prefill": True, "transfer_id": transfer_id,
"remote_engine_id": src_eid,
"remote_bootstrap_addr": src_bootstrap_addr})
return {
"input_tokens": input_tokens,
"t_transfer_s": t_xfer,
"t_step2_start_unix": t_start_unix,
"t_step2_end_unix": time.time(),
"kv_bytes": input_tokens * KV_PER_TOK,
"eff_gbps": input_tokens * KV_PER_TOK / 1e9 / t_xfer if t_xfer > 0 else 0,
}
async def main_async(a):
sizes = [int(s) for s in a.sizes.split(",")]
bg_loads = [int(s) for s in a.bg_loads.split(",")]
src_host, dst_host = a.src_host, a.dst_host
limits = httpx.Limits(max_connections=256, max_keepalive_connections=256)
async with httpx.AsyncClient(limits=limits, trust_env=False) as client:
src_eid = await get_engine_id(client, src_host, a.src_bp)
src_bootstrap_addr = f"http://{src_host}:{a.src_bp}"
print(f"[mb6] src eid={src_eid[:16]}... endpoints A={src_host}:{a.src_port} "
f"B={dst_host}:{a.dst_port}")
endpoints = [(src_host, a.src_port), (dst_host, a.dst_port)]
results = []
for bg in bg_loads:
loader = None
if bg > 0:
loader = BackgroundLoad(client, endpoints, bg,
prompt_tokens=a.bg_prompt,
out_tokens=a.bg_out)
loader.start()
# wait for instances to actually be busy
print(f"[mb6] bg={bg}: ramping background load ...")
for _ in range(40):
await asyncio.sleep(1.0)
na = await num_running(client, src_host, a.src_port)
nb = await num_running(client, dst_host, a.dst_port)
if na >= 1 and nb >= 1:
print(f"[mb6] bg={bg}: busy (A running={na} B running={nb})")
break
else:
print(f"[mb6] bg=0: idle baseline")
# ensure idle
await asyncio.sleep(2.0)
for sz in sizes:
for rep in range(a.repeats):
na = await num_running(client, src_host, a.src_port)
nb = await num_running(client, dst_host, a.dst_port)
row = await measure_transfer(
client, src_host, a.src_port, dst_host, a.dst_port,
src_eid, src_bootstrap_addr, sz, seed=sz * 100 + rep + bg * 7)
row["bg_load"] = bg
row["A_running_at_measure"] = na
row["B_running_at_measure"] = nb
results.append(row)
kv_mib = sz * KV_PER_TOK / 2**20
print(f" bg={bg:>3} size={sz:>6} ({kv_mib:6.0f}MiB) rep={rep} "
f"A_run={na:>2} B_run={nb:>2} "
f"transfer={row['t_transfer_s']*1000:7.0f}ms "
f"eff={row['eff_gbps']:5.2f}GB/s")
if loader:
await loader.stop()
# let the instances drain before next bg level
print(f"[mb6] bg={bg}: draining ...")
for _ in range(60):
await asyncio.sleep(1.0)
na = await num_running(client, src_host, a.src_port)
nb = await num_running(client, dst_host, a.dst_port)
if na <= 0 and nb <= 0:
break
# summary per (bg, size)
print("\n=== summary: effective transfer bandwidth vs background load ===")
print(f"{'bg':>4} {'size':>7} {'n':>3} {'xfer_p50_ms':>12} {'eff_p50_GBps':>13} "
f"{'eff_mean':>9}")
summary = []
for bg in bg_loads:
for sz in sizes:
rs = [r for r in results if r["bg_load"] == bg and r["input_tokens"] == sz]
if not rs:
continue
xfer = sorted(r["t_transfer_s"] for r in rs)
eff = sorted(r["eff_gbps"] for r in rs)
p50x = xfer[len(xfer) // 2]
p50e = eff[len(eff) // 2]
meane = statistics.mean(eff)
summary.append({"bg": bg, "size": sz, "n": len(rs),
"xfer_p50_ms": p50x * 1000,
"eff_p50_gbps": p50e, "eff_mean_gbps": meane})
print(f"{bg:>4} {sz:>7} {len(rs):>3} {p50x*1000:>12.0f} "
f"{p50e:>13.2f} {meane:>9.2f}")
Path(a.out).write_text(json.dumps(
{"model": MODEL_PATH, "kv_bytes_per_token": KV_PER_TOK,
"label": a.label, "raw": results, "summary": summary}, indent=2))
print(f"\n[mb6] wrote {a.out}")
def main():
p = argparse.ArgumentParser()
p.add_argument("--src-host", default="127.0.0.1")
p.add_argument("--dst-host", default="127.0.0.1")
p.add_argument("--src-port", type=int, default=8000)
p.add_argument("--dst-port", type=int, default=8001)
p.add_argument("--src-bp", type=int, default=8998)
p.add_argument("--dst-bp", type=int, default=8999)
p.add_argument("--sizes", default="16384,65536")
p.add_argument("--bg-loads", default="0,8,24")
p.add_argument("--repeats", type=int, default=4)
p.add_argument("--bg-prompt", type=int, default=2000)
p.add_argument("--bg-out", type=int, default=6000)
p.add_argument("--label", default="main-venv")
p.add_argument("--out", default="mb6_result.json")
args = p.parse_args()
asyncio.run(main_async(args))
if __name__ == "__main__":
main()

109
microbench/fresh_setup/run_mb6.sh Executable file
View File

@@ -0,0 +1,109 @@
#!/usr/bin/env bash
# MB6 launcher: 2 vLLM instances (kv_both, Mooncake) + transfer-under-load
# sweep. Parameterized by VENV so it runs on either the patched main venv
# or the fresh upstream venv, to test whether the bandwidth degradation is
# our patch or inherent to upstream mooncake.
#
# Usage:
# VENV=/home/admin/cpfs/wjh/agentic-kv/.venv bash run_mb6.sh # main
# VENV=/home/admin/cpfs/wjh/agentic-kv-fresh/.venv bash run_mb6.sh # fresh
set -uo pipefail
PROJ_DIR="${PROJ_DIR:-/home/admin/cpfs/wjh/agentic-kv}"
VENV="${VENV:-$PROJ_DIR/.venv}"
LABEL="${LABEL:-$(basename $(dirname $VENV))}"
MODEL="${MODEL:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-Coder-30B-A3B-Instruct}"
GPUS="${GPUS:-0 1}"
SIZES="${SIZES:-16384,65536}"
BG_LOADS="${BG_LOADS:-0,8,24}"
REPEATS="${REPEATS:-4}"
DATE="$(date +%Y%m%d_%H%M)"
OUTDIR="${OUTDIR:-$PROJ_DIR/outputs/mb6_${LABEL}_${DATE}}"
PYTHON="$VENV/bin/python"
MC_INSTR="$PROJ_DIR/microbench/fresh_setup/instrument_mooncake.py"
DRIVER="$PROJ_DIR/microbench/fresh_setup/mb6_transfer_under_load.py"
MC_FILE="$VENV/lib/python3.12/site-packages/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py"
mkdir -p "$OUTDIR/logs"
XFER_LOG_DIR="$OUTDIR/xfer_log"; mkdir -p "$XFER_LOG_DIR"
echo "=== MB6 transfer-under-load ($LABEL) ==="
echo "VENV : $VENV"
echo "Out : $OUTDIR"
echo ""
PORTS=(8000 8001); BPS=(8998 8999)
gpu_arr=($GPUS)
cleanup() {
pkill -9 -f "vllm serve" 2>/dev/null || true
pkill -9 -f "EngineCore" 2>/dev/null || true
sleep 4
"$PYTHON" "$MC_INSTR" --venv "$MC_FILE" --revert 2>/dev/null || true
}
trap cleanup EXIT
cleanup
echo "[0] apply MB2 mooncake instrument to $LABEL venv"
"$PYTHON" "$MC_INSTR" --venv "$MC_FILE" --apply
"$PYTHON" "$MC_INSTR" --venv "$MC_FILE" --check
echo "[1] launch 2 instances"
i=0
for gpu in ${gpu_arr[@]:0:2}; do
port=${PORTS[$i]}; bp=${BPS[$i]}; master=$((29600 + i))
PYTHONHASHSEED=42 \
VLLM_MOONCAKE_BOOTSTRAP_PORT=$bp \
MB2_LOG_DIR="$XFER_LOG_DIR" \
CUDA_VISIBLE_DEVICES=$gpu \
MASTER_PORT=$master \
nohup "$VENV/bin/vllm" serve "$MODEL" \
--host 0.0.0.0 --port "$port" \
--tensor-parallel-size 1 --trust-remote-code --enable-prefix-caching \
--dtype auto --gpu-memory-utilization 0.9 --max-model-len 200000 \
--kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_both"}' \
--enable-prompt-tokens-details \
> "$OUTDIR/logs/vllm_${i}_gpu${gpu}.log" 2>&1 &
disown
sleep 2
i=$((i + 1))
done
echo "[2] wait for health"
for i in 0 1; do
port=${PORTS[$i]}; tries=0
while ! curl -sf "http://127.0.0.1:$port/health" >/dev/null 2>&1; do
tries=$((tries + 1))
if [ $tries -gt 180 ]; then echo "FATAL inst_$i not healthy"; exit 1; fi
sleep 2
done
echo " inst_$i ready"
done
# bootstrap /query reachable?
for i in 0 1; do
bp=${BPS[$i]}; tries=0
while ! curl -sf "http://127.0.0.1:$bp/query" >/dev/null 2>&1; do
tries=$((tries + 1))
if [ $tries -gt 60 ]; then echo "WARN bootstrap $bp not ready"; break; fi
sleep 2
done
done
echo "[3] run MB6 driver"
"$PYTHON" "$DRIVER" \
--src-port "${PORTS[0]}" --dst-port "${PORTS[1]}" \
--src-bp "${BPS[0]}" --dst-bp "${BPS[1]}" \
--sizes "$SIZES" --bg-loads "$BG_LOADS" --repeats "$REPEATS" \
--label "$LABEL" --out "$OUTDIR/mb6_result.json" \
2>&1 | tee "$OUTDIR/mb6_run.txt"
echo "[4] teardown + revert"
pkill -9 -f "vllm serve" 2>/dev/null || true
pkill -9 -f "EngineCore" 2>/dev/null || true
sleep 4
"$PYTHON" "$MC_INSTR" --venv "$MC_FILE" --revert
echo ""
echo "Done. Artifacts in $OUTDIR/"
echo " mb6_result.json mb6_run.txt xfer_log/ logs/"