chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
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
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/mooncake/conn.py
+13 -11
View File
@@ -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]