Expose CP shared-KV transfer page-count mismatches before decode

Mooncake previously truncated source prefill pages when the selected decode destination page list was shorter. That made an invalid prefill/decode page mapping continue as an incomplete KV transfer, which can surface later as decode garbage instead of the original mapping error.\n\nThis changes the transfer contract to require exact source/destination page-count equality and records compact page summaries in the fail-fast error. The helper is shared so the page-count contract can be unit-tested without constructing the transfer worker thread.\n\nConstraint: CP shared-KV page ownership requires a one-to-one prefill source page to decode destination page mapping.\nRejected: Keep warning-and-truncate | masks page-map corruption and can silently drop KV.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not restore source-page truncation; fix the mapping producer if this fail-fast triggers.\nTested: git diff --check; python -m py_compile python/sglang/srt/disaggregation/utils.py python/sglang/srt/disaggregation/mooncake/conn.py test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py\nNot-tested: pytest blocked locally by missing orjson; remote g0034 was unavailable during this pass.

python/sglang/srt/disaggregation/mooncake/conn.py

@@ -32,6 +32,7 @@ from sglang.srt.disaggregation.mooncake.utils import (
 from sglang.srt.disaggregation.utils import (
     DisaggregationMode,
     filter_kv_indices_for_cp_rank,
+    validate_transfer_page_count_or_raise,
 )
 from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine
 from sglang.srt.environ import envs
@@ -1052,17 +1053,18 @@ class MooncakeKVManager(CommonKVManager):
                                 kv_chunk.is_last_chunk,
                             )
 
-                        # NOTE: This is temporarily a workaround to deal with the case where the prefill_kv_indices
-                        # is mismatched with the dst_kv_indices when page size > 1, this should never happen.
-                        if len(chunked_dst_kv_indice) < len(
-                            kv_chunk.prefill_kv_indices
-                        ):
-                            logger.warning(
-                                f"len(chunked_dst_kv_indice) = {len(chunked_dst_kv_indice)}, len(kv_chunk.prefill_kv_indices) = {len(kv_chunk.prefill_kv_indices)}"
-                            )
-                            kv_chunk.prefill_kv_indices = kv_chunk.prefill_kv_indices[
-                                : len(chunked_dst_kv_indice)
-                            ]
+                        validate_transfer_page_count_or_raise(
+                            prefill_indices=kv_chunk.prefill_kv_indices,
+                            dst_indices=chunked_dst_kv_indice,
+                            room=kv_chunk.room,
+                            cp_rank=self.attn_cp_rank,
+                            logical_page_positions=kv_chunk.logical_page_positions,
+                            index_slice=kv_chunk.index_slice,
+                            is_cp_shared_kv=(
+                                kv_chunk.logical_page_positions is not None
+                            ),
+                            path="mooncake_kv",
+                        )
 
                         target_rank_registration_info: KVArgsRegisterInfo = (
                             self.decode_kv_args_table[req.mooncake_session_id]

python/sglang/srt/disaggregation/utils.py

@@ -668,6 +668,72 @@ def select_pages_by_request_positions(
     return pages[request_positions]
 
 
+def _transfer_indices_summary(indices: Any) -> str:
+    if indices is None:
+        return "None"
+    arr = np.asarray(indices)
+    if arr.size == 0:
+        return f"shape={arr.shape} dtype={arr.dtype} size=0"
+    head = arr.reshape(-1)[: min(8, arr.size)].tolist()
+    tail = arr.reshape(-1)[-min(8, arr.size) :].tolist()
+    return (
+        f"shape={arr.shape} dtype={arr.dtype} size={arr.size} "
+        f"head={head} tail={tail}"
+    )
+
+
+def _slice_summary(index_slice: Optional[slice]) -> str:
+    if index_slice is None:
+        return "None"
+    return (
+        f"slice(start={index_slice.start}, stop={index_slice.stop}, "
+        f"step={index_slice.step})"
+    )
+
+
+def validate_transfer_page_count_or_raise(
+    *,
+    prefill_indices: Any,
+    dst_indices: Any,
+    room: Optional[int],
+    cp_rank: Optional[int],
+    logical_page_positions: Optional[Any],
+    index_slice: Optional[slice],
+    is_cp_shared_kv: bool,
+    path: str,
+) -> None:
+    """Fail fast if a transfer chunk would drop or invent pages.
+
+    Mooncake used to truncate source pages when the destination page selection
+    was shorter.  Under CP shared-KV this hides a broken logical-page mapping
+    and lets decode continue with incomplete KV.  Keep the check as a small
+    shared helper so tests can cover the contract without constructing the
+    transfer worker thread.
+    """
+
+    prefill_len = len(prefill_indices) if prefill_indices is not None else 0
+    dst_len = len(dst_indices) if dst_indices is not None else 0
+    if prefill_len == dst_len:
+        return
+
+    prefix = (
+        "[CP_SHARED_KV_FAIL_FAST]"
+        if is_cp_shared_kv or logical_page_positions is not None
+        else "[KV_TRANSFER_FAIL_FAST]"
+    )
+    raise RuntimeError(
+        f"{prefix}[mooncake_transfer_page_count_mismatch] "
+        "source and destination page counts must match exactly; "
+        "continuing would silently drop or mis-map KV pages. "
+        f"path={path} room={room} cp_rank={cp_rank} "
+        f"prefill_pages={prefill_len} dst_pages={dst_len} "
+        f"logical_positions={_transfer_indices_summary(logical_page_positions)} "
+        f"index_slice={_slice_summary(index_slice)} "
+        f"prefill_indices={_transfer_indices_summary(prefill_indices)} "
+        f"dst_indices={_transfer_indices_summary(dst_indices)}"
+    )
+
+
 #########################
 # Misc
 #########################