fix(disagg): cover notifier-missed layers via post-forward fallback (lever A)

E2E diagnostic was precise: "finish TIMEOUT processed=78/79", submit_failed=0 —
the per-layer notifier fires 78x but kv_data_ptrs has 79 layers (the 79th is the
MTP/nextn EAGLE buffer: present in kv_data_ptrs so the monolithic send moves it,
but it doesn't fire the per-layer hook). The old completion required all num_layers,
so it both hung to the timeout AND would silently miss that layer's KV (corruption).

Redesign: gate completion on the ACTUAL enqueued count (note_enqueued), and in
finish() SYNCHRONOUSLY transfer any layers the notifier didn't fire for (KV is fully
written post-forward, no event needed). Net: the per-layer set == kv_data_ptrs,
byte-identical to the monolithic send; robust to any model firing fewer hooks than
KV buffers. The fired layers stay overlapped with the forward.

Unit tests updated (25 pass): fallback transfers the missed layers; submit failures
still report -1.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

python/sglang/srt/disaggregation/cp_per_layer_transfer.py

@@ -60,12 +60,18 @@ _FINISH_TIMEOUT_S = 30.0
 
 
 class PerLayerTransferContext:
-    """Tracks the per-layer async transfers for one prefill request.
+    """Per-request overlapped per-layer KV transfer.
 
-    num_layers = the number of per-layer notifier fires expected (this pp stage's local
+    During the forward, the notifier submits each FIRED layer's transfer after its
-    layer count). finish() blocks until all num_layers layers have been processed
+    write event (overlap). finish() (post-forward) waits those, then SYNCHRONOUSLY
-    (submitted or skipped) before waiting the batch_ids — otherwise it could race ahead
+    transfers any layers the notifier did NOT fire for — e.g. an MTP/nextn buffer that
-    of the worker threads and silently miss layers still in flight (corrupt KV)."""
+    is present in kv_data_ptrs (so the monolithic send moves it) but does not fire the
+    per-layer hook. By then the whole forward is done, so the missed layers are fully
+    written and need no event. Net: the set of layers moved == kv_data_ptrs, byte-
+    identical to the monolithic path. num_layers = len(kv_data_ptrs).
+
+    Completion is gated on the ACTUAL enqueued count (note_enqueued), not num_layers,
+    so it is robust to the model firing fewer hooks than there are KV buffers."""
 
     def __init__(self, engine, session_id: str, get_blocks: BlocksFn, num_layers: int):
         self.engine = engine
@@ -74,90 +80,102 @@ class PerLayerTransferContext:
         self.num_layers = num_layers
         self._batch_ids: List[int] = []
         self._submitted_layers: set[int] = set()
+        self._enqueued = 0
         self._processed = 0
         self._failed = False
-        self._lock = threading.Lock()
+        self._cond = threading.Condition()
-        self._all_processed = threading.Event()
-        if num_layers <= 0:
-            self._all_processed.set()
 
     @property
     def failed(self) -> bool:
-        with self._lock:
+        with self._cond:
             return self._failed
 
-    def _bump_processed_locked(self) -> None:
+    def note_enqueued(self) -> None:
-        self._processed += 1
+        """Called by the manager when it enqueues a (ctx, layer) item — the true count
-        if self._processed >= self.num_layers:
+        of notifier fires for this request, which finish() waits on."""
-            self._all_processed.set()
+        with self._cond:
+            self._enqueued += 1
 
