# 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.