fix(cp_per_layer_transfer): use absolute chunk_page_start for chunked KV dst mapping

The per-layer KV transfer registration hardcoded chunk_page_start=0 when
filtering CP-shared-KV owned pages. The CP filter's second return (`positions`)
are absolute full-sequence page positions built from chunk_page_start, and the
transfer indexes the FULL-request dst_kv_indices by those absolute positions
(mirroring the monolithic send(), which passes chunk_page_start=index_slice.start
— the cumulative page offset). With start=0, chunk N>0's positions were
chunk-local, so its KV was written onto chunk 0's decode pages, corrupting the
decode output. Non-chunked requests (single chunk, start=0) were unaffected,
matching the observed symptom (non-chunked byte-identical, chunked garbage).

Fix: chunk_page_start = chunk_key // page_size, where chunk_key is the chunk's
start_send_idx (page-aligned), making it exactly the monolithic index_slice.start.

Verified: opus first-principles code audit; empirical mapping-invariant on the
deployed modules (per-chunk == whole-request for all 8 CP ranks; old start=0
sends chunk1 to chunk0's dst); 2 new regression tests (TestChunkedDstMapping).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-06-07 08:53:10 +00:00
parent beffd3a82f
commit e9354c41bc
2 changed files with 95 additions and 1 deletions

View File

@@ -417,11 +417,19 @@ class MooncakeKVManager(CommonKVManager):
cp_rank=self.attn_cp_rank,
)
pages = np.asarray(page_indices, dtype=np.int32)
# chunk_key is this chunk's start_send_idx in TOKENS (page-aligned for any
# non-first chunk). The CP filter's `positions` (second return) are absolute
# full-sequence page positions built from chunk_page_start, and the transfer
# indexes the FULL-request dst_kv_indices by those absolute positions (mirrors
# send(): chunk_page_start=index_slice.start, worker: dst_kv_indices[logical_
# 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=0
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(