agentic-kvc/microbench/connector_tax/layerwise/DESIGN.md

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

• worktree + connector backup + design
• modified connector (LAYERWISE.py, +193/-4 lines, env-gated)
• correctness microbench (mb7_layerwise.py) + launcher (run_mb7.sh)
• correctness run on dash0 — PASS (KV lands; cached == prompt)
• 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.