diff --git a/python/sglang/srt/disaggregation/cp_per_layer_transfer.py b/python/sglang/srt/disaggregation/cp_per_layer_transfer.py index 2c0155903..8053647b9 100644 --- a/python/sglang/srt/disaggregation/cp_per_layer_transfer.py +++ b/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() diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 4de1bfdb7..84bc45f87 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/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. diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 79b024215..17c837ec5 100644 --- a/python/sglang/srt/environ.py +++ b/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) 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 d9e2eb31a..ee8e556de 100644 --- a/test/registered/unit/disaggregation/test_cp_per_layer_transfer.py +++ b/test/registered/unit/disaggregation/test_cp_per_layer_transfer.py @@ -118,13 +118,15 @@ class TestPerLayerTransferContext(unittest.TestCase): 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. + # missing layers 1,2 -> all 3 moved, success. The fallback BATCHES the missing + # layers into one submit. eng = _FakeEngine() ctx = _ctx(eng, num_layers=3) - ctx.submit_layer(0, None) # notifier fired layer 0 only + ctx.submit_layer(0, None) # notifier fired layer 0 only -> 1 submit 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]) + self.assertEqual(len(eng.submits), 2) # layer 0 + fallback batch [1,2] + self.assertEqual(eng.waits[-1], [101, 102]) + self.assertEqual(eng.submits[1], ("sess", [1000, 2000], [6000, 7000], [64, 64])) def test_failed_submit_still_reported_after_fallback(self): # layer 2 submit fails during the overlap; finish reports -1 even though it @@ -163,21 +165,24 @@ class TestBuildLayerBlocks(unittest.TestCase): class _MockCtx: + num_layers = 8 + def __init__(self): - self.submitted = [] + self.submitted = [] # (start, end, event) self.enqueued = 0 self.finished = False self.failed = False + self._enq = set() - def note_enqueued(self, layer_id): - if layer_id in getattr(self, "_enq_layers", set()): + def note_enqueued(self, key): + if key in self._enq: return False - self._enq_layers = getattr(self, "_enq_layers", set()) | {layer_id} + self._enq.add(key) self.enqueued += 1 return True - def submit_layer(self, layer_id, event): - self.submitted.append((layer_id, event)) + def submit_group(self, start, end, event): + self.submitted.append((start, end, event)) def finish(self): self.finished = True @@ -199,11 +204,14 @@ class _RecEvent: self.synced = True -def _manager(num_workers=0): +def _manager(num_workers=0, group_size=1): from sglang.srt.disaggregation.cp_per_layer_transfer import PerLayerTransferManager return PerLayerTransferManager( - num_workers=num_workers, event_factory=_RecEvent, current_stream=lambda: "STREAM" + num_workers=num_workers, + event_factory=_RecEvent, + current_stream=lambda: "STREAM", + group_size=group_size, ) @@ -221,23 +229,33 @@ def _drain(q): class TestPerLayerTransferManager(unittest.TestCase): def test_on_layer_end_enqueues_per_active_ctx_with_recorded_event(self): - m = _manager() + m = _manager(group_size=1) # per-layer: every layer is a group boundary c1, c2 = _MockCtx(), _MockCtx() m.register("r1", c1) m.register("r2", c2) m.on_layer_end(5) items = _drain(m._q) self.assertEqual(len(items), 2) # one per active context - for ctx, layer_id, ev in items: - self.assertEqual(layer_id, 5) - self.assertEqual(ev.recorded_stream, "STREAM") # event recorded on compute stream + for ctx, start, end, ev in items: + self.assertEqual((start, end), (5, 5)) + self.assertEqual(ev.recorded_stream, "STREAM") - def test_worker_step_calls_submit_layer(self): + def test_grouping_only_enqueues_at_boundaries(self): + m = _manager(group_size=4) + c = _MockCtx() + m.register("r1", c) + for L in range(8): + m.on_layer_end(L) + items = _drain(m._q) + # boundaries at L=3 (group [0,3]) and L=7 (group [4,7]) -> 2 groups, not 8 + self.assertEqual([(s, e) for _, s, e, _ in items], [(0, 3), (4, 7)]) + + def test_worker_step_calls_submit_group(self): m = _manager() c = _MockCtx() ev = _RecEvent() - m._worker_step((c, 3, ev)) - self.assertEqual(c.submitted, [(3, ev)]) + m._worker_step((c, 0, 3, ev)) + self.assertEqual(c.submitted, [(0, 3, ev)]) def test_worker_step_marks_failed_on_exception(self): m = _manager() @@ -246,14 +264,14 @@ class TestPerLayerTransferManager(unittest.TestCase): def __init__(self): self.failed = False - def submit_layer(self, layer_id, event): + def submit_group(self, start, end, event): raise RuntimeError("boom") def mark_failed(self): self.failed = True b = _Boom() - m._worker_step((b, 0, None)) + m._worker_step((b, 0, 0, None)) self.assertTrue(b.failed) def test_finish_pops_and_calls_ctx_finish(self): @@ -267,7 +285,7 @@ class TestPerLayerTransferManager(unittest.TestCase): self.assertEqual(m.finish("r1"), 0) # idempotent after pop def test_on_layer_end_no_active_is_noop(self): - m = _manager() + m = _manager(group_size=1) m.on_layer_end(0) self.assertEqual(_drain(m._q), [])