diff --git a/python/sglang/srt/disaggregation/cp_per_layer_transfer.py b/python/sglang/srt/disaggregation/cp_per_layer_transfer.py index ec3c87e84..2c0155903 100644 --- a/python/sglang/srt/disaggregation/cp_per_layer_transfer.py +++ b/python/sglang/srt/disaggregation/cp_per_layer_transfer.py @@ -81,6 +81,7 @@ class PerLayerTransferContext: self.num_layers = num_layers self._batch_ids: List[int] = [] self._submitted_layers: set[int] = set() + self._enqueued_layers: set[int] = set() self._enqueued = 0 self._processed = 0 self._failed = False @@ -91,11 +92,19 @@ class PerLayerTransferContext: with self._cond: return self._failed - 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.""" + def note_enqueued(self, layer_id: int) -> bool: + """Reserve layer_id for enqueue. Returns True only the FIRST time this layer is + seen, so the manager enqueues + counts each layer ONCE. Critical: the notifier + re-fires every layer on every subsequent forward while this ctx is still active + (awaiting finish); without dedup, _enqueued would grow unbounded (target>>layers) + and finish() could never reach it -> 30s timeouts. The first fire for a layer is + from THIS request's own forward, so its event is the correct one.""" with self._cond: + if layer_id in self._enqueued_layers: + return False + self._enqueued_layers.add(layer_id) self._enqueued += 1 + return True def _submit_one(self, layer_id: int) -> None: """Submit layer_id's transfer (KV assumed written). Updates state; logs + @@ -222,6 +231,8 @@ class PerLayerTransferManager: def register(self, room, ctx: PerLayerTransferContext) -> None: with self._active_lock: + if room in self._active: + return # already registered for its forward; don't overwrite (leak) self._active[room] = ctx def has_active(self) -> bool: @@ -242,8 +253,8 @@ class PerLayerTransferManager: else: event.record() for ctx in ctxs: - ctx.note_enqueued() - self._q.put((ctx, layer_id, event)) + if ctx.note_enqueued(layer_id): # enqueue each layer ONCE per ctx + self._q.put((ctx, layer_id, event)) def has_room(self, room) -> bool: with self._active_lock: 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 73e6a503b..d9e2eb31a 100644 --- a/test/registered/unit/disaggregation/test_cp_per_layer_transfer.py +++ b/test/registered/unit/disaggregation/test_cp_per_layer_transfer.py @@ -102,6 +102,20 @@ class TestPerLayerTransferContext(unittest.TestCase): ctx.submit_layer(0, None) self.assertEqual(ctx.finish(), -9) + def test_repeated_enqueues_dedup_so_finish_cannot_hang(self): + # the notifier re-fires every layer on every subsequent forward while the ctx + # is still active; note_enqueued must dedup so target stays capped at the layer + # count and finish() completes (the high-cache-hit hang was target>>layers). + eng = _FakeEngine() + ctx = _ctx(eng, num_layers=3) + for _forward in range(5): # 5 forwards all re-fire layers 0,1,2 + for L in range(3): + self.assertEqual(ctx.note_enqueued(L), _forward == 0) # True only 1st + if _forward == 0: + ctx.submit_layer(L, None) + self.assertEqual(len(eng.submits), 3) # 3 unique despite 15 note_enqueued + self.assertEqual(ctx.finish(timeout=0.5), 0) # completes, no 30s timeout + 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. @@ -155,8 +169,12 @@ class _MockCtx: self.finished = False self.failed = False - def note_enqueued(self): + def note_enqueued(self, layer_id): + if layer_id in getattr(self, "_enq_layers", set()): + return False + self._enq_layers = getattr(self, "_enq_layers", set()) | {layer_id} self.enqueued += 1 + return True def submit_layer(self, layer_id, event): self.submitted.append((layer_id, event))