Skip to content

[KV Connector][Mooncake] Renew prefill KV leases from the decoder - #64193

Open
CJJ1008 wants to merge 3 commits into
vllm-project:mainfrom
CJJ1008:mooncake-kv-lease
Open

CJJ1008 wants to merge 3 commits into
vllm-project:mainfrom
CJJ1008:mooncake-kv-lease

Conversation

@CJJ1008

@CJJ1008 CJJ1008 commented Oct 28, 2026

Copy link
Copy Markdown

Overview

Today the Mooncake prefiller holds a finished request's KV until the decoder pulls it, or until VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT (480 s) runs out. This PR adds an opt-in kv_lease_duration that replaces that fixed timeout with a lease the decoder renews, following NixlConnector's lease (#41383, design). It addresses the "480s expire fallback" in item 6 of RFC #60523.

Claims

  • D crash: P frees the dead decoder's KV within one lease of D's last renewal, instead of after 480 s.
  • D overload: requests queued on a busy decoder keep their KV for as long as they wait. A short fixed timeout fails them.
  • Late pulls: a pull for KV that P has already freed fails at once, instead of waiting out the timeout. This also applies when the lease is off.
  • Without kv_lease_duration, behavior is unchanged.

Validation

Unit tests: pytest tests/v1/kv_connector/unit/test_mooncake_connector.py → 144 passed (A800, huggingface_hub==1.31.0 as in CI). The new tests use the real scheduler, real ZMQ sockets and a stubbed bootstrap service. They cover:

  • renewal is off without the option, and throttled to every lease / 6 with it;
  • D renews until the pull settles or the request aborts;
  • renewals reach every paired P worker, querying the bootstrap first if needed;
  • on P, renewals extend only known transfers and never shorten a deadline, and a P without the lease rejects them;
  • pulls on the same channel still work;
  • a pull for freed KV gets an immediate error.

pre-commit passes.

End to end: 2× A800 with 200G HDR IB, Mooncake TE 0.3.13.post1, DeepSeek-MoE-16B, TP=1 per side. Both sides run main 10e8e921b, with only mooncake_connector.py swapped (baseline vs this PR).

  • Load: 128 requests, 3,000-token prompts, 1,000 output tokens. That is 2.4× D's KV capacity, so D always has a queue.
  • Probe: a 512-token request sent straight to P every 10 s.
Case Unmodified main This PR
A. D killed at t≈23 s; 480 s timeout vs 30 s lease P's KV stays at 99% until t≈484 s; the probe sent at t≈10 s takes 475.5 s P's KV is below 5% at t≈103 s; that probe takes 61.3 s, later probes 36–46 ms
B. No crash, requests wait on D up to ~80 s; 20 s timeout vs 20 s lease 79 / 128 requests fail 128 / 128 succeed, no lease expires

In case A, requests still queued on P when D died are prefilled afterwards. The baseline holds them for another 480 s, and P is still at 37% KV when the run ends at t=660 s. The lease holds each for one 30 s lease.

Reproduction
# Same on D with "kv_role":"kv_consumer" and --port 8020. Drop kv_lease_duration for the baseline;
# case B's baseline sets VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT=20 on both sides.
vllm serve deepseek-ai/deepseek-moe-16b-chat --port 8010 --dtype bfloat16 --trust-remote-code \
    --max-model-len 4096 --gpu-memory-utilization 0.85 --no-enable-prefix-caching \
    --kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_producer","kv_connector_extra_config":{"device_name":"mlx5_102","kv_lease_duration":30}}'
python examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py \
    --port 8000 --prefill http://<P>:8010 8998 --decode http://<D>:8020
  • Load: 128 streaming completions sent at once through the proxy, each with random prompt token ids, max_tokens=1000 and ignore_eos.
  • Probe: a plain completion with max_tokens=1, sent to P without kv_transfer_params.
  • Case A sends SIGKILL to D's process group 20 s after the load starts.

Overhead:

  • P checks each side-channel message for a renewal: 0.7 µs for a steady pull, 1.3 µs for a pull carrying the layout, 21 µs for an 18 KB pull of 32 requests.
  • D sends one message of about 1 KB per 128 transfer ids, to each paired P worker every lease / 6.

Details

Flow

  1. D starts tracking a request in on_new_request, while it may still be queued.
  2. Every lease / 6, build_connector_meta groups the tracked transfer ids by P engine.
  3. start_load_kv sends them to each P worker that D would pull from.
  4. P sets expire_time = now + lease when prefill ends. On each renewal it sets expire_time = max(expire_time, now + 2/3 · lease).
  5. D stops renewing on finished_recving or request_finished.
  6. When renewals stop, P frees the blocks on its next scan and remembers the transfer id, so a later pull fails at once.

Design notes

  • Opt-in, unlike NIXL. Mooncake has no handshake compatibility check, so an old decoder paired with a lease-enabled prefiller would silently lose the KV of requests queued longer than the lease. The docs say to set the same value on both sides and to upgrade prefillers first. An old prefiller rejects renewals, and D warns once.
  • Transport. Renewals use P's existing ZMQ ROUTER, so there is no new port and the pull format is unchanged.
    • A pull decodes as an empty renewal, because msgspec skips unknown fields.
    • P's listener answers renewals itself, since all sender workers may be blocked waiting for prefills.
    • D waits for the reply, because a linger=0 socket drops unsent messages.
    • Renewals go to the same P workers as the pull, using the address selection now shared with receive_kv. A P engine that D has not contacted yet gets a bootstrap query first.
  • Expiry scan. P keeps scanning every entry, because renewals reorder deadlines; NIXL's ordered reaper hit this in [Bug] NIXL lease reaper strands expired KV blocks behind a heartbeated head entry #58222. Deadlines never cross processes, so there is no clock skew (cf. [PD][Bugfix] Rebase KV lease deadlines onto worker clock #50326).
  • Limitation. Backpressure is unchanged: P holds KV for as long as a slow decoder needs it. In case B with the lease, P sits at 99% for about 40 s and a probe waits 35 s. The 20 s timeout frees KV sooner (probe 17 s) only by failing requests. Fixing that needs D-side staging (rd2h / host receive).
  • [RFC]: A failure spine for the Mooncake connectors #57530 failure spine. With the lease on, P's deadline is the lease. VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT still bounds the wait for an unfinished prefill and how long freed transfer ids are remembered.

Not a duplicate. I searched for Mooncake lease, heartbeat, ABORT_REQUEST_TIMEOUT and expire_time. #63420 and #59522 are leases in the Mooncake Store connector. #63804 is an external connector, and #57529 is a standalone failure ledger. Open PRs on this file (#62657, #62658, #62255, #61974) do not change P's expiry.

AI assistance: written with Claude Code. I reviewed every changed line, ran all tests and experiments above, and can answer for the design.


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.

🤖 Generated with Claude Code

The prefiller holds a finished request's KV blocks until the decoder
pulls them or VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT (480 s) runs out. A
crashed decoder pins those blocks for the whole timeout, while a shorter
timeout frees the KV of requests that are only queued on a busy decoder.

Add an opt-in kv_lease_duration, following NixlConnector's lease
(vllm-project#41383). The prefiller grants a short lease when prefill ends. The
decoder renews it over the existing ZMQ side channel from the moment a
request arrives until its pull settles, and the prefiller frees blocks
whose lease runs out. The prefiller also answers a pull for KV it has
already freed at once, instead of waiting for a prefill that already ran.

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: CJJ1008 <843515223@qq.com>

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@mergify

mergify Bot commented Oct 28, 2026

Copy link
Copy Markdown
Contributor

Documentation preview: https://vllm--64193.org.readthedocs.build/en/64193/

@mergify mergify Bot added documentation Improvements or additions to documentation kv-connector mooncake Mooncake KV-transfer / EC-transfer labels Oct 28, 2026
@CJJ1008 CJJ1008 mentioned this pull request Oct 28, 2026
2 of 6 tasks

@wangxiyuan wangxiyuan left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The change LGTM. just a nit. Feel free to update if need follow-up change.

# The requests may still be queued. Query now so this renewal,
# and the pull later, can reach P's workers.
await self._query_bootstrap(heartbeat.remote_bootstrap_addr)
if remote_engine_id not in self._remote_agents:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it's good to add a log here before return for better debug.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added a warning before the return in fbf1e7f. Thanks for the review!

@wangxiyuan wangxiyuan added the ready ONLY add when PR is ready to merge/full CI is needed label Oct 28, 2026
@github-actions

Copy link
Copy Markdown

✅ @CJJ1008, CI is now available for this PR.

  • /ci run starts upstream CI; /amd-ci run starts AMD CI only.
  • Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
  • /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
  • /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
  • /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: CJJ1008 <843515223@qq.com>
@CJJ1008

CJJ1008 commented Oct 28, 2026

Copy link
Copy Markdown
Author

/ci run

@github-actions

Copy link
Copy Markdown

✅ Triggered Buildkite CI #97935 for commit 55c29ff24b9b.

ℹ️ This PR is 24 commits behind upstream main. Merged upstream main into this branch as 55c29ff24b9b before starting CI. Pull the branch before pushing more changes.

@wangxiyuan

Copy link
Copy Markdown
Contributor

Follow-up idea (not blocking): turn P/D lease misconfiguration into an explicit error instead of a docs-level constraint.

@CJJ1008

CJJ1008 commented Oct 28, 2026

Copy link
Copy Markdown
Author

Thanks, good point! I'll handle it in a follow-up PR after this one lands. Rough idea for now (haven't validated it yet):

  • P advertises its kv_lease_duration via the bootstrap server, and D checks it when it looks P up. That catches D renewing against a P that has the lease off or a different value.
  • D also sends its kv_lease_duration with each pull (an optional field, so older P just ignores it), and P checks that. That catches P having the lease on while D isn't renewing, including older D builds.
  • If they don't match, the affected requests fail with a clear error showing both values, instead of relying on the docs.

If you have any other ideas, feel free to let me know. I'd be glad to keep contributing here.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation kv-connector mooncake Mooncake KV-transfer / EC-transfer ready ONLY add when PR is ready to merge/full CI is needed

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants