Repository navigation
Conversation
maybe_send_cached_prefix_chunk records a completion event on forward_stream so the transfer worker can wait for the prior step's prefill writes. torch.cuda.Event raises "Tried to instantiate dummy base class Event" wherever torch has no CUDA build, killing the prefill scheduler as soon as a request reuses a cached prefix under overlap scheduling. self.device_module is resolved once from self.device in init_overlap(), and is the same module forward_stream was built from, so the event and the stream cannot come from different accelerators. Resolving the module from the process default instead would reinstate that mismatch whenever --device names something other than the default accelerator. Reproduced on Intel XPU with PD (nixl transfer) + hierarchical cache: the shared prefix in the second request of a group crashes the prefill scheduler; with this fix a 64-request shared-prefix run completes with 0 KV-transfer failures. The CPU unit test injects a stub device module and asserts the event comes from it and records against forward_stream, so it fails on the pre-fix source, which never touches the stub. It pins SGLANG_DISAGG_PREFILL_EARLY_SEND_CACHED_PREFIX with override() rather than relying on the default, so a deployment that turns early send off cannot make the case fail for an unrelated reason. Signed-off-by: KP, Jeeja <jeeja.kp@intel.com>
maybe_send_cached_prefix_chunk records a completion event on forward_stream and attaches it to the sender so the transfer worker can wait for the prior step's prefill writes before reading those prefix pages. sgl-project#41395 made mooncake wait on it; nixl's add_transfer_request still took no wait_event, so the event was recorded, attached and dropped, leaving the RDMA read racing the writes it was meant to order against. SGLANG_DISAGG_PREFILL_EARLY_SEND_CACHED_PREFIX defaults to True, so this is the default PD path. The failure is silent: the transfer did move the bytes, they just were not final, so decode reads partially written KV and only an output-quality check can catch it. The nixl worker now waits the way mooncake's does: before the failed-room check and before the staging gather, skipped for a room that already failed, and cleared so a chunk re-enqueued on a watermark defer does not wait again. The sender side moves to a CommonKVSender._take_early_send_wait_event() helper, with _early_send_wait_event initialised in __init__, so mooncake, nixl and mori read the one-shot event the same way instead of through getattr. Adds tests that the mooncake worker waits before the staging gather, and that the nixl worker waits before reading KV and before the staging gather, skips the wait for a failed room, and forwards the event onto the queued chunk.
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
|
/tag-run-ci-label |
|
CI is complete. Failures are not related to this PR: Test: test/registered/layers/mamba/test_mamba_slot_fused.py Test: test/registered/npu/basic_function/HiCache/test_npu_hicache_mla.py Test: test/registered/npu/basic_function/pd_disaggregation/test_npu_pd_disaggregation.py |
Lzy17
left a comment
Lzy17
left a comment
There was a problem hiding this comment.
LGTM. test_add_transfer_request_forwards_the_wait_event_to_the_queued_chunk passes chunk_id positionally (False, 0, wait_event=event). #41140 removes chunk_id and makes wait_event keyword-only, so whichever lands second needs to adjust this call.
This branch has not been deployed
3 participants
Motivation
Under overlap scheduling,
maybe_send_cached_prefix_chunkcan start the KV transfer for a cached prefix while the previous step's prefill forward is still writing those pages onforward_stream. It records a completion event and attaches it to the KV sender, so the transfer worker can wait for those writes before reading the pages. Two bugs broke this:torch.cuda.Event(), which raisesTried to instantiate dummy base class Eventon torch builds without CUDA. On Intel XPU with PD + hierarchical cache, the prefill scheduler died on the first request that reused a cached prefix.add_transfer_requesthad nowait_eventparameter. The event was recorded, attached and then dropped, so the RDMA read raced the writes it was meant to wait for.SGLANG_DISAGG_PREFILL_EARLY_SEND_CACHED_PREFIXdefaults toTrue, so this is the default nixl PD path. Nothing fails visibly: the bytes do arrive, they just may not be final yet, and only an output-quality check catches it.Modifications
disaggregation/prefill.py: create the event withself.device_module.Event().device_moduleis set once fromself.deviceininit_overlap(), the same moduleforward_streamwas built from, so the event and the stream are always on the same accelerator. Using the process-default accelerator would bring that mismatch back whenever--deviceis not the default.disaggregation/nixl/conn.py:add_transfer_requestnow acceptswait_eventand puts it on the queuedTransferKVChunk. The transfer worker waits on it the same way mooncake's does:disaggregation/common/conn.py: add aCommonKVSender._take_early_send_wait_event()helper (returns the event and clears it), and set_early_send_wait_eventin__init__. mooncake, mori and nixl now all read the event through this helper instead ofgetattr.Tests
test_prefill_abort_result_cleanup.py: injects a stub device module and checks that the event comes from it and is recorded onforward_stream. This fails on the old code, which never touches the stub. The env flag is pinned withoverride(True).test_nixl_backend_basic.py: the nixl worker waits before reading KV and before the staging gather, skips the wait for a failed room,add_transfer_requestputs the event on the queued chunk, and taking the event clears it.test_mooncake_transfer_batching.py: the mooncake worker waits before the staging gather; the existing early-send test uses the new helper.Validation
Intel XPU, PD with nixl transfer + hierarchical cache, overlap scheduling:
Checklist
CI States
Latest PR Test (Base): 🚫 Run #38246004209
Latest PR Test (Extra): ❌ Run #38246003869
Latest PR Test (AMD ROCm 10): ❌ Run #38246004352