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

MB2: working end-to-end intra-node KV transfer microbench

Browse Source 
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>
This commit is contained in:
Gahow Wang
2026-05-27 18:53:25 +08:00
parent 622e0bc04c
commit 91673f1fb8
3 changed files with 112 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files

133
microbench/fresh_setup/instrument_mooncake.py
View File

@@ -1,26 +1,26 @@
#!/usr/bin/env python3
"""Instrument mooncake_connector_v1.py with per-stage timing for MB2.
"""Instrument vLLM's shipped MooncakeConnector for per-stage MB2 timing.
Adds structured JSONL logging at two key call sites so we can recover
the full breakdown (queueing / prefill / kv transfer / decode) when
combined with the client-side TTFT + finish timestamps:
vLLM 0.18.1 ships its OWN MooncakeConnector at
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py
which is the one --kv-transfer-config '{"kv_connector":"MooncakeConnector"}'
loads. The mooncake-package's mooncake_connector_v1.py is a separate
implementation and not used by vLLM 0.18.1.
send_blocks (P-side, the actual `batch_transfer_sync_write`):
{event=send_blocks, remote_session, total_bytes,
duration_s, t_start_unix, ret, tp_rank}
We add structured JSONL events at two key sites:
receive_kv (D-side, the ZMQ-driven pull request):
{event=receive_kv, path, local_req_ids, remote_req_ids,
duration_s, t_start_unix, tp_rank}
_send_blocks (P-side, the actual batch_transfer_sync_write):
{event=send_blocks, remote_session, total_bytes, duration_s,
t_start_unix, ret, tp_rank, t_log_unix}
Logs go to ``${MB2_LOG_DIR}/mb2_transfer_pid<pid>.jsonl`` (default
``MB2_LOG_DIR=/tmp``). One file per vLLM worker PID so concurrent send +
receive on the same machine don't fight for fsync ordering.
receive_kv_from_single_worker (D-side, async ZMQ-driven pull):
{event=receive_kv_enter / receive_kv_finish, worker_addr,
req_ids, duration_s, t_start_unix, tp_rank, t_log_unix}
Apply / revert pattern follows microbench/connector_tax/cache_sweep/
apply_direct_read_fix.py. All inserted code is bracketed with
``# MB2_INSTRUMENT_START`` / ``# MB2_INSTRUMENT_END`` so revert is a
single text scan.
Logs land in ``$MB2_LOG_DIR/mb2_transfer_pid<pid>.jsonl`` (default
MB2_LOG_DIR=/tmp). One file per PID so multi-worker vllm doesn't fight.
Apply/revert markers: ``# MB2_INSTRUMENT_START/END``.
Usage:
python instrument_mooncake.py --apply [--venv PATH]
@@ -34,11 +34,16 @@ import re
from pathlib import Path
DEFAULT_VENV = Path("/home/admin/cpfs/wjh/agentic-kv-fresh/.venv")
TARGET_REL = "lib/python3.12/site-packages/mooncake/mooncake_connector_v1.py"
TARGET_REL = (
"lib/python3.12/site-packages/vllm/"
"distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py"
)
START_MARK = "# MB2_INSTRUMENT_START"
END_MARK = "# MB2_INSTRUMENT_END"
HEADER_ANCHOR = "import time\n"
HEADER_INSERT = f"""
{START_MARK}
@@ -47,7 +52,13 @@ import os as _mb2_os
import threading as _mb2_threading
import time as _mb2_time
_MB2_LOG_DIR = _mb2_os.environ.get("MB2_LOG_DIR", "/tmp")
_MB2_LOG_PATH = _mb2_os.path.join(_MB2_LOG_DIR, f"mb2_transfer_pid{{_mb2_os.getpid()}}.jsonl")
try:
_mb2_os.makedirs(_MB2_LOG_DIR, exist_ok=True)
except Exception:
pass
_MB2_LOG_PATH = _mb2_os.path.join(
_MB2_LOG_DIR, f"mb2_transfer_pid{{_mb2_os.getpid()}}.jsonl"
)
_MB2_LOG_FILE = None
_MB2_LOG_LOCK = _mb2_threading.Lock()
def _mb2_log_event(d):
@@ -60,8 +71,7 @@ def _mb2_log_event(d):
{END_MARK}
"""
# Insert HEADER right after the `import time` near the top.
HEADER_ANCHOR = "import time\n"
# ---- Patch 2: _send_blocks ----
SEND_TARGET = """ def _send_blocks(
self,
@@ -97,54 +107,60 @@ SEND_REPLACE = f""" def _send_blocks(
"duration_s": time.perf_counter() - start_time,
"t_start_unix": _mb2_t_start_unix,
"ret": ret_value,
"tp_rank": self.tp_rank}})
"tp_rank": getattr(self, "tp_rank", -1)}})
{END_MARK}"""
RECV_TARGET = """ async def receive_kv(
self, path: str, req_blocks: list[tuple[str, str, list[int]]]
):
local_req_ids, remote_req_ids, block_ids = map(list, zip(*req_blocks))"""
# ---- Patch 3: receive_kv_from_single_worker entry ----
RECV_REPLACE = f""" async def receive_kv(
self, path: str, req_blocks: list[tuple[str, str, list[int]]]
RECV_ENTRY_TARGET = """ async def receive_kv_from_single_worker(
self,
worker_addr: str,
pull_metas: dict[ReqId, PullReqMeta],
):
local_req_ids, remote_req_ids, block_ids = map(list, zip(*req_blocks))
req_ids = set(pull_metas)"""
RECV_ENTRY_REPLACE = f""" async def receive_kv_from_single_worker(
self,
worker_addr: str,
pull_metas: dict[ReqId, PullReqMeta],
):
req_ids = set(pull_metas)
{START_MARK}
_mb2_t_recv_start = time.perf_counter()
_mb2_t_recv_start_unix = time.time()
_mb2_recv_t_start = time.perf_counter()
_mb2_recv_t_start_unix = time.time()
_mb2_log_event({{"event": "receive_kv_enter",
"worker_addr": worker_addr,
"req_ids": [str(r) for r in req_ids],
"t_start_unix": _mb2_recv_t_start_unix,
"tp_rank": getattr(self, "tp_rank", -1)}})
{END_MARK}"""
RECV_DONE_TARGET = """ self.finished_recving_reqs.update(local_req_ids)
# ---- Patch 4: receive_kv_from_single_worker FINISH ----
logger.debug(
"pulling kv_caches for %s finished (local requests: %s)",
remote_req_ids, local_req_ids)"""
RECV_FINISH_TARGET = """ if response.status == MooncakeXferResponseStatus.FINISH:
break"""
RECV_DONE_REPLACE = f""" {START_MARK}
_mb2_log_event({{"event": "receive_kv",
"path": path,
"local_req_ids": list(local_req_ids),
"remote_req_ids": list(remote_req_ids),
"duration_s": time.perf_counter() - _mb2_t_recv_start,
"t_start_unix": _mb2_t_recv_start_unix,
"tp_rank": self.tp_rank}})
{END_MARK}
self.finished_recving_reqs.update(local_req_ids)
logger.debug(
"pulling kv_caches for %s finished (local requests: %s)",
remote_req_ids, local_req_ids)"""
RECV_FINISH_REPLACE = f""" if response.status == MooncakeXferResponseStatus.FINISH:
{START_MARK}
_mb2_log_event({{"event": "receive_kv_finish",
"worker_addr": worker_addr,
"req_ids": [str(r) for r in req_ids],
"duration_s": time.perf_counter() - _mb2_recv_t_start,
"t_start_unix": _mb2_recv_t_start_unix,
"tp_rank": getattr(self, "tp_rank", -1)}})
{END_MARK}
break"""
PATCHES = [
("header", HEADER_ANCHOR, HEADER_ANCHOR + HEADER_INSERT),
("send_blocks", SEND_TARGET, SEND_REPLACE),
("receive_kv (entry)", RECV_TARGET, RECV_REPLACE),
("receive_kv (done)", RECV_DONE_TARGET, RECV_DONE_REPLACE),
("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),
]
def find_target(path: Path) -> Path:
candidates = [path, DEFAULT_VENV / TARGET_REL]
def find_target(venv_or_path: Path) -> Path:
candidates = [venv_or_path, DEFAULT_VENV / TARGET_REL]
for c in candidates:
if c.is_file():
return c
@@ -152,7 +168,7 @@ def find_target(path: Path) -> Path:
sub = c / TARGET_REL
if sub.is_file():
return sub
raise FileNotFoundError(f"cannot find mooncake_connector_v1.py near {path}")
raise FileNotFoundError(f"cannot find vllm MooncakeConnector at {venv_or_path}")
def is_patched(text: str) -> bool:
@@ -173,7 +189,7 @@ def apply(target: Path) -> None:
)
new = new.replace(src, dst, 1)
target.write_text(new)
print(f"[mb2-instr] applied 4 patches -> {target}")
print(f"[mb2-instr] applied {len(PATCHES)} patches -> {target}")
def revert(target: Path) -> None:
@@ -181,14 +197,11 @@ def revert(target: Path) -> None:
if not is_patched(text):
print(f"[mb2-instr] not patched (nothing to revert): {target}")
return
# Remove everything between START_MARK and END_MARK, including the markers
# and the surrounding whitespace they introduced.
pat = re.compile(
r"[ \t]*" + re.escape(START_MARK) + r".*?" + re.escape(END_MARK) + r"\n?",
flags=re.DOTALL,
)
new = pat.sub("", text)
# Header insert leaves a leading blank line — collapse runs of 3+ newlines.
new = re.sub(r"\n{3,}", "\n\n", new)
target.write_text(new)
print(f"[mb2-instr] reverted: {target}")