From 0cd5c38a9e3b5a0e3deef4ea82eea1da435945d4 Mon Sep 17 00:00:00 2001 From: leavelet Date: Sat, 6 Jun 2026 21:59:43 +0000 Subject: [PATCH] fix(disagg/mooncake): clip CP/NSA state indices to the shorter side 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) --- python/sglang/srt/disaggregation/mooncake/conn.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) 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