[PD] Support pipeline-parallel prefill with Mooncake staging buffer (#33807)

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

test/registered/unit/disaggregation/test_disaggregation_wire.py

@@ -1,12 +1,16 @@
 import struct
+import threading
 import unittest
 from types import SimpleNamespace
-from unittest.mock import patch
+from unittest.mock import Mock, patch
 
 import numpy as np
 import torch
 
 from sglang.srt.disaggregation.base.conn import KVArgs, StateType
+from sglang.srt.disaggregation.common.staging_handler import (
+    handle_staging_req,
+)
 from sglang.srt.disaggregation.common.utils import (
     group_concurrent_contiguous,
     pack_int_lists,
@@ -15,7 +19,8 @@ from sglang.srt.disaggregation.common.utils import (
     unpack_list_of_buffers,
 )
 from sglang.srt.disaggregation.mooncake.conn import (
-    KVArgsRegisterInfo as MooncakeKVArgsRegisterInfo,
+    KVArgsRegisterInfo,
+    MooncakeKVManager,
 )
 from sglang.srt.disaggregation.utils import (
     MetadataBuffers,
@@ -57,7 +62,7 @@ class TestDisaggregationWire(unittest.TestCase):
             b"2",
         ]
 
-        info = MooncakeKVArgsRegisterInfo.from_zmq(msg)
+        info = KVArgsRegisterInfo.from_zmq(msg)
 
         self.assertEqual(info.staging_base_ptr, 0x3000)
         self.assertEqual(info.staging_total_size, 4096)
@@ -134,6 +139,84 @@ class TestGroupConcurrentContiguous(unittest.TestCase):
             group_concurrent_contiguous(self._arr([1, 2, 3]), self._arr([1, 2]))
 
 
+class TestMooncakePPStaging(unittest.TestCase):
+    def test_staging_response_targets_requesting_pp_rank(self):
+        sock = Mock()
+        receiver = SimpleNamespace(
+            chunk_staging_infos=[],
+            _connect_to_bootstrap_server=Mock(return_value=(sock, threading.Lock())),
+        )
+        allocator = SimpleNamespace(
+            assign=Mock(return_value=(3, 128, 0)), total_size=1 << 20
+        )
+        kv_args = SimpleNamespace(
+            page_size=64,
+            kv_item_lens=[4096, 4096],
+            total_kv_head_num=4,
+            engine_rank=0,
+        )
+        target = {"pp_rank": 3}
+
+        handle_staging_req(
+            [b"STAGING_REQ", b"7", b"0", b"1", b"peer", b"3"],
+            allocator,
+            kv_args,
+            attn_tp_size=16,
+            prefill_attn_tp_size=1,
+            kv_buffer_tensors=None,
+            room_receivers={7: receiver},
+            room_bootstrap={7: [{"pp_rank": 2}, target]},
+        )
+
+        receiver._connect_to_bootstrap_server.assert_called_once_with(target)
+        sock.send_multipart.assert_called_once()
+
+    @patch(
+        "sglang.srt.disaggregation.common.staging_buffer.gather_all_layers_to_staging"
+    )
+    def test_pp_stage_writes_its_global_layer_slots(self, gather):
+        manager = object.__new__(MooncakeKVManager)
+        tensor = SimpleNamespace(shape=(1, 1, 8), element_size=lambda: 2)
+        manager.kv_buffer_tensors = {
+            "k_buffers": [tensor],
+            "v_buffers": [tensor],
+            "page_size": 2,
+        }
+        manager.attn_tp_size = 1
+        manager.pp_size = 16
+        manager.kv_args = SimpleNamespace(
+            engine_rank=0,
+            gpu_id=0,
+            total_kv_head_num=4,
+            kv_head_num=4,
+            kv_layer_ids=[7, 7],
+        )
+        manager._transfer_data = Mock(return_value=0)
+        staging = SimpleNamespace(fits=lambda size: True, get_ptr=lambda: 0x9000)
+
+        ret = manager.send_kvcache_staged(
+            "peer",
+            np.array([1, 2], dtype=np.int32),
+            dst_staging_ptr=0x100000,
+            dst_staging_size=1 << 20,
+            dst_tp_rank=0,
+            dst_attn_tp_size=16,
+            dst_kv_item_len=128,
+            dst_layer_ids=[3, 7, 11, 3, 7, 11],
+            staging_buffer=staging,
+        )
+
+        self.assertEqual(ret, 0)
+        gather.assert_called_once()
+        manager._transfer_data.assert_called_once_with(
+            "peer",
+            [
+                (0x9000, 0x100000 + 64, 64),
+                (0x9000 + 64, 0x100000 + 4 * 64, 64),
+            ],
+        )
+
+
 class TestEagleDsaSeedTransfer(unittest.TestCase):
     @staticmethod
     def _make_req(seed, metadata_buffer_index=0):