[Core] Sleep mode only with KV connectors that follow the KV cache mapping

[aoshen02]

Problem

Sleep mode (and a KV cache discard with the cumem allocator) unmaps the KV cache, and the wake-up maps new physical pages at the same virtual addresses. A KV connector that registered that memory with a transport (RDMA memory regions in NIXL and Mooncake, IPC handles in others) keeps addressing the old pages. P/D disaggregation then returns wrong output with HTTP 200, and the pinned pages are not freed.

Principle

• A registration must live exactly as long as the mapping it covers. The allocator, which owns the mapping, signals unmap and map; the connector that registered releases and restores. Worker, executor and engine contain no connector code.
• Whether a connector can do that is a static property of its class. The base class cannot know it, so the default refuses, and the check runs once, up front.

Design

• vllm/device_allocator: register_tag_hooks(tag, before_unmap, after_map) returns the function that unregisters them. CuMem and XPU allocators run before_unmap before they unmap any allocation of the tag (sleep, discard), so a failing hook leaves everything mapped, and after_map after every wake-up that leaves all allocations of the tag mapped. Restores are idempotent, so the next wake-up retries a failed one.
• KVConnectorBase_V1.supports_sleep_mode is False by default; MultiConnector supports it only if every child does and forwards to them. A supporting connector implements release_kv_caches (drain, then release) and restore_kv_caches.
• ensure_kv_transfer_initialized subscribes the worker connector for the kv_cache tag exactly when it declares support; ensure_kv_transfer_shutdown unsubscribes it. Declaring support therefore means being called.
• EngineCore refuses an unsupported connector at startup when sleep mode is enabled. With the cumem allocator alone (supported on main with any connector as long as nothing is unmapped), sleep(level >= 1) and release_kv_cache_memory() refuse it before anything is unmapped. After a wake-up it republishes the workers' handshake metadata, since the KV cache was registered anew.
• A failing release raises. So that the engine can still recover with a wake-up:
  • Executor.sleep counts as asleep even if a rank failed (as discard already does);
  • MultiprocExecutor.collective_rpc drains every rank's response before raising, so the next RPC is answered by its own responses;
  • a rank enters the collective communicator suspend/resume (enable_nccl_comm_suspend) even when its own allocator step raised.

NIXL support: #59635 (on this PR and #59713). Mooncake over RDMA: #59625.

Behaviour change

With sleep mode enabled, an engine whose KV connector does not declare support fails at startup with <Connector> does not support sleep mode; with the cumem allocator alone, the sleep or KV cache release fails instead. A connector opts in by declaring support and implementing release and restore. One that only copies KV blocks (offload connectors) releases by waiting for its queued and in-flight copies and restores nothing; none is switched on here before its own hardware validation.

Tests

Mutation check: 22 mutations, all killed (each hook call, the mapped/asleep filters, subscription, unsubscription, MultiConnector all/forwarding, each refusal, the republish, the executor finally, the RPC drain, both comm finallys).

• test_cumem.py::test_tag_hooks_bracket_the_mapping_of_their_tag (real allocator): hooks bracket sleep, discard and selective wake for every subscriber; a failed restore is retried by a wake-up that maps nothing; a failing before_unmap leaves everything mapped.
• test_kv_connector_lifecycle.py: through the real ensure_kv_transfer_initialized/shutdown, a supporting connector is released and restored around the mapping, an unsupported one never is, and a shut down one is not called.
• test_engine_core.py: startup refusal on a real engine; refusal before unmapping for sleep levels 1 and 2 and the KV cache release, never for a pause; republish after wake-up.
• test_executor.py, test_multiproc_executor.py, test_sleep_mode_backend.py (comm suspension, also when the backend step raises), test_multi_connector.py.
• Hardware (NIXL over rc_mlx5, see [KV Connector] Support sleep mode with the NIXL pull connector #59635): decode, prefill and both sleeping, 2 cycles each: 803 MiB asleep, 30/30 identical after every wake-up.

Not covered

• The XPU allocator gets the same hooks, but they are not exercised on XPU hardware.
• Ray and uniproc executors are unchanged.

AI assistance: written with Claude Code and reviewed by the submitter.

🤖 Generated with Claude Code

[mergify]

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

[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

[aoshen02]

Superseded by #59625, which includes this connector API (with connectors supporting sleep mode by default and only the ones that cannot follow the KV cache remap refused). Closing.