fix(glm-5.2-nvfp4): bound Mooncake synchronous transfer batches (#32758)

test/registered/unit/disaggregation/test_mooncake_transfer_batching.py

@@ -0,0 +1,133 @@
+import concurrent.futures
+import unittest
+from types import SimpleNamespace
+from unittest.mock import MagicMock, call
+
+import numpy as np
+
+from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+class TestMooncakeTransferBatching(unittest.TestCase):
+    @staticmethod
+    def _make_manager(
+        side_effect=None, enable_custom_mem_pool=False, max_batch_indices=0
+    ):
+        engine = MagicMock()
+        if side_effect is None:
+            engine.batch_transfer_sync.return_value = 0
+        else:
+            engine.batch_transfer_sync.side_effect = side_effect
+        manager = SimpleNamespace(
+            engine=engine,
+            is_mla_backend=True,
+            is_hybrid_mla_backend=False,
+            pp_size=1,
+            enable_custom_mem_pool=enable_custom_mem_pool,
+            enable_deferred_decode_kv_release=False,
+            max_transfer_batch_indices=max_batch_indices,
+            get_mla_kv_ptrs_with_pp=MagicMock(
+                return_value=([1000, 2000], [5000, 6000], 2)
+            ),
+        )
+        manager._transfer_data = lambda session, blocks: (
+            MooncakeKVManager._transfer_data(manager, session, blocks)
+        )
+        manager._await_transfer_futures = lambda futures: (
+            MooncakeKVManager._await_transfer_futures(manager, futures)
+        )
+        return manager
+
+    @staticmethod
+    def _send(
+        manager,
+        dst_device_data_indices=None,
+        dst_device_data_ptrs=None,
+    ):
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            return MooncakeKVManager._send_kvcache_generic(
+                manager,
+                mooncake_session_id="session",
+                src_data_ptrs=[1000, 2000],
+                dst_data_ptrs=[5000, 6000],
+                item_lens=[10, 20],
+                prefill_data_indices=np.array([0, 1, 2, 3, 4], dtype=np.int32),
+                dst_data_indices=np.array([10, 11, 12, 13, 14], dtype=np.int32),
+                executor=executor,
+                dst_device_data_indices=dst_device_data_indices,
+                dst_device_data_ptrs=dst_device_data_ptrs,
+            )
+
+    def test_slices_index_arrays_before_forming_transfer_ranges(self):
+        manager = self._make_manager(max_batch_indices=2)
+        ret = self._send(manager)
+
+        self.assertEqual(ret, 0)
+        self.assertEqual(
+            manager.engine.batch_transfer_sync.call_args_list,
+            [
+                call("session", [1000, 2000], [5100, 6200], [20, 40]),
+                call("session", [1020, 2040], [5120, 6240], [20, 40]),
+                call("session", [1040, 2080], [5140, 6280], [10, 20]),
+            ],
+        )
+
+    def test_preserves_legacy_single_batch_path_for_short_transfers(self):
+        for max_batch_indices in (0, 5, 6):
+            with self.subTest(max_batch_indices=max_batch_indices):
+                manager = self._make_manager(max_batch_indices=max_batch_indices)
+                ret = self._send(manager)
+
+                self.assertEqual(ret, 0)
+                manager.engine.batch_transfer_sync.assert_called_once_with(
+                    "session",
+                    [1000, 2000],
+                    [5100, 6200],
+                    [50, 100],
+                )
+
+    def test_stops_after_first_failed_index_batch(self):
+        manager = self._make_manager(side_effect=[0, -1], max_batch_indices=2)
+        ret = self._send(manager)
+
+        self.assertEqual(ret, -1)
+        self.assertEqual(manager.engine.batch_transfer_sync.call_count, 2)
+
+    def test_uses_device_page_indices_in_batched_path(self):
+        manager = self._make_manager(max_batch_indices=2)
+        ret = self._send(
+            manager,
+            dst_device_data_indices=np.array([20, 21, 22, 23, 24], dtype=np.int32),
+            dst_device_data_ptrs={6000},
+        )
+
+        self.assertEqual(ret, 0)
+        self.assertEqual(
+            manager.engine.batch_transfer_sync.call_args_list,
+            [
+                call("session", [1000, 2000], [5100, 6400], [20, 40]),
+                call("session", [1020, 2040], [5120, 6440], [20, 40]),
+                call("session", [1040, 2080], [5140, 6480], [10, 20]),
+            ],
+        )
+
+    def test_preserves_one_transfer_per_layer_for_custom_mem_pool(self):
+        manager = self._make_manager(enable_custom_mem_pool=True, max_batch_indices=2)
+        ret = self._send(manager)
+
+        self.assertEqual(ret, 0)
+        self.assertEqual(manager.engine.batch_transfer_sync.call_count, 2)
+        manager.engine.batch_transfer_sync.assert_has_calls(
+            [
+                call("session", [1000], [5100], [50]),
+                call("session", [2000], [6200], [100]),
+            ],
+            any_order=True,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/disaggregation/test_pp_hybrid_kv_transfer.py

@@ -75,6 +75,7 @@ class _RecordingKVManager:
         self.is_mla_backend = False
         self.is_hybrid_mla_backend = False
         self.enable_custom_mem_pool = False
+        self.max_transfer_batch_indices = 0
         self.pp_size = pp_size
         self.kv_args = SimpleNamespace(prefill_start_layer=prefill_start_layer)
         self.blocks = []

test/registered/unit/mem_cache/test_hisparse_allocator.py

@@ -187,6 +187,7 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase):
         manager.is_mla_backend = True
         manager.is_hybrid_mla_backend = False
         manager.enable_custom_mem_pool = False
+        manager.max_transfer_batch_indices = 0
         manager._transfer_data = MagicMock(return_value=0)
 
         with ThreadPoolExecutor(max_workers=1) as executor: