[KV Connector] Support sleep mode with MooncakeStoreConnector

[aoshen02]

Purpose

Stacked on #59625; only the last commit is new.

MooncakeStoreConnector registers the GPU KV cache with the Mooncake store (register_buffer). Over RDMA, that registration pins the physical pages. Sleep mode unmaps the KV cache, and wake maps new pages at the same virtual addresses. The store keeps using the old pages, so puts and gets read and write memory the model no longer uses.

Evidence on main (GB300, Qwen3-8B, two instances sharing one store; instance A sleeps at level 1, instance B never sleeps):

• After the wake, every store hit on A gives wrong output: 0/8 for KV that B put, and 0/8 for A's own post-wake puts.
• A's puts after the wake store stale bytes. B reads them and also returns wrong output (0/8), even though B never slept.
• While asleep, A still holds 67.8 GB of GPU memory, because the KV cache stays pinned. After the wake it holds 150.8 GB (85.4 GB before the sleep).
• There are no errors and no log lines.

With #59625 alone, this combination is refused at startup ("MooncakeStoreConnector does not support sleep mode").

This PR implements the #59625 hooks for MooncakeStoreConnector:

• release_kv_caches() runs before the KV cache is unmapped. It waits until the put queue and the shared get queue have no queued or running tasks, then calls unregister_buffer on each device KV region. All store puts and gets run on those queues; lookups only query keys.
  • The wait is bounded by VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT (480 s by default), the same bound MooncakeConnector uses for its release in [KV Connector] Support sleep mode with MooncakeConnector over RDMA #59625. That is well above Mooncake's 60 s per-transfer timeout, so a slow queue does not fail the sleep. On timeout it raises TimeoutError and keeps the registration. The wait shares one helper with close().
• restore_kv_caches() runs after the KV cache is mapped again. It registers the same addresses and lengths.
• Both calls are all or nothing and idempotent. If an (un)register call fails partway, the calls already made are undone and the error is raised. A failed sleep leaves the engine awake and fully registered; a failed wake leaves the regions released, so a retry starts clean.
• Host-resident (HiSparse) regions are host memory, which sleep does not touch, so they stay registered.
• supports_sleep_mode() reads the protocol from the Mooncake config file (MOONCAKE_CONFIG_PATH). The workers use the same file to set up the store. It accepts rdma and tcp and refuses other configured protocols. The check runs only when sleep mode is enabled, and it does not import mooncake. If the config is missing, startup fails with the existing "MOONCAKE_CONFIG_PATH is not set" error.
  • Why RDMA and TCP: no peer holds a key to our KV cache. Puts and gets use our local lkey against peers' host-DRAM segments, so registering again locally is enough.
  • TCP registration only records the address range. Data moves with cudaMemcpy on the virtual address, and nothing is pinned.
  • Other configured protocols (e.g. nvlink) can export handles to peers, so they stay refused.
  • It also refuses when MC_FORCE_MNNVL or MC_INTRANODE_NVLINK is set. Mooncake then installs an nvlink transport even with protocol: rdma. Like Mooncake, it checks only that the variable is set, not its value.

Test Plan

• Unit tests (CPU, fake store), one per invariant:
  • Release and restore round-trip exactly the device regions, with the same addresses and lengths (dense layout, and HiSparse with one host group plus one device group). Both calls are idempotent across two cycles.
  • Release waits for a queued or running put or get. A real thread finishes the task when a threading.Event is set. A timeout keeps the registration.
  • A failed unregister or register is rolled back, and a retry then succeeds.
  • The config gate: rdma/tcp accepted; nvlink, MC_FORCE_MNNVL (also =0) and MC_INTRANODE_NVLINK refused; missing config. Mooncake is blocked from import.
• Mutation check: each piece of the change was altered on its own (gate and each env var, region recording, host skip, idempotence, state, lengths, rollback, each drained queue, the bound, the deadline). 23 mutants, all killed, and every test kills at least one.
• pytest tests/v1/kv_connector/unit/{test_config,test_mooncake_store_*,test_mooncake_connector,test_multi_connector}.py tests/v1/worker/test_sleep_mode_backend.py. The store tests were also run 5 times. pre-commit was run on the changed files.
• Hardware (GB300, Mooncake 0.3.13.post1, RDMA, standalone mooncake_master, Qwen3-8B):
  • Two servers share one store: A sleeps and B never sleeps. GPU prefix caching is off, so every cached token comes from the store.
  • 8 random 2048-token prompts per phase, greedy decoding, 32 output tokens.
  • After every cycle: A gets KV that B put, B gets KV that A put after the wake, and A gets its own post-wake put back.
  • Cycles: 3× /sleep?level=1 + /wake_up; 1× level-1 sleep then /wake_up?tags=weights and /wake_up?tags=kv_cache; 1× /pause + /release_kv_cache_memory + /wake_up?tags=kv_cache.
  • Control run: /reset_prefix_cache?reset_external=true at the same points. This performs the same store wipe as sleep, without unmapping anything.

Test Result

