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

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>

microbench/fresh_setup/mb2_kv_transfer.py

@@ -29,8 +29,10 @@ import httpx
 MODEL_PATH = "/home/admin/cpfs/wjh/models/Qwen/Qwen3-Coder-30B-A3B-Instruct"
 
 
-async def get_engine_id(client: httpx.AsyncClient, port: int) -> str:
-    r = await client.get(f"http://127.0.0.1:{port}/query")
+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"]
@@ -74,36 +76,39 @@ 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)."""
+    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,
-            "remote_block_ids": None,
-            "remote_engine_id": src_eid,
-            "remote_host": "127.0.0.1",
-            "remote_port": src_port,
+            "transfer_id": transfer_id,
         },
     )
-    src_kvp = prefill_resp.get("kv_transfer_params") or {}
     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,
-            "remote_block_ids": src_kvp.get("remote_block_ids"),
+            "transfer_id": transfer_id,
             "remote_engine_id": src_eid,
-            "remote_host": "127.0.0.1",
-            "remote_port": src_kvp.get("remote_port", src_port),
+            "remote_bootstrap_addr": src_bootstrap_addr,
         },
     )
     t_step2_end_client = time.time()
@@ -144,18 +149,28 @@ async def main_async(args: argparse.Namespace) -> None:
     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_port)
-        dst_eid = await get_engine_id(client, dst_port)
-        print(f"[mb2] src_eid={src_eid[:16]}...  dst_eid={dst_eid[:16]}...")
+        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,
                 )
@@ -206,8 +221,14 @@ async def main_async(args: argparse.Namespace) -> None:
 
 def main() -> None:
     p = argparse.ArgumentParser()
-    p.add_argument("--src-port", type=int, default=8000)
-    p.add_argument("--dst-port", type=int, default=8001)
+    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",