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

Refactor staging registration metadata fields (#33910)

Browse Source 
Co-authored-by: Shangming Cai <csmthu@gmail.com>
This commit is contained in:
YAMY
2026-08-10 17:53:32 +08:00
committed by GitHub
co-authored by Shangming Cai
parent 0b6189d0e8
commit c971d7ac9c
5 changed files with 74 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/common/staging_handler.py
-26
View File
@@ -10,7 +10,6 @@ from __future__ import annotations
import dataclasses
import logging
import struct
import threading
import time
from typing import TYPE_CHECKING, List, Optional, Tuple
@@ -529,31 +528,6 @@ class StagingTransferInfo:
self.ends[idx] = end
@dataclasses.dataclass
class StagingRegisterInfo:
"""Staging buffer registration info attached to a KVArgsRegisterInfo."""
base_ptr: int = 0
total_size: int = 0
@classmethod
def from_zmq_fields(
cls, msg: list, msg_start_offset: int
) -> Optional[StagingRegisterInfo]:
i = msg_start_offset
base_ptr = (
struct.unpack("Q", msg[i])[0] if len(msg) > i and len(msg[i]) == 8 else 0
)
total_size = (
int(msg[i + 1].decode("ascii"))
if len(msg) > i + 1 and len(msg[i + 1]) > 0
else 0
)
if base_ptr == 0 and total_size == 0:
return None
return cls(base_ptr=base_ptr, total_size=total_size)
class PrefillStagingStrategy:
"""Prefill-side staging transfer: readiness check + gather-RDMA execution.
python/sglang/srt/disaggregation/mooncake/conn.py
+16 -9
View File
@@ -27,7 +27,6 @@ from sglang.srt.disaggregation.common.staging_handler import (
STAGING_WATERMARK_WAIT_S,
DecodeStagingContext,
PrefillStagingContext,
StagingRegisterInfo,
StagingTransferInfo,
)
from sglang.srt.disaggregation.common.utils import (
@@ -144,8 +143,8 @@ class KVArgsRegisterInfo:
dst_dcp_rank: int = 0
requires_dcp_relayout: bool = False
dcp_token_item_lens: Optional[List[int]] = None
# Note: always put the staging field at the final (since the staging field is optional and contains multiple inputs)
staging: Optional[StagingRegisterInfo] = None
staging_base_ptr: int = 0
staging_total_size: int = 0
@classmethod
def from_zmq(cls, msg: List[bytes]):
@@ -176,15 +175,20 @@ class KVArgsRegisterInfo:
if len(msg) > 13 and msg[13] != b""
else []
),
# msg[14:16] belong to the staging field below; DCP trails it.
staging_base_ptr=(
struct.unpack("Q", msg[14])[0]
if len(msg) > 14 and len(msg[14]) == 8
else 0
),
staging_total_size=(
int(msg[15].decode("ascii")) if len(msg) > 15 and msg[15] != b"" else 0
),
dst_dcp_size=(
int(msg[16].decode("ascii")) if len(msg) > 16 and msg[16] != b"" else 1
),
dst_dcp_rank=(
int(msg[17].decode("ascii")) if len(msg) > 17 and msg[17] != b"" else 0
),
# Note: always put the staging field at the final
staging=StagingRegisterInfo.from_zmq_fields(msg, 14),
)
@@ -467,8 +471,8 @@ class MooncakeKVManager(CommonKVManager):
ret = staging_strategy.transfer(
req.mooncake_session_id,
kv_chunk.prefill_kv_indices,
target_info.staging.base_ptr + c_offset,
target_info.staging.total_size - c_offset,
target_info.staging_base_ptr + c_offset,
target_info.staging_total_size - c_offset,
target_info,
)
if ret == -1:
@@ -1693,7 +1697,10 @@ class MooncakeKVManager(CommonKVManager):
elif (
self.enable_staging
and staging_strategy is not None
and target_rank_registration_info.staging is not None
and (
target_rank_registration_info.staging_base_ptr != 0
or target_rank_registration_info.staging_total_size != 0
)
):
ret, deferred = self._do_staging_transfer(
staging_strategy,
python/sglang/srt/disaggregation/nixl/conn.py
+17 -11
View File
@@ -25,10 +25,7 @@ from sglang.srt.disaggregation.common.conn import (
CommonKVSender,
KVTransferError,
)
from sglang.srt.disaggregation.common.staging_handler import (
STAGING_WATERMARK_WAIT_S,
StagingRegisterInfo,
)
from sglang.srt.disaggregation.common.staging_handler import STAGING_WATERMARK_WAIT_S
from sglang.srt.disaggregation.common.utils import (
FastQueue,
TransferKVChunk,
@@ -233,9 +230,8 @@ class KVArgsRegisterInfo:
dst_state_layer_ids: List[List[int]] = dataclasses.field(default_factory=list)
dst_homogeneous_mem_kind: Optional[str] = None
kv_xfer_segments: Optional[List[_KVXferPreparedSegment]] = None
# Keep last: optional, parsed from a variable-length tail of the ZMQ
# frame in from_zmq() below, so positional construction stays stable.
staging: Optional[StagingRegisterInfo] = None
staging_base_ptr: int = 0
staging_total_size: int = 0
@classmethod
def from_zmq(cls, msg: List[bytes]):
@@ -304,7 +300,14 @@ class KVArgsRegisterInfo:
dst_state_item_lens=dst_state_item_lens,
dst_state_dim_per_tensor=dst_state_dim_per_tensor,
dst_state_layer_ids=dst_state_layer_ids,
staging=StagingRegisterInfo.from_zmq_fields(msg, 14),
staging_base_ptr=(
struct.unpack("Q", msg[14])[0]
if len(msg) > 14 and len(msg[14]) == 8
else 0
),
staging_total_size=(
int(msg[15].decode("ascii")) if len(msg) > 15 and msg[15] != b"" else 0
),
)
@@ -1203,7 +1206,10 @@ class NixlKVManager(CommonKVManager):
and not self.is_mla_backend
and not self.is_hybrid_mla_backend
and decode_tp_size != self.attn_tp_size
and dst_info.staging is not None
and (
dst_info.staging_base_ptr != 0
or dst_info.staging_total_size != 0
)
)
kv_xfer_handle = None
@@ -1972,8 +1978,8 @@ class NixlKVManager(CommonKVManager):
handle = self.send_kvcache_staged(
req.agent_name,
src_prefill_kv_indices,
dst_info.staging.base_ptr + c_offset,
dst_info.staging.total_size - c_offset,
dst_info.staging_base_ptr + c_offset,
dst_info.staging_total_size - c_offset,
dst_info.gpu_id,
dst_info.decode_tp_rank,
dst_info.decode_tp_size,