diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 2bbfbd6de..5796cbf7c 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -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 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