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

[PD] Fix clip logic when state indices lens are mismatch (#23323)

Browse Source
This commit is contained in:
Shangming Cai
2026-04-21 13:22:20 +08:00
committed by GitHub
parent 5595f6e988
commit a58c7f381e
1 changed files with 10 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/mooncake/conn.py
+10 -6
View File
@@ -1003,16 +1003,20 @@ class MooncakeKVManager(CommonKVManager):
raise RuntimeError( raise RuntimeError(
f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {state_type.upper()} hybrid models yet." f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {state_type.upper()} hybrid models yet."
) )
if len(prefill_state_indices) < len(req.dst_state_indices): dst_state_indices = req.dst_state_indices
if len(prefill_state_indices) > len(dst_state_indices):
logger.warning( logger.warning(
f"len(prefill_state_indices) = {len(prefill_state_indices)}, len(dst_state_indices) = {len(req.dst_state_indices)}" f"len(prefill_state_indices) = {len(prefill_state_indices)}, len(dst_state_indices) = {len(dst_state_indices)}"
) )
prefill_state_indices = prefill_state_indices[ prefill_state_indices = prefill_state_indices[: len(dst_state_indices)]
: len(req.dst_state_indices) elif len(prefill_state_indices) < len(dst_state_indices):
] logger.warning(
f"len(prefill_state_indices) = {len(prefill_state_indices)}, len(dst_state_indices) = {len(dst_state_indices)}"
)
dst_state_indices = dst_state_indices[: len(prefill_state_indices)]
# Reuse _send_kvcache_generic interface to send extra pool data # Reuse _send_kvcache_generic interface to send extra pool data
prefill_state_indices = np.array(prefill_state_indices, dtype=np.int32) prefill_state_indices = np.array(prefill_state_indices, dtype=np.int32)
dst_state_indices = np.array(req.dst_state_indices, dtype=np.int32) dst_state_indices = np.array(dst_state_indices, dtype=np.int32)
return self._send_kvcache_generic( return self._send_kvcache_generic(
mooncake_session_id=req.mooncake_session_id, mooncake_session_id=req.mooncake_session_id,
src_data_ptrs=self.kv_args.state_data_ptrs, src_data_ptrs=self.kv_args.state_data_ptrs,