feat(disagg): per-layer transfer manager (lever A, A2 orchestration)

Add PerLayerTransferManager: owns a worker thread-pool + the active per-request
contexts for the current forward batch, registered as a KV-pool notifier.
on_layer_end (forward thread) records ONE CUDA event on the compute stream and
enqueues (ctx, layer, event) per active context; workers do the event-wait +
async submit OFF the forward thread. finish(room) waits the request's transfers.
event_factory/current_stream injected for unit-testability (no CUDA needed).

Unit-tested (5 manager cases, 11 total in test_cp_per_layer_transfer.py):
per-active-ctx enqueue with the event recorded on the stream, worker-step submit +
mark-failed-on-exception, finish pop + idempotency, no-op when idle. The A3
scheduler wiring (notifier registration + setup-before-forward + finish-after +
no-double-send reconcile) is the remaining hot-path step; plan locked in
lever-a-implementation-plan.md.

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

python/sglang/srt/disaggregation/cp_per_layer_transfer.py

@@ -76,7 +76,76 @@ class PerLayerTransferContext:
         wait_status = self.engine.wait_batch_transfers(batch_ids)
         return -1 if failed else wait_status
 
+    def mark_failed(self) -> None:
+        with self._lock:
+            self._failed = True
+
     @property
     def num_submitted(self) -> int:
         with self._lock:
             return len(self._batch_ids)
+
+
+class PerLayerTransferManager:
+    """Owns the per-layer transfer worker pool + the active per-request contexts for
+    the current forward batch, and is registered as a notifier on the KV pool.
+
+    on_layer_end(layer_id) runs on the FORWARD thread and is cheap: it records ONE
+    CUDA event on the compute stream (all requests in the batch are written by the
+    same forward) and enqueues (ctx, layer_id, event) for each active context. The
+    worker threads do the event-wait + RDMA submit OFF the forward thread.
+
+    event_factory/current_stream are injected so the class is unit-testable without
+    CUDA; num_workers=0 starts no threads (drain manually via _worker_step in tests).
+    """
+
+    def __init__(self, num_workers=4, event_factory=None, current_stream=None):
+        import queue as _queue
+
+        self._event_factory = event_factory
+        self._current_stream = current_stream
+        self._q = _queue.SimpleQueue()
+        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:
+        with self._active_lock:
+            self._active[room] = ctx
+
+    def has_active(self) -> bool:
+        with self._active_lock:
+            return bool(self._active)
+
+    def on_layer_end(self, layer_id: int) -> None:
+        with self._active_lock:
+            if not self._active:
+                return
+            ctxs = list(self._active.values())
+        event = None
+        if self._event_factory is not None:
+            event = self._event_factory()
+            stream = self._current_stream() if self._current_stream is not None else None
+            if stream is not None:
+                event.record(stream)
+            else:
+                event.record()
+        for ctx in ctxs:
+            self._q.put((ctx, layer_id, event))
+
+    def finish(self, room) -> int:
+        with self._active_lock:
+            ctx = self._active.pop(room, None)
+        return ctx.finish() if ctx is not None else 0
+
+    def _worker_step(self, item) -> None:
+        ctx, layer_id, event = item
+        try:
+            ctx.submit_layer(layer_id, event)
+        except Exception:
+            ctx.mark_failed()
+
+    def _worker(self) -> None:
+        while True:
+            self._worker_step(self._q.get())

test/registered/unit/disaggregation/test_cp_per_layer_transfer.py

@@ -103,5 +103,107 @@ class TestPerLayerTransferContext(unittest.TestCase):
         self.assertEqual(ctx.finish(), -9)
 
 
+class _MockCtx:
+    def __init__(self):
+        self.submitted = []
+        self.finished = False
+        self.failed = False
+
+    def submit_layer(self, layer_id, event):
+        self.submitted.append((layer_id, 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):
+    from sglang.srt.disaggregation.cp_per_layer_transfer import PerLayerTransferManager
+
+    return PerLayerTransferManager(
+        num_workers=num_workers, event_factory=_RecEvent, current_stream=lambda: "STREAM"
+    )
+
+
+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()
+        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, layer_id, ev in items:
+            self.assertEqual(layer_id, 5)
+            self.assertEqual(ev.recorded_stream, "STREAM")  # event recorded on compute stream
+
+    def test_worker_step_calls_submit_layer(self):
+        m = _manager()
+        c = _MockCtx()
+        ev = _RecEvent()
+        m._worker_step((c, 3, ev))
+        self.assertEqual(c.submitted, [(3, ev)])
+
+    def test_worker_step_marks_failed_on_exception(self):
+        m = _manager()
+
+        class _Boom:
+            def __init__(self):
+                self.failed = False
+
+            def submit_layer(self, layer_id, event):
+                raise RuntimeError("boom")
+
+            def mark_failed(self):
+                self.failed = True
+
+        b = _Boom()
+        m._worker_step((b, 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()
+        m.on_layer_end(0)
+        self.assertEqual(_drain(m._q), [])
+
+
 if __name__ == "__main__":
     unittest.main()