chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
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>
This commit is contained in:
孙远航
2026-06-07 02:02:46 +00:00
parent 246dbddac0
commit 64f767bb42
2 changed files with 97 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files

141
python/sglang/srt/disaggregation/cp_per_layer_transfer.py
View File

@@ -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:

30
test/registered/unit/disaggregation/test_cp_per_layer_transfer.py
View File

@@ -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))