[PD] Centralize drain-aware abort acknowledgements

[livingshade]

Motivation

When decode aborts or fails a request, prefill may still be RDMA-writing into the decode's KV pages; releasing them early lets a late write corrupt the request that reuses them. Deferred decode KV release (default-on since #41023) holds the pages until every prefill rank ACKs that its writes drained.

This PR adds the drain-aware ACK to Mori and moves the ACK protocol into the common layer shared by Mooncake, NIXL and Mori.

Built on #40645, #41402 and #41404 and resolved potental conflicts.
Related to #33861 and #34510.

Modifications

Mori

• Count in-flight chunks per room and send the ABORT_ACK only after the room is Failed and the count drains to zero.
• Transfers still in progress when the shard worker gives up on them are handed to a per-shard drain worker, which waits for them before uncounting.
• Opt Mori into deferred release (supports_deferred_decode_kv_release = True).

Common

The common layer is extracted from what Mori established:

• Shared AbortNotification / AbortAck wire format. ABORT carries a per-room generation armed before sending, and decode only counts ACKs of the current generation.
• ACK fan-out: prefill keeps one ACK target per decode endpoint and ACKs all of them on drain, instead of [PD] Fan drain abort ACKs out to every decode peer of the room #41402's transfer_infos snapshot, which has no generation.
• Mooncake and NIXL ABORT handling moves onto the common helpers.

Tests

• Refactored related tests into Common and backend-specific test.

Speed Tests and Profiling

No change on the normal transfer path beyond per-chunk counter bookkeeping.

Accuracy Tests

No change to model forward or sampling. 1P1D H100 (Mooncake), TestDisaggregationAccuracy: gsm8k 0.83.

Unit tests (CPU): test_deferred_decode_kv_release.py (common protocol, tracker, fan-out, decode release), plus backend-specific test_mooncake_deferred_kv_release.py, test_nixl_deferred_kv_release.py and test_mori_deferred_kv_release.py.

1P1D e2e on H100 (TP1 / TP1 over IB), Mooncake and NIXL. test_disaggregation_chunked_prefill_abort, TestDisaggregationMooncakeFailure

Mori: AMD CI stage-b-test-large-8-gpu-mi35x-disaggregation-amd (includes test_mori_transfer_engine_e2e.py).

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ✅ Run #37089529256
Latest PR Test (Extra): ⚠️ Not enabled -- add run-ci-extra label to opt in.
Latest PR Test (AMD ROCm 10): ❌ Run #37089529181

[inkcherry]

Thanks @livingshade! LGTM, verified on MI355*2 node and left minor comments.

[inkcherry]

I validated it on MI355X E2E: 45 focused tests passed, the registered Mori E2E passed for TP1→1 and TP2→4, and a cross-node 1P1D abort test remained healthy with a successful follow-up request.

[inkcherry]

we can userun-mori-pd label if the AMD CI still has some issue. see
#37077

[jambow0320]

Two thoughts, both about scope rather than correctness.

1. Can we move the pre-arm into CommonKVReceiver._send_abort_notification

_send_abort_notification has a second caller: _check_waiting_timeout (common/conn.py:1617, reached from poll() in all three backends). It sends the ABORT and sets abort_notified = True without going through Scheduler._abort_request, so register_deferred_abort_room is never called for it.

This PR already fixes this for Mori. The same logic in the base method would fix mooncake and nixl too:

# CommonKVReceiver
def _send_abort_notification(self):
    if self.kv_mgr.enable_deferred_decode_kv_release and self.init_time is not None:
        self.kv_mgr.register_deferred_abort_room(self.bootstrap_room)
    for bootstrap_info in self.bootstrap_infos:
        ...

self.init_time already encodes what metadata_published encodes — it stays None until the end of send_metadata in all three receivers (mooncake:2600, mori:2055, nixl:2968) and _check_waiting_timeout already uses it as that marker — so the new field metadata_published may not be needed. (However, adding a field for readability is also reasonable.)

With putting 'register_deferred_abort_room' into '_send_abort_notification', scheduler.py:4918-4930 (the re-arm right after receiver.abort()) can be deleted. That in turn makes reset_deferred_abort_room and the setdefault override unnecessary.

2. Consider nixl's "no ack" fallback instead of the drain queue

nixl handles the same situation — writes possibly still in flight when completion can no longer be confirmed — by not acking at all:

# nixl/conn.py:1390-1393
self.record_failure(room, str(e))
self.update_status(room, KVPoll.Failed)
# No ack here on purpose: the DONE barrier bails on the first
# ERR, so siblings may still be writing; fall back to the timeout.

It needs no extra threads or queue.

The current drain path can stop prefill-side transfers: _enqueue_transfer_drain calls put without block=False, so a transfer worker blocks once the queue is full, while status.Wait() has no timeout parameter and _drain_worker retries indefinitely. _num_shards (default 8) sizes the transfer workers, the drain workers and the queue, so roughly 16 transfers that never reach a terminal status are enough to block every transfer worker — and shards are keyed by room % _num_shards, so healthy rooms stop too.

[ShangmingCai]

/tag-and-rerun-ci

[livingshade]

/rerun-failed-ci

[ShangmingCai]

/rerun-failed-ci

[ShangmingCai]

CI has passed: