From 246dbddac0cd6d2c877196ff2fc50b20ed00e294 Mon Sep 17 00:00:00 2001 From: leavelet Date: Sun, 7 Jun 2026 01:44:00 +0000 Subject: [PATCH] fix(disagg): per-layer completion can't hang + worker CUDA device + diagnostics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit E2E (clean run) showed finished ret=-1 at exactly the 30s timeout for every request: finish() hung because _worker_step swallowed submit_layer exceptions without counting the layer toward completion (so _processed never reached num_layers). Fixes: - submit_layer: try/finally that ALWAYS counts the layer (completion can never hang on a per-layer error) and LOGS the actual exception. - PerLayerTransferManager worker_init: torch.cuda.set_device on each worker thread (likely cause — event.synchronize() needs the device set on these fresh threads, unlike the transfer_worker where A1's engine call worked). - finish() logs processed/num_layers on timeout to separate exception-failure from notifier-undercount. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../disaggregation/cp_per_layer_transfer.py | 60 ++++++++++++++----- python/sglang/srt/disaggregation/prefill.py | 2 + 2 files changed, 47 insertions(+), 15 deletions(-) diff --git a/python/sglang/srt/disaggregation/cp_per_layer_transfer.py b/python/sglang/srt/disaggregation/cp_per_layer_transfer.py index b4315fb05..bf8b71ef0 100644 --- a/python/sglang/srt/disaggregation/cp_per_layer_transfer.py +++ b/python/sglang/srt/disaggregation/cp_per_layer_transfer.py @@ -17,9 +17,12 @@ prefill/conn integration. Engine + events are injected so this is unit-testable. """ from __future__ import annotations +import logging import threading from typing import Callable, List, Optional, Tuple +logger = logging.getLogger(__name__) + # get_blocks(layer_id) -> (src_addrs, dst_addrs, lengths) for THIS rank's owned pages # of layer_id, or None to skip (no owned pages / dummy). BlocksFn = Callable[[int], Optional[Tuple[List[int], List[int], List[int]]]] @@ -91,30 +94,40 @@ class PerLayerTransferContext: 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; counts every unique layer - toward completion even when skipped/failed so finish() can't hang.""" + 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) - if self._failed: - self._bump_processed_locked() # already failing: count + skip submit + skip = self._failed # already failing: count it but don't submit + 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 - if event is not None: - event.synchronize() # wait the GPU write kernel for layer_id to finish - blocks = self.get_blocks(layer_id) - batch_id = None - 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) ) - with self._lock: - if batch_id is not None and batch_id < 0: + with self._lock: + if batch_id < 0: + self._failed = True + else: + self._batch_ids.append(batch_id) + except Exception as e: + with self._lock: self._failed = True - elif batch_id is not None: - self._batch_ids.append(batch_id) - self._bump_processed_locked() + logger.warning( + "[CP_PER_LAYER_TRANSFER] submit_layer %d failed: %r", layer_id, e + ) + finally: + with self._lock: + self._bump_processed_locked() def finish(self, timeout: float = _FINISH_TIMEOUT_S) -> int: """Wait for ALL layers to be processed, then wait the submitted transfers. @@ -124,6 +137,15 @@ class PerLayerTransferContext: 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_status = self.engine.wait_batch_transfers(batch_ids) if failed or not completed: return -1 @@ -152,11 +174,14 @@ class PerLayerTransferManager: 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): + def __init__( + self, num_workers=4, event_factory=None, current_stream=None, worker_init=None + ): import queue as _queue self._event_factory = event_factory self._current_stream = current_stream + self._worker_init = worker_init # called once per worker thread (e.g. set CUDA device) self._q = _queue.SimpleQueue() self._active = {} self._active_lock = threading.Lock() @@ -212,5 +237,10 @@ class PerLayerTransferManager: ctx.mark_failed() def _worker(self) -> None: + if self._worker_init is not None: + try: + self._worker_init() # e.g. torch.cuda.set_device for event.synchronize + except Exception as e: + logger.warning("[CP_PER_LAYER_TRANSFER] worker_init failed: %r", e) while True: self._worker_step(self._q.get()) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index b27c4ec71..3d541d522 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -289,9 +289,11 @@ class PrefillBootstrapQueue: PerLayerTransferManager, ) + device = torch.cuda.current_device() manager = PerLayerTransferManager( event_factory=torch.cuda.Event, current_stream=torch.cuda.current_stream, + worker_init=lambda: torch.cuda.set_device(device), ) pool = self.token_to_kv_pool if hasattr(pool, "register_layer_backup_notifier"):