From e9354c41bcfbf98d07d290efc7eefdd2e5c902e1 Mon Sep 17 00:00:00 2001 From: leavelet Date: Sun, 7 Jun 2026 08:53:10 +0000 Subject: [PATCH] fix(cp_per_layer_transfer): use absolute chunk_page_start for chunked KV dst mapping MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../srt/disaggregation/mooncake/conn.py | 10 ++- .../test_cp_per_layer_transfer.py | 86 +++++++++++++++++++ 2 files changed, 95 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 399c8f738..e93a3b2e7 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -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( diff --git a/test/registered/unit/disaggregation/test_cp_per_layer_transfer.py b/test/registered/unit/disaggregation/test_cp_per_layer_transfer.py index 010bc3d15..020edf465 100644 --- a/test/registered/unit/disaggregation/test_cp_per_layer_transfer.py +++ b/test/registered/unit/disaggregation/test_cp_per_layer_transfer.py @@ -309,5 +309,91 @@ class TestPerLayerTransferManager(unittest.TestCase): self.assertEqual(_drain(m._q), []) +class TestChunkedDstMapping(unittest.TestCase): + """Regression for the chunked per-layer dst-offset bug. + + register_per_layer_transfer maps a chunk's owned source pages to decode pages + via dst_kv_indices[positions], where positions come from the CP filter built on + chunk_page_start. The monolithic send() uses chunk_page_start=index_slice.start + (the absolute, cumulative page offset) and the worker indexes the FULL-request + dst_kv_indices by those absolute positions. So a per-chunk registration MUST pass + chunk_page_start = start_send_idx // page_size; passing 0 for every chunk makes + chunk N>0's positions chunk-local and silently writes its KV onto chunk 0's + decode pages (observed as corrupted decode output).""" + + def _layout(self, cp_rank, cp_size=8, page_size=64): + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + return CpSharedKVLayout( + page_size=page_size, cp_size=cp_size, cp_rank=cp_rank + ) + + def _map_chunk(self, layout, pages, dst_kv_indices, chunk_page_start): + """Mirror register_per_layer_transfer's source->dst mapping for one chunk.""" + import numpy as np + + from sglang.srt.disaggregation.utils import ( + filter_kv_pages_for_cp_shared_kv, + ) + + owned, positions = filter_kv_pages_for_cp_shared_kv( + layout=layout, + logical_pages=np.asarray(pages, dtype=np.int32), + chunk_page_start=chunk_page_start, + ) + dst = np.asarray(dst_kv_indices, dtype=np.int32)[positions] + return list(owned), list(dst) + + def test_correct_offset_matches_whole_request_per_rank(self): + import numpy as np + + page_size = 64 + cp_size = 8 + total_pages = 24 # logical pages 0..23 + all_pages = np.arange(total_pages, dtype=np.int32) + # distinct decode page per absolute logical position + dst_kv_indices = (90000 + np.arange(total_pages)).astype(np.int32) + # 2 chunks split at logical page 16 (token 1024 with page_size 64) + chunk0 = all_pages[:16] + chunk1 = all_pages[16:] + for cp_rank in range(cp_size): + layout = self._layout(cp_rank, cp_size, page_size) + # whole-request ground truth (single shot, offset 0) + whole_src, whole_dst = self._map_chunk( + layout, all_pages, dst_kv_indices, chunk_page_start=0 + ) + # chunked with CORRECT absolute offsets -> identical pairing + s0, d0 = self._map_chunk(layout, chunk0, dst_kv_indices, 0) + s1, d1 = self._map_chunk(layout, chunk1, dst_kv_indices, 16) + self.assertEqual(whole_src, s0 + s1, f"src mismatch rank={cp_rank}") + self.assertEqual(whole_dst, d0 + d1, f"dst mismatch rank={cp_rank}") + + def test_zero_offset_corrupts_second_chunk(self): + """The pre-fix behaviour: chunk_page_start=0 for chunk 1 maps its owned + pages onto chunk 0's decode region (dst < the chunk's true start).""" + import numpy as np + + page_size = 64 + cp_size = 8 + total_pages = 24 + all_pages = np.arange(total_pages, dtype=np.int32) + dst_kv_indices = (90000 + np.arange(total_pages)).astype(np.int32) + chunk1 = all_pages[16:] + # find a rank that actually owns at least one page in chunk 1 + for cp_rank in range(cp_size): + layout = self._layout(cp_rank, cp_size, page_size) + _, d_bug = self._map_chunk(layout, chunk1, dst_kv_indices, 0) + _, d_fixed = self._map_chunk(layout, chunk1, dst_kv_indices, 16) + if not d_fixed: + continue + # fixed dst pages are all in chunk1's region (>= 90000+16) + self.assertTrue(all(x >= 90000 + 16 for x in d_fixed)) + # buggy dst pages land in chunk0's region (< 90000+16) -> corruption + self.assertTrue(all(x < 90000 + 16 for x in d_bug)) + self.assertNotEqual(d_bug, d_fixed) + return + self.fail("no rank owned a chunk-1 page; adjust fixture") + + if __name__ == "__main__": unittest.main()