diff --git a/python/sglang/srt/disaggregation/cp_per_layer_transfer.py b/python/sglang/srt/disaggregation/cp_per_layer_transfer.py index bf8b71ef0..d6a1a318c 100644 --- a/python/sglang/srt/disaggregation/cp_per_layer_transfer.py +++ b/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 - layer count). finish() blocks until all num_layers layers have been processed - (submitted or skipped) before waiting the batch_ids — otherwise it could race ahead - of the worker threads and silently miss layers still in flight (corrupt KV).""" + During the forward, the notifier submits each FIRED layer's transfer after its + write event (overlap). finish() (post-forward) waits those, then SYNCHRONOUSLY + transfers any layers the notifier did NOT fire for — e.g. an MTP/nextn buffer that + 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._all_processed = threading.Event() - if num_layers <= 0: - self._all_processed.set() + self._cond = threading.Condition() @property def failed(self) -> bool: - with self._lock: + with self._cond: return self._failed - def _bump_processed_locked(self) -> None: - self._processed += 1 - if self._processed >= self.num_layers: - self._all_processed.set() + def note_enqueued(self) -> None: + """Called by the manager when it enqueues a (ctx, layer) item — the true count + of notifier fires for this request, which finish() waits on.""" + with self._cond: + self._enqueued += 1 - def submit_layer(self, layer_id: int, event) -> None: - """Wait the layer-L write event, then async-submit layer L's transfer. - Called on a BACKGROUND worker thread (event.synchronize() must not run on - the compute/forward thread). Idempotent per layer; ALWAYS counts a unique - 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 + def _submit_one(self, layer_id: int) -> None: + """Submit layer_id's transfer (KV assumed written). Updates state; logs + + marks failed on any error.""" + batch_id = None 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]): - return - src_addrs, dst_addrs, lengths = blocks - batch_id = self.engine.batch_transfer_async_submit( - 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) + if blocks and blocks[0]: + src_addrs, dst_addrs, lengths = blocks + batch_id = self.engine.batch_transfer_async_submit( + self.session_id, list(src_addrs), list(dst_addrs), list(lengths) + ) 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: - self._bump_processed_locked() + with self._cond: + 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. - Returns 0 on success; non-zero on failure (a submit failed, the wait reported - failure, or layers never completed within the timeout).""" - completed = self._all_processed.wait(timeout) - with self._lock: - failed = self._failed - 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, + """Wait the notifier-fired transfers, synchronously cover any un-fired layers, + then wait all submitted. Returns 0 on success, -1 on any failure.""" + with self._cond: + target = self._enqueued + completed = self._cond.wait_for( + lambda: self._processed >= target, timeout ) + 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: 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 5f02cc148..73e6a503b 100644 --- a/test/registered/unit/disaggregation/test_cp_per_layer_transfer.py +++ b/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): - # only 1 of 3 layers ever processed -> finish must NOT silently succeed + def test_finish_falls_back_for_notifier_missed_layers(self): + # 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) - self.assertEqual(ctx.finish(timeout=0.05), -1) # incomplete -> failure + ctx.submit_layer(0, None) # notifier fired layer 0 only + 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): - # layer 1 has no blocks (skip), layer 2 submit fails; all 3 still complete - eng = _FakeEngine(submit_rets=[101, -1]) - - def blocks(L): - return None if L == 1 else ([L], [L + 10], [8]) - - ctx = _ctx(eng, blocks, num_layers=3) + def test_failed_submit_still_reported_after_fallback(self): + # layer 2 submit fails during the overlap; finish reports -1 even though it + # also runs the fallback for any unfired layers. + eng = _FakeEngine(submit_rets=[101, -1, 103]) + ctx = _ctx(eng, 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))