chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

fix(disagg/mooncake): clip CP/NSA state indices to the shorter side

Browse Source 
Port upstream sgl-project/sglang #23323. Our state_type in [swa, nsa]
transfer path only handled len(prefill_state_indices) < len(dst), and even
then clipped prefill_state_indices (a no-op when prefill is shorter), and
never handled the prefill > dst case. Now clip prefill when it is longer and
clip the dst state indices when it is shorter, so the extra-pool transfer
never reads past prefill_state_indices or writes past dst_state_indices.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-06 21:59:43 +00:00
parent f75ffff8d9
commit 0cd5c38a9e
1 changed files with 12 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
python/sglang/srt/disaggregation/mooncake/conn.py
View File

@@ -822,13 +822,24 @@ class MooncakeKVManager(CommonKVManager):
if dst_state_indices is not None
else np.asarray(req.dst_state_indices, dtype=np.int32)
)
if len(prefill_state_indices) < len(effective_dst_state_indices):
# Clip src/dst state indices to the shorter side so the transfer never
# reads past prefill_state_indices or writes past dst_state_indices.
# Ported from upstream sgl-project/sglang #23323 (the prior logic only
# handled prefill<dst and clipped the wrong array, a no-op).
if len(prefill_state_indices) > len(effective_dst_state_indices):
logger.warning(
f"len(prefill_state_indices) = {len(prefill_state_indices)}, len(dst_state_indices) = {len(effective_dst_state_indices)}"
)
prefill_state_indices = prefill_state_indices[
: len(effective_dst_state_indices)
]
elif len(prefill_state_indices) < len(effective_dst_state_indices):
logger.warning(
f"len(prefill_state_indices) = {len(prefill_state_indices)}, len(dst_state_indices) = {len(effective_dst_state_indices)}"
)
effective_dst_state_indices = effective_dst_state_indices[
: len(prefill_state_indices)
]
src_state_data_ptrs = self.kv_args.state_data_ptrs
dst_state_ptrs = dst_state_data_ptrs
state_item_lens = self.kv_args.state_item_lens