kvcache-ai / kvcache-ai/Mooncake

[RFC]: Deadline-safe bounded remote transfer lifecycle

Open
#3,793 1 comment 0 reactions 1 assignee Claimed by @alogfans View on GitHub
Dominant language
C++
Stars
6.6k
Forks
1.2k
Avg merge
3d 5h
Merged PRs (30d)
312

Description

## Changes proposed

Target baseline: `kvcache-ai/Mooncake` `main` at `88230cb08b3c15d41b4176aa9821ac8ff4648ee0`.

Implementation companion: #3792

Related research: https://arxiv.org/abs/2608.17826

## Summary

Mooncake's scatter API already owns a transfer batch, opened segments,
per-fragment callbacks, and physical completion polling. It also correctly
documents that `waitFor(timeout)` does not cancel a transfer and that the
operation and buffers must remain alive. The missing abstraction is a
deadline-safe lifecycle result that tells integrations whether a logical
request failed, whether the transport is physically terminal, and whether the
local buffer and memory registration may be reused.

This RFC proposes two additive changes:

1. A Transfer Engine lifecycle API for absolute-deadline waits, best-effort
cancellation followed by bounded drain, explicit quarantine, idempotent
late completion, and operation-local observability.
2. A later Mooncake Store bounded-retrieve API that submits a large object in
windows limited by bytes, fragments, and operations under one total
deadline.

The first pull request is intentionally limited to item 1. It does not add a
cross-rank coordinator and does not change Store behavior.

## Current API and evidence

At the baseline above:

- `TransferEngine::ScatterTransferOperation` owns the backend, batch,
segments, requests, callbacks, and completion polling.
- `wait()` blocks until physical completion and object destruction calls it.
- `waitFor(duration)` returns a clock error at the logical timeout without
cancelling; its public contract requires the operation and local buffers to
remain alive.
- TENT exposes best-effort task cancellation, but cancellation support is
transport-specific. The classic path has no generic cancellation API.
- Store ranged reads ultimately create one scatter operation. Lease refreshes
use repeated relative waits; there is no user-visible total transfer
deadline or bounded submission window.
- Store `TransferFuture` has a fixed logical timeout, but its result does not
express physical completion or buffer reuse safety.

Relevant upstream work:

- [#3000](https://github.com/kvcache-ai/Mooncake/pull/3000) introduced the
framework-neutral scatter API.
- [#3310](https://github.com/kvcache-ai/Mooncake/pull/3310) made the scatter
operation the owner of submission, waiting, callbacks, segments, and batch
release, including draining already-published work after submission errors.
- [#2887](https://github.com/kvcache-ai/Mooncake/issues/2887) and
[#2881](https://github.com/kvcache-ai/Mooncake/pull/2881) establish that the
framework owns compute scheduling while Store exposes range/session
transfer primitives.
- Closed draft [#2552](https://github.com/kvcache-ai/Mooncake/pull/2552)
explored chunked reads but is not a safe base: review identified wakeup,
synchronization, complexity, and ownership problems, and the author closed
it after the newer range/scatter APIs landed.
- Open [#2916](https://github.com/kvcache-ai/Mooncake/issues/2916) /
[#2917](https://github.com/kvcache-ai/Mooncake/pull/2917) addresses
transport-specific NVMe-oF cancellation/drain/quarantine. The generic API
must compose with it rather than duplicate transport cleanup.
- Open [#3523](https://github.com/kvcache-ai/Mooncake/issues/3523) documents a
TENT RDMA case where a posted work request may never produce a CQE. This is
why a bounded caller path must report quarantine instead of falsely
reporting physical termination.
- Observability should align with tracing RFC
[#1850](https://github.com/kvcache-ai/Mooncake/issues/1850), especially its
persistent execution span and deduplicated terminal markers.

This RFC is the design companion to Draft PR #3792. The implementation PR is intentionally kept in Draft while the API and responsibility boundary are discussed here.

## Scope and responsibility boundary

### Transfer Engine owns

- accurate physical transfer status;
- best-effort transport cancellation dispatch;
- draining a submitted batch to a safe terminal state when possible;
- preventing duplicate callbacks and duplicate terminal accounting;
- retaining caller-provided lifetime anchors until physical completion;
- reporting whether buffer/MR reuse is safe;
- per-operation lifecycle counters and timings.

### Mooncake Store owns in PR 2

- splitting a range retrieve into bounded windows;
- limits on in-flight bytes, fragments, and operations;
- one absolute deadline across all windows;
- exact transferred-byte validation;
- Store-facing metrics and failure propagation.

### LMCache/vLLM/integration owns

- TP rank membership and request epochs;
- submitting each rank's shard;
- the all-rank barrier and validation;
- commit only after every rank succeeds;
- aborting the complete restore after any rank failure;
- discarding successful temporary shards after a peer fails;
- scheduler transitions out of `WAITING_FOR_REMOTE_KVS`;
- local recomputation policy.

Mooncake should provide the states required to implement this barrier, but it
should not contain vLLM request IDs, rank constants, scheduler logic, or a
model-specific connector.

## Proposed PR 1 API

The change is additive to `ScatterTransferOperation`:

```cpp
enum class ScatterTransferState : uint8_t {
SUBMITTED,
IN_PROGRESS,
SUCCEEDED,
FAILED,
TIMED_OUT,
CANCEL_REQUESTED,
DRAINING,
DRAINED,
QUARANTINED,
};

struct ScatterTransferSnapshot {
ScatterTransferState state;
Status first_failure;
size_t total_bytes;
size_t in_flight_bytes;
size_t quarantined_bytes;
size_t total_fragments;
size_t completed_fragments;
size_t deadline_timeouts;
size_t cancel_requests;
size_t drain_attempts;
size_t quarantine_events;
size_t late_completions;
std::chrono::nanoseconds submit_to_physical_completion;
bool physical_complete;
bool buffer_reusable;
};

Status waitUntil(std::chrono::steady_clock::time_point deadline);
Status cancelAndDrainUntil(std::chrono::steady_clock::time_point deadline);
ScatterTransferSnapshot snapshot() const;
```

`waitFor(duration)` remains source- and behavior-compatible. It is a legacy
relative wait and does not change the operation's logical outcome. New code
that needs one request-level deadline uses `waitUntil`.

An additive `ScatterTransferOptions` overload accepts opaque
`std::shared_ptr` lifetime anchors. The engine does not interpret them;
it merely retains the owners of local allocations and/or MRs until physical
completion. Keeping anchors in a separate options object avoids changing the
layout of `ScatterTransferRange`. Existing callers that use externally managed
buffers remain source-compatible.

## State machine

```text
SUBMITTED -> IN_PROGRESS -> SUCCEEDED
-> FAILED

SUBMITTED/IN_PROGRESS -> TIMED_OUT
TIMED_OUT -> CANCEL_REQUESTED -> DRAINING -> DRAINED
`-> QUARANTINED
QUARANTINED -> DRAINING -> DRAINED (later retry/poll)
```

`SUCCEEDED`, `FAILED`, and `DRAINED` imply physical completion and safe buffer
reuse. `TIMED_OUT`, `CANCEL_REQUESTED`, `DRAINING`, and `QUARANTINED` do not.

`DRAINED` is a resource-safety terminal state, not a success result. The first
logical failure remains available in the snapshot and is returned by later
waits. This prevents a late successful CQE from converting a request-level
timeout into success.

## Timeout, cancellation, drain, and quarantine

1. `waitUntil(deadline)` polls against an absolute steady-clock deadline.
2. On expiry, pending callbacks are resolved exactly once with the timeout;
physical requests remain tracked.
3. `cancelAndDrainUntil(deadline)` requests cancellation where the backend
supports it, then polls to physical terminal under the supplied drain
deadline.
4. Unsupported cancellation is not an error in itself; the operation drains
naturally.
5. If the drain deadline expires, the state becomes `QUARANTINED`, the return
value is a clock error, and `buffer_reusable` remains false.
6. A quarantined operation must be retained by the integration. Dropping it
keeps the existing safe behavior: destruction waits for physical
completion. A caller that needs a bounded request thread moves the operation
into its quarantine registry.
7. If a lifetime anchor was supplied, it is retained through quarantine and
released only at physical completion.

PR 1 deliberately does not introduce an engine-global detached reaper. Such a
reaper would complicate engine shutdown and could form ownership cycles. A
separate follow-up may add a bounded-capacity quarantine registry after the
operation contract is accepted.

## Completion and callback rules

- Logical callbacks fire at most once per fragment.
- Timeout/cancel may publish a logical failure callback before physical
completion.
- A later physical completion updates counters and releases resources but does
not call the fragment callback again.
- The first failure is immutable.
- Terminal accounting is idempotent.
- Physical completion time is measured from submission to batch release, not
merely to logical timeout.

These rules prevent callbacks captured from a returned integration stack frame
from being invoked by a late completion.

## Proposed PR 2: bounded Store retrieve

An additive options object should carry:

```cpp
struct BoundedRetrieveOptions {
size_t max_in_flight_bytes;
size_t max_in_flight_fragments;
size_t max_in_flight_operations;
std::chrono::steady_clock::time_point deadline;
std::chrono::nanoseconds drain_timeout;
};
```

The Store session/range layer submits only enough fragments to fill the
window. A staging allocation is not recycled until its operation snapshot says
`buffer_reusable`. The same total deadline is passed to every window; it is not
reset for each fragment. On failure, Store stops submitting new windows,
cancel/drains submitted operations, quarantines unresolved resources, validates
the exact number of transferred bytes, and returns a non-success result.

The first implementation should be sequential-window safe. Parallel windows
up to `max_in_flight_operations` can be added without changing the API.

## TP atomic restore contract

The integration creates one bounded retrieve per rank shard and treats the
collection as a transaction:

```text
prepare(rank 0..N-1)
-> transfer and validate every rank
-> all succeeded: commit all temporary shard views
-> any failed: abort all, drain/quarantine, local recompute
```

No rank publishes restored KV into the active scheduler before the barrier.
Mooncake returns precise per-rank state; the integration owns the barrier and
commit token. This keeps the API framework-neutral while preventing partial TP
restore.

## Compatibility and risks

- Existing scatter callers compile unchanged.
- Existing `waitFor` retains its non-cancelling relative-wait behavior.
- A caller that does not supply a lifetime anchor must still own its buffers
until `buffer_reusable` is true.
- Cancellation remains transport-dependent. `QUARANTINED` is required for
never-terminal hardware/driver paths.
- Per-operation snapshots are not a substitute for process-wide metrics; they
are an integration point for the tracing/metrics design.
- A quarantine registry needs explicit memory limits and shutdown semantics;
it is intentionally outside the minimum PR 1.
- Store range callbacks currently capture request-local state. PR 2 must adopt
lifetime anchors or heap-owned callback state before bounded return.

## Test plan

PR 1 uses a hardware-free scripted transport to cover:

- success and failure terminal states;
- absolute deadline expiration;
- logical callback exactly once across late physical completion;
- cancel/drain timeout to quarantine;
- lifetime anchor retention through quarantine and release after drain;
- first-failure preservation;
- in-flight/quarantined byte accounting;
- zero-duration and already-expired deadlines;
- sanitizer-friendly destruction after a late completion.

PR 2 should add Store fault injection for partial windows, short transfer,
timeout, cancel unsupported, late completion, large-object window bounds, and
multi-window total-deadline enforcement.

## Benchmark method

The minimum public benchmark is hardware-free and reports control-plane cost,
not RDMA throughput:

1. Script immediate success for 1, 16, 256, and 4096 fragments.
2. Run submit + wait + snapshot for a fixed iteration count.
3. Report median and p95 operation latency and allocations.
4. Compare baseline scatter with lifecycle snapshot enabled.

RDMA bandwidth, 512K-prefix TTFT, and deployment-specific performance are out
of scope until a public hardware matrix and exact reproducible images exist.

## Suggested commits and pull requests

PR 1, title prefix `[TransferEngine]`:

1. Add the lifecycle state/snapshot and absolute-deadline wait.
2. Add cancel/drain/quarantine and lifetime anchors.
3. Add scripted fault-injection tests and the control-plane benchmark.
4. Add API documentation and AI-assistance disclosure in the PR body.

PR 2, title prefix `[Store]`:

1. Add bounded retrieve options and a sequential window executor.
2. Integrate Store session/range retrieve and exact byte validation.
3. Add Store metrics, fault injection, and large-object tests.
4. Document integration responsibility for TP atomic commit.

No private connector code, host identifiers, production logs, deployment
addresses, model assets, or registry paths are needed for either pull request.

## Before submitting a new issue...

- [x] I searched the related issues and pull requests referenced above and read the Mooncake contribution and design documentation.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.