[RFC]: Support Layerwise KV Cache Transfer on User-Defined Stream under Graph-Mode for Mooncake Store

### Changes proposed

## Overview

We propose an **additive, non-breaking `*_on_stream` Store path** that threads a user-supplied `stream_ptr` from the Store layer all the way down to the transport, so the transfer executes **on the caller's stream** and its completion is naturally observable by a device event (`ExternalEvent.record` / `cudaEventRecord`), with **no host blocking**. This requires an explicit **pin/lease** for the async window. Detailed motivation and analysis are in the accompanying design note.

The real ask is not only "add a `stream_ptr` parameter to `batch_get`." It is: **convert an asynchronous remote-KV transfer's completion signal from host-visible to device-visible**, so that a full-graph forward (Prefill or Decode) can wait for each layer's KV to arrive *without host intervention*.

## Motivation

### The application scenario

In pooled-KV disaggregated (Prefill/Decode) architectures, long-context LLM, Generative Recommendation (GR), and Agent workloads are increasingly scaling context lengths from $K$ to $M$ tokens. Under this setup, KV Cache lifecycle management across workers follows a distinct pattern:

- **KV Save Phase:** The kv produce worker flushes request historical/prefix KV cache to a **Mooncake DRAM KV pool** using `batch_put_into_multi_buffers()`.
- **KV Retrieval Phase:** The downstream worker must load each layer's required KV cache back from the shared pool prior to executing that layer's attention computation (i.e. decode workers fetching historical KV, or prefill workers loading prefix KV upon a prefix cache hit).

The cost of pooling is that "getting KV cache from the pool" lands on the forward's critical path: each layer's compute is preceded by loading that layer's KV from the Store DRAM pool to HBM, and this load is the same order of magnitude as the layer's compute. To keep pooling profitable, the framework adopts a **layerwise pipeline** — placing each layer group's KV transfer on a separate (load) execution stream so its latency is masked by the previous group's compute.

To further reduce host scheduling overhead, the pooled inference framework can capture the model's execution into a compute graph (**full graph**). Under the layerwise pipeline, every Layer's forward must gate on its KV load completing; once the whole forward is captured into one graph, that gating sync must become a *device-side* graph node rather than a host wait — otherwise the host is forced back in between groups and the graph breaks. **The collision between "forward must wait per layer" (layerwise pipeline) and "host must be absent between layers" (full graph) is where the difficulty lies.**

The target scenario this RFC serves is:

> **`full Graph + Store DRAM KV Pool + layerwise KV load/compute fully overlapped`** — applicable to both Prefill and Decode, anchored on the shared step of "getting KV cache from the pool."

```mermaid
flowchart LR
subgraph Pool["Mooncake Store · DRAM KV Pool"]
KV0["KV[layer0]"]
KV1["KV[layer1]"]
KVN["KV[layerN]"]
end

subgraph Xfer["Transfer Stream (user-owned)"]
L0["Load KV[0]"]
L1["Load KV[1]"]
LN["Load KV[N]"]
L0 --> L1 --> LN
end

subgraph Graph["Compute Stream · full graph (one replay)"]
direction TB
W0["wait(E0)"]
A0["Attention 0"]
W1["wait(E1)"]
A1["Attention 1"]
WN["wait(EN)"]
AN["Attention N"]
W0 --> A0 --> W1 --> A1 --> WN --> AN
end

KV0 -. "RH2D" .-> L0
KV1 -. "RH2D" .-> L1
KVN -. "RH2D" .-> LN
L0 -. "record E0" .-> W0
L1 -. "record E1" .-> W1
LN -. "record EN" .-> WN
```

The compute stream is captured into a **single full graph**: one `replay()` runs all layers back-to-back with no host code between them. Each layer's KV load runs on a separate **user-owned transfer stream**; when it finishes it `record`s an event, and the graph's `wait(Ei)` node — a true device-side dependency — gates `Attention i` on it. The host submits all loads and the single replay, then steps away; the device autonomously overlaps transfer and compute.

### What this is not: distinct from Mooncake's existing layerwise API (#3120)

Mooncake already has a layerwise-flavored API — #3120, session-based ranged transfer:

```python
batch_get_session_start()
for layer:
batch_get_into_multi_buffer_ranges(...)
batch_get_session_end()
```

It and this RFC **both carry the "layerwise" label but solve different layers of different problems — they are orthogonal and composable, not substitutes.** The contrast that matters for this scenario:

