[Perf][KV Connector][Mooncake] Send the KV layout once per peer

[staryxchen]

Overview

MooncakeConnector now sends the decode KV layout to each prefill worker once, instead of on every pull. Later pulls carry only block ids and a layout_id, and decode stops attaching the layout only after prefill echoes that id back. This is the exchange-once, cache-per-peer model NixlConnector already uses for its handshake. Follows the packed-region work in #57952.

Claims

• On DeepSeek-V4-Flash 4P4D (146 registered regions), each later pull is 9184 B smaller per rank: 9500 B → 316 B at ~8k tokens. The saving is the same at 508 and 8008 tokens.
• Prefill-side control-plane time per pull drops from 0.078 ms to 0.014 ms (steady-state P50). The absolute saving is ~0.06 ms, so TTFT does not change (1.10 s at ~8k tokens on both sides).
• Mixed versions stay correct in both directions. A prefill without this change never receives a layout-less pull; a decode without this change keeps today's behavior.
• No accuracy change on gsm8k (Qwen3-0.6B 1P1D; DeepSeek-V4-Flash 4P4D vs. single node).

Validation

Unit tests: 123 passed.

.venv/bin/python -m pytest -v \
  tests/v1/kv_connector/unit/test_mooncake_connector.py \
  tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py

The environment was uv venv --system-site-packages, with the default test model facebook/opt-125m. The cu129 precompiled wheel for the base commit returned 404, so this is not a VLLM_USE_PRECOMPILED=1 install.

New tests cover:

• layout-less pulls after the echo;
• a LAYOUT_MISS resend;
• a prefill that never echoes (old peer);
• a reused layout arriving under a new id;
• LRU eviction.

pre-commit passes: ruff-check, ruff-format, typos, and mypy-3.12 --hook-stage manual.

Version matrix (Qwen2.5-0.5B, greedy, TP=1, RDMA). "Old" is origin/main (4e4d75c459). The MooncakeXferMetadata to <addr> is N bytes (layout|blocks) line is a DEBUG log (VLLM_LOGGING_LEVEL=DEBUG). These runs used an earlier revision that logged the same line once per address at INFO.

Prefill	Decode	PP (P/D)	Decode log	Result
new	new	1/1	1615 B (layout), then 256 B (blocks)	same output, 100% external prefix hits
old	new	1/1	(layout) only, never (blocks)	same output
old	new	2/1	(layout) only for both prefill addresses; no no matching consumer on prefill	same output
new	old	1/1	(old decode has no size log)	same output

The old/new-PP 1/1 and new/old rows were run on the pre-rebase commit; the Mooncake code is identical. The LRU commit does not touch the code those two rows run: an old prefill never executes it, and an old decode sends layout_id=0, which the new prefill never stores.

Bytes and TTFT (DeepSeek-V4-Flash, TP=4, EP, fp8, block 256, 4P4D on one node). Each TTFT row is 8 streamed requests with max_tokens=1, after warmup.

	origin/main	this PR
Metadata per pull, per rank, ~8k tokens	9483 B every pull	9500 B first pull, then 316 B
Prefill control plane, steady-state P50 / P99	0.078 / 0.099 ms (n=44)	0.014 / 0.029 ms (n=24)
TTFT, ~8k tokens	1.092–1.097 s	1.096–1.099 s

The control-plane time was measured from send_kv_to_decode entry to the end of region-length validation, on all 4 ranks, dropping the first pull. The timing code was temporary and is not part of this PR.

Accuracy (gsm8k, exact_match,strict-match, 1319 questions).

Setup	Score
Qwen3-0.6B 1P1D, tests/v1/kv_connector/mooncake_integration/run_accuracy_test.sh	0.422 (expected 0.41 ± 0.05)
DeepSeek-V4-Flash 4P4D	0.951 (1255/1319)
DeepSeek-V4-Flash single node, TP=4, same eager / fp8 / max_num_seqs=8	0.951 (1254/1319)

Both P/D runs had 100% external prefix hits and 0 failed transfers.

Not covered end to end. Prefill side-channel ports come from bind_to_random_port, and decode RPC ports from the transfer engine, so a restarted peer always comes back at a new address. Three paths therefore have unit coverage only:

• the LAYOUT_MISS resend;
• the same address reusing a layout under a new id;
• LRU eviction followed by a resend.

Prefill and decode restarts were run end to end; they re-sent the layout to the new address, with unchanged output.

Details

