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