From 7ac8067db157b75859a8f3e279ba48be5e171596 Mon Sep 17 00:00:00 2001 From: leavelet Date: Sun, 7 Jun 2026 00:35:19 +0000 Subject: [PATCH] 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) --- .../disaggregation/cp_per_layer_transfer.py | 69 ++++++++++++ .../test_cp_per_layer_transfer.py | 102 ++++++++++++++++++ 2 files changed, 171 insertions(+) diff --git a/python/sglang/srt/disaggregation/cp_per_layer_transfer.py b/python/sglang/srt/disaggregation/cp_per_layer_transfer.py index 54abe49e8..301a1a901 100644 --- a/python/sglang/srt/disaggregation/cp_per_layer_transfer.py +++ b/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()) 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 afe5e97c1..1802ea54f 100644 --- a/test/registered/unit/disaggregation/test_cp_per_layer_transfer.py +++ b/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()