[RFC]: Layerwise KV Cache Transfer for MooncakeStore with Session API Optimization

[ClockOfDestiny]

Motivation.

Background

As model context windows continue to scale (128K → 1M+ tokens), the KV cache has become one of the dominant memory and bandwidth consumers in LLM serving. vLLM already supports layerwise KV cache transfer through connectors such as LMCacheConnector (PR #16625) and OffloadingConnector (RFC #19854), which use per-layer hooks to overlap KV cache transfer with model forward computation.

In the distributed KV cache ecosystem, MooncakeStoreConnector provides a shared KV cache backend across GPU nodes via RDMA. However, the current MooncakeStore connector transfers all layers' KV data in a single bulk operation per request — it performs the full load before the first attention layer and submits the full save after the last layer. This design misses the opportunity to pipeline transfer with per-layer computation and suffers from increasing metadata overhead as models grow deeper.

The Key Explosion Problem

A straightforward layerwise transfer approach would extend each KV block key with the layer index to distinguish data belonging to different layers. With N layers and M blocks per layer, this transforms the key space from M keys to N × M keys.

This explosion in key count directly impacts:

• Lookup performance: Each prefix cache hit check must scan N × M candidates per block, linearly increasing latency.
• Master server overhead: Each layer’s block-keys lookup requires a master RPC.
• Backend scheduling cost: The store's internal metadata index grows proportionally.

The same problem has been noted by RFC #33398 (Layerwise KV cache offloading) in its "Further Optimization", where it describes how "individual layer operations generate excessive keys" and proposes block-level aggregation as a future optimization.

Opportunity: Mooncake Session API

MooncakeStore recently introduced a Session API in kvcache-ai/Mooncake#2881, that enables per-layer byte-offset addressing within a block-level page. This allows us to:

1. Keep a single key per KV block — eliminate layer index suffixes, restoring the O(M) key space.
2. Address per-layer data via byte offsets — each block is stored as a contiguous page; reading/writing layer's data is just a byte offset into that page.
3. Simplify metadata management — only M keys need to be stored in lookup tables, scheduled, and tracked for error recovery.

By combining the layerwise pipeline with the Session API, we get both the latency benefits of pipelined transfer and the efficiency of a single-key-per-block addressing model.

Proposed Change.

1. Leverage Existing vLLM Layerwise Infrastructure

vLLM's KV connector framework already provides the necessary infrastructure for layerwise transfer: @maybe_transfer_kv_layer decorator wraps each attention layer's forward pass, calling wait_for_layer_load before computation and save_kv_layer after computation. The MooncakeStore connector simply needs to implement these callbacks.

Our work implements the missing per-layer dispatch inside MooncakeStoreConnector, reusing these framework hooks without any changes to the scheduler, engine, or model execution path.

2. Per-Layer Task Construction and Dispatch

When start_load_kv() is called at the beginning of a scheduler step, we enumerate every active request and build per-layer LayerTransferTask objects offline. These tasks are stored in per-layer arrays and consumed one layer at a time during the subsequent forward pass.

During forward, the @maybe_transfer_kv_layer decorator drives the pipeline:

Layer 0:  wait_for_layer_load(L0)  →  attention forward  →  save_kv_layer(L0)
Layer 1:  wait_for_layer_load(L1)  →  attention forward  →  save_kv_layer(L1)
...
Layer N-1 (last):  ...  →  save_kv_layer(N-1)  →  _wait_for_all_layer_saves()

The recv threads load KV blocks before each layer's forward pass; the send thread saves KV blocks after each layer's forward pass. Coordination uses per-layer threading.Event objects — one per layer for both load and save completion. The last layer's save_kv_layer() call triggers a barrier that waits for all layers' saves to complete before proceeding to the next scheduler step.

3. Mooncake Session API Integration

To solve the key explosion problem, we integrate the Mooncake Session API for per-layer byte-offset addressing.

Block-Level Key Strategy

Instead of encoding the layer index into the key, each KV block uses a single unified key derived from its content hash. The block's data in MooncakeStore is organized as a page covering all layers:

Block page:
├── Layer 0 data  (byte offset 0 to page_size - 1)
├── Layer 1 data  (byte offset page_size to 2*page_size - 1)
├── ...
└── Layer N-1 data

When writing or reading layer L's data, we compute offset = L * page_size + segment_offset and pass it to the range-based API. The layer index disappears from the key entirely.

Session Lifecycle

Phase	Save Path	Load Path
Session Start (step start)	batch_put_session_start(block_keys) — reserve space for all blocks	batch_get_session_start(block_keys) — open read session
Per Layer (layer L forward)	batch_put_from_multi_buffer_ranges(keys, addrs, sizes, offsets@L)	batch_get_into_multi_buffer_ranges(keys, addrs, sizes, offsets@L)
Session End (last layer)	batch_put_session_end(committed_keys) — make data visible to readers	(session ends when load completes)

Failure Recovery

• Put failures during a per-layer range write trigger batch_put_session_revoke(failed_keys) to atomically discard partially-written data. Failed keys are excluded from subsequent layers' transfers.
• Get failures during a per-layer range read record the failed block IDs into an error set. The scheduler retrieves this set via get_block_ids_with_load_errors() and uses it to recompute from the correct token position.

Feedback Period.

Two weeks from submission.

CC List.

Any Other Things.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[Yusufchy]

The single-key-per-block design looks like the right trade-off for scaling metadata. One question worth pinning down is how partial-session visibility behaves when a later layer fails: can readers safely consume earlier committed ranges, or is the block revoked as a unit? An explicit invariant plus a failure-injection test would make the recovery contract easier to validate.