Refactor staging registration metadata fields (#33910)

Co-authored-by: Shangming Cai <csmthu@gmail.com>

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.
 

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_base_ptr: int = 0
-    staging: Optional[StagingRegisterInfo] = None
+    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_base_ptr + c_offset,
-            target_info.staging.total_size - 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

@@ -25,10 +25,7 @@ from sglang.srt.disaggregation.common.conn import (
     CommonKVSender,
     KVTransferError,
 )
-from sglang.srt.disaggregation.common.staging_handler import (
+from sglang.srt.disaggregation.common.staging_handler import STAGING_WATERMARK_WAIT_S
-    STAGING_WATERMARK_WAIT_S,
-    StagingRegisterInfo,
-)
 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
+    staging_base_ptr: int = 0
-    # frame in from_zmq() below, so positional construction stays stable.
+    staging_total_size: int = 0
-    staging: Optional[StagingRegisterInfo] = None
 
     @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_base_ptr + c_offset,
-            dst_info.staging.total_size - c_offset,
+            dst_info.staging_total_size - c_offset,
             dst_info.gpu_id,
             dst_info.decode_tp_rank,
             dst_info.decode_tp_size,

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]]),

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.base_ptr, staging_ptr)
+        self.assertEqual(info.staging_total_size, 1048576)
-        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 = (