"""Unit tests for A2 core: PerLayerTransferContext (per-request overlapped transfer). Key correctness properties: - submit_layer waits the layer's CUDA write event BEFORE issuing the RDMA (so the transfer never reads a layer before its write kernel finished); - idempotent per layer; skips layers with no owned blocks; - a submit failure marks the context failed and stops further submits; - finish waits all accumulated batch_ids and returns 0 / non-zero correctly. Engine and CUDA events are mocked — no GPU/RDMA. """ import unittest from unittest import mock from sglang.srt.disaggregation.cp_per_layer_transfer import PerLayerTransferContext class _FakeEvent: def __init__(self): self.synced = False def synchronize(self): self.synced = True class _FakeEngine: def __init__(self, submit_rets=None, wait_ret=0): self.submits = [] self.waits = [] self._rets = submit_rets self._wait = wait_ret self._n = 0 self.submit_order = [] # (layer-derived addr, whether event was synced first) def batch_transfer_async_submit(self, sid, src, dst, lens): self.submits.append((sid, list(src), list(dst), list(lens))) if self._rets is not None: r = self._rets[self._n] self._n += 1 return r return 100 + len(self.submits) def wait_batch_transfers(self, ids): self.waits.append(list(ids)) return self._wait def _blocks(layer_id): base = 1000 * layer_id return ([base], [base + 5000], [64]) def _ctx(engine, blocks=_blocks, num_layers=8): return PerLayerTransferContext(engine, "sess", blocks, num_layers=num_layers) class TestPerLayerTransferContext(unittest.TestCase): def test_submit_waits_event_then_submits(self): eng = _FakeEngine() ev = _FakeEvent() _ctx(eng).submit_layer(3, ev) self.assertTrue(ev.synced) # write event waited BEFORE the RDMA submit self.assertEqual(eng.submits, [("sess", [3000], [8000], [64])]) def test_idempotent_per_layer(self): eng = _FakeEngine() ctx = _ctx(eng) ctx.submit_layer(0, _FakeEvent()) ctx.submit_layer(0, _FakeEvent()) self.assertEqual(len(eng.submits), 1) def test_skip_layer_with_no_blocks(self): eng = _FakeEngine() def blocks(L): return None if L == 1 else ([L], [L + 1], [8]) ctx = _ctx(eng, blocks) for L in range(3): ctx.submit_layer(L, None) self.assertEqual(len(eng.submits), 2) # layer 1 (no blocks) skipped def test_finish_waits_all_and_returns_zero(self): eng = _FakeEngine() ctx = _ctx(eng, num_layers=3) for L in range(3): ctx.submit_layer(L, _FakeEvent()) self.assertEqual(ctx.finish(), 0) self.assertEqual(eng.waits, [[101, 102, 103]]) def test_submit_failure_marks_failed_and_stops(self): eng = _FakeEngine(submit_rets=[101, -1, 103]) ctx = _ctx(eng, num_layers=3) for L in range(3): ctx.submit_layer(L, None) self.assertTrue(ctx.failed) self.assertEqual(len(eng.submits), 2) # stopped submitting after the -1 self.assertEqual(ctx.finish(), -1) def test_finish_propagates_wait_failure(self): eng = _FakeEngine(wait_ret=-9) ctx = _ctx(eng, num_layers=1) ctx.submit_layer(0, None) self.assertEqual(ctx.finish(), -9) def test_repeated_enqueues_dedup_so_finish_cannot_hang(self): # the notifier re-fires every layer on every subsequent forward while the ctx # is still active; note_enqueued must dedup so target stays capped at the layer # count and finish() completes (the high-cache-hit hang was target>>layers). eng = _FakeEngine() ctx = _ctx(eng, num_layers=3) for _forward in range(5): # 5 forwards all re-fire layers 0,1,2 for L in range(3): self.assertEqual(ctx.note_enqueued(L), _forward == 0) # True only 1st if _forward == 0: ctx.submit_layer(L, None) self.assertEqual(len(eng.submits), 3) # 3 unique despite 15 note_enqueued self.assertEqual(ctx.finish(timeout=0.5), 0) # completes, no 30s timeout def test_finish_falls_back_for_notifier_missed_layers(self): # only layer 0 fired by the notifier; finish must SYNCHRONOUSLY transfer the # missing layers 1,2 -> all 3 moved, success. The fallback BATCHES the missing # layers into one submit. eng = _FakeEngine() ctx = _ctx(eng, num_layers=3) ctx.submit_layer(0, None) # notifier fired layer 0 only -> 1 submit self.assertEqual(ctx.finish(timeout=0.5), 0) self.assertEqual(len(eng.submits), 2) # layer 0 + fallback batch [1,2] self.assertEqual(eng.waits[-1], [101, 102]) self.assertEqual(eng.submits[1], ("sess", [1000, 2000], [6000, 7000], [64, 64])) def test_successful_finish_does_not_emit_hot_path_info_log(self): # Successful per-request transfer completion is hot-path and can fire once per # request/rank/chunk. It must not emit INFO logs by default. eng = _FakeEngine() ctx = _ctx(eng, num_layers=1) ctx.submit_layer(0, None) with mock.patch( "sglang.srt.disaggregation.cp_per_layer_transfer.logger.isEnabledFor", return_value=True, ), mock.patch( "sglang.srt.disaggregation.cp_per_layer_transfer.logger.info" ) as info: self.assertEqual(ctx.finish(timeout=0.5), 0) info.assert_not_called() def test_failed_submit_still_reported_after_fallback(self): # layer 2 submit fails during the overlap; finish reports -1 even though it # also runs the fallback for any unfired layers. eng = _FakeEngine(submit_rets=[101, -1, 103]) ctx = _ctx(eng, num_layers=3) for L in range(3): ctx.submit_layer(L, None) self.assertEqual(ctx.finish(timeout=0.5), -1) class TestBuildLayerBlocks(unittest.TestCase): def test_addresses_and_lengths(self): from sglang.srt.disaggregation.cp_per_layer_transfer import build_layer_blocks src, dst, lens = build_layer_blocks( 1000, 5000, 64, [[0, 1, 2], [5, 6]], [[10, 11, 12], [20, 21]] ) self.assertEqual(src, [1000 + 0 * 64, 1000 + 5 * 64]) self.assertEqual(dst, [5000 + 10 * 64, 5000 + 20 * 64]) self.assertEqual(lens, [64 * 3, 64 * 2]) # item_len * run length def test_empty(self): from sglang.srt.disaggregation.cp_per_layer_transfer import build_layer_blocks self.assertEqual(build_layer_blocks(1, 2, 8, [], []), ([], [], [])) def test_only_base_ptr_changes_across_layers(self): from sglang.srt.disaggregation.cp_per_layer_transfer import build_layer_blocks s0, d0, l0 = build_layer_blocks(1000, 2000, 64, [[3]], [[7]]) s1, d1, l1 = build_layer_blocks(9000, 8000, 64, [[3]], [[7]]) self.assertEqual(s0, [1000 + 3 * 64]) self.assertEqual(s1, [9000 + 3 * 64]) self.assertEqual(l0, l1) # lengths identical across layers (the invariant) class _MockCtx: num_layers = 8 def __init__(self): self.submitted = [] # (start, end, event) self.enqueued = 0 self.finished = False self.failed = False self._enq = set() def note_enqueued(self, key): if key in self._enq: return False self._enq.add(key) self.enqueued += 1 return True def submit_group(self, start, end, event): self.submitted.append((start, end, event)) def finish(self): self.finished = True return 0 def mark_failed(self): self.failed = True class _RecEvent: def __init__(self): self.recorded_stream = "UNSET" self.synced = False def record(self, stream=None): self.recorded_stream = stream def synchronize(self): self.synced = True def _manager(num_workers=0, group_size=1): from sglang.srt.disaggregation.cp_per_layer_transfer import PerLayerTransferManager return PerLayerTransferManager( num_workers=num_workers, event_factory=_RecEvent, current_stream=lambda: "STREAM", group_size=group_size, ) def _drain(q): import queue as _q items = [] while True: try: items.append(q.get(block=False)) except _q.Empty: break return items class TestPerLayerTransferManager(unittest.TestCase): def test_on_layer_end_enqueues_per_active_ctx_with_recorded_event(self): m = _manager(group_size=1) # per-layer: every layer is a group boundary c1, c2 = _MockCtx(), _MockCtx() m.register("r1", c1) m.register("r2", c2) m.on_layer_end(5) items = _drain(m._q) self.assertEqual(len(items), 2) # one per active context for ctx, start, end, ev in items: self.assertEqual((start, end), (5, 5)) self.assertEqual(ev.recorded_stream, "STREAM") def test_grouping_only_enqueues_at_boundaries(self): m = _manager(group_size=4) c = _MockCtx() m.register("r1", c) for L in range(8): m.on_layer_end(L) items = _drain(m._q) # boundaries at L=3 (group [0,3]) and L=7 (group [4,7]) -> 2 groups, not 8 self.assertEqual([(s, e) for _, s, e, _ in items], [(0, 3), (4, 7)]) def test_chunked_fifo_separate_contexts_and_dedup(self): m = _manager(group_size=1) c1, c2 = _MockCtx(), _MockCtx() m.register("r1", c1, chunk_key=0) # chunk 1 m.register("r1", c1, chunk_key=0) # SAME chunk re-register -> deduped m.register("r1", c2, chunk_key=64) # chunk 2 (different start_send_idx) self.assertTrue(m.has_chunk("r1", 0)) self.assertTrue(m.has_chunk("r1", 64)) self.assertFalse(m.has_chunk("r1", 128)) m.on_layer_end(5) # both chunk contexts active -> both enqueued self.assertEqual(len(_drain(m._q)), 2) # finish pops FIFO: chunk 1 (c1) first, then chunk 2 (c2) self.assertEqual(m.finish("r1"), 0) self.assertTrue(c1.finished) self.assertFalse(c2.finished) self.assertEqual(m.finish("r1"), 0) self.assertTrue(c2.finished) self.assertFalse(m.has_room("r1")) def test_worker_step_calls_submit_group(self): m = _manager() c = _MockCtx() ev = _RecEvent() m._worker_step((c, 0, 3, ev)) self.assertEqual(c.submitted, [(0, 3, ev)]) def test_worker_step_marks_failed_on_exception(self): m = _manager() class _Boom: def __init__(self): self.failed = False def submit_group(self, start, end, event): raise RuntimeError("boom") def mark_failed(self): self.failed = True b = _Boom() m._worker_step((b, 0, 0, None)) self.assertTrue(b.failed) def test_finish_pops_and_calls_ctx_finish(self): m = _manager() c = _MockCtx() m.register("r1", c) self.assertTrue(m.has_active()) self.assertEqual(m.finish("r1"), 0) self.assertTrue(c.finished) self.assertFalse(m.has_active()) self.assertEqual(m.finish("r1"), 0) # idempotent after pop def test_on_layer_end_no_active_is_noop(self): m = _manager(group_size=1) m.on_layer_end(0) 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") class TestRegisterGuardSingleInfo(unittest.TestCase): """register_per_layer_transfer must register a context only when there is exactly one non-dummy decode info. The transfer worker calls finish() once per non-dummy info for the room while the per-layer path registers a single context, so >1 non-dummy info (decode attn_tp < prefill attn_tp) must fall back to the monolithic transfer.""" def _fake_self(self, infos, recorder): import types mgr = types.SimpleNamespace( register=lambda room, ctx, chunk_key=0: recorder.append( (room, chunk_key) ), ) return types.SimpleNamespace( per_layer_transfer_manager=mgr, server_args=types.SimpleNamespace(enable_nsa_prefill_cp_shared_kv=True), transfer_infos={7: infos}, kv_args=types.SimpleNamespace(page_size=64), attn_cp_size=8, attn_cp_rank=0, build_per_layer_context=lambda sid, owned, dst: ("ctx", sid), ) def _info(self, is_dummy): import types import numpy as np return types.SimpleNamespace( is_dummy=is_dummy, dst_kv_indices=np.arange(64, dtype=np.int32), mooncake_session_id="s", ) def _call(self, infos, recorder): try: from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager except Exception as e: # pragma: no cover - env without mooncake self.skipTest(f"mooncake conn not importable: {e}") s = self._fake_self(infos, recorder) # logical pages incl. page 1 which cp_rank 0 owns ((1-1)%8==0) return MooncakeKVManager.register_per_layer_transfer( s, 7, [0, 1, 2, 3], chunk_key=0 ) def test_two_non_dummy_infos_falls_back(self): rec = [] infos = {"a": self._info(False), "b": self._info(False)} self.assertFalse(self._call(infos, rec)) self.assertEqual(rec, []) # nothing registered def test_single_non_dummy_info_registers(self): rec = [] infos = {"a": self._info(True), "b": self._info(False)} # 1 dummy + 1 real self.assertTrue(self._call(infos, rec)) self.assertEqual(len(rec), 1) def test_all_dummy_falls_back(self): rec = [] infos = {"a": self._info(True), "b": self._info(True)} self.assertFalse(self._call(infos, rec)) self.assertEqual(rec, []) if __name__ == "__main__": unittest.main()