| | **#3120 — session-based ranged read** | **This RFC — stream-bound, device-visible completion** |
| :--- | :--- | :--- |
| **Problem layer** | Metadata / call overhead | Completion signal |
| **Solves** | Session-scoped per-byte-offset layered read; **0 Master RPC per layer** | Transfer runs on the caller's stream; completion **observable by a device event** |
| **Completion semantics** | Host-synchronous — range read path still does `future->get()` (`Client::TransferReadInternal → submitRangeRead`) | Device-asynchronous — `ExternalEvent.record(stream_ptr)` is the "KV ready" signal |
| **Under full graph** | **Cannot** — `future->get()` is a host sync primitive the graph cannot wait on; host must block between layers → graph breaks | **Can** — graph's `wait(Ei)` is a device-side dependency, host absent between layers |
| **What it omits for this scenario** | The critical blocker: no device-visible completion signal | — |

**One-line positioning:** #3120 = session-scoped layered read (host sync); this RFC = stream binding + device-visible completion (device async). #3120 lowers the per-layer metadata cost and is a **useful prerequisite**, but it does **not** contain the blocker this scenario requires — a device-visible completion the graph can wait on. The two are designed to **compose** (e.g. a `*_ranges_on_stream` variant that gets #3120's 0-Master-RPC-per-layer benefit *and* device-visible completion in one path; see Open questions), not to replace each other. This RFC supplies the one layer #3120 deliberately leaves out.

### Three constraints that make the scenario inevitable

The motivation rests on three physical constraints. Each is individually uncontroversial; **their collision is the gap this RFC asks Mooncake to close.**

**Constraint 1 — Full graph is the best execution model for both Prefill and Decode, but it demands host absence between layers.**
Graph mode captures the entire forward into one replayable graph, eliminating per-step kernel-launch host overhead — the single largest host-side tax on memory-bound decode (and a meaningful one for prefill's dense attention over long history). The trade-off is hard, not preferential: **after `replay()`, the host must not intervene between layers.** Any host code inserted between layers is either a forbidden sync primitive (capture fails) or silently absent on replay (graph break).

**Constraint 2 — "Getting KV from the pool" must be masked → the layerwise-overlap requirement is *derived* from Constraint 1, not assumed.**
Full graph + KV-cache-bearing decode, combined, are the root cause of the memory-bound regime, and this is what *derives* the layerwise-masking requirement. With pooling, each layer's compute is preceded by loading that layer's KV from the Store DRAM pool to HBM; this load is on the forward's critical path and is the same order of magnitude as the layer's compute ($T_{load,i}$ vs $T_{compute,i}$ for prefill's dense attention / decode's single-step attention), so overlapping the two is what makes pooling worthwhile rather than latency-neutral. Layerwise masking requires $T_{load,i} \le T_{compute,i}$; once overlapped, end-to-end latency converges to:

$$T \approx T_{load,first} + \sum_i \max(T_{load,i},\, T_{compute,i})$$

If KV transfer sits serially in front of attention ($T = T_{load,first} + \sum_i (T_{load,i} + T_{compute,i})$), the pooling benefit is eaten by transfer latency. **Overlap is not a nice-to-have; it is the precondition for the pooled-KV architecture to pay off at all — and it holds for both prefill and decode, wherever "getting KV from the pool" lands on the forward's critical path.**

**Constraint 3 — Remote KV "arrival" is inherently a device-side event, but the Mooncake Store API forces it host-visible.**
The NIC completing an RMA is a device/NIC-side fact that a stream event could observe directly. But `store.batch_get_into_multi_buffers()` pulls completion back to the host: inside `Client::BatchGet` it blocks the calling thread on `future.get()`, and completion is only ever observed by a **host-side** polling thread (`query_thread_` in `AscendDirectTransport`'s `AsyncTransferExecutor`, looping on `GetTransferStatus`). `adxl::AdxlEngine::TransferAsync` has **no stream parameter and no device-event callback** — completion can only be learned by host polling. There is **no `stream_ptr` anywhere** in the Store API surface — not in `PyClient`/`RealClient`/`Client::BatchGet`/`TransferSubmitter`, not even in the `Slice{ptr,size}` struct.

> **The core contradiction: Constraint 1 (host must be absent) collides with Constraint 3 (completion is forced host-visible).** A graph-mode forward that reads pooled KV cannot, with the current Store API, both keep the graph intact *and* know when each layer's KV has arrived. The community is forced to choose: drop full graph (lose the latency win), or drop pooling-on-demand (lose the capacity win). vLLM MoRIIO's READ mode is forced to piecewise CUDA graphs for exactly this class of reason — see the same-problem RFC [vLLM RFC #49643: Enable full CUDA graphs for MoRIIO READ mode via async KV-load gating](https://github.com/vllm-project/vllm/issues/49643).

### What needs to change

Resolving the contradiction means converting Mooncake Store's asynchronous remote KV transfer completion into an **XPU-visible synchronization dependency**, so layerwise KV load and forward compute overlap on the device side (for both prefill and decode). This decomposes across three layers:

```
Layer 1: Storage Layer 2: Transport Layer 3: Execution
Mooncake Store Transfer Engine CUDA/ACL Graph
│ │ │
└── DRAM KV pool └── RDMA async KV load ├── Layer i KV wait
│ └── Layer i Attention
▼
completion signal ──► must be device-visible
(currently host-visible)
```

- **Layer 1 (Storage) + Layer 2 (Transport) are Mooncake's responsibility** — and this is where the gap lives. Today the completion signal exiting Layer 2 is host-visible (`future.get()` / `query_thread_`), with no stream binding. **No amount of Layer 3 orchestration can manufacture a device-visible completion signal that the transport never emits** — there is no signal source to wait on. This is why the fix must land in Mooncake, not in the inference framework.
- **Layer 3 (Execution) is the inference framework's responsibility** — inserting `wait(Ei)`/`reset(Ei)` graph nodes and the capture/replay timeline. The inference-framework side already scopes per-layer transfer/attention overlap as framework orchestration. But that orchestration is *blocked* until Layers 1–2 expose a stream-bound, device-visible completion.

**Conclusion for maintainers:** the pooled-KV + full-graph scenario (Prefill and Decode alike) is not speculative — it is the direction long-context serving is already moving, and the constraints above make the gap unavoidable. The single thing only Mooncake can provide is a Store read API that runs the transfer on a caller-supplied stream and lets a device event observe its completion. That is what this RFC proposes.

---

## Goals / Non-Goals

**Goals**

1. **Device-side completion.** After the call returns, `stream_ptr`'s progress reflects transfer completion, so a user `ExternalEvent.record(stream_ptr)` / `cudaEventRecord(stream_ptr)` is a valid "KV ready" signal a compute stream can wait on — with **no host blocking**.
2. **Stream-aware Store read API.** An additive, non-breaking method such as `store.batch_get_into_multi_buffers_on_stream(keys, all_buffers, all_sizes, stream_ptr, ...)`, mirroring the existing API, defaulting `stream_ptr = 0` to the current blocking behavior.
3. **Preserve Store semantics.** Replica selection, lease/pin, and address validity during the transfer window must still hold — the stream-aware path must **not** degrade to a bare-address query (`batch_get_replica_desc` + direct TE) that drops lifecycle protection.
4. **Cross-platform.** Work on both CUDA and Ascend (NPU), covering the **cross-process RMA** case that dominates PD pooling (Standalone/embedded deploy: real-store client ↔ pool is always cross-process).

**Non-Goals**

- Changing the existing blocking `batch_get_into_multi_buffers` semantics.
- Re-implementing graph capture or event primitives — those remain the caller's responsibility (`torch.npu.ExternalEvent`, `cudaEvent`).

---

## Proposed Design: direct stream binding

We propose a **layered, opt-in** stream-aware path that threads a user-supplied `stream_ptr` from the Store layer **all the way down to the transport**, so the transfer executes on the caller's stream. Because the transfer itself runs on `stream_ptr`, a subsequent `record_event(stream_ptr)` is a true "KV ready" signal — no host callback, no CPU busy-wait, true pipeline overlap.

**A. Store layer — thread `stream_ptr` through the existing chain (additive, non-breaking).**

Add an overload (or parallel `*_on_stream` method) at each Store layer, `stream_ptr = 0` defaulting to current blocking behavior:

- `PyClient::batch_get_into_multi_buffers_on_stream(..., uintptr_t stream_ptr)` — new virtual (`pyclient.h`)
- `RealClient::batch_get_into_multi_buffers_on_stream` — delegates internally; instead of blocking on `future.get()`, submits the batch with the bound `stream_ptr` and returns a lightweight async handle (`real_client.cpp`)
- `Client::BatchGetOnStream` — same `TransferSubmitter::submit` submission, but **skips the `future.get()` wait loop** and passes `stream_ptr` down (`client_service.cpp`)
- pybind: `store_py.cpp` exposes `batch_get_into_multi_buffers_on_stream`

**B. Transport layer — bind the transfer onto the user stream.**

Thread `stream_ptr` from `TransferSubmitter` → `TransferEngine::submitTransfer` → the selected `Transport`, so the device-side copy/RMA is issued on the caller's stream rather than an internal one:

- **CUDA / same-process local copy:** issue `cuMemcpyAsync`/`cudaMemcpyAsync` on `stream_ptr` instead of the engine's internal stream.
- **Ascend same-process local copy:** `LocalCopyEngine` already has an async path (`CopyWithAsync`); open it to accept an external `aclrtStream` instead of its private internal `stream_`.
- **Ascend cross-process RMA (the PD-pooling case):** `adxl::AdxlEngine::TransferAsync` currently has **no stream parameter** (it is an external SDK whose ABI Mooncake does not own). Here the stream cannot be threaded into the RMA itself; instead, the transport **submits the RMA from the user stream's context and records completion on `stream_ptr`** (e.g. via the transport's existing completion notification wired to a stream event), so `ExternalEvent.record(stream_ptr)` still marks transfer completion. The exact hook point is to be confirmed with the ADXL/HiXL maintainers — this is the one place where direct binding needs vendor cooperation.

> The key invariant across all paths: **the user `stream_ptr` is the execution carrier of the transfer (for local copy) or the completion-recording carrier (for cross-process RMA), not merely a trigger.** This is what makes `record_event(stream_ptr)` a valid device-visible completion signal.

**C. Lifecycle — explicit pin for the async window (correctness-critical).**

Because the host no longer blocks on `future.get()`, the Store's **implicit lease protection is decoupled from transfer completion**. The stream-aware path must **pin** (or refresh the lease for the duration of) the selected replicas before returning, and **unpin** via a stream-completion callback (or a bounded TTL sized to the expected transfer). Without this, an eviction mid-transfer yields a silent **"valid address, Wrong content"** read — the single most important correctness requirement and the main divergence from the existing blocking API.

---

## Known limitations / trade-offs

1. **Ascend cross-process RMA needs a vendor hook.** `adxl::AdxlEngine::TransferAsync` has no stream param and is an external SDK; binding completion onto `stream_ptr` for cross-process RMA requires confirming a hook point with the ADXL/HiXL maintainers. This is the main open dependency of the direct-binding approach.
2. **Process-exit on failure must be replaced.** Any existing failure path that terminates the process (since a host callback cannot propagate errors) must be replaced by a queryable failed-event the caller can inspect.
3. **Pin/lease bookkeeping.** The caller must be able to query which replicas were pinned and ensure they are released; a leaky pin path degrades the pool.

---

## Requested from maintainers

1. **Agreement on direct stream binding** as the mechanism for stream-aware Store transfers — threading `stream_ptr` from Store to transport, with `record_event(stream_ptr)` as the device-visible completion signal.
2. **Guidance on the pin/lease primitive** to expose for the async window — does Mooncake already have an internal pin we can reuse, or is a new `pin_for_transfer(keys, ttl)` needed?
3. **Acceptance of an additive, non-breaking `*_on_stream` API surface** on `PyClient`/`RealClient`/`Client`, mirroring the existing `batch_get_into_multi_buffers`.
4. **Confirmation of the Ascend cross-process RMA hook point** — how to bind `adxl::AdxlEngine::TransferAsync` completion onto a user `aclrtStream` (or whether a vendor-side stream param is feasible).

## Open questions

- Should `*_on_stream` return a future/handle the caller can await outside the graph, or is "fire-and-forget + device event" sufficient?
- How to bound the pin TTL for very large KV transfers without knowing the transfer duration a priori?
- Does the existing `USE_EVENT_DRIVEN_COMPLETION` path in `TransferTask` offer a lower-CPU-overhead completion signal reusable here?
- Relationship to #3120: should `*_on_stream` compose with the session-based ranged API (`batch_get_session_start` / `batch_get_into_multi_buffer_ranges` / `batch_get_session_end`), i.e. a `*_ranges_on_stream` variant that gets the 0-Master-RPC-per-layer benefit *and* device-visible completion in one path?

### Before submitting a new issue...

- [x] Make sure you already searched for relevant issues and read the [documentation](https://kvcache-ai.github.io/Mooncake/)

Start by tracing the existing path through PyClient, RealClient, Client::BatchGet, and TransferSubmitter, then inspect AscendDirectTransport's AsyncTransferExecutor and AdxlEngine::TransferAsync. The RFC's target is a user-supplied stream_ptr propagated through this path, with transfer completion exposed to a device event without host blocking; the design note and open questions need resolution before implementation.

Written by the indexing model from the issue text.