fix(disagg): per-layer completion can't hang + worker CUDA device + diagnostics

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) <noreply@anthropic.com>

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

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"):