minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

refactor(disagg): hoist staging helper imports out of the bootstrap loops (#35980)

Browse Source
This commit is contained in:
Shangming Cai
2026-08-23 01:24:58 +08:00
committed by GitHub
parent 7d22b7a875
commit cce0a1244b
2 changed files with 4 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/mooncake/conn.py
+2 -8
View File
@@ -29,6 +29,8 @@ from sglang.srt.disaggregation.common.staging_handler import (
PrefillStagingContext, PrefillStagingContext,
StagingManagerMixin, StagingManagerMixin,
StagingTransferInfo, StagingTransferInfo,
handle_staging_rsp,
handle_watermark_msg,
) )
from sglang.srt.disaggregation.common.utils import ( from sglang.srt.disaggregation.common.utils import (
AuxDataCodec, AuxDataCodec,
@@ -1944,18 +1946,10 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
room = waiting_req_bytes[0].decode("ascii") room = waiting_req_bytes[0].decode("ascii")
# Staging: decode reports consumption watermark back to prefill # Staging: decode reports consumption watermark back to prefill
if room == "WATERMARK": if room == "WATERMARK":
from sglang.srt.disaggregation.common.staging_handler import (
handle_watermark_msg,
)
handle_watermark_msg(self._staging_ctx, waiting_req_bytes) handle_watermark_msg(self._staging_ctx, waiting_req_bytes)
continue continue
# Staging: decode replies with allocated staging offset # Staging: decode replies with allocated staging offset
if room == "STAGING_RSP": if room == "STAGING_RSP":
from sglang.srt.disaggregation.common.staging_handler import (
handle_staging_rsp,
)
handle_staging_rsp(waiting_req_bytes, self.transfer_infos) handle_staging_rsp(waiting_req_bytes, self.transfer_infos)
continue continue
# Decode-side abort notification: mark room as failed and ACK # Decode-side abort notification: mark room as failed and ACK
python/sglang/srt/disaggregation/nixl/conn.py
+2 -8
View File
@@ -28,6 +28,8 @@ from sglang.srt.disaggregation.common.conn import (
from sglang.srt.disaggregation.common.staging_handler import ( from sglang.srt.disaggregation.common.staging_handler import (
STAGING_WATERMARK_WAIT_S, STAGING_WATERMARK_WAIT_S,
StagingManagerMixin, StagingManagerMixin,
handle_staging_rsp,
handle_watermark_msg,
) )
from sglang.srt.disaggregation.common.utils import ( from sglang.srt.disaggregation.common.utils import (
FastQueue, FastQueue,
@@ -2674,20 +2676,12 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager):
# Staging: decode reports consumption watermark back to prefill # Staging: decode reports consumption watermark back to prefill
if waiting_req_bytes[0] == b"WATERMARK": if waiting_req_bytes[0] == b"WATERMARK":
if self.enable_staging: if self.enable_staging:
from sglang.srt.disaggregation.common.staging_handler import (
handle_watermark_msg,
)
handle_watermark_msg(self._staging_ctx, waiting_req_bytes) handle_watermark_msg(self._staging_ctx, waiting_req_bytes)
continue continue
# Staging: decode replies with allocated staging offset # Staging: decode replies with allocated staging offset
if waiting_req_bytes[0] == b"STAGING_RSP": if waiting_req_bytes[0] == b"STAGING_RSP":
if self.enable_staging: if self.enable_staging:
from sglang.srt.disaggregation.common.staging_handler import (
handle_staging_rsp,
)
handle_staging_rsp(waiting_req_bytes, self.transfer_infos) handle_staging_rsp(waiting_req_bytes, self.transfer_infos)
continue continue