feat(disagg): unify per-layer transfer over chunked prefill (lever A)

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) <noreply@anthropic.com>

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

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

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.