minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
HZY
2026-09-08 22:14:33 +08:00
committed by GitHub
parent 4df5df911b
commit a6b542813f
6 changed files with 205 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/docs/advanced_features/pd_disaggregation.mdx
+5
View File
@@ -180,6 +180,11 @@ The `SGLANG_MOONCAKE_CUSTOM_MEM_POOL` environment variable enables the custom me
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>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 <code>1</code>, then we transfer requests one by one according to fcfs strategy</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`4`</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>**`SGLANG_MOONCAKE_MAX_TRANSFER_BATCH_INDICES`**</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>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.</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>0</code> (disabled)</td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>**`SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT`**</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Timeout (seconds) for receiving destination KV indices during request initialization</td>
python/sglang/srt/disaggregation/mooncake/conn.py
+61 -2
View File
@@ -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,
python/sglang/srt/environ.py
+4
View File
@@ -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)
test/registered/unit/disaggregation/test_mooncake_transfer_batching.py
+133
View File
@@ -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
+1
View File
@@ -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
+1
View File
@@ -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: