gahow/agentic-kvc
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

PD-sep server-side profiling: vLLM patches + per-request breakdown

Browse Source 
Instrumentation patches (microbench/patches/):
  - pd_profile.py: shared event emitter (VLLM_PD_PROFILE_LOG env var)
  - apply_patches.py: idempotent patch installer for mooncake_connector.py
    and scheduler.py, marks insertions with # PD_PROFILE_PATCH
  - analyze_events.py: joins per-process JSONL event logs by transfer_id
    into per-request phase durations

Seven events captured per request:
  D_get_num_matched → P_zmq_received → P_prefill_done →
  P_rdma_start → P_rdma_end → D_recv_complete → D_request_promoted

Driver fix (microbench/lifecycle/driver.py):
  seed_prefix_cache now sends via the proxy URL so P and D both cache
  the seeded prefix with matching block hashes. Previously seeding D
  directly produced different block hashes than the proxy-routed
  measurement requests, making incremental transfer impossible.

Real breakdown (fig_breakdown_real.png, server_breakdown.csv, n=93):
  prefill_compute  620 ms median (95% of overhead)
  rdma_transfer     42 ms median (~71 Gbps effective)
  other overhead    10 ms median (dispatch + params + signal + promote)

Mooncake transfer is NOT the bottleneck. Even with bulk RDMA the
transfer cost is <10% of prefill cost for Qwen3-30B-A3B on H20.
This commit is contained in:
Gahow Wang
2026-05-26 13:59:09 +08:00
parent d76eb02637
commit 72790ae6c1
12 changed files with 1099 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

273
microbench/patches/analyze_events.py Normal file
View File

@@ -0,0 +1,273 @@
#!/usr/bin/env python3
"""
Merge PD-sep event logs from P and D into per-request breakdown.
Reads JSONL event logs produced by the patched mooncake_connector +
scheduler, joins events by request_id / transfer_id, and emits a
per-request CSV with the lifecycle phase durations:
prefill_compute : P_prefill_done.t - <P start time>
(we don't have explicit P start; use min of D_get_num_matched - 0
or use prefill duration ≈ P_zmq_received - P_prefill_done? NO.
P_prefill_done = when prefill finished, blocks ready.
We approximate prefill_compute = P_prefill_done - D_get_num_matched
because D and P receive the request simultaneously from proxy.)
zmq_handshake : P_rdma_start - P_zmq_received
(time from D's pull request reaching P to RDMA write start)
rdma_transfer : P_rdma_end - P_rdma_start
(pure RDMA write duration on P side)
completion_signal : D_recv_complete - P_rdma_end
(RDMA completion event back to D side)
D_promote_delay : D_request_promoted - D_recv_complete
(D scheduler step delay to wake the blocked request)
full_pdsep_overhead: D_request_promoted - D_get_num_matched
(total server-side overhead from request arrival to schedulable)
Usage:
python analyze_events.py --events-dir LOGDIR --out breakdown.csv
"""
import argparse
import json
from collections import defaultdict
from pathlib import Path
import csv
def load_events(paths):
"""Yield event dicts from multiple JSONL files."""
for p in paths:
with open(p) as f:
for line in f:
line = line.strip()
if not line:
continue
try:
yield json.loads(line)
except json.JSONDecodeError:
continue
def group_events(events):
"""Group events by transfer_id (preferred) or req_id."""
by_key = defaultdict(dict) # key -> {event_name: event}
# First pass: figure out req_id <-> transfer_id mapping
req_to_transfer = {}
for ev in events:
tid = ev.get("transfer_id", "")
rid = ev.get("req_id", "")
if tid and rid:
req_to_transfer[rid] = tid
# Second pass: assign each event to its transfer_id key
# Need to re-iterate; load again
return req_to_transfer
def _merge_event(slot: dict, ev: dict) -> None:
"""Add event to per-request slot. For repeating events:
- 'end'/'complete'/'promoted' → keep the LATEST
- everything else → keep the EARLIEST."""
name = ev["event"]
existing = slot.get(name)
if existing is None:
slot[name] = ev
return
is_end_like = any(s in name for s in ("rdma_end", "recv_complete", "promoted", "prefill_done"))
if is_end_like:
if ev["t_ns"] > existing["t_ns"]:
slot[name] = ev
else:
if ev["t_ns"] < existing["t_ns"]:
slot[name] = ev
def build_per_request(event_files):
"""Walk all events, group by transfer_id, compute breakdown."""
# Collect all events into memory (these logs are small enough)
all_events = list(load_events(event_files))
# Map req_id <-> transfer_id
req_to_transfer = {}
for ev in all_events:
tid = ev.get("transfer_id", "")
rid = ev.get("req_id", "")
if tid and rid:
req_to_transfer[rid] = tid
# Pre-pass: assign transfer_ids to P_zmq_received events from their `data.transfer_ids` field
# Also collect P_zmq_received timestamps with their transfer_ids so we can link
# P_rdma_start/end events that happen nearby.
zmq_records = [] # list of (t_ns, [transfer_ids...])
for ev in all_events:
if ev["event"] == "P_zmq_received":
tids = ev.get("data", {}).get("transfer_ids", []) or []
if tids:
# Synthetically tag this event with the first transfer_id for primary key.
ev["transfer_id"] = tids[0]
ev["_tids_in_batch"] = tids
zmq_records.append((ev["t_ns"], tids))
# Sort zmq_records by time so we can binary-search later
zmq_records.sort()
def find_zmq_batch(t_ns):
"""Find the most recent ZMQ batch whose timestamp <= t_ns and within 5 seconds."""
best = None
for ts, tids in zmq_records:
if ts <= t_ns and (t_ns - ts) < 5e9:
best = tids # take the latest qualifying
elif ts > t_ns:
break
return best
# Tag P_rdma_start / P_rdma_end with the transfer_ids from the nearest preceding ZMQ batch
for ev in all_events:
if ev["event"] in ("P_rdma_start", "P_rdma_end") and not ev.get("transfer_id"):
tids = find_zmq_batch(ev["t_ns"])
if tids:
ev["transfer_id"] = tids[0]
ev["_tids_in_batch"] = tids
# Group events by transfer_id
by_xfer = defaultdict(dict)
orphans = []
for ev in all_events:
tid = ev.get("transfer_id", "")
rid = ev.get("req_id", "")
# find key
if tid:
key = tid
elif rid in req_to_transfer:
key = req_to_transfer[rid]
else:
orphans.append(ev)
continue
# Also handle events that belong to multiple transfers in a single batch:
# we replicate the event under each transfer_id key for fan-out.
tids_in_batch = ev.get("_tids_in_batch", [tid] if tid else [])
if len(tids_in_batch) > 1:
for t in tids_in_batch:
_merge_event(by_xfer[t], ev)
else:
_merge_event(by_xfer[key], ev)
print(f"Loaded {len(all_events)} events, grouped into {len(by_xfer)} requests, "
f"{len(orphans)} orphans")
# Build per-request rows
rows = []
for tid, evmap in by_xfer.items():
def t(name):
e = evmap.get(name)
return e["t_ns"] if e else None
def d(name, field, default=None):
e = evmap.get(name)
return e["data"].get(field, default) if e else default
row = {
"transfer_id": tid,
"n_events": len(evmap),
"events_seen": ",".join(sorted(evmap.keys())),
# data fields
"num_local_cached": d("D_get_num_matched", "num_local_cached"),
"prompt_tokens": d("D_get_num_matched", "prompt_tokens"),
"remote_total": d("D_get_num_matched", "remote_total"),
"delta_to_pull": d("D_get_num_matched", "delta_to_pull"),
"num_send_blocks": d("P_prefill_done", "num_send_blocks"),
"num_prompt_tokens_P": d("P_prefill_done", "num_prompt_tokens"),
"rdma_num_ops": d("P_rdma_end", "num_ops"),
"rdma_bytes": d("P_rdma_end", "bytes_total"),
}
# timestamps (ns → ms)
ts = {
"t_D_get_num_matched": t("D_get_num_matched"),
"t_P_prefill_done": t("P_prefill_done"),
"t_P_zmq_received": t("P_zmq_received"),
"t_P_rdma_start": t("P_rdma_start"),
"t_P_rdma_end": t("P_rdma_end"),
"t_D_recv_complete": t("D_recv_complete"),
"t_D_request_promoted": t("D_request_promoted"),
}
row.update(ts)
# phase durations in ms
def dur(a, b):
ta, tb = ts.get(a), ts.get(b)
if ta is None or tb is None:
return None
return (tb - ta) / 1e6
row["d_to_p_dispatch_ms"] = dur("t_D_get_num_matched", "t_P_zmq_received")
row["prefill_compute_ms"] = dur("t_P_zmq_received", "t_P_prefill_done")
row["build_params_ms"] = dur("t_P_prefill_done", "t_P_rdma_start")
row["rdma_transfer_ms"] = dur("t_P_rdma_start", "t_P_rdma_end")
row["completion_sig_ms"] = dur("t_P_rdma_end", "t_D_recv_complete")
row["D_promote_ms"] = dur("t_D_recv_complete", "t_D_request_promoted")
row["full_overhead_ms"] = dur("t_D_get_num_matched", "t_D_request_promoted")
# transfer bandwidth
rt = row["rdma_transfer_ms"]
bytes_ = row["rdma_bytes"]
if rt and bytes_ and rt > 0:
row["rdma_bandwidth_gbps"] = (bytes_ * 8 / (rt / 1000)) / 1e9
else:
row["rdma_bandwidth_gbps"] = None
rows.append(row)
return rows
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--events", nargs="+", required=True,
help="One or more JSONL event log files")
ap.add_argument("--out", default="breakdown.csv")
args = ap.parse_args()
paths = [Path(p) for p in args.events]
for p in paths:
if not p.exists():
raise SystemExit(f"Not found: {p}")
rows = build_per_request(paths)
if not rows:
print("No grouped requests found.")
return
# Write CSV
fieldnames = sorted({k for r in rows for k in r.keys()})
with open(args.out, "w", newline="") as f:
w = csv.DictWriter(f, fieldnames=fieldnames)
w.writeheader()
w.writerows(rows)
print(f"Wrote {len(rows)} request breakdowns → {args.out}")
# Quick summary
complete = [r for r in rows if r.get("full_overhead_ms") is not None]
print(f"\nComplete-event requests: {len(complete)}/{len(rows)}")
if complete:
import statistics as st
for k in ("d_to_p_dispatch_ms", "prefill_compute_ms", "build_params_ms",
"rdma_transfer_ms", "completion_sig_ms", "D_promote_ms",
"full_overhead_ms", "rdma_bandwidth_gbps",
"delta_to_pull", "num_local_cached", "rdma_bytes"):
vals = [r[k] for r in complete if r.get(k) is not None]
if vals:
med = st.median(vals)
print(f" {k:<24} median={med:.1f} n={len(vals)}")
if __name__ == "__main__":
main()

285
microbench/patches/apply_patches.py Normal file
View File

@@ -0,0 +1,285 @@
#!/usr/bin/env python3
"""
Apply (or revert) PD-sep profiling instrumentation to vLLM.
Inserts time.perf_counter_ns() emit calls into mooncake_connector.py at:
- get_num_new_matched_tokens() [D: cache hit, delta]
- update_state_after_alloc() [D: blocks allocated]
- send_kv_to_decode() [P: zmq from D arrived]
- record_send_reqs() [P: prefill ready event set]
- _send_blocks() [P: RDMA start/end + bytes]
- process_pulling_result() [D: RDMA recv complete]
- request_finished() [common: request lifecycle end]
And into scheduler.py at:
- _try_promote_blocked_waiting_request() [D: request promoted]
Marker comment "# PD_PROFILE_PATCH" is added so revert can locate inserts.
Usage:
python apply_patches.py [--apply | --revert] [--vllm-root PATH]
The patches are idempotent: --apply on already-patched code is a no-op.
"""
import argparse
import re
import sys
from pathlib import Path
MARKER = "# PD_PROFILE_PATCH"
# Default location: venv vLLM 0.18.1 on dash0
DEFAULT_VLLM_ROOT = Path.home() / "agentic-kv/.venv/lib/python3.12/site-packages/vllm"
def _find_block(text: str, signature: str) -> tuple[int, int] | None:
"""Find the start/end line numbers of a function or method definition."""
lines = text.splitlines()
for i, line in enumerate(lines):
if signature in line:
# find indent
stripped = line.lstrip()
indent = len(line) - len(stripped)
# find function end: next line with <= indent and not blank
for j in range(i + 1, len(lines)):
next_line = lines[j]
if next_line.strip() == "":
continue
next_indent = len(next_line) - len(next_line.lstrip())
if next_indent <= indent and not next_line.lstrip().startswith("#"):
return i, j
return i, len(lines)
return None
def _insert_after_line(text: str, line_no: int, snippet: str) -> str:
"""Insert snippet after line_no (1-indexed). snippet should not include trailing newline."""
lines = text.splitlines()
lines.insert(line_no, snippet)
return "\n".join(lines) + ("\n" if text.endswith("\n") else "")
def _already_patched(text: str) -> bool:
return MARKER in text
def _revert(text: str) -> str:
"""Remove all lines containing the marker."""
lines = text.splitlines()
out = [l for l in lines if MARKER not in l]
return "\n".join(out) + ("\n" if text.endswith("\n") else "")
# ─── Patch definitions ──────────────────────────────────────────────────────
def patch_mooncake_connector(text: str) -> str:
"""Apply patches to mooncake_connector.py."""
if _already_patched(text):
print(" mooncake_connector.py already patched, skipping")
return text
# 1. Add import at top (after first 'import' block)
import_snippet = (
"from vllm.distributed.kv_transfer.kv_connector.v1.mooncake "
"import _pd_profile as _pdp " + MARKER
)
# Insert after the last 'import' statement near the top
lines = text.splitlines()
last_import_line = 0
for i, line in enumerate(lines[:80]):
if line.startswith("import ") or line.startswith("from "):
last_import_line = i
lines.insert(last_import_line + 1, import_snippet)
text = "\n".join(lines) + "\n"
# 2. Patch get_num_new_matched_tokens (D side, scheduler)
# Inject right before "return count, True" and "return 0, False"
text = re.sub(
r"(\n if count > 0:\n return count, True\n)",
r"\n if count > 0:\n _pdp.emit('D_get_num_matched', "
r"req_id=request.request_id, role='kv_consumer', "
r"num_local_cached=num_computed_tokens, prompt_tokens=len(token_ids), "
r"remote_total=remote_total, delta_to_pull=count) " + MARKER + "\n"
r" return count, True\n",
text, count=1,
)
# 3. Patch update_state_after_alloc (D side) — add emit for blocks allocated
# The function is around line 348. Find "self._reqs_need_recv[" assignment.
text = re.sub(
r"(\n self\._reqs_need_recv\[request_id\] = PullReqMeta\()",
r"\n _pdp.emit('D_alloc_blocks', req_id=request_id, "
r"role='kv_consumer', num_local_blocks=len(local_block_ids), "
r"num_external_tokens=num_external_tokens) " + MARKER +
r"\n self._reqs_need_recv[request_id] = PullReqMeta(",
text, count=1,
)
# 4. Patch send_kv_to_decode entry (P side) — ZMQ message received
text = re.sub(
r"( async def send_kv_to_decode\(\n"
r" self, identity: bytes, sock: zmq\.asyncio\.Socket, meta: MooncakeXferMetadata\n"
r" \):\n)",
r"\1 _pdp.emit('P_zmq_received', role='kv_producer', "
r"num_reqs=len(meta.req_blocks), remote_host=meta.remote_hostname, "
r"transfer_ids=[tid for tid, _ in meta.req_blocks.values()]) " + MARKER + "\n",
text, count=1,
)
# 5. Patch _send_blocks (P side) — wrap RDMA write
# Inject before and after self.engine.batch_transfer_sync_write
text = re.sub(
r"( start_time = time\.perf_counter\(\)\n"
r" ret_value = self\.engine\.batch_transfer_sync_write\(\n"
r" remote_session, src_ptrs, dst_ptrs, lengths\n"
r" \))",
r" _pdp.emit('P_rdma_start', role='kv_producer', "
r"num_ops=len(src_ptrs), bytes_total=sum(lengths), remote=str(remote_session)) " + MARKER + "\n"
r"\1\n"
r" _pdp.emit('P_rdma_end', role='kv_producer', "
r"num_ops=len(src_ptrs), bytes_total=sum(lengths), ret=ret_value) " + MARKER,
text, count=1,
)
# 6. Patch process_pulling_result (D side) — RDMA recv complete (success path)
# Match the specific pattern inside `if pull_meta.pull_tasks_count == 0:`
text = re.sub(
r"(pull_meta\.pull_tasks_count -= 1\n"
r" if pull_meta\.pull_tasks_count == 0:\n"
r")( self\.finished_recving_reqs\.add\(pull_meta\.d_req_id\))",
r"\1 _pdp.emit('D_recv_complete', req_id=pull_meta.d_req_id, "
r"transfer_id=pull_meta.transfer_id, role='kv_consumer') " + MARKER + "\n"
r"\2",
text, count=1,
)
# 7. Patch request_finished (P side) — when prefill blocks are marked ready to send
# Find: self._reqs_need_send[request.request_id] = (request, send_block_ids)
text = re.sub(
r"(\n if delay_free_blocks:\n"
r" self\._reqs_need_send\[request\.request_id\] = \(request, send_block_ids\))",
r"\n if delay_free_blocks:\n"
r" _pdp.emit('P_prefill_done', req_id=request.request_id, "
r"transfer_id=params.get('transfer_id', ''), role='kv_producer', "
r"num_send_blocks=len(send_block_ids), num_prompt_tokens=request.num_prompt_tokens) "
+ MARKER + "\n"
r" self._reqs_need_send[request.request_id] = (request, send_block_ids)",
text, count=1,
)
return text
def patch_scheduler(text: str) -> str:
"""Patch v1/core/sched/scheduler.py for D-side request promotion."""
if _already_patched(text):
print(" scheduler.py already patched, skipping")
return text
# Add import at the top of the file
lines = text.splitlines()
last_import_line = 0
for i, line in enumerate(lines[:100]):
if line.startswith("import ") or line.startswith("from "):
last_import_line = i
import_snippet = (
"try:\n"
" from vllm.distributed.kv_transfer.kv_connector.v1.mooncake import _pd_profile as _pdp\n"
"except Exception:\n"
" class _pdp: # fallback no-op\n"
" @staticmethod\n"
" def emit(*a, **kw): pass\n"
" @staticmethod\n"
" def enabled(): return False\n"
f"{MARKER}"
)
lines.insert(last_import_line + 1, import_snippet)
text = "\n".join(lines) + "\n"
# Patch _update_waiting_for_remote_kv — match exact `request: Request` (no quotes)
text = re.sub(
r"( def _update_waiting_for_remote_kv\(self, request: Request\) -> None:\n)",
r"\1 _pdp.emit('D_request_promoted', req_id=request.request_id, "
r"role='kv_consumer', num_computed_tokens=request.num_computed_tokens) " + MARKER + "\n",
text, count=1,
)
return text
# ─── Driver ────────────────────────────────────────────────────────────────
def apply_to_file(path: Path, patcher) -> bool:
if not path.exists():
print(f" SKIP {path} (not found)")
return False
original = path.read_text()
patched = patcher(original)
if patched == original:
print(f" unchanged: {path}")
return False
path.write_text(patched)
n_marks = patched.count(MARKER)
print(f" patched ({n_marks} marks): {path}")
return True
def revert_file(path: Path) -> bool:
if not path.exists():
return False
original = path.read_text()
reverted = _revert(original)
if reverted == original:
print(f" no marks found: {path}")
return False
path.write_text(reverted)
print(f" reverted: {path}")
return True
def install_profile_module(vllm_root: Path) -> None:
"""Copy pd_profile.py to mooncake/_pd_profile.py inside vLLM."""
src = Path(__file__).parent / "pd_profile.py"
dst = vllm_root / "distributed/kv_transfer/kv_connector/v1/mooncake/_pd_profile.py"
dst.parent.mkdir(parents=True, exist_ok=True)
dst.write_text(src.read_text())
print(f" installed: {dst}")
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--apply", action="store_true")
ap.add_argument("--revert", action="store_true")
ap.add_argument("--vllm-root", type=Path, default=DEFAULT_VLLM_ROOT)
args = ap.parse_args()
if not (args.apply or args.revert) or (args.apply and args.revert):
ap.error("Specify exactly one of --apply or --revert")
root = args.vllm_root
if not root.exists():
print(f"ERROR: vLLM root not found: {root}")
sys.exit(1)
mc_path = root / "distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py"
sched_path = root / "v1/core/sched/scheduler.py"
if args.apply:
print(f"Applying PD-sep profile patches to {root}")
install_profile_module(root)
apply_to_file(mc_path, patch_mooncake_connector)
apply_to_file(sched_path, patch_scheduler)
else:
print(f"Reverting PD-sep profile patches from {root}")
revert_file(mc_path)
revert_file(sched_path)
# also remove the module
prof_module = root / "distributed/kv_transfer/kv_connector/v1/mooncake/_pd_profile.py"
if prof_module.exists():
prof_module.unlink()
print(f" removed: {prof_module}")
if __name__ == "__main__":
main()