• Unit tests (test_config, test_mooncake_store_worker, test_mooncake_connector, test_multi_connector, test_sleep_mode_backend): 407 passed, 2 failed. The 2 failures are in test_multi_connector.py and are unrelated to this change: one needs the gated meta-llama/Llama-3.2-1B-Instruct, the other needs ray, which is not installed. They fail the same way on [KV Connector] Support sleep mode with MooncakeConnector over RDMA #59625 alone. All test_mooncake_store_* files: 332 passed in each of the 5 runs.
• Hardware, sleep run vs awake control: 224/224 outputs are token-for-token identical across all 28 phases. Every hit read 2032 tokens from the store.
  • Before this change, each post-wake hit phase was 0/8.
  • Each run also compares hits with the cold output of the same prompt. A few phases there show 7/8, and the control run shows exactly the same mismatches. These come from near-tie prompts where a store hit and a cold prefill differ, not from sleep.

GPU memory on A	this PR	main (earlier run)
awake	85.1 GB	85.4 GB
asleep, level 1 (3 cycles)	2.1 GB each time	67.8 GB
after wake	85.1 GB each time	150.8 GB
selective wake: weights only	19.4 GB	—
after release_kv_cache_memory	19.7 GB	—

• After the review changes (rollback, env check, shared drain), a shorter rerun on the final commit with 2 level-1 cycles gave 104/104 outputs identical to its control. Every hit read 2032 tokens from the store. A held 2.1 GB asleep and 85.1 GB awake in both cycles.
• Transfer errors: none in either run (non-zero error counters, glog E lines, ERROR lines, and tracebacks were all 0 on A and B).

Not covered:

• TCP on hardware. On this cluster, TCP puts and gets fail with TRANSFER_FAIL even without sleep, with both P2PHANDSHAKE and the HTTP metadata server. An earlier H200 run of main with TCP and sleep returned 8/8 with memory fully freed, but it was not repeated with this commit.
• The HiSparse path on hardware (unit-tested only).
• TP > 1 and multi-GPU runs.
• The in-process engine (VLLM_ENABLE_V1_MULTIPROCESSING=0). The drain inside release_kv_caches is what covers it there.

Known limitations:

• Mooncake builds with NVLink support (such as the cuda13 wheel) also install an nvlink transport when they find no RDMA NIC. That can't be detected at config time, so it is not refused.
• Pre-existing: sleep at level ≥ 1 and release_kv_cache_memory reset the connector before the release. On rank 0 that reset waits for the put queue with no bound, so a stuck put hangs there, not in the bounded release.
• Pre-existing: a load task that raises before the try in KVCacheStoreRecvingThread._handle_request never calls task_done. The drain then never sees an empty queue, and every later sleep fails with TimeoutError. We found no scheduler input that reaches it. A fix needs a load-error report too, so it is left for a follow-up.
• Pre-existing, not specific to this connector: /sleep?mode=wait (and /pause?mode=wait) while connector loads are in flight returns 500 "Failed to reset KV cache ...", and the scheduler stays paused until /resume. Under PAUSED_NEW the scheduler counts only running requests, so requests holding blocks for an async load look idle. Use mode=abort, or call /resume and retry.

Known separate issue, not fixed here: sleep at level ≥ 1 resets the connector, which runs remove_all(force=True) on the Mooncake master. One instance's sleep therefore deletes the KV of every instance that shares the master. It should be scoped by tenant.

Additional Test Result: stress run

• Stress run (GB300, 1 node, Qwen3-8B, Mooncake RDMA, one shared store; the image plus per-file mounts of this stack, sha-checked):
  • Three servers: TP2 and TP1 sleep with staggered, overlapping windows, and a TP1 server never sleeps; store_tp_size=2, so TP2 and TP1 share keys. 20 cycles. After each wake: get of keys another server put while it was asleep, own put read back, and cross-server read-back of its post-wake puts.
    • 2080/2080 outputs are token-identical to an awake control run (/reset_prefix_cache?reset_external=true at the same points). Every hit read 2032/2048 tokens from the store.
  • TP2 plus a never-sleep server, 10 cycles: full sleep, selective wake (weights, then kv_cache), and release_kv_cache_memory. 400/400 identical to control.
  • Memory had no drift: TP2 at 3583 MiB/GPU asleep and 86467 woken, TP1 at 2077 / 85101, the same in every cycle. Host RSS of each server was unchanged from cycle 1 to cycle 20. 0 transfer errors.
  • 64 concurrent requests up to /sleep?mode=wait, with puts in flight: 64/64 complete, 0 timeouts, and requests not yet scheduled run after the wake. mode=abort: all aborted, and sleep takes 0.5 s.
  • HiSparse (GLM-5.2-FP8 TP4, HiSparse + MooncakeStore through MultiConnector): 78 host regions stay registered and the device region is released and registered again. 12 release_kv_cache_memory cycles: device memory returns to the same value every cycle, and store hits read 4160+ tokens. Exact-retrieval accuracy on store hits after a release equals an interleaved control on the same prompts (72/96 vs 72/96, with 140/144 per-request outcomes equal). Level-1 sleep was not run on this model: it would back up about 704 GB of weights to host RAM.

---

This PR contains AI-assisted code (Claude Code); commits carry a Co-authored-by trailer.

🤖 Generated with Claude Code

[mergify]

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

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @aoshen02.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork