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
python/sglang/srt/disaggregation/decode.py
+27 -1
View File
@@ -1249,6 +1249,32 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
page_indices = kv_to_page_indices(kv_indices, kv_transfer_page_size).astype(
np.int32
)
device_page_indices = None
if (
self.scheduler.enable_hisparse
and isinstance(self.token_to_kv_pool, DeepSeekV4TokenToKVPool)
and not _is_fake_transfer(decode_req.req, self.scheduler.server_args)
):
# alloc_logical_only() already allocated the shared logical pages
# used by C4 indexer and C128 KV. These device buffers do not use
# the C4 sparse physical-slot mapping; carry their logical page IDs
# alongside the independently allocated C4 host page IDs.
full_kv_indices = self.req_to_token_pool.req_to_token[
decode_req.req.req_pool_idx,
prefix_len:origin_input_len,
]
device_page_indices = kv_to_page_indices(
full_kv_indices,
page_size,
).astype(np.int32)
if self.transfer_backend != TransferBackend.MOONCAKE:
raise NotImplementedError(
"DSV4 HiSparse direct PD transfer currently requires "
"the Mooncake backend"
)
metadata_kwargs = {"decode_prefix_len": total_prefix_len}
if device_page_indices is not None:
metadata_kwargs["device_kv_indices"] = device_page_indices
if (
self.transfer_queue.enable_staging
and hasattr(decode_req.kv_receiver, "require_staging")
@@ -1263,7 +1289,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
page_indices,
decode_req.metadata_buffer_index,
state_indices,
decode_prefix_len=total_prefix_len,
**metadata_kwargs,
)
if decode_req.is_rebootstrap:
self.kv_manager.submit_prefill_recompute(
python/sglang/srt/disaggregation/mooncake/conn.py
+67 -3
View File
@@ -85,6 +85,7 @@ class TransferInfo:
required_dst_info_num: int
is_dummy: bool
decode_prefix_len: Optional[int] = None
dst_device_kv_indices: Optional[npt.NDArray[np.int32]] = None
# Note: always put the optional staging field at the final (it will be set through 'STAGING_RSP' pkg when needed)
staging: Optional[StagingTransferInfo] = None
@@ -113,6 +114,11 @@ class TransferInfo:
decode_prefix_len=(
int(msg[8].decode("ascii")) if len(msg) > 8 and msg[8] != b"" else None
),
dst_device_kv_indices=(
np.frombuffer(msg[9], dtype=np.int32)
if len(msg) > 9 and msg[9] != b""
else None
),
)
@@ -623,6 +629,8 @@ class MooncakeKVManager(CommonKVManager):
force_flat: bool = False,
src_layer_ids: Optional[List[int]] = None,
dst_layer_ids: Optional[List[int]] = None,
dst_device_data_indices: Optional[npt.NDArray[np.int32]] = None,
dst_device_data_ptrs: Optional[set[int]] = None,
) -> int:
"""
Generic KV cache transfer supporting both MHA and MLA architectures.
@@ -632,10 +640,18 @@ class MooncakeKVManager(CommonKVManager):
even on a non-MLA backend, for K-only state buffers (e.g. MiniMax sparse
index) whose per-layer list must not be half-split into K/V.
"""
# Group by indices for optimization
# Host and device buffers may use different destination page spaces.
# Build both transfer plans once, then select per destination buffer.
prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous(
prefill_data_indices, dst_data_indices
)
device_prefill_kv_blocks = device_dst_kv_blocks = None
if dst_device_data_indices is not None:
device_prefill_kv_blocks, device_dst_kv_blocks = (
group_concurrent_contiguous(
prefill_data_indices, dst_device_data_indices
)
)
layers_params = None
@@ -701,7 +717,18 @@ class MooncakeKVManager(CommonKVManager):
src_ptr: int, dst_ptr: int, item_len: int
) -> List[Tuple[int, int, int]]:
transfer_blocks = []
for prefill_index, decode_index in zip(prefill_kv_blocks, dst_kv_blocks):
if dst_device_data_ptrs and int(dst_ptr) in dst_device_data_ptrs:
assert (
device_prefill_kv_blocks is not None
and device_dst_kv_blocks is not None
)
src_blocks, dst_blocks = (
device_prefill_kv_blocks,
device_dst_kv_blocks,
)
else:
src_blocks, dst_blocks = prefill_kv_blocks, dst_kv_blocks
for prefill_index, decode_index in zip(src_blocks, dst_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)
@@ -750,7 +777,19 @@ class MooncakeKVManager(CommonKVManager):
dst_kv_indices: npt.NDArray[np.int32],
executor: concurrent.futures.ThreadPoolExecutor,
dst_layer_ids: Optional[List[int]] = None,
dst_device_kv_indices: Optional[npt.NDArray[np.int32]] = None,
):
dst_device_kv_ptrs = None
if dst_device_kv_indices is not None:
compression_ratios = self.kv_args.mla_compression_ratios
assert compression_ratios is not None
if len(dst_kv_ptrs) == len(self.kv_args.kv_data_ptrs):
start = self.kv_args.prefill_start_layer
end = self.kv_args.prefill_end_layer
assert end is not None
compression_ratios = compression_ratios[start:end]
c4_layer_num = sum(ratio == 4 for ratio in compression_ratios)
dst_device_kv_ptrs = set(dst_kv_ptrs[c4_layer_num:])
return self._send_kvcache_generic(
mooncake_session_id=mooncake_session_id,
src_data_ptrs=self.kv_args.kv_data_ptrs,
@@ -761,6 +800,8 @@ class MooncakeKVManager(CommonKVManager):
executor=executor,
src_layer_ids=self.kv_args.kv_layer_ids,
dst_layer_ids=dst_layer_ids,
dst_device_data_indices=dst_device_kv_indices,
dst_device_data_ptrs=dst_device_kv_ptrs,
)
def send_kvcache_dcp(
@@ -1557,12 +1598,22 @@ class MooncakeKVManager(CommonKVManager):
is_dcp_transfer = (
target_rank_registration_info.requires_dcp_relayout
)
chunked_dst_device_kv_indice = None
if is_dcp_transfer:
if req.dst_device_kv_indices is not None:
raise RuntimeError(
"HiSparse destination device indices are not "
"supported by PD DCP relayout"
)
chunked_dst_kv_indice = req.dst_kv_indices
else:
chunked_dst_kv_indice = req.dst_kv_indices[
kv_chunk.index_slice
]
if req.dst_device_kv_indices is not None:
chunked_dst_device_kv_indice = (
req.dst_device_kv_indices[kv_chunk.index_slice]
)
# NOTE: This is temporarily a workaround to deal with the case where the prefill_kv_indices
# is mismatched with the dst_kv_indices when page size > 1, this should never happen.
@@ -1577,6 +1628,12 @@ class MooncakeKVManager(CommonKVManager):
: len(chunked_dst_kv_indice)
]
)
if chunked_dst_device_kv_indice is not None:
chunked_dst_device_kv_indice = (
chunked_dst_device_kv_indice[
: len(kv_chunk.prefill_kv_indices)
]
)
skip_kv, skip_state = self._get_dsa_cache_transfer_skip_flags(
target_rank_registration_info
@@ -1620,7 +1677,8 @@ class MooncakeKVManager(CommonKVManager):
target_rank_registration_info.dst_kv_ptrs,
chunked_dst_kv_indice,
executor,
target_rank_registration_info.dst_kv_layer_ids,
dst_layer_ids=target_rank_registration_info.dst_kv_layer_ids,
dst_device_kv_indices=chunked_dst_device_kv_indice,
)
elif (
self.enable_staging
@@ -2244,6 +2302,7 @@ class MooncakeKVReceiver(CommonKVReceiver):
aux_index: Optional[int] = None,
state_indices: Optional[List] = None,
decode_prefix_len: Optional[int] = None,
device_kv_indices: Optional[npt.NDArray[np.int32]] = None,
):
if self.bootstrap_infos is None:
self.kv_mgr.record_failure(
@@ -2282,6 +2341,11 @@ class MooncakeKVReceiver(CommonKVReceiver):
),
str(self.required_dst_info_num).encode("ascii"),
str(decode_prefix_len or 0).encode("ascii"),
(
np.asarray(device_kv_indices, dtype=np.int32).tobytes()
if not is_dummy and device_kv_indices is not None
else b""
),
]
)
except zmq.ZMQError: