Layerwise KV transfer on Mooncake: PoC + microbench (worktree exploration)

Implements per-layer KV push during prefill (write mode) on vLLM's
MooncakeConnector, env-gated by MOONCAKE_LAYERWISE=1. 2-instance microbench
(mb7) shows correctness (KV lands, cached==prompt) and that the transfer is
hidden behind prefill compute: critical-path overhead drops from O(KV size)
(123/202/529ms for 8k/16k/32k) to a flat ~58ms (2-9x), with no prefill
slowdown, on idle instances. Caveats: idle-only, chunked-prefill disabled,
single concurrent transfer — see DESIGN.md.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

microbench/connector_tax/layerwise/DESIGN.md

@@ -0,0 +1,124 @@
+# Layer-wise KV transfer on Mooncake — exploration
+
+Goal: make vLLM's `MooncakeConnector` push KV **per-layer during prefill**
+(write mode) instead of the current **post-hoc full-request transfer**, then
+microbench correctness + whether it hides the transfer behind prefill compute
+(the thing MoRIIO's write mode does on AMD; no NVIDIA connector ships it).
+
+Everything here is isolated in worktree `worktree-mooncake-layerwise`. The
+dash0 venv connector is backed up at `mooncake_connector.py.ORIG_BACKUP`;
+revert = copy the backup back. Opt-in via env `MOONCAKE_LAYERWISE=1`, so with
+the env unset the connector behaves exactly as upstream.
+
+## Baseline flow (post-hoc, what we have)
+
+1. Proxy: prefill on src (`do_remote_decode`, max_tokens=1) → **await done** →
+   decode on dst (`do_remote_prefill`) which pulls.
+2. dst `start_load_kv`→`receive_kv` sends ZMQ `MooncakeXferMetadata` (its block
+   addrs) to src bootstrap.
+3. src `send_kv_to_decode`: waits `send_meta.ready` (set at `request_finished`,
+   i.e. **after full prefill**) → `_build_transfer_params` (all layers) →
+   `_send_blocks` (one big `batch_transfer_sync_write`) → FINISH response.
+
+Measured: this full transfer is on the critical path, runs at ~3 GB/s under
+load (vs ~10 GB/s idle), dominating migration TTFT.
+
+## Layer-wise flow (write mode, this exploration)
+
+Key idea: keep all RDMA + completion on the `sender_loop` thread (clean), but
+issue **one `batch_transfer_sync_write` per layer**, each fired as soon as that
+layer's KV is computed — so writes overlap the remaining prefill compute.
+
+Signaling: `save_kv_layer(layer_name, ...)` (called by vLLM's attention hook
+after each layer's forward, on the main worker thread) records "layer L
+computed" and wakes the sender_loop. `send_kv_to_decode` loops L=0..N-1,
+waits until L is computed, writes layer L's blocks, then sends FINISH.
+
+### Edits to `mooncake_connector.py` (all gated by `_lw_enabled`)
+
+1. **Worker `__init__`**: `_lw_enabled` (env), layer-name→position map,
+   `_lw_computed: dict[transfer_id,int]`, `_lw_active: set[transfer_id]`,
+   wake event, lock.
+2. **`register_kv_caches`**: build `_lw_layer_pos[layer_name]` (0..N-1) and
+   `_lw_addr_idx[pos]` = indices into `kv_caches_base_addr` (×2 if
+   `split_k_and_v`).
+3. **Scheduler `update_state_after_alloc`** (`do_remote_decode` branch): in
+   layer-wise mode capture `blocks.get_block_ids()[0]` and store non-empty in
+   `_reqs_need_send` so the worker learns local block_ids + sets `ready`
+   **before** prefill finishes.
+4. **Worker `note_layer_computed(layer_name)`** (new) called from
+   `MooncakeConnector.save_kv_layer`: bump `_lw_computed[tid]` for active
+   producers, `call_soon_threadsafe(wake.set)`.
+5. **Worker `send_kv_to_decode`**: in layer-wise mode, mark transfer active,
+   loop layers: await `_lw_computed[tid] >= L`, `_send_blocks` for layer L
+   only (subset of `_build_transfer_params`), then send FINISH.
+6. **Worker `_build_layer_transfer_params`** (new): like
+   `_build_transfer_params` but only the addr indices for one layer position.
+
+### Microbench requirements
+
+- Disable chunked prefill (`--max-num-batched-tokens` ≥ prompt) so prefill is a
+  single forward and `save_kv_layer` fires once per layer in order.
+- Dispatch the dst (`do_remote_prefill`) request **first/concurrently** so the
+  ZMQ handshake reaches src during prefill.
+- Correctness: dst follow-up `cached_tokens == prompt_len` (KV landed),
+  identical to baseline.
+- Perf: src prefill wall-clock (does layer-wise slow it?) and dst TTFT (does
+  transfer leave the critical path?), swept over KV size, vs baseline.
+
+## Status
+
+- [x] worktree + connector backup + design
+- [x] modified connector (LAYERWISE.py, +193/-4 lines, env-gated)
+- [x] correctness microbench (mb7_layerwise.py) + launcher (run_mb7.sh)
+- [x] correctness run on dash0 — PASS (KV lands; cached == prompt)
+- [x] perf run + verdict — POSITIVE (transfer hidden behind prefill)
+
+## Results (2-instance, idle, chunked-prefill off, Qwen3-30B-A3B, 48 layers)
+
+Metric: `overhead = total − prefill_only` = the transfer cost left on the
+critical path (TTFT). Baseline = post-hoc full pull (sequential).
+
+| KV size | baseline overhead | **layerwise overhead** | reduction |
+|--------:|------------------:|-----------------------:|----------:|
+| 8192 (0.75 GiB)  | 123 ms | **58 ms** | 2.1× |
+| 16384 (1.5 GiB)  | 202 ms | **58 ms** | 3.5× |
+| 32768 (3.0 GiB)  | 529 ms | **57 ms** | 9.3× |
+
+Key signatures:
+- **Layerwise overhead is ~constant (~58 ms)** regardless of KV size, while
+  baseline grows O(KV size). The 58 ms is handshake + last-layer tail + 1
+  decode; the bulk transfer is hidden behind prefill compute.
+- **Prefill did NOT slow down**: layerwise `t_A` (575/1495/4440 ms) ==
+  `prefill_only` (574/1492/4440 ms). The concurrent RDMA was "free" on idle
+  GPUs — no measurable HBM contention with prefill compute here.
+- Producer logs confirm the transfer itself took 0.39/0.55/4.37 s (grows with
+  size) yet ran *inside* the prefill window, so it left the critical path.
+- **Correctness PASS**: B's follow-up cached == prompt for all sizes; the
+  48-layer / 96-base-addr (split K&V) per-layer addressing is correct.
+
+## Caveats (why this is a proof-of-concept, not a verdict for production)
+
+1. **Idle instances only.** Real migration happens between *busy* instances.
+   Under load both prefill and transfer slow; transfer (even at ~3 GB/s) is
+   still < prefill for big contexts so it should still hide, but receive-side
+   (B) and HBM contention during prefill are untested here. NEXT: rerun with
+   background load on both A and B.
+2. **Chunked prefill disabled.** The monotonic layer counter assumes one
+   forward, layers in order. Production uses chunked prefill (multi-step),
+   which needs per-(chunk,layer) tracking — not implemented.
+3. **Single concurrent producer transfer.** Global counter; real migration is
+   concurrent. Would need per-transfer state.
+4. **Microbench dispatch.** mb7 fires B then A with a 50 ms head start to get
+   the handshake to A before its forward. The real proxy path
+   (`_handle_combined_pd_sep_v2`) dispatches sequentially and would need the
+   write-mode (concurrent) restructure.
+
+## Verdict
+
+The mechanism **works and delivers the predicted benefit**: layer-wise push
+turns migration's KV-transfer cost from O(KV size) on the critical path into a
+near-constant tail, by overlapping it with prefill compute — exactly what
+MoRIIO's write mode does on AMD, now demonstrated on NVIDIA/Mooncake. Whether
+it flips agentic *migration* to net-positive still depends on the busy-instance
+behavior (caveat 1) and is the next experiment.