-    def submit_layer(self, layer_id: int, event) -> None:
+    def _submit_one(self, layer_id: int) -> None:
-        """Wait the layer-L write event, then async-submit layer L's transfer.
+        """Submit layer_id's transfer (KV assumed written). Updates state; logs +
-        Called on a BACKGROUND worker thread (event.synchronize() must not run on
+        marks failed on any error."""
-        the compute/forward thread). Idempotent per layer; ALWAYS counts a unique
+        batch_id = None
-        layer toward completion (the finally), even when skipped/failed/raising, so
-        finish() can never hang on a per-layer error."""
-        with self._lock:
-            if layer_id in self._submitted_layers:
-                return  # duplicate fire — don't double-count
-            self._submitted_layers.add(layer_id)
-            skip = self._failed  # already failing: count it but don't submit
         try:
-            if skip:
-                return
-            if event is not None:
-                event.synchronize()  # wait the GPU write kernel for layer_id
             blocks = self.get_blocks(layer_id)
-            if not (blocks and blocks[0]):
+            if blocks and blocks[0]:
-                return
+                src_addrs, dst_addrs, lengths = blocks
-            src_addrs, dst_addrs, lengths = blocks
+                batch_id = self.engine.batch_transfer_async_submit(
-            batch_id = self.engine.batch_transfer_async_submit(
+                    self.session_id, list(src_addrs), list(dst_addrs), list(lengths)
-                self.session_id, list(src_addrs), list(dst_addrs), list(lengths)
+                )
-            )
-            with self._lock:
-                if batch_id < 0:
-                    self._failed = True
-                else:
-                    self._batch_ids.append(batch_id)
         except Exception as e:
-            with self._lock:
-                self._failed = True
             logger.warning(
                 "[CP_PER_LAYER_TRANSFER] submit_layer %d failed: %r", layer_id, e
             )
+            with self._cond:
+                self._failed = True
+            return
+        with self._cond:
+            if batch_id is not None and batch_id < 0:
+                self._failed = True
+            elif batch_id is not None:
+                self._batch_ids.append(batch_id)
+
+    def submit_layer(self, layer_id: int, event) -> None:
+        """Worker thread: wait the layer-L write event, then submit. Idempotent per
+        layer; always advances _processed (the finally) so finish() can't hang."""
+        with self._cond:
+            if layer_id in self._submitted_layers:
+                return  # duplicate fire
+            self._submitted_layers.add(layer_id)
+            skip = self._failed
+        try:
+            if not skip:
+                if event is not None:
+                    event.synchronize()  # wait the GPU write kernel for layer_id
+                self._submit_one(layer_id)
         finally:
-            with self._lock:
+            with self._cond:
-                self._bump_processed_locked()
+                self._processed += 1
+                self._cond.notify_all()
 
     def finish(self, timeout: float = _FINISH_TIMEOUT_S) -> int:
-        """Wait for ALL layers to be processed, then wait the submitted transfers.
+        """Wait the notifier-fired transfers, synchronously cover any un-fired layers,
-        Returns 0 on success; non-zero on failure (a submit failed, the wait reported
+        then wait all submitted. Returns 0 on success, -1 on any failure."""
-        failure, or layers never completed within the timeout)."""
+        with self._cond:
-        completed = self._all_processed.wait(timeout)
+            target = self._enqueued
-        with self._lock:
+            completed = self._cond.wait_for(
-            failed = self._failed
+                lambda: self._processed >= target, timeout
-            batch_ids = list(self._batch_ids)
-            processed = self._processed
-        if not completed:
-            logger.warning(
-                "[CP_PER_LAYER_TRANSFER] finish TIMEOUT processed=%d/%d submitted=%d failed=%s",
-                processed,
-                self.num_layers,
-                len(batch_ids),
-                failed,
             )
+            missing = [
+                L for L in range(self.num_layers) if L not in self._submitted_layers
+            ]
+            failed = self._failed
+        if missing:
+            logger.debug(
+                "[CP_PER_LAYER_TRANSFER] finish: %d notifier-missed layer(s) sent sync",
+                len(missing),
+            )
+            for layer_id in missing:
+                with self._cond:
+                    if layer_id in self._submitted_layers:
+                        continue
+                    self._submitted_layers.add(layer_id)
+                self._submit_one(layer_id)  # post-forward: written, no event needed
+        with self._cond:
+            failed = failed or self._failed
+            batch_ids = list(self._batch_ids)
         wait_status = self.engine.wait_batch_transfers(batch_ids)
         if failed or not completed:
             return -1
         return wait_status
 
     def mark_failed(self) -> None:
-        with self._lock:
+        with self._cond:
             self._failed = True
 
     @property
     def num_submitted(self) -> int:
-        with self._lock:
+        with self._cond:
             return len(self._batch_ids)
 
 
@@ -210,6 +228,7 @@ class PerLayerTransferManager:
             else:
                 event.record()
         for ctx in ctxs:
+            ctx.note_enqueued()
             self._q.put((ctx, layer_id, event))
 
     def has_room(self, room) -> bool:

test/registered/unit/disaggregation/test_cp_per_layer_transfer.py

@@ -102,24 +102,24 @@ class TestPerLayerTransferContext(unittest.TestCase):
         ctx.submit_layer(0, None)
         self.assertEqual(ctx.finish(), -9)
 
-    def test_finish_times_out_to_failure_when_layers_missing(self):
+    def test_finish_falls_back_for_notifier_missed_layers(self):
-        # only 1 of 3 layers ever processed -> finish must NOT silently succeed
+        # only layer 0 fired by the notifier; finish must SYNCHRONOUSLY transfer the
+        # missing layers 1,2 (e.g. an MTP buffer not hooked) -> all 3 moved, success.
         eng = _FakeEngine()
         ctx = _ctx(eng, num_layers=3)
-        ctx.submit_layer(0, None)
+        ctx.submit_layer(0, None)  # notifier fired layer 0 only
-        self.assertEqual(ctx.finish(timeout=0.05), -1)  # incomplete -> failure
+        self.assertEqual(ctx.finish(timeout=0.5), 0)
+        self.assertEqual(len(eng.submits), 3)  # layer 0 (overlap) + 1,2 (fallback)
+        self.assertEqual(eng.waits[-1], [101, 102, 103])
 
-    def test_skipped_and_failed_layers_still_count_toward_completion(self):
+    def test_failed_submit_still_reported_after_fallback(self):
-        # layer 1 has no blocks (skip), layer 2 submit fails; all 3 still complete
+        # layer 2 submit fails during the overlap; finish reports -1 even though it
-        eng = _FakeEngine(submit_rets=[101, -1])
+        # also runs the fallback for any unfired layers.
-
+        eng = _FakeEngine(submit_rets=[101, -1, 103])
-        def blocks(L):
+        ctx = _ctx(eng, num_layers=3)
-            return None if L == 1 else ([L], [L + 10], [8])
-
-        ctx = _ctx(eng, blocks, num_layers=3)
         for L in range(3):
             ctx.submit_layer(L, None)
-        self.assertEqual(ctx.finish(timeout=0.5), -1)  # completes (no hang), reports failure
+        self.assertEqual(ctx.finish(timeout=0.5), -1)
 
 
 class TestBuildLayerBlocks(unittest.TestCase):
@@ -151,9 +151,13 @@ class TestBuildLayerBlocks(unittest.TestCase):
 class _MockCtx:
     def __init__(self):
         self.submitted = []
+        self.enqueued = 0
         self.finished = False
         self.failed = False
 
+    def note_enqueued(self):
+        self.enqueued += 1
+
     def submit_layer(self, layer_id, event):
         self.submitted.append((layer_id, event))