Why. Every pull sends the whole registered layout in MooncakeXferMetadata: base addresses, block lengths, layer names, group ids and row offsets. On a 146-region model that is ~9 KB per rank per pull, mostly layer names. Prefill then builds and hashes a tuple of all of it on every batch, just to hit _prepared_transfer_regions. Both costs grow with the number of registered regions, not with the request.

Prior art: NixlConnector. NIXL already exchanges the KV layout once per remote engine and keeps the per-request message small:

	NixlConnector (pull mode)	MooncakeConnector, this PR
Data direction	Decode reads from prefill, so decode caches the prefill layout	Prefill writes into decode, so prefill caches the decode layout
When the layout is exchanged	Once per remote engine, in _ensure_handshake; concurrent requests to the same engine share one handshake future	Once per decode worker address, piggybacked on the first pull
Where it is cached	add_remote_agent stores kv_caches_base_addr[engine_id][tp_rank], block lengths, and region metadata	_layout_by_peer, keyed by decode host, port, ranks, and layout_id
Per-request message	RemoteMeta: block ids, host, port, engine_id, request id	req_blocks and layout_id
How a new layout is detected	The cache is keyed by engine_id; a restarted engine has a new id and is handshaken anew, and the replaced engine is cleaned up (_track_remote_engine_replacement)	Each register_kv_caches mints a new layout_id; an unknown id gets LAYOUT_MISS and a resend

Mooncake has no separate handshake round trip; decode simply sends its pull to the prefill side channel. So instead of adding a handshake, this PR makes the first pull carry the layout. Prefill's echo of layout_id then plays the part of NIXL's handshake response: decode treats the layout as delivered only after it sees the echo, so a peer that never echoes keeps getting the full layout. NIXL's compatibility hash (NixlHandshakePayload.compatibility_hash) is out of scope here.

Protocol.

• register_kv_caches draws a random nonzero 63-bit layout_id. A counter would restart at 1 in every process, so a restarted decode at the same address could match a stale prefill entry and be written with the old layout.
• Prefill stores aligned regions by (decode host, port, tp_rank, pp_size, tp_size, layout_id). It stores them on both a fresh alignment and a _prepared_transfer_regions hit. A hit can carry a new id for a layout prefill has already seen.
• Prefill echoes layout_id in MooncakeXferResponse once it holds that layout. Decode omits the layout for a worker address only after seeing the echo.
  • A plain FINISH is not enough. A prefill without this change ignores the new field and answers FINISH. With equal PP it would then reject every layout-less pull. With different PP, allow_partial_layers would align zero regions, transfer nothing, and report success.
• A layout-less pull whose id prefill does not hold gets the new LAYOUT_MISS status. Decode resends with the layout once; a second miss fails the request normally. Only layout-less pulls can receive LAYOUT_MISS, and only a new decode sends those, so an old decode never sees the new enum value.
• _layout_by_peer is an LRU capped at 1024 entries, because every decode re-registration mints a new id. Eviction is always safe: the evicted decode gets LAYOUT_MISS and resends.

Tradeoff. About 130 net lines in the connector, for ~9 KB and ~0.06 ms per pull, with no TTFT change on DeepSeek-V4-Flash. The case for the change is that this overhead is fixed per pull and grows with the region count and layer-name length, so it grows with deeper and multi-cache models. It is not a latency win at today's scale.

Limitations (pre-existing, unchanged here).

• _prepared_transfer_regions is still unbounded. Its key includes decode base addresses, so it grows on every decode restart. Its entries could take the same LRU in a follow-up.
• After a prefill restart, the proxy must be restarted too. Otherwise decode keeps the old engine_id and requests stall instead of failing.
• A pull with empty req_blocks gets no response.

Related PRs.

• [Bugfix][KV Connector] Mooncake: reject mismatched block sizes in the handshake #59592 adds remote_block_size to MooncakeXferMetadata, in code this PR moves into _build_pull_metadata. Whichever lands second should set it on both full and layout-less pulls.
• [KV Connector] Support sleep mode with MooncakeConnector over RDMA #59625 (draft, sleep mode over RDMA) re-registers KV caches on wake-up. Here that mints a new layout_id and clears the caches, so a woken peer re-sends or re-receives the layout. That PR also restructures send_kv_to_decode, so expect a textual conflict.

Duplicate check. No open PR sends the Mooncake layout once or caches it per peer. Searched mooncake layout, mooncake handshake and mooncake metadata with gh pr list --repo vllm-project/vllm --state open --search ....

AI assistance. This PR was written with AI assistance (Cursor). I reviewed every changed line and ran the validation above.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.