From 15a439832054c1809a2bf59f7b94bf9dd71de282 Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Sat, 22 Aug 2026 15:35:55 +0800 Subject: [PATCH] refactor(disagg): hoist duplicated _handle_staging_req into a mixin (#35948) --- .../disaggregation/common/staging_handler.py | 40 ++++++++++++++++++ .../srt/disaggregation/mooncake/conn.py | 37 +---------------- python/sglang/srt/disaggregation/nixl/conn.py | 41 +++---------------- 3 files changed, 47 insertions(+), 71 deletions(-) diff --git a/python/sglang/srt/disaggregation/common/staging_handler.py b/python/sglang/srt/disaggregation/common/staging_handler.py index bccfac37d..74316ca05 100644 --- a/python/sglang/srt/disaggregation/common/staging_handler.py +++ b/python/sglang/srt/disaggregation/common/staging_handler.py @@ -824,6 +824,46 @@ def handle_staging_req( pass +class StagingManagerMixin: + """Shared STAGING_REQ handling for KV managers that support staging. + + Mixed into the managers whose decode thread receives STAGING_REQ messages + (currently Mooncake and NIXL). Expects the concrete manager to provide + ``_staging_handler``, ``_staging_ctx``, ``kv_args``, ``attn_tp_size`` and + optionally ``kv_buffer_tensors``. + """ + + def _handle_staging_req(self, msg): + room = int(msg[1].decode("ascii")) + session_id = msg[4].decode("ascii") + handler = self._staging_handler + assert ( + handler is not None + ), "STAGING_REQ received before staging handler initialized" + decode_req = handler._room_to_decode_req.get(room) + if decode_req is None: + logger.warning( + "STAGING_REQ received for unregistered room=%s, skipping", + room, + ) + return + prefill_tp = decode_req.kv_receiver.prefill_info.attn_tp_size + handle_staging_req( + msg, + self._staging_ctx.allocator, + self.kv_args, + self.attn_tp_size, + prefill_tp, + getattr(self, "kv_buffer_tensors", None), + self._staging_ctx.room_receivers, + self._staging_ctx.room_bootstrap, + ) + + receiver = self._staging_ctx.room_receivers.get(room) + if receiver is not None: + handler.register_wm_subscriber(receiver, session_id) + + def prefetch_staging_reqs( room: int, transfer_infos: dict, diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index ab158e8be..4b34b3f43 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -27,6 +27,7 @@ from sglang.srt.disaggregation.common.staging_handler import ( STAGING_WATERMARK_WAIT_S, DecodeStagingContext, PrefillStagingContext, + StagingManagerMixin, StagingTransferInfo, ) from sglang.srt.disaggregation.common.utils import ( @@ -192,7 +193,7 @@ class KVArgsRegisterInfo: ) -class MooncakeKVManager(CommonKVManager): +class MooncakeKVManager(StagingManagerMixin, CommonKVManager): AUX_DATA_HEADER = b"AUX_DATA" def __init__( @@ -361,40 +362,6 @@ class MooncakeKVManager(CommonKVManager): ) self.kv_buffer_tensors = None - def _handle_staging_req(self, msg): - from sglang.srt.disaggregation.common.staging_handler import ( - handle_staging_req, - ) - - room = int(msg[1].decode("ascii")) - session_id = msg[4].decode("ascii") - handler = self._staging_handler - assert ( - handler is not None - ), "STAGING_REQ received before staging handler initialized" - decode_req = handler._room_to_decode_req.get(room) - if decode_req is None: - logger.warning( - "STAGING_REQ received for unregistered room=%s, skipping", - room, - ) - return - prefill_tp = decode_req.kv_receiver.prefill_info.attn_tp_size - handle_staging_req( - msg, - self._staging_ctx.allocator, - self.kv_args, - self.attn_tp_size, - prefill_tp, - getattr(self, "kv_buffer_tensors", None), - self._staging_ctx.room_receivers, - self._staging_ctx.room_bootstrap, - ) - - receiver = self._staging_ctx.room_receivers.get(room) - if receiver is not None: - handler.register_wm_subscriber(receiver, session_id) - def _is_watermark_ready( self, session_id: str, alloc_round: int, alloc_end: int ) -> bool: diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index ec642b84a..eae55f6cd 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -25,7 +25,10 @@ from sglang.srt.disaggregation.common.conn import ( CommonKVSender, KVTransferError, ) -from sglang.srt.disaggregation.common.staging_handler import STAGING_WATERMARK_WAIT_S +from sglang.srt.disaggregation.common.staging_handler import ( + STAGING_WATERMARK_WAIT_S, + StagingManagerMixin, +) from sglang.srt.disaggregation.common.utils import ( FastQueue, TransferKVChunk, @@ -390,7 +393,7 @@ class TransferStatus: return True -class NixlKVManager(CommonKVManager): +class NixlKVManager(StagingManagerMixin, CommonKVManager): def __init__( self, args: KVArgs, @@ -618,40 +621,6 @@ class NixlKVManager(CommonKVManager): threading.Thread(target=decode_listener_thread, daemon=True).start() - def _handle_staging_req(self, msg): - from sglang.srt.disaggregation.common.staging_handler import ( - handle_staging_req, - ) - - room = int(msg[1].decode("ascii")) - session_id = msg[4].decode("ascii") - handler = self._staging_handler - assert ( - handler is not None - ), "STAGING_REQ received before staging handler initialized" - decode_req = handler._room_to_decode_req.get(room) - if decode_req is None: - logger.warning( - "STAGING_REQ received for unregistered room=%s, skipping", - room, - ) - return - prefill_tp = decode_req.kv_receiver.prefill_info.attn_tp_size - handle_staging_req( - msg, - self._staging_ctx.allocator, - self.kv_args, - self.attn_tp_size, - prefill_tp, - getattr(self, "kv_buffer_tensors", None), - self._staging_ctx.room_receivers, - self._staging_ctx.room_bootstrap, - ) - - receiver = self._staging_ctx.room_receivers.get(room) - if receiver is not None: - handler.register_wm_subscriber(receiver, session_id) - def _prefetch_staging_reqs(self, room: int): """Send STAGING_REQ for all chunks before the prefill forward starts.