[PD] Consolidate sender and receiver implementations

[jambow0320]

Motivation

Mooncake, NIXL, and Mori each implement their own KV sender / receiver lifecycle (sending, registration, metadata publication, polling, timeouts, cleanup, failure handling). The code is duplicated, and the backends are inconsistent in details such as failure-reason recording, terminal-state latching, and timeout trigger conditions.

NIXL determines transfer completion differently from the other two backends:

• Mooncake / Mori: after all chunks have been written, prefill sends a status message to decode over ZMQ, and decode determines completion by counting prefill ranks.
• NIXL: prefill attaches a notif string to every RDMA WRITE ({room}_kv_*, {room}_aux, {room}_state_*, {room}_stg_*, etc.). Decode calls agent.get_new_notifs() in NixlKVReceiver.poll() and uses TransferStatus to count, per PP rank, the arrival of chunks, parts, aux, and state to determine completion. Decode has to maintain a separate set of counting logic, and notifs are only read inside the receiver's poll().

This PR consolidates the sender / receiver lifecycle of the three backends into a common implementation, and switches NIXL completion notification to the same ZMQ mechanism used by Mooncake / Mori, no longer using NIXL notifs.

Refs: RFC #33861, Tracking #34510.

Modifications

1. NIXL switches to the same ZMQ completion mechanism as Mooncake and removes notifs

Prefill side

• For each chunk, the transfer worker waits until all of its handles reach DONE before processing the next chunk. When the last chunk completes, it calls conclude_transfer(Success), which sends a KV_STATUS message to all non-dummy decode endpoints of the room. The failure path is the same as upstream: it sends KV_STATUS=Failed with the failure reason.
• Staging path: after the handle of each staged write reaches DONE, a ZMQ CHUNK_READY is sent (sharing send_chunk_ready with Mooncake). Before the next peer reuses the prefill-side packing buffer, the worker likewise waits for that handle to complete.
• NIXL transfer calls no longer pass a notif and no longer call send_notif.

Decode side

• The listener thread handles four message types: KV_STATUS, CHUNK_READY, STAGING_REQ, and ABORT_ACK. KV_STATUS is counted per prefill rank via apply_prefill_status, and the room is set to Success once required_prefill_response_num is reached; CHUNK_READY is passed to the staging handler via handle_chunk_ready (shared with Mooncake).
• The receiver's poll() only reads request_status and no longer calls get_new_notifs().

Removed

• TransferStatus, update_transfer_status() and its notif parsing functions (_track_kv_arrival, _track_kv_part_arrival, _handle_stg_notification, _handle_aux_notification, _maybe_submit_last_scatter), and check_transfer_done().
• transfer_source_rank, TransferKVChunk.chunk_id (along with the counter in the sender and the chunk_id parameter of add_transfer_request), the NIXL manager's exceptions dict, and the receiver's started_transfer / expects_state.
• The notif parameter of each NIXL send function.

Unchanged

• The NIXL agent configuration is kept as in upstream.
• The wire format of NIXL registration and metadata messages is unchanged.

2. Sender / receiver lifecycle consolidated into the common implementation

• send, poll, clear, failure_exception, abort (sender) and _register_kv_args, send_metadata, clear, failure_exception, abort (receiver) are implemented in CommonKVSender / CommonKVReceiver. The receiver's poll() uses the upstream common implementation.
• Backend subclasses:
  • NixlKVSender, NixlKVReceiver, MooncakeKVReceiver, and MoriKVReceiver are empty subclasses.
  • MooncakeKVSender only keeps the trace decorator (mooncake_trace_func) on send.
  • MoriKVSender only keeps clear, which holds the Mori-specific _abort_ack_lock when deferred KV release is enabled.
• Backend differences are provided through hooks on the manager:
  • build_receiver_registration_message() / build_receiver_metadata_message(): per-backend message encoding; the wire format is unchanged. Mooncake's device_kv_indices (DSV4 HiSparse) is forwarded through send_metadata(**backend_kwargs).
  • prepare_send_state_indices(): returns the input unchanged by default; Mori uses it to skip CP-replicated state and to normalize per component.
  • sender_trace_module: set to "mooncake" by Mooncake, used to create the sender's trace context.
  • add_transfer_request(): unified across the three backends as 4 positional parameters plus keyword parameters wait_event / trace_ctx; each backend ignores the parameters it does not use (wait_event is used only by Mori, trace_ctx only by Mooncake).
• send_chunk_ready() / handle_chunk_ready() are moved to CommonKVManager and shared by Mooncake and NIXL.
• The decode-side prefill_response_tracker changes from defaultdict(set) to a plain dict, with the entry created when the receiver is constructed. The listener thread only reads the entry, so a Success arriving after clear() does not re-create it.
• _staging_ctx is declared on CommonKVManager with a default value of None.

Behavior Changes After Unification

NIXL protocol

• Source of the completion signal: decode completion changes from "counting notifs" to "counting KV_STATUS=Success sent by prefill ranks", the same as Mooncake.

• When Success is sent: Success is sent when all handles of the last chunk have completed. If an earlier staging chunk has been deferred due to insufficient staging space and is still queued:

  • the prefill sender returns Transferring while _staging_outstanding > 0, so the source KV is not released;
  • the decode staging handler only marks the request complete after all allocations have been scattered.

  This matches how Mooncake handles it. Before this change, NIXL decode relied on notif counting to wait for all chunks to arrive.

• Staging deferral retry: when a chunk is deferred, the worker first waits for the writes already posted in the current round to complete, and then retries (before this change it retried immediately without waiting for the posted handles). The retry still reprocesses all decode peers of the chunk, the same as before this change and as Mooncake.

• Failure exception: the sender's failure_exception() raises KVTransferError, with the reason text taken from failure_records (the worker writes str(e) on the failure path). Before this change, NIXL re-raised the original exception object caught in the worker, so the exception type and text format in scheduler logs change.

• Abort received only by prefill: if an abort reaches only the prefill side while the last chunk's writes are in flight, prefill does not send Success, and decode does not complete the request. When prefill has already cleared the room, it sends no status, and decode fails the request after SGLANG_DISAGGREGATION_WAITING_TIMEOUT (300 s by default). This matches Mooncake. Before this change, NIXL decode counted all notifs once the writes landed and completed the request.

Sender (all three backends)

• send() returns immediately once the sender has reached a terminal state (Success / Failed) and no longer enqueues. Before this change, Mooncake / Mori relied on the status check inside add_transfer_request, and NIXL relied on _send_failed, which was only set in failure_exception().
• Transfer latency metric: Mooncake and Mori start reporting the sender-measured transfer_latency_s, timed from the first enqueue until poll() observes Success. Before this change these two backends did not set the field, and the metric used completion_time - prefill_transfer_queue_entry_time (covering only the last chunk). With chunked prefill, the new value includes the compute time of subsequent chunks, so latency_ms is larger and speed_gb_s is smaller. NIXL's timing is unchanged.
• poll() returns Failed when the room is not in request_status. Before this change, Mooncake raised KeyError in this case and NIXL returned WaitingForInput. poll() writes conclude_state once it observes a terminal state; before this change NIXL did not latch it.
• clear() saves the terminal state before removing manager state; cleanup of staging prefetch records (prefetched_rooms / prefetch_requested) is extended from NIXL to Mooncake.
• aux_index and state_indices are only passed with the last chunk. Before this change NIXL passed them with every chunk, but the worker only read them on the last chunk, so behavior is unchanged.
• When no failure reason was recorded locally, the default text of failure_exception() is unified as Failed due to an unknown reason from another rank (before this change it was NIXL KVSender Exception for NIXL and KV transfer failed for Mori).

Receiver (all three backends)

• When bootstrap_infos is None, send_metadata() uniformly records a failure reason and sets the status to Failed. Before this change, NIXL only logged without recording a reason (so the failure was classified as coming from another rank), and Mori returned directly.
• send_metadata() raises NotImplementedError when given KVTransferDestination.HOST. None of the three backends implements host destinations; before this change, passing this argument raised TypeError due to a signature mismatch.
• clear() writes conclude_state before removing manager state (a terminal state keeps its value; a non-terminal state is recorded as Failed).
• abort() can be called repeatedly and does not overwrite an already recorded failure reason when already Failed. Mori's abort() no longer immediately calls clear() and removes the failure reason: cleanup is done by failure_exception(), which the scheduler calls after the next poll() returns Failed, so reasons such as Aborted by AbortReq. are preserved. Before this change, Mori's failure reason was removed early, and failure_exception() could only report unknown reason from another rank.
• When no failure reason was recorded locally, the default text of failure_exception() is unified (before this change it was NIXL KVReceiver Exception for NIXL and KV transfer failed for Mori).
• The NIXL and Mori receivers switch to the upstream common poll(): after metadata is published, the waiting timeout is checked in the WaitingForInput state. Before this change, NIXL conditioned this on started_transfer and read notifs inside poll().
• The bootstrap_room is None checks in the Mori receiver are removed. The scheduler already rejects disaggregation requests without a bootstrap_room on intake (except for the FAKE backend, which uses FakeKVReceiver).

Mooncake decode CHUNK_READY handling

• Switches to the common handle_chunk_ready(): malformed messages are dropped after logging a warning; late messages for cleared rooms are dropped; a warning is logged when the staging handler is not initialized. Before this change these cases raised an exception and triggered an assert, respectively.

Wire compatibility

• The registration, metadata, and status message formats of Mooncake and Mori are unchanged, and the registration and metadata message formats of NIXL are unchanged. Unit tests assert the registration and metadata messages of all three backends frame by frame.

Accuracy Tests

• CPU unit tests test/registered/unit/disaggregation/ (based on upstream 0bdd0cc8dc): 637 passed, 11 skipped, 303 subtests passed.
• Added and adjusted tests:
  • New test_common_sender_lifecycle.py:
    • the add_transfer_request argument contract of the three backends with two chunks (aux_index / state_indices are only passed with the last chunk);
    • Mori's early-send event is handed to only one chunk, and CP-replicated state is skipped;
    • the sender's Success is held while a staging chunk is deferred;
    • the Mooncake trace context is passed to the worker and finished correctly.
  • Added to test_receiver_connection_pool.py:
    • frame-by-frame assertions of the registration and metadata messages of the three backends, and forwarding of Mooncake device_kv_indices;
    • deduplicated counting of Success from multiple prefill ranks;
    • an encode/parse round trip of CHUNK_READY;
    • a late Success concurrent with clear() does not re-create the rank tracker.
  • test_nixl_backend_basic.py:
    • new tests for the NIXL ZMQ completion path: Success is sent only after the KV, aux, and state handles have all completed; Success is not sent when any handle fails; a deferred staging chunk still completes its transfer after the last chunk has sent Success;
    • a new test for the decode listener thread handling CHUNK_READY and KV_STATUS;
    • the _do_staging_transfer test now checks the returned chunk_idx;
    • notif-dependent tests are removed (TestNixlTransferStatus, TestNixlNotifications, and the old receiver poll and staging notif tests).
  • Adaptations of upstream tests:
    • the helpers in test_mooncake_deferred_kv_release.py and test_nixl_deferred_kv_release.py drop chunk_id;
    • the fake agent in test_nixl_dsv4_kv_slots.py adds a default value for the notif parameter;
    • test_nixl_sender_failure_cleanup.py now asserts KVTransferError and its failure reason.
• Single-node GPU end-to-end validation (8× H200; NIXL over UCX cuda_ipc and TCP, Mooncake over TCP; Qwen3-8B and Qwen3.6-27B): GSM8K accuracy differs from the baseline within run-to-run variation; at concurrency 1 the output tokens are bitwise identical to a non-PD deployment; the servers stay healthy with no KV leak after abort stress and failure injection. Cross-node RDMA end-to-end validation has not been completed.

Speed Tests and Profiling

Setup: a single node with 8× H200, Qwen3-8B, one prefill (TP1) and one decode (TP1) instance, NIXL backend (UCX). The baseline is upstream 0bdd0cc8dc; the baseline and this PR were run alternately, restarting the servers for every run, with random-token inputs and fixed output lengths.

• TCP transport (UCX_TLS=tcp,cuda_copy): across loads with 2048–32768 input tokens and concurrency 1–128 (3 runs per version), no detectable difference in TTFT, ITL, or throughput.
• NVLink transport (UCX cuda_ipc): at concurrency 1 with 2048 input and 4 output tokens (15 runs per version), mean TTFT increases by 0.99 ms (69.59 → 70.58 ms, +1.4%) and p99 TTFT by 1.21 ms. The added time is in transfer-completion detection: decode-side decode_transferred increases by about 1.1 ms and prefill-side prefill_transfer_kv_cache by about 0.6 ms, while prefill forward time is unchanged; the cause has not been identified. At concurrency 16–256, no detectable difference was observed.

Checklist

• pre-commit checks passed (isort, ruff format, ruff check).
• Added tests for the sender / receiver lifecycle, wire format, and NIXL completion ordering.
• Re-run the two-process NIXL validation on the current revision.
• GPU / RDMA end-to-end validation and performance testing.

---

CI States

Latest PR Test (Base): ❌ Run #37925410328
Latest PR Test (Extra): ❌ Run #37925409795
Latest PR Test (AMD ROCm 10): ❌ Run #37925410225

[Lzy17]

Thanks for the PR. Items:

1. Rebase. Main's Mooncake worker already calls wait_event.synchronize(). Keep that path in the unified send() and correct the "only Mori uses wait_event" statement.
2. Agree the merge order with #38959 (poll_pp_consensus), #37358 (four new mori registration fields) and #40861 (NIXL wait_event accessor).
3. Add cross-node RDMA, staging and mori/AMD validation; only single-node H200 is covered.
4. Mooncake/mori transfer_latency_s is now measured from the first enqueue, which shifts nightly baselines. Please flag this to the nightly owners.

[jambow0320]

@Lzy17 Thanks for the review! I have rebased onto the latest main and resolved the conflicts. Once #38959, #37358, and #40861 are merged, I will update this PR accordingly. For performance testing, I currently only have a single H200 node locally, so I will ask around whether someone can help run the tests. Finally, regarding transfer_latency_s, who specifically should I flag this to?