From 2524fe4c9dbc592eb5ced7b5fc3239daf6c65886 Mon Sep 17 00:00:00 2001 From: leavelet Date: Sat, 6 Jun 2026 22:03:52 +0000 Subject: [PATCH] fix(disagg/mooncake): mark abort Failed and skip aborted chunks in worker Port the abort-correctness pair from upstream sgl-project/sglang: - #24522: abort() now calls kv_mgr.update_status(room, Failed) so the transfer worker and other pollers observe the abort, not only the local conclude_state. - #27372 (guard half): the transfer worker skips a chunk whose room has already failed or been aborted, so it never transfers into KV pages that may have been reclaimed. This matters more once per-layer async transfer widens the abort-vs-in-flight window. The decode->prefill ABORT/ABORT_ACK notification half of #27372 is deferred: it interleaves with staging/tracing infra this branch lacks and needs the PD test harness to verify, so porting it now would be unverified guesswork. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../srt/disaggregation/mooncake/conn.py | 23 +++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 5796cbf7c..472b92671 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -1000,6 +1000,21 @@ class MooncakeKVManager(CommonKVManager): while True: try: kv_chunk: TransferKVChunk = queue.get() + # Skip a chunk whose room has already failed or been aborted, so we never + # transfer into KV pages that may have been reclaimed. Port of upstream + # sgl-project/sglang #27372 (abort KV-cache corruption guard). The + # decode->prefill ABORT/ABORT_ACK notification half of that PR is deferred + # until the PD test harness exists (it interleaves with staging/tracing + # infra this branch does not carry). + if ( + kv_chunk.room not in self.request_status + or self.check_status(kv_chunk.room) == KVPoll.Failed + ): + logger.debug( + "Skipping chunk for room %s because it has already failed or been aborted", + kv_chunk.room, + ) + continue reqs_to_be_processed = ( self.transfer_infos[kv_chunk.room].values() if kv_chunk.room in self.transfer_infos @@ -1580,6 +1595,10 @@ class MooncakeKVSender(CommonKVSender): self.bootstrap_room, "Aborted by AbortReq.", ) + # Mark the manager status Failed (not only the local conclude_state) so the + # transfer worker and other pollers observe the abort. Port of upstream + # sgl-project/sglang #24522. + self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) # Explicitly set the status to failure since this request has been aborted self.conclude_state = KVPoll.Failed @@ -1760,6 +1779,10 @@ class MooncakeKVReceiver(CommonKVReceiver): self.bootstrap_room, "Aborted by AbortReq.", ) + # Mark the manager status Failed (not only the local conclude_state) so the + # transfer worker and other pollers observe the abort. Port of upstream + # sgl-project/sglang #24522. + self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) # Explicitly set the status to failure since this request has been aborted self.conclude_state = KVPoll.Failed