[KV Connector] Support sleep mode with the NIXL pull connector

[aoshen02]

Purpose

Stacked on #59625; only the last commit is new. This replaces the earlier lease-based revision of this PR (and the
closed #59713): there is no lease and no new protocol between engines.

#59625 refuses sleep mode with NIXL because peers keep the stale KV cache registration. Sleep keeps the virtual
addresses and maps new physical pages, so everything NIXL bound to the old pages goes stale: the memory registration
(rkeys), the local prepped dlists, the handshake payload served to peers, and each peer's loaded remote agent.

Evidence on #59625 with the NIXL refusal forced off (GB300, cross-node 1P1D, Qwen3-8B, UCX RDMA over RoCE):

• After a decode sleep and wake: 0/30 outputs match the awake control. After a prefill sleep and wake: 0/30.
• No errors and no failed transfers: on GB300 there is no nvidia_peermem, so UCX registers GPU memory through
  dmabuf, and the stale MR pins the old pages; peers read old data without any error.
• Sleep frees nothing of the KV cache (67763 MiB asleep instead of 2077), and after the wake the GPU holds both
  (150787 MiB instead of 85101).

This PR implements the #59625 hooks for the NIXL pull connector (NixlConnector / NixlPullConnector):

• release_kv_caches() (before the unmap): waits until no own read is in flight and no handshake is pending (bounded,
  TimeoutError keeps the registration), then releases every local prepped dlist, drops every remote engine (their
  split handles reference local memory; they are re-handshaked on demand), and deregisters the KV cache.
• restore_kv_caches() (after the remap): registers the same descriptors again, re-preps the local dlist, and
  rebuilds the handshake payload with the new agent metadata.
• Peer refresh through the existing handshake only:
  • After a wake-up that restored the KV cache (tags is None or contains kv_cache), EngineCore hands the
    workers' payloads to the scheduler (the startup code, extracted).
  • The scheduler-side connector owns the registration epoch: each republish bumps it and stamps the same epoch on
    every rank's payload, so ranks cannot advertise different epochs, even after a partly failed release.
  • The side-channel listener serves the new payloads; request_finished adds remote_registration_epoch to
    kv_transfer_params.
  • The reader records the epoch of the payloads it loaded (the minimum over the remote ranks of one handshake, in
    case it straddled the republish). A request naming a newer epoch drops that remote engine and handshakes again
    through the existing path; while reads through the old registration are still in flight, the request is
    re-queued instead. Older epochs never trigger a re-handshake.
  • The reader checks whether the agent is loaded once per request (_start_read), so a handshake that completes
    concurrently cannot install an old agent between the epoch check and the read.
• Under sleep mode the worker sets UCX_RCACHE_ENABLE=n and UCX_TLS=^cuda_ipc right before creating the agent,
  each only if the user did not set it, and logs once which defaults it applied (as UCX_RCACHE_MAX_UNRELEASED in
  nixl_utils.py does). They are process-wide: other UCX users in the worker process (the NIXL EPLB communicator,
  the EC NIXL connector, NVSHMEM with a UCX transport) and child processes inherit them.
  • UCX's registration cache defers the real deregistration. On GB300 the cached dmabuf MR kept the released pages
    pinned in about half of the sleeps, and release_kv_cache_memory freed nothing. With the cache off, every sleep
    freed the KV cache.
  • A peer's cuda_ipc import (intra-node NVLink or MNNVL) pins the exporter's pages inside the peer's UCX cache. Over
    MNNVL the outputs stayed correct, but the prefiller's KV cache stayed resident asleep and was doubled after the
    wake. Without cuda_ipc in the exporter's UCX context no peer can import it; reads use RDMA (on one node: RDMA
    loopback through the NIC). Cost measured below: none on GB300.
• Static policy, checked once at startup: NixlPullConnector.supports_sleep_mode is True only with the UCX backend
  (backends == ["UCX"]), the only one validated. NixlPushConnector stays refused: the producer writes through the
  decoder's registration and is not refreshed.
• NIXL_CONNECTOR_VERSION 13 -> 14 (payload field and kv_transfer_params key). Routers must forward
  kv_transfer_params unchanged, as they already must for remote_block_ids.

Test Plan

• Unit tests (CPU-capable, fake NIXL agent), organized around invariants:
  • test_decoder_reads_from_republished_prefiller[all-ranks|one-rank|straddled] (scheduler level, real side
    channel): whichever ranks re-registered, a handshake after the republish loads the epoch the prefiller
    advertises, so the decoder reads instead of handshaking again; a handshake that straddled the republish is redone
    exactly once. The one-rank row (one rank's restore was a no-op) livelocks on the previous revision of this PR.
  • test_sleep_cycle_releases_and_restores_registration: with a tracking agent, nothing built on the registration
    survives a release (registrations, local and remote dlists, remote agents), and a restore registers the same memory
    with exactly one local dlist and publishes the new agent metadata; both calls idempotent over two cycles.
  • test_release_waits_for_reads_and_handshakes: an in-flight read or a pending handshake holds the release; a wedged
    one times out and keeps the registration; shutdown while released does not deregister again.
  • TestPeerReplacement::test_rehandshakes_re_registered_peer: same epoch reads through the loaded agent; a newer
    epoch drops it and handshakes again; with reads in flight the request is deferred until they finish.
  • test_sleep_mode_ucx_defaults: the UCX defaults under sleep mode only, and a user setting wins per variable.
  • test_config.py: NIXL pull supported, NIXL with a non-UCX backend and NIXL push refused, MultiConnector rows.
  • test_engine_core.py::test_wake_up_republishes_kv_connector_handshake_metadata[all|kv_cache|weights]: republish
    after the executor woke up, and only when the KV cache was restored.
• Mutation check: every piece of the change altered on its own, including the previous revision's worker-owned epoch.
• pytest tests/v1/kv_connector/unit/{test_nixl_connector,test_config,test_mooncake_connector,test_multi_connector}.py tests/v1/engine/test_engine_core.py tests/v1/worker/test_sleep_mode_backend.py; the NIXL and config files 5 times;
  pre-commit on the changed files.
• Hardware: NVIDIA GB300 Slurm cluster, NUMA-local to the only active NIC (mlx5_4, RoCE, GID 3). Image: vLLM nightly
  aarch64 + this branch as per-file mounts with sha256 gates (all 12 mounted files equal to this commit). NIXL 1.4.1
  (bundled UCX). Qwen3-8B, prefix caching off, --gpu-memory-utilization 0.3, level-1 sleep, UCX settings left to
  the branch defaults. Toy proxy from tests/v1/kv_connector/nixl_integration. Each check sends 30 prompts of
  1.1k-1.7k tokens (greedy, 32 tokens) and compares them with the awake control; transfers are counted from NIXL's KV
  transfer metrics and the logs are scanned for transfer failures and UCX errors. The UCX protocol in use is taken
  from UCX_PROTO_INFO=y.
  • Cross-node: P and D on two nodes, one GPU each.
  • Single node: P on GPUs 0-1 and D on GPUs 2-3 (TP2 each), with and without sleep mode, plus a transfer benchmark
    (30 sequential requests of ~1.3k and ~5k tokens, twice; 128 requests at concurrency 32), with transfer time and
    bytes from D's vllm:nixl_xfer_time_seconds and vllm:nixl_bytes_transferred.

Test Result

Cross-node 1P1D TP1 (final commit; UCX picked rc_mlx5 zero-copy GPU-to-GPU reads). "Asleep" and "woken" are MiB on
the slept GPU; awake is 85101 MiB.

Case	Cycles	After every wake	Asleep	Woken	Failed transfers
Sleep D	3	30/30	2077	85101	0
Sleep P	3	30/30	2077	85101	0
Sleep both	3	30/30	2077 / 2077	85101	0
Selective wake, weights then kv_cache (both)	1	30/30	2077 -> 19415 -> 85101		0
Selective wake, kv_cache then weights (both)	1	30/30	2077 -> 68809 -> 85101		0
release_kv_cache_memory on both (pause wait, release, wake kv_cache, resume)	1	30/30	P 20983, D 19415	85101	0
P sleep (mode=wait) under load: 90 requests at concurrency 30, wake after 5 s	1	0 errors; 30/30 after	2077	85101	0

• The decoder re-handshakes exactly once per wake that restored the KV cache, including the selective wakes (the
  weights-only step no longer republishes); the first request after a wake takes 0.21-0.24 s versus 0.18 s.
• Under load, 18/90 outputs differ from the sequential control, as in a concurrent run without any sleep (batching
  noise); every post-wake check is 30/30.

Single node, P on GPUs 0-1, D on GPUs 2-3, TP2 each (sleep mode with the branch defaults, so reads use RDMA loopback):

Case	Cycles	After every wake	Asleep (2 GPUs)	Woken	Failed transfers
Sleep D / P / both	3 each	30/30	7166	172934	0
release_kv_cache_memory on both	1	30/30	P 26054, D 25194	172934	0

Single-node transfer cost of the UCX_TLS=^cuda_ipc default (medians of two launches each; "transfer" is NIXL's
per-transfer time, one transfer per TP rank and request):

Workload	Without sleep mode: cuda_ipc (NVLink)	Sleep mode: RDMA loopback	Without sleep mode, UCX_TLS=^cuda_ipc
~1.3k tokens, sequential: request p50	59 ms	54 ms	52 ms
~1.3k tokens: transfer time / rate (90 MB)	8.0-8.9 ms / 10-11 GB/s	4.0-4.6 ms / 20-22 GB/s	4.1 ms / 22 GB/s
~5k tokens, sequential: request p50	101-119 ms	85 ms	84 ms
~5k tokens: transfer time / rate (355 MB)	28-36 ms / 10-13 GB/s	15.1 ms / 23.5 GB/s	15.0 ms / 23.7 GB/s
128 requests at concurrency 32, ~1.3k / ~5k tokens	91-96 / 30.8 req/s	117.5 / 32.8 req/s	120.7 / 32.4 req/s

• No NVLink cost was measured on GB300: with UCX's defaults, RDMA loopback through the NIC was faster than
  cuda_ipc for these transfers (cuda_ipc spent 5-29 ms per transfer in posting). This is UCX 1.4.1-bundled cuda_ipc
  between two instances with disjoint CUDA_VISIBLE_DEVICES; a tuned cuda_ipc setup, or a node whose NIC is slower
  relative to NVLink, may differ. Users who want cuda_ipc can set UCX_TLS themselves; then a peer's import keeps the
  slept KV cache resident (measured over MNNVL: 67763 MiB asleep, 150787 MiB after the wake; UCX_CUDA_IPC_CACHE=n on
  the reader frees it at +140 ms per request).
• Baseline ([KV Connector] Support sleep mode with MooncakeConnector over RDMA #59625 head with the refusal forced off) under the cross-node harness: decode sleep 0/30, prefill sleep
  0/30, 0 errors, KV cache pinned asleep (67763 MiB) and doubled after the wake (150787 MiB).
• Earlier revision of this PR at scale (same peer-refresh path): V4-Flash TP4+TP4 cross-node 30/30 cycles, 900/900
  probes, 23,792 transfers, 0 failed; Qwen3-8B 100/100 cycles, 0 failed.
• Unit: 391 passed across the files above. The failures are tests that need models this node cannot download (gated
  meta-llama/Llama-3.2-1B-Instruct, and test_engine_core_tp, _concurrent_batches and
  _invalid_request_id_type); run alone on [KV Connector] Support sleep mode with MooncakeConnector over RDMA #59625's head, the same 4 fail. The NIXL and config files: 5 repeats,
  208 passed each time. Mutation check: 37 mutants, all killed; every
  new test kills at least one. The reviewer's fix probes pass (28/28).

Known limitations:

• Host memory grows by about 1.3 MiB per sleep cycle on both P and D (measured on the previous revision with
  smaps_rollup Rss at Qwen3-8B TP1, num_threads=4: 1.25-1.34 MiB/cycle, linear over 500 cycles; Mooncake and sleep
  without a connector are flat). Each re-handshake creates and closes UCX RC endpoints, and UCX does not return their
  send-queue buffers (about 640 KiB pinned per side per re-handshake, plus allocator fragmentation). The cost scales
  with NIXL threads and TP ranks (num_threads=1 cuts it to 2/5). This revision keeps exactly one re-handshake per
  wake. Follow-up: keep the remote agent across a wake and reload only the peer's memory sections (NIXL partial
  metadata without connection info), which needs NIXL to replace the old sections of a loaded agent; that is not
  verified for NIXL 1.4.1. The UCX endpoint leak itself is to be reported upstream.
• NIXL backends other than UCX (refused), push mode (refused), and cuda_ipc/MNNVL transfers under sleep mode (disabled
  by default; a user who sets UCX_TLS with cuda_ipc keeps the peers' imports pinned, as measured above).
• A release that fails on one TP rank (timeout or deregistration error) leaves that worker released while
  Executor.sleep reports the engine awake, so nothing restores it. Partial-failure handling belongs to the executor
  and is shared with [KV Connector] Support sleep mode with MooncakeConnector over RDMA #59625's other connectors; it is not changed here. The registration epoch no longer diverges
  across ranks in that case, because the scheduler stamps it.
• With enforce_handshake_compat: false, a version-13 decoder ignores the epoch and never refreshes after a
  version-14 prefiller wakes; with the rcache-off default that shows up as failed transfers.

---

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

🤖 Generated with Claude Code

[mergify]

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

[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

[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

[claude]

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.