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:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user