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

[HiSparse]Fix DeepSeek V4 HiSparse PD Transfers with Separate Host and Device KV Indices (#31901)

Browse Source 
Co-authored-by: jackyYang6 <82102811+jackyYang6@users.noreply.github.com>
This commit is contained in:
huangtingwei
2026-08-03 20:35:04 +08:00
committed by GitHub
co-authored by jackyYang6
parent d48ab2d386
commit 3953788596
3 changed files with 201 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_hisparse_allocator.py
+107 -9
View File
@@ -1,7 +1,9 @@
import unittest
from concurrent.futures import ThreadPoolExecutor
from types import SimpleNamespace
from unittest.mock import MagicMock
import numpy as np
import torch
from sglang.srt.mem_cache.allocator.hisparse import (
@@ -78,8 +80,8 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase):
fill_len = 512
swa_tail_len = 128
kv_loc = torch.arange(fill_len, dtype=torch.int64)
host_indices = torch.arange(1000, 1000 + fill_len, dtype=torch.int64)
kv_loc = torch.arange(512, 512 + fill_len, dtype=torch.int64)
host_indices = torch.arange(1000, 1128, dtype=torch.int64)
req = SimpleNamespace(
rid="req-0",
@@ -108,18 +110,17 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase):
req_to_token_pool = ReqToTokenPool()
allocator = SimpleNamespace(
device=torch.device("cpu"),
page_size=64,
page_size=256,
available_size=MagicMock(return_value=fill_len),
alloc_extend_swa_tail=MagicMock(return_value=kv_loc),
alloc_logical_only=MagicMock(return_value=kv_loc),
)
regular_host_alloc = MagicMock(return_value=host_indices)
coordinator = SimpleNamespace(
mem_pool_host=SimpleNamespace(
alloc_paged_token_slots=MagicMock(return_value=host_indices)
),
mem_pool_host=SimpleNamespace(alloc_paged_token_slots=regular_host_alloc),
req_to_host_pool=object(),
req_to_host_pool_allocated_len=object(),
host_token_len=MagicMock(side_effect=lambda length: length),
host_token_len=MagicMock(side_effect=lambda token_len: token_len // 4),
)
queue = DecodePreallocQueue.__new__(DecodePreallocQueue)
queue.req_to_token_pool = req_to_token_pool
@@ -138,7 +139,7 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase):
result = queue._pre_alloc(req)
self.assertIs(result, host_indices)
self.assertTrue(torch.equal(result, host_indices))
allocator.alloc_extend_swa_tail.assert_called_once()
allocator.alloc_logical_only.assert_not_called()
_, kwargs = allocator.alloc_extend_swa_tail.call_args
@@ -149,7 +150,104 @@ class TestDeepSeekV4HiSparseAllocator(CustomTestCase):
self.assertEqual(req.kv_committed_len, fill_len)
self.assertEqual(req.extend_range.length, fill_len)
self.assertEqual(len(req_to_token_pool.writes), 1)
coordinator.mem_pool_host.alloc_paged_token_slots.assert_called_once()
coordinator.host_token_len.assert_called_once_with(fill_len)
regular_host_alloc.assert_called_once_with(
coordinator.req_to_host_pool,
coordinator.req_to_host_pool_allocated_len,
req.req_pool_idx,
0,
len(host_indices),
)
self.assertTrue(torch.equal(req_to_token_pool.writes[0][1], kv_loc))
# C4 indexer/C128 use the logical allocator's full-page IDs. They do not
# use either the independently allocated host pages or C4 sparse slots.
np.testing.assert_array_equal(
np.unique(kv_loc.numpy() // allocator.page_size),
np.array([2, 3]),
)
def test_mooncake_uses_separate_host_and_device_page_indices(self):
from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager
manager = object.__new__(MooncakeKVManager)
manager.is_mla_backend = True
manager.is_hybrid_mla_backend = False
manager.enable_custom_mem_pool = False
manager._transfer_data = MagicMock(return_value=0)
with ThreadPoolExecutor(max_workers=1) as executor:
ret = manager._send_kvcache_generic(
mooncake_session_id="session",
src_data_ptrs=[1000, 2000, 3000],
dst_data_ptrs=[10000, 20000, 30000],
item_lens=[100, 100, 100],
prefill_data_indices=np.array([1, 2], dtype=np.int32),
dst_data_indices=np.array([7, 8], dtype=np.int32),
executor=executor,
dst_device_data_indices=np.array([21, 22], dtype=np.int32),
dst_device_data_ptrs={20000, 30000},
)
self.assertEqual(ret, 0)
manager._transfer_data.assert_called_once_with(
"session",
[
(1100, 10700, 200),
(2100, 22100, 200),
(3100, 32100, 200),
],
)
def test_mooncake_derives_device_buffers_from_local_pp_layout(self):
from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager
manager = object.__new__(MooncakeKVManager)
manager.kv_args = SimpleNamespace(
kv_data_ptrs=[1000, 2000, 3000],
kv_item_lens=[100, 100, 100],
kv_layer_ids=[],
mla_compression_ratios=[4, 128, 4, 128],
prefill_start_layer=0,
prefill_end_layer=2,
)
manager._send_kvcache_generic = MagicMock(return_value=0)
executor = MagicMock()
manager.send_kvcache(
"session",
np.array([1], dtype=np.int32),
[10000, 20000, 30000],
np.array([7], dtype=np.int32),
executor,
dst_device_kv_indices=np.array([21], dtype=np.int32),
)
kwargs = manager._send_kvcache_generic.call_args.kwargs
self.assertEqual(kwargs["dst_device_data_ptrs"], {20000, 30000})
def test_mooncake_transfer_metadata_carries_device_page_indices(self):
from sglang.srt.disaggregation.mooncake.conn import TransferInfo
host_pages = np.array([7, 8], dtype=np.int32)
device_pages = np.array([21, 22], dtype=np.int32)
info = TransferInfo.from_zmq(
[
b"9",
b"127.0.0.1",
b"12345",
b"session",
host_pages.tobytes(),
b"0",
b"",
b"1",
b"0",
device_pages.tobytes(),
]
)
np.testing.assert_array_equal(info.dst_kv_indices, host_pages)
np.testing.assert_array_equal(info.dst_device_kv_indices, device_pages)
if __name__ == "__main__":