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.
This commit is contained in:
laoyao0822
2026-06-04 23:14:22 +08:00
parent f50e2b1e00
commit e200091638
4 changed files with 243 additions and 11 deletions
@@ -8,6 +8,7 @@ from sglang.srt.disaggregation.utils import (
append_cp_draft_state_buffers,
filter_kv_pages_for_cp_shared_kv,
select_pages_by_request_positions,
validate_transfer_page_count_or_raise,
)
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
from sglang.test.ci.ci_register import register_cpu_ci
@@ -79,6 +80,34 @@ class TestCPSharedKVTransferMapping(unittest.TestCase):
self.assertEqual(logical_positions.tolist(), [2, 6])
self.assertEqual(selected_decode_pages.tolist(), [43, 47])
def test_cp_shared_kv_transfer_page_count_mismatch_fail_fast(self):
with self.assertRaisesRegex(
RuntimeError,
r"\[CP_SHARED_KV_FAIL_FAST\]\[mooncake_transfer_page_count_mismatch\].*room=17.*cp_rank=3.*prefill_pages=3.*dst_pages=2",
):
validate_transfer_page_count_or_raise(
prefill_indices=np.array([11, 12, 13], dtype=np.int32),
dst_indices=np.array([21, 22], dtype=np.int32),
room=17,
cp_rank=3,
logical_page_positions=np.array([0, 1, 2], dtype=np.int32),
index_slice=None,
is_cp_shared_kv=True,
path="mooncake_kv",
)
def test_cp_shared_kv_transfer_page_count_match_passes(self):
validate_transfer_page_count_or_raise(
prefill_indices=np.array([11, 12], dtype=np.int32),
dst_indices=np.array([21, 22], dtype=np.int32),
room=18,
cp_rank=1,
logical_page_positions=np.array([4, 8], dtype=np.int32),
index_slice=None,
is_cp_shared_kv=True,
path="mooncake_kv",
)
def test_append_cp_draft_state_buffers_registers_nsa_state_as_target_payload(self):
kv_args = SimpleNamespace(
state_type="nsa",