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

This commit is contained in:
Shangming Cai
2026-04-21 13:22:20 +08:00
committed by GitHub
parent 5595f6e988
commit a58c7f381e
@@ -1003,16 +1003,20 @@ class MooncakeKVManager(CommonKVManager):
raise RuntimeError(
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(
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[
: len(req.dst_state_indices)
]
prefill_state_indices = prefill_state_indices[: len(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
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(
mooncake_session_id=req.mooncake_session_id,
src_data_ptrs=self.kv_args.state_data_ptrs,