From c971d7ac9c9c8d677e7feba3b73a2c4b428d0b86 Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Mon, 10 Aug 2026 02:53:32 -0700 Subject: [PATCH] Refactor staging registration metadata fields (#33910) Co-authored-by: Shangming Cai --- .../disaggregation/common/staging_handler.py | 26 --------------- .../srt/disaggregation/mooncake/conn.py | 25 +++++++++----- python/sglang/srt/disaggregation/nixl/conn.py | 28 +++++++++------- .../test_disaggregation_wire.py | 33 +++++++++++++++++++ .../disaggregation/test_nixl_backend_basic.py | 14 ++++---- 5 files changed, 74 insertions(+), 52 deletions(-) diff --git a/python/sglang/srt/disaggregation/common/staging_handler.py b/python/sglang/srt/disaggregation/common/staging_handler.py index 823bd9460..0b7860210 100644 --- a/python/sglang/srt/disaggregation/common/staging_handler.py +++ b/python/sglang/srt/disaggregation/common/staging_handler.py @@ -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. diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 353e50bbe..5780fdf7e 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -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, diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 447f88935..97e625193 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -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, diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index 746ea0132..0b4639661 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -1,3 +1,4 @@ +import struct import unittest from types import SimpleNamespace from unittest.mock import patch @@ -13,6 +14,9 @@ from sglang.srt.disaggregation.common.utils import ( unpack_int_lists, unpack_list_of_buffers, ) +from sglang.srt.disaggregation.mooncake.conn import ( + KVArgsRegisterInfo as MooncakeKVArgsRegisterInfo, +) from sglang.srt.disaggregation.utils import ( MetadataBuffers, get_dsv4_c128_state_indices, @@ -31,6 +35,35 @@ register_cpu_ci(est_time=2, suite="base-a-test-cpu") class TestDisaggregationWire(unittest.TestCase): + def test_mooncake_registration_staging_fields(self): + msg = [ + b"room", + b"127.0.0.1", + b"1234", + b"session", + struct.pack("Q", 0x1000), + struct.pack("Q", 0x2000), + b"", + b"0", + b"1", + b"128", + b"", + b"", + b"", + b"", + struct.pack("Q", 0x3000), + b"4096", + b"4", + b"2", + ] + + info = MooncakeKVArgsRegisterInfo.from_zmq(msg) + + self.assertEqual(info.staging_base_ptr, 0x3000) + self.assertEqual(info.staging_total_size, 4096) + self.assertEqual(info.dst_dcp_size, 4) + self.assertEqual(info.dst_dcp_rank, 2) + def test_int_lists_roundtrip(self): cases = [ ("Q", [[1, 2, 3], [4]]), diff --git a/test/registered/unit/disaggregation/test_nixl_backend_basic.py b/test/registered/unit/disaggregation/test_nixl_backend_basic.py index 763842ea2..9136395ed 100644 --- a/test/registered/unit/disaggregation/test_nixl_backend_basic.py +++ b/test/registered/unit/disaggregation/test_nixl_backend_basic.py @@ -244,9 +244,8 @@ class TestNixlKVArgsRegisterInfo(CustomTestCase): self.assertEqual(info.dst_dcp_rank, 3) self.assertEqual(info.dst_state_layer_ids, [[4], [4, 5]]) self.assertEqual(info.dst_kv_layer_ids, [2, 7]) - self.assertIsNotNone(info.staging) - self.assertEqual(info.staging.base_ptr, staging_ptr) - self.assertEqual(info.staging.total_size, 1048576) + self.assertEqual(info.staging_base_ptr, staging_ptr) + self.assertEqual(info.staging_total_size, 1048576) def test_from_zmq_allows_missing_state_and_staging_fields(self): msg = [ @@ -272,7 +271,8 @@ class TestNixlKVArgsRegisterInfo(CustomTestCase): self.assertEqual(info.dst_kv_item_lens, [256]) self.assertEqual(info.dst_dcp_size, 1) self.assertEqual(info.dst_dcp_rank, 0) - self.assertIsNone(info.staging) + self.assertEqual(info.staging_base_ptr, 0) + self.assertEqual(info.staging_total_size, 0) class TestNixlTransferStatus(CustomTestCase): @@ -452,7 +452,8 @@ class TestNixlTransferWorker(CustomTestCase): dst_kv_ptrs=[0], dst_aux_ptrs=[0], gpu_id=0, - staging=None, + staging_base_ptr=0, + staging_total_size=0, kv_xfer_segments=None, dst_homogeneous_mem_kind="VRAM", # Non-DCP peer. Without this the worker raises AttributeError @@ -936,7 +937,8 @@ class TestNixlStaging(CustomTestCase): decode_tp_rank=0, dst_kv_item_len=128, dst_kv_item_lens=[], - staging=SimpleNamespace(base_ptr=0x8000, total_size=4096), + staging_base_ptr=0x8000, + staging_total_size=4096, ) calls = [] mgr.send_kvcache_staged = (