[PD] Don't release KV pages while Mooncake transfers are in flight

[KastanDay]

The bug

An abort or timeout during a prefill→decode transfer frees the request's KV pages back to the allocator, but it does not cancel the transfer. The next request can then be allocated those same pages and have them overwritten by the earlier, still-in-flight write, causing the next request to hold the first request's KV, with no error raised anywhere.

@ShangmingCai already anticipated this in #27372, which prevents the prefill starting new work after an abort and this PR implements what he mentioned as a TODO, which is handling transfers that are already in flight.

The bug's window is narrow: the transfer has to still be in flight and the pages have to be reallocated before it lands. But in our setup, we often abort requests that are too slow, so I believe we are hitting this scenario. The only symptom is a slightly wrong answer to our users.

No cost on the happy-path: no extra collective and no wait unless a request has failed.

Only ~995 lines of implementation code, the rest is tests and doc strings.

Reproduction script

This script reproduces the bug on current main (CPU only, no GPU):

[T1] transfer worker is inside a write to decode page 0xdead0000
[T2] request aborted while that write is still in flight
[T3] scheduler was told: Failed
RESULT: REPRODUCED

Reporting a terminal state is what lets the allocator hand those pages out, so Failed at [T3] is the bug. Asserting the ordering makes it deterministic. With this PR, [T3] reports Transferring and terminal arrives only after the write returns.

The core change

No new states and no new message types. KVPoll keeps its five states (Failed, Bootstrapping, WaitingForInput, Transferring, Success); what changes is which one the scheduler is told, and when:

• Before: the first failure reported Failed immediately. A terminal state is what authorizes the allocator to hand the request's pages out — so pages could be reused while a transfer was still writing to them.
• After: until every rank proves the pages are idle, the poll path reports Transferring — a state the scheduler already knows to wait on. Failed arrives only with proof. (Skipped entirely when nothing failed, so the happy path is unchanged.)

Correspondingly, ABORT_ACK changes meaning: from "I heard your abort" to "your pages are safe to reuse." Prefill now sends it only after that room's transfer work has drained, and it must echo a per-request nonce.

Each bootstrap room gets a lease-counted lifetime. Transfer work holds a lease; a failed request closes the room; pages are released once it is closed with zero leases and, on decode, once every peer has acknowledged.

Required for the barrier to hold

Fixed these three paths from bypassing the abort:

• The lease would drop early. send_kvcache returns on the first failing future, but Future.cancel() is a no-op once a future is running — so siblings keep writing while the worker reports the chunk done.
• HiCache bypassed the gate. A failed restore releases pages regardless of poll state, including while the barrier is holding.
• Room identity is ambiguous. bootstrap_room is recycled, so a queued chunk can lease the wrong request's lifetime.

Tests

• Integrated 3-node H200 validation passed. With two TP8 prefill and one TP8 decode (24 H200s, STRICT policy), 192/192 immediate victim reuse probes succeeded after 192 in-flight donor aborts. There were zero transport, semantic, hard, or critical failures; zero foreign canaries; a decode transfer-queue peak of 20; and zero model restarts. Raw summary, matrix, telemetry, and the admission-control disclosure.

python -m pytest -q test/registered/unit/disaggregation/test_transfer_ownership_barrier.py

71 tests named by property. Every fix has a regression test verified to fail with that fix reverted. unit/disaggregation 72 → 143 passed; unit/managers + unit/mem_cache unchanged apart from additions.

Related PRs

• PP is unvalidated. The added collectives cover TP and CP. [PD+PP] Honor PP consensus for bootstrap and prealloc #31869 looks complementary.
• Overlaps Fix race on shared ZMQ sockets in mooncake KV manager senders #29984 (socket serialization) and [PD] Prevent decode scheduler from blocking on ZMQ sends to a stalled prefill peer #31144 (non-blocking sends). I deliberately used the same function names and placements so the PR diffs collapse and merge cleanly.

Glad for this to land as a completion of #27372. @AsterWang identified this race first in #24580 but I think my solution is a bit safer.

---

CI States

Latest PR Test (Base): ❌ Run #31566610945
Latest PR Test (Extra): ✅ Run #31620191653

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[ShangmingCai]

Hi @KastanDay , thank you for the digging and the PR, but I think this PR is quite complicated since it introduces many behavior changes. I will explore a more lightweight and universal fix for all transfer backends in the following week.

[KastanDay]

Mostly resolved now.