Gahow Wang
fec50fa45d
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>
125 lines
6.5 KiB
Markdown
125 lines
6.5 KiB
Markdown
# 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.
|