[Core][KV Connector] Release transfer-held KV before a cache reset

[aoshen02]

Purpose

For RL with prefill/decode disaggregation, the trainer pauses P and D, updates the weights, and resumes. With PD, every operation that resets the KV cache failed on D with "Failed to reset KV cache even when all the running requests are preempted…". That covers:

• pause(wait|keep, clear_cache=True);
• sleep(level>=1);
• release_kv_cache_memory.

The reset failed because remote KV transfers still held blocks:

• D requests were still waiting for remote KV.
• P's finished requests were holding KV that D had not pulled yet.

A pause that keeps the KV had the opposite problem: P waited for those exports, up to 480 s when D was paused too and would not pull them.

The fix starts from two facts:

• A pause only stops compute.
• A cache reset needs every block back, including blocks held by remote transfers.

So:

• Release, then reset. Once the engine is idle, an operation that resets the cache calls Scheduler.release_transfer_kv() once. That call:

  • asks the connector to end its transfers, through the new KVConnectorBase_V1.abort_transfers() hook: P stops holding finished requests' KV for remote readers, and D asks P to drop loads still in flight. MultiConnector fans it out; this PR implements it for Mooncake.
  • keeps every request whose load is cut short. It is marked as a failed load with 0 computed tokens, so it recomputes its KV after resume through the existing failed-load recovery, under either kv_load_failure_policy. Nothing is aborted (like SGLang's retract, with D recomputing).

  The engine then waits for the transfers to give their blocks back, and resets. All of this lives in one place, EngineCoreProc._when_idle. It covers pause with clear_cache, sleep and release_kv_cache_memory. A resume fails any such operation still pending.

• Keep pauses do not wait. A paused engine does not wait for KV that finished requests hold for remote readers, unless it is releasing it. So a pause that keeps the KV returns at once, and the exports finish after resume.

• Waiting requests recompute. The reset recomputes waiting requests that hold KV, including remote loads that arrived but were not scheduled yet. It shares the preemption path, and the queue order is kept.

• No raise on held blocks. reset_prefix_cache(reset_running_requests=True) returns False while transfers hold blocks, as documented, instead of raising.

• Docs. The PD note in async_rl.md: pause both P and D, and use kv_load_failure_policy="recompute".

Not covered:

• NIXL keeps the default no-op hook. D reads P's blocks one-sided and cannot see a revoked lease, so P waits out its 30 s lease TTL. That wait is bounded and safe.
• The in-process engine cannot wait for transfers. Its cache reset still fails while transfers hold blocks, as on main.

Stacked on vllm-project#59625 and the Mooncake transfer PR (this fork PR's base branch).

Test Plan

• Unit: scheduler, async scheduler, engine core, Mooncake, MultiConnector and NIXL suites.

  • test_release_transfer_kv_frees_blocks_for_reset checks:
    • loads in flight are retracted, not aborted, and recompute after the load ends (written, failed blocks, or a failed request), also when a resume cancels the reset first;
    • held sends are dropped;
    • loaded and arrived KV is recomputed in order;
    • a streaming session keeps its KV;
    • a retracted load keeps its blocks until its transfer ends.
  • test_cache_reset_waits_for_transfer_kv covers every pause mode and release_kv_cache_memory, each idle or draining, with or without held transfers. It drives the real input loop, which must never block on input while work is pending.
  • test_resume_cancels_a_pending_cache_reset.

• Mutation: each key line removed fails its test (31 of 31 across both PRs).

• Hardware: GB300, 1P1D over Mooncake RDMA. The lifecycle matrix covers:

  • pause abort/wait/keep on P, D and both, including keep with clear_cache;
  • sleep L0/L1/L2 in abort and wait;
  • abort_requests;
  • reset_prefix_cache;
  • release_kv_cache_memory;
  • in-place weight reload under abort/wait/keep;
  • client disconnects.

  All cases run under load.

Test Result

	main	main + vllm-project#59625	+ both PRs
lifecycle matrix (Mooncake, 27 cases)	6/22 pass	most abort and clear cases fail	27/27 on every revision (v20–v26)
NIXL matrix (19 cases, no sleep)	–	–	19/19; P pause ≤ 30 s (lease)

🤖 Generated with Claude Code