From beffd3a82fa1d0ac74cb74700e84ff21c4f133fd Mon Sep 17 00:00:00 2001 From: leavelet Date: Sun, 7 Jun 2026 07:33:46 +0000 Subject: [PATCH] feat(disagg): unify per-layer transfer over chunked prefill (lever A) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per code review (HiCache load is layer-by-layer & correctly ordered into the compute stream before the backup hook; current-reuse is a within-forward read that doesn't rewrite pool pages): unify registration to the EXACT range this forward's send_kv_chunk transmits — req_to_token[start_send_idx:end_idx], page-floored for a non-last chunk. Non-chunked = one full range; chunked = one range per chunk. Drop the is_chunked/start_send_idx skip. To avoid the review's collision risk (chunk N still finishing when chunk N+1 registers), the manager keys contexts per (room, start_send_idx): _active[room] is a FIFO deque of (chunk_key, ctx); register dedups the same chunk but appends a new one; finish(room) pops the FRONT (chunks finish in send order — no chunk key needed in the transfer_worker); drop drains all the room's chunks; on_layer_end enqueues for all active chunk contexts (per-ctx note_enqueued dedup keeps each chunk's own events). 28 unit tests pass incl. chunked FIFO + per-chunk dedup. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../disaggregation/cp_per_layer_transfer.py | 49 ++++++++++++++----- .../srt/disaggregation/mooncake/conn.py | 7 +-- python/sglang/srt/disaggregation/prefill.py | 36 +++++++++----- .../test_cp_per_layer_transfer.py | 19 +++++++ 4 files changed, 82 insertions(+), 29 deletions(-) diff --git a/python/sglang/srt/disaggregation/cp_per_layer_transfer.py b/python/sglang/srt/disaggregation/cp_per_layer_transfer.py index 8053647b9..557384f80 100644 --- a/python/sglang/srt/disaggregation/cp_per_layer_transfer.py +++ b/python/sglang/srt/disaggregation/cp_per_layer_transfer.py @@ -245,19 +245,33 @@ class PerLayerTransferManager: self._group_size = max(1, int(group_size)) # layers per RDMA submit (overhead) self._num_layers = 0 # learned from the first registered ctx (model layer count) self._q = _queue.SimpleQueue() + # room -> FIFO deque of (chunk_key, ctx). One ctx per chunk (chunked prefill + # sends per chunk); FIFO because the transfer_worker finishes chunks in the + # order they were sent, so finish(room) pops the front without needing a key. self._active = {} self._active_lock = threading.Lock() for _ in range(max(0, num_workers)): threading.Thread(target=self._worker, daemon=True).start() - def register(self, room, ctx: PerLayerTransferContext) -> None: + def register(self, room, ctx: PerLayerTransferContext, chunk_key=0) -> None: + import collections + with self._active_lock: - if room in self._active: - return # already registered for its forward; don't overwrite (leak) - self._active[room] = ctx + dq = self._active.get(room) + if dq is not None and any(k == chunk_key for k, _ in dq): + return # this chunk already registered (bs>1 batch-forming re-iterates) + if dq is None: + dq = collections.deque() + self._active[room] = dq + dq.append((chunk_key, ctx)) if self._num_layers == 0: self._num_layers = ctx.num_layers + def has_chunk(self, room, chunk_key) -> bool: + with self._active_lock: + dq = self._active.get(room) + return dq is not None and any(k == chunk_key for k, _ in dq) + def has_active(self) -> bool: with self._active_lock: return bool(self._active) @@ -274,7 +288,9 @@ class PerLayerTransferManager: with self._active_lock: if not self._active: return - ctxs = list(self._active.values()) + # every active ctx across all rooms/chunks; per-ctx note_enqueued dedups so + # a still-finishing chunk's ctx ignores a later chunk's re-fired layers. + ctxs = [ctx for dq in self._active.values() for _, ctx in dq] event = None if self._event_factory is not None: event = self._event_factory() @@ -290,20 +306,27 @@ class PerLayerTransferManager: def has_room(self, room) -> bool: with self._active_lock: - return room in self._active + return bool(self._active.get(room)) def finish(self, room) -> int: + # Pop the FRONT chunk (FIFO): the transfer_worker finishes chunks in send order. with self._active_lock: - ctx = self._active.pop(room, None) - return ctx.finish() if ctx is not None else 0 + dq = self._active.get(room) + if not dq: + return 0 + _, ctx = dq.popleft() + if not dq: + del self._active[room] + return ctx.finish() def drop(self, room) -> None: - """Abort path: drain + discard a room's context so any outstanding RDMA - completes before its KV pages are freed (never write into reclaimed pages).""" + """Abort path: drain + discard ALL of a room's chunk contexts so any + outstanding RDMA completes before its KV pages are freed.""" with self._active_lock: - ctx = self._active.pop(room, None) - if ctx is not None: - ctx.finish() + dq = self._active.pop(room, None) + if dq: + for _, ctx in dq: + ctx.finish() def _worker_step(self, item) -> None: ctx, start, end, event = item diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 1d4b03d61..399c8f738 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -395,7 +395,7 @@ class MooncakeKVManager(CommonKVManager): self.engine, mooncake_session_id, get_blocks, num_layers=n_layers ) - def register_per_layer_transfer(self, room, page_indices) -> bool: + def register_per_layer_transfer(self, room, page_indices, chunk_key=0) -> bool: """Lever A: before the forward, build + register a per-layer transfer context for `room` so the per-layer notifier overlaps its main-KV transfer with the forward. Reuses send()'s CP filter exactly (no re-derivation). Scoped to @@ -428,10 +428,11 @@ class MooncakeKVManager(CommonKVManager): info.mooncake_session_id, owned_pages, dst_indices ) if ctx is not None: - mgr.register(room, ctx) + mgr.register(room, ctx, chunk_key=chunk_key) logger.info( - "[CP_PER_LAYER_TRANSFER] registered room=%s owned_pages=%d", + "[CP_PER_LAYER_TRANSFER] registered room=%s chunk=%s owned_pages=%d", room, + chunk_key, len(owned_pages), ) return True diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 84bc45f87..15881b4f0 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -667,23 +667,33 @@ class SchedulerDisaggregationPrefillMixin: if getattr(req, "disagg_kv_sender", None) is None: continue room = getattr(req, "bootstrap_room", None) - if room is None or mgr.has_room(room): - continue # already registered (bs>1 batch-forming re-iterates the - # same reqs many times) — skip the CP filter + build entirely - if getattr(req, "is_chunked", 0) != 0 or getattr(req, "start_send_idx", 0) != 0: - continue # chunked / partially-sent: this forward isn't the full range - # Cover the FULL sequence (cached prefix + new tokens). The per-layer - # backup hook fires after each layer's full processing, by which point - # layer L's HiCache-loaded prefix KV and forward-written new KV are both - # final — so transferring all of layer L's pages then is correct, and it - # overlaps the (dominant, at high cache-hit) prefix transfer with the - # forward. Verified by output-equality. + if room is None: + continue + start_idx = getattr(req, "start_send_idx", 0) + if mgr.has_chunk(room, start_idx): + continue # this chunk already registered (bs>1 batch-forming re-iterates) + # Register the EXACT range this forward's send_kv_chunk will transmit: + # req_to_token[start_send_idx : end_idx], page-floored for a NON-last chunk + # (mirrors send_kv_chunk). UNIFIED: non-chunked = one full range [0:end]; + # chunked = one range per chunk. The per-layer backup hook fires after each + # layer's full processing (HiCache-loaded prefix KV + forward-written new KV + # both final), so transferring layer L's pages then is correct and overlaps + # the transfer with the forward. Multiple chunks of a room get separate + # contexts (keyed by start_send_idx), finished FIFO. end_idx = min(len(req.fill_ids), len(req.origin_input_ids)) + if getattr(req, "is_chunked", 0) > 0: # not the last chunk + end_idx -= end_idx % page_size + if end_idx <= start_idx: + continue page_indices = _kv_locs_to_page_indices_cpu( - self.req_to_token_pool.req_to_token[req.req_pool_idx, 0:end_idx], + self.req_to_token_pool.req_to_token[ + req.req_pool_idx, start_idx:end_idx + ], page_size, ) - kv_manager.register_per_layer_transfer(room, page_indices) + kv_manager.register_per_layer_transfer( + room, page_indices, chunk_key=start_idx + ) except Exception as e: # Never let lever-A setup crash the scheduler; fall back to the # monolithic post-forward transfer for this request. 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 ee8e556de..010bc3d15 100644 --- a/test/registered/unit/disaggregation/test_cp_per_layer_transfer.py +++ b/test/registered/unit/disaggregation/test_cp_per_layer_transfer.py @@ -250,6 +250,25 @@ class TestPerLayerTransferManager(unittest.TestCase): # 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()