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

fix(cp_per_layer_transfer): gate per-layer path on a single non-dummy decode info

Browse Source 
The transfer worker iterates every non-dummy decode info for a room and calls
per_layer_mgr.finish() once per info, but register_per_layer_transfer registers
exactly one context per room/chunk (built for one info's dst_kv_indices). This is
only sound when there is exactly one non-dummy info (required_dst_info_num == 1).
With decode attn_tp < prefill attn_tp a single prefill rank holds >1 non-dummy
infos; finishing once-per-info would over-pop chunk contexts and under-deliver KV
to the other infos. Make the assumption explicit: register only when there is one
non-dummy info, otherwise fall back to the monolithic post-forward transfer (which
fans out to all infos correctly). Found by an independent first-principles audit.

Adds TestRegisterGuardSingleInfo (2-info fallback, 1-info register, all-dummy
fallback) exercising the real MooncakeKVManager.register_per_layer_transfer.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-07 08:59:02 +00:00
parent e9354c41bc
commit 243f1b964c
2 changed files with 96 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
python/sglang/srt/disaggregation/mooncake/conn.py
View File

@@ -408,6 +408,20 @@ class MooncakeKVManager(CommonKVManager):
infos = self.transfer_infos.get(room)
if not infos:
return False
# The per-layer path registers exactly ONE context per room/chunk, but the
# transfer worker iterates every non-dummy decode info for the room and calls
# finish() per info (conn.py reqs_to_be_processed loop). That is only sound
# when there is exactly one non-dummy info (required_dst_info_num == 1). For
# decode attn_tp < prefill attn_tp a single prefill rank holds >1 non-dummy
# infos; finishing once-per-info would over-pop chunk contexts and the single
# ctx only carries one info's dst_kv_indices. Fall back to the monolithic
# post-forward transfer (which fans out to all infos) in that case.
non_dummy = [
info for info in infos.values() if not getattr(info, "is_dummy", False)
]
if len(non_dummy) != 1:
return False
info = non_dummy[0]
from sglang.srt.disaggregation.utils import filter_kv_pages_for_cp_shared_kv
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
@@ -425,27 +439,23 @@ class MooncakeKVManager(CommonKVManager):
# page_positions]). Must offset by the chunk's absolute page start, else chunk
# N>0 writes its KV onto chunk 0's decode pages. page_size divides chunk_key.
chunk_page_start = int(chunk_key) // self.kv_args.page_size
for info in infos.values():
if getattr(info, "is_dummy", False):
continue
owned_pages, positions = filter_kv_pages_for_cp_shared_kv(
layout=layout, logical_pages=pages, chunk_page_start=chunk_page_start
)
dst_indices = np.asarray(info.dst_kv_indices, dtype=np.int32)[positions]
ctx = self.build_per_layer_context(
info.mooncake_session_id, owned_pages, dst_indices
)
if ctx is not None:
mgr.register(room, ctx, chunk_key=chunk_key)
logger.info(
"[CP_PER_LAYER_TRANSFER] registered room=%s chunk=%s owned_pages=%d",
room,
chunk_key,
len(owned_pages),
)
return True
owned_pages, positions = filter_kv_pages_for_cp_shared_kv(
layout=layout, logical_pages=pages, chunk_page_start=chunk_page_start
)
dst_indices = np.asarray(info.dst_kv_indices, dtype=np.int32)[positions]
ctx = self.build_per_layer_context(
info.mooncake_session_id, owned_pages, dst_indices
)
if ctx is None:
return False
return False
mgr.register(room, ctx, chunk_key=chunk_key)
logger.info(
"[CP_PER_LAYER_TRANSFER] registered room=%s chunk=%s owned_pages=%d",
room,
chunk_key,
len(owned_pages),
)
return True
def _send_kvcache_generic(
self,