From a6b542813ff46de8ba3856e0236f1af35f4fca54 Mon Sep 17 00:00:00 2001 From: HZY Date: Tue, 8 Sep 2026 16:14:33 +0200 Subject: [PATCH] fix(glm-5.2-nvfp4): bound Mooncake synchronous transfer batches (#32758) --- .../advanced_features/pd_disaggregation.mdx | 5 + .../srt/disaggregation/mooncake/conn.py | 63 ++++++++- python/sglang/srt/environ.py | 4 + .../test_mooncake_transfer_batching.py | 133 ++++++++++++++++++ .../test_pp_hybrid_kv_transfer.py | 1 + .../unit/mem_cache/test_hisparse_allocator.py | 1 + 6 files changed, 205 insertions(+), 2 deletions(-) create mode 100644 test/registered/unit/disaggregation/test_mooncake_transfer_batching.py diff --git a/docs/docs/advanced_features/pd_disaggregation.mdx b/docs/docs/advanced_features/pd_disaggregation.mdx index d2a781042..685cbd956 100644 --- a/docs/docs/advanced_features/pd_disaggregation.mdx +++ b/docs/docs/advanced_features/pd_disaggregation.mdx @@ -180,6 +180,11 @@ The `SGLANG_MOONCAKE_CUSTOM_MEM_POOL` environment variable enables the custom me Sets the number of parallel transfer queues. KVCache transfer requests from multiple decode instances will be sharded into these queues so that they can share the threads and the transfer bandwidth at the same time. If it is set to 1, then we transfer requests one by one according to fcfs strategy `4` + + **`SGLANG_MOONCAKE_MAX_TRANSFER_BATCH_INDICES`** + Opt-in limit for the number of KV cache indices represented by one synchronous all-layer Mooncake transfer batch. Set it to a positive value to slice larger index arrays into ordered sub-batches before contiguous address ranges are formed. The custom-memory-pool layerwise path is unchanged. + 0 (disabled) + **`SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT`** Timeout (seconds) for receiving destination KV indices during request initialization diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index d8914b736..0f9e9ce38 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -219,6 +219,9 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager): self.init_engine() self.register_buffer_to_engine() self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get() + self.max_transfer_batch_indices = ( + envs.SGLANG_MOONCAKE_MAX_TRANSFER_BATCH_INDICES.get() + ) self.enable_trace = get_observability().enable_trace if self.disaggregation_mode == DisaggregationMode.PREFILL: self.session_failures = defaultdict(int) @@ -796,8 +799,64 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager): return self._await_transfer_futures(futures) else: # Combining all layers' params in one batch transfer is more efficient - # compared to using multiple threads - return process_layers(layers_params) + # compared to using multiple threads. Preserve this legacy path unless + # users explicitly opt in to bounded index batches. + max_batch_indices = self.max_transfer_batch_indices + if max_batch_indices <= 0 or prefill_data_indices.size <= max_batch_indices: + return process_layers(layers_params) + + def process_index_batch( + prefill_blocks, + dst_blocks, + device_prefill_blocks=None, + device_dst_blocks=None, + ) -> int: + transfer_blocks = [] + for src_ptr, dst_ptr, item_len in layers_params: + if dst_device_data_ptrs and int(dst_ptr) in dst_device_data_ptrs: + assert ( + device_prefill_blocks is not None + and device_dst_blocks is not None + ) + src_blocks, target_blocks = ( + device_prefill_blocks, + device_dst_blocks, + ) + else: + src_blocks, target_blocks = prefill_blocks, dst_blocks + for prefill_index, decode_index in zip(src_blocks, target_blocks): + src_addr = src_ptr + int(prefill_index[0]) * item_len + dst_addr = dst_ptr + int(decode_index[0]) * item_len + length = item_len * len(prefill_index) + transfer_blocks.append((src_addr, dst_addr, length)) + return self._transfer_data(mooncake_session_id, transfer_blocks) + + for start in range( + 0, + prefill_data_indices.size, + max_batch_indices, + ): + batch_prefill_blocks, batch_dst_blocks = group_concurrent_contiguous( + prefill_data_indices[start : start + max_batch_indices], + dst_data_indices[start : start + max_batch_indices], + ) + batch_device_prefill_blocks = batch_device_dst_blocks = None + if dst_device_data_indices is not None: + batch_device_prefill_blocks, batch_device_dst_blocks = ( + group_concurrent_contiguous( + prefill_data_indices[start : start + max_batch_indices], + dst_device_data_indices[start : start + max_batch_indices], + ) + ) + ret = process_index_batch( + batch_prefill_blocks, + batch_dst_blocks, + batch_device_prefill_blocks, + batch_device_dst_blocks, + ) + if ret != 0: + return ret + return 0 def _validate_envelope_kv_layout( self, diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 1239dd34d..0b9b87903 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -752,6 +752,10 @@ class Envs: # staging_buffer.py once Triton kernels are fully validated in production. SGLANG_STAGING_USE_TORCH = EnvBool(False) SGLANG_MOONCAKE_CUSTOM_MEM_POOL = EnvStr(None) + # Opt-in limit for the number of KV cache indices represented by one + # synchronous all-layer Mooncake batch. Set to a positive value to split + # larger transfers; 0 preserves the legacy single-batch behavior. + SGLANG_MOONCAKE_MAX_TRANSFER_BATCH_INDICES = EnvInt(0) ENABLE_ASCEND_TRANSFER_WITH_MOONCAKE = EnvBool(False) ASCEND_NPU_PHY_ID = EnvInt(-1) SGLANG_MOONCAKE_SEND_AUX_TCP = EnvBool(False) diff --git a/test/registered/unit/disaggregation/test_mooncake_transfer_batching.py b/test/registered/unit/disaggregation/test_mooncake_transfer_batching.py new file mode 100644 index 000000000..55b80323d --- /dev/null +++ b/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() diff --git a/test/registered/unit/disaggregation/test_pp_hybrid_kv_transfer.py b/test/registered/unit/disaggregation/test_pp_hybrid_kv_transfer.py index 7c7cb40f7..a065e1cdb 100644 --- a/test/registered/unit/disaggregation/test_pp_hybrid_kv_transfer.py +++ b/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 = [] diff --git a/test/registered/unit/mem_cache/test_hisparse_allocator.py b/test/registered/unit/mem_cache/test_hisparse_allocator.py index dd1337261..d0fc86156 100644 --- a/test/registered/unit/mem_cache/test_hisparse_allocator.py +++ b/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: