perf(disagg): coarse-grained per-layer transfer + skip re-registration (lever A)

For large bs (production target: bs~10 x ~100k tokens), per-layer granularity does
10x79 = 790 submitTransfer calls + CUDA events + enqueues per forward on the forward
thread. Two overhead cuts:
- Group SGLANG_CP_SHARED_KV_PER_LAYER_GROUP (default 8) consecutive layers into ONE
  RDMA submit: ~num_layers/K submits + events + enqueues instead of per-layer; same
  bytes (page index lists are identical across layers). on_layer_end is O(1) at
  non-boundary layers. The last partial group enqueues via the num_layers boundary;
  any misses fall back to one batched sync submit.
- Scheduler hook skips reqs already registered (bs>1 batch-forming re-iterates the
  same reqs ~9x -> was rebuilding the CP filter + context every time).

27 unit tests pass incl. grouping-boundary + batched-fallback.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

python/sglang/srt/disaggregation/cp_per_layer_transfer.py

@@ -106,43 +106,52 @@ class PerLayerTransferContext:
             self._enqueued += 1
             return True
 
-    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
+    def _submit_layers(self, layer_ids) -> None:
+        """Submit the COMBINED blocks of several (consecutive) layers as ONE RDMA
+        batch — coarser granularity means far fewer submitTransfer calls + batch_ids,
+        the key CPU/overhead reduction at large batch sizes. The page index lists are
+        identical across layers (only the per-layer base ptr differs), so this is the
+        same bytes as per-layer, just batched. Updates state; logs + fails closed."""
+        src_all: List[int] = []
+        dst_all: List[int] = []
+        len_all: List[int] = []
         try:
-            blocks = self.get_blocks(layer_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)
-                )
+            for layer_id in layer_ids:
+                blocks = self.get_blocks(layer_id)
+                if blocks and blocks[0]:
+                    src_all.extend(blocks[0])
+                    dst_all.extend(blocks[1])
+                    len_all.extend(blocks[2])
         except Exception as e:
-            logger.warning(
-                "[CP_PER_LAYER_TRANSFER] submit_layer %d failed: %r", layer_id, e
-            )
+            logger.warning("[CP_PER_LAYER_TRANSFER] get_blocks failed: %r", e)
             with self._cond:
                 self._failed = True
             return
+        if not src_all:
+            return
+        batch_id = self.engine.batch_transfer_async_submit(
+            self.session_id, src_all, dst_all, len_all
+        )
         with self._cond:
-            if batch_id is not None and batch_id < 0:
+            if batch_id < 0:
                 self._failed = True
-            elif batch_id is not None:
+            else:
                 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."""
+    def submit_group(self, start: int, end: int, event) -> None:
+        """Worker thread: wait the group's last-layer write event, then submit layers
+        [start..end] as one RDMA batch. Idempotent; always advances _processed (the
+        finally) so finish() can't hang. (Per-layer is the special case start==end.)"""
         with self._cond:
-            if layer_id in self._submitted_layers:
-                return  # duplicate fire
-            self._submitted_layers.add(layer_id)
+            new = [L for L in range(start, end + 1) if L not in self._submitted_layers]
+            for L in new:
+                self._submitted_layers.add(L)
             skip = self._failed
         try:
-            if not skip:
+            if not skip and new:
                 if event is not None:
-                    event.synchronize()  # wait the GPU write kernel for layer_id
-                self._submit_one(layer_id)
+                    event.synchronize()  # last layer of the group -> all written
+                self._submit_layers(new)
         finally:
             with self._cond:
                 self._processed += 1
@@ -164,12 +173,13 @@ class PerLayerTransferContext:
             submitted_at_entry = len(self._batch_ids)
         t1 = time.perf_counter()  # workers drained
         if 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:
+                still = [L for L in missing if L not in self._submitted_layers]
+                for L in still:
+                    self._submitted_layers.add(L)
+            if still:
+                # post-forward: all written, no event needed; one batch for the tail
+                self._submit_layers(still)
         t2 = time.perf_counter()  # fallback submitted
         with self._cond:
             failed = failed or self._failed
@@ -192,6 +202,10 @@ class PerLayerTransferContext:
             return -1
         return wait_status
 
+    def submit_layer(self, layer_id: int, event) -> None:
+        """Per-layer convenience (group of one). The manager uses submit_group."""
+        self.submit_group(layer_id, layer_id, event)
+
     def mark_failed(self) -> None:
         with self._cond:
             self._failed = True
@@ -216,13 +230,20 @@ class PerLayerTransferManager:
     """
 
     def __init__(
-        self, num_workers=4, event_factory=None, current_stream=None, worker_init=None
+        self,
+        num_workers=4,
+        event_factory=None,
+        current_stream=None,
+        worker_init=None,
+        group_size=8,
     ):
         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._group_size = max(1, int(group_size))  # layers per RDMA submit (overhead)
+        self._num_layers = 0  # learned from the first registered ctx (model layer count)
         self._q = _queue.SimpleQueue()
         self._active = {}
         self._active_lock = threading.Lock()
@@ -234,12 +255,22 @@ class PerLayerTransferManager:
             if room in self._active:
                 return  # already registered for its forward; don't overwrite (leak)
             self._active[room] = ctx
+            if self._num_layers == 0:
+                self._num_layers = ctx.num_layers
 
     def has_active(self) -> bool:
         with self._active_lock:
             return bool(self._active)
 
     def on_layer_end(self, layer_id: int) -> None:
+        # Group K consecutive layers into ONE submit. Only do work at a group boundary
+        # (every K layers, or the last layer); other layers are O(1). This is the CPU/
+        # overhead reduction at large bs: ~num_layers/K submits + events + enqueues
+        # instead of one per layer.
+        K = self._group_size
+        is_last = self._num_layers > 0 and layer_id == self._num_layers - 1
+        if (layer_id + 1) % K != 0 and not is_last:
+            return  # not a group boundary
         with self._active_lock:
             if not self._active:
                 return
@@ -252,9 +283,10 @@ class PerLayerTransferManager:
                 event.record(stream)
             else:
                 event.record()
+        start = (layer_id // K) * K  # first layer of this group
         for ctx in ctxs:
-            if ctx.note_enqueued(layer_id):  # enqueue each layer ONCE per ctx
-                self._q.put((ctx, layer_id, event))
+            if ctx.note_enqueued(layer_id):  # one group per (ctx, boundary) — dedup
+                self._q.put((ctx, start, layer_id, event))
 
     def has_room(self, room) -> bool:
         with self._active_lock:
@@ -274,9 +306,9 @@ class PerLayerTransferManager:
             ctx.finish()
 
     def _worker_step(self, item) -> None:
-        ctx, layer_id, event = item
+        ctx, start, end, event = item
         try:
-            ctx.submit_layer(layer_id, event)
+            ctx.submit_group(start, end, event)
         except Exception:
             ctx.mark_failed()
 

python/sglang/srt/disaggregation/prefill.py

@@ -294,6 +294,7 @@ class PrefillBootstrapQueue:
             event_factory=torch.cuda.Event,
             current_stream=torch.cuda.current_stream,
             worker_init=lambda: torch.cuda.set_device(device),
+            group_size=envs.SGLANG_CP_SHARED_KV_PER_LAYER_GROUP.get(),
         )
         pool = self.token_to_kv_pool
         if hasattr(pool, "register_layer_backup_notifier"):
@@ -657,13 +658,18 @@ class SchedulerDisaggregationPrefillMixin:
         prefix requests (whose prefix is HiCache-loaded, not forward-written) are left
         on the monolithic post-forward path. No-op unless the flag is on."""
         kv_manager = self.disagg_prefill_bootstrap_queue.kv_manager
-        if getattr(kv_manager, "per_layer_transfer_manager", None) is None:
+        mgr = getattr(kv_manager, "per_layer_transfer_manager", None)
+        if mgr is None:
             return
         page_size = self.token_to_kv_pool_allocator.page_size
         for req in batch.reqs:
             try:
                 if getattr(req, "disagg_kv_sender", None) is None:
                     continue
+                room = getattr(req, "bootstrap_room", None)
+                if room is None or mgr.has_room(room):
+                    continue  # already registered (bs>1 batch-forming re-iterates the
+                    # same reqs many times) — skip the CP filter + build entirely
                 if getattr(req, "is_chunked", 0) != 0 or getattr(req, "start_send_idx", 0) != 0:
                     continue  # chunked / partially-sent: this forward isn't the full range
                 # Cover the FULL sequence (cached prefix + new tokens). The per-layer
@@ -677,7 +683,7 @@ class SchedulerDisaggregationPrefillMixin:
                     self.req_to_token_pool.req_to_token[req.req_pool_idx, 0:end_idx],
                     page_size,
                 )
-                kv_manager.register_per_layer_transfer(req.bootstrap_room, page_indices)
+                kv_manager.register_per_layer_transfer(room, page_indices)
             except Exception as e:
                 # Never let lever-A setup crash the scheduler; fall back to the
                 # monolithic post-forward transfer for this request.

python/sglang/srt/environ.py

@@ -215,6 +215,10 @@ class Envs:
     # asynchronously (pipelined batch_transfer_async) instead of one monolithic
     # all-layers batch_transfer_sync. Default off until validated.
     SGLANG_CP_SHARED_KV_PER_LAYER_TRANSFER = EnvBool(False)
+    # Lever A granularity: group this many consecutive layers into one RDMA submit.
+    # Larger = fewer submitTransfer calls / events / enqueues (less CPU overhead at
+    # large bs) but coarser overlap. 1 = per-layer.
+    SGLANG_CP_SHARED_KV_PER_LAYER_GROUP = EnvInt(8)
     SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE = EnvBool(False)
     SGLANG_CP_SHARED_KV_FUSED_MLA_STORE = EnvBool(False)
     SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE = EnvBool(False)