Refactor staging registration metadata fields (#33910)

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

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.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 = (