57
microbench/patches/pd_profile.py Normal file
View File

@@ -0,0 +1,57 @@
"""
PD-Sep Lifecycle Profiler — patch for mooncake_connector.py + scheduler.py
Activated by env var: VLLM_PD_PROFILE_LOG=/path/to/events.jsonl
Each line is one event:
{"t_ns": <perf_counter_ns>, "event": "<name>", "req_id": "<id>",
"transfer_id": "<id>", "role": "<kv_producer|kv_consumer>",
"data": {...event-specific fields...}}
This module is imported by the patched mooncake_connector.py and scheduler.py
to centralize event emission.
"""
import os
import json
import time
import threading
from typing import Any
_LOG_FILE = os.environ.get("VLLM_PD_PROFILE_LOG", "")
_LOCK = threading.Lock()
_HANDLE = None
_ENABLED = bool(_LOG_FILE)
def _get_handle():
global _HANDLE
if _HANDLE is None and _ENABLED:
_HANDLE = open(_LOG_FILE, "a", buffering=1) # line-buffered
return _HANDLE
def emit(event: str, req_id: str = "", transfer_id: str = "",
role: str = "", **data: Any) -> None:
"""Record a profile event. No-op if VLLM_PD_PROFILE_LOG is not set."""
if not _ENABLED:
return
record = {
"t_ns": time.perf_counter_ns(),
"event": event,
"req_id": str(req_id),
"transfer_id": str(transfer_id),
"role": role,
"data": data,
}
try:
h = _get_handle()
if h:
with _LOCK:
h.write(json.dumps(record, default=str) + "\n")
except Exception:
pass # never let profiling break vLLM
def enabled() -> bool:
return _ENABLED