From ae18e3adc8100d5a8af68532f34059294e30dce9 Mon Sep 17 00:00:00 2001 From: leavelet Date: Sun, 7 Jun 2026 01:01:40 +0000 Subject: [PATCH] feat(disagg): wire per-layer overlapped transfer end-to-end (lever A, A3-step3) Complete the lever-A hot-path integration behind SGLANG_CP_SHARED_KV_PER_LAYER_TRANSFER: - PerLayerTransferContext: add num_layers + a completion event so finish() waits until ALL layers are processed before wait_batch_transfers (never races ahead of the worker threads and silently drops in-flight layers); times out to FAILURE. - PerLayerTransferManager: has_room (for the swap) + drop (abort/failure drain so outstanding RDMA finishes before pages are reclaimed). - MooncakeKVManager.register_per_layer_transfer: build + register a context before the forward, reusing send()'s exact CP filter (no re-derivation). - transfer_worker: when a room is per-layer-active, wait those transfers (finish) instead of the monolithic send_kvcache -- no double-send; aux/state/completion unchanged. The skip path drops the context on abort/failure. - prefill scheduler: _register_per_layer_transfers(batch) before run_batch, scoped (first impl) to single-forward, no-cached-prefix requests (the notifier transfers forward-written pages; chunked/cached-prefix are HiCache-loaded -> lever B). Unit-tested (25 cases). e2e output-equality + TTFT verification next. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../disaggregation/cp_per_layer_transfer.py | 80 ++++++++++++++----- .../srt/disaggregation/mooncake/conn.py | 69 +++++++++++++++- python/sglang/srt/disaggregation/prefill.py | 27 +++++++ .../test_cp_per_layer_transfer.py | 31 +++++-- 4 files changed, 179 insertions(+), 28 deletions(-) diff --git a/python/sglang/srt/disaggregation/cp_per_layer_transfer.py b/python/sglang/srt/disaggregation/cp_per_layer_transfer.py index 4766ca1bc..b4315fb05 100644 --- a/python/sglang/srt/disaggregation/cp_per_layer_transfer.py +++ b/python/sglang/srt/disaggregation/cp_per_layer_transfer.py @@ -50,56 +50,84 @@ def build_layer_blocks( return src_addrs, dst_addrs, lengths -class PerLayerTransferContext: - """Tracks the per-layer async transfers for one prefill request.""" +# finish() default cap waiting for the forward's per-layer events to all be processed. +# Post-forward all events are already enqueued, so this is normally hit in <1s; the cap +# is only a safety net against a lost notifier (treated as failure, never silent loss). +_FINISH_TIMEOUT_S = 30.0 - def __init__(self, engine, session_id: str, get_blocks: BlocksFn): + +class PerLayerTransferContext: + """Tracks the per-layer async transfers for one prefill request. + + 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).""" + + def __init__(self, engine, session_id: str, get_blocks: BlocksFn, num_layers: int): self.engine = engine self.session_id = session_id self.get_blocks = get_blocks + self.num_layers = num_layers self._batch_ids: List[int] = [] self._submitted_layers: set[int] = set() + self._processed = 0 self._failed = False self._lock = threading.Lock() + self._all_processed = threading.Event() + if num_layers <= 0: + self._all_processed.set() @property def failed(self) -> bool: with self._lock: return self._failed + def _bump_processed_locked(self) -> None: + self._processed += 1 + if self._processed >= self.num_layers: + self._all_processed.set() + 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; no-op after failure.""" + the compute/forward thread). Idempotent per layer; counts every unique layer + toward completion even when skipped/failed so finish() can't hang.""" with self._lock: - if self._failed or layer_id in self._submitted_layers: - return + 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 + return if event is not None: event.synchronize() # wait the GPU write kernel for layer_id to finish blocks = self.get_blocks(layer_id) - if not blocks: - return - src_addrs, dst_addrs, lengths = blocks - if not src_addrs: - return - batch_id = self.engine.batch_transfer_async_submit( - self.session_id, list(src_addrs), list(dst_addrs), list(lengths) - ) + 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 < 0: + if batch_id is not None and batch_id < 0: self._failed = True - else: + elif batch_id is not None: self._batch_ids.append(batch_id) + self._bump_processed_locked() - def finish(self) -> int: - """Wait for all submitted layer transfers. Returns 0 on success, non-zero on - failure (a submit failed, or wait reported failure). Drains regardless.""" + 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) wait_status = self.engine.wait_batch_transfers(batch_ids) - return -1 if failed else wait_status + if failed or not completed: + return -1 + return wait_status def mark_failed(self) -> None: with self._lock: @@ -159,11 +187,23 @@ class PerLayerTransferManager: for ctx in ctxs: self._q.put((ctx, layer_id, event)) + def has_room(self, room) -> bool: + with self._active_lock: + return room in self._active + def finish(self, room) -> int: with self._active_lock: ctx = self._active.pop(room, None) return ctx.finish() if ctx is not None else 0 + def drop(self, room) -> None: + """Abort path: drain + discard a room's context so any outstanding RDMA + completes before its KV pages are freed (never write into reclaimed pages).""" + with self._active_lock: + ctx = self._active.pop(room, None) + if ctx is not None: + ctx.finish() + def _worker_step(self, item) -> None: ctx, layer_id, event = item try: diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index d7d47b3ec..1d4b03d61 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -391,7 +391,52 @@ class MooncakeKVManager(CommonKVManager): dst_kv_blocks, ) - return PerLayerTransferContext(self.engine, mooncake_session_id, get_blocks) + return PerLayerTransferContext( + self.engine, mooncake_session_id, get_blocks, num_layers=n_layers + ) + + def register_per_layer_transfer(self, room, page_indices) -> bool: + """Lever A: before the forward, build + register a per-layer transfer context + for `room` so the per-layer notifier overlaps its main-KV transfer with the + forward. Reuses send()'s CP filter exactly (no re-derivation). Scoped to + CP-shared-KV (the target scenario). Returns True iff a context was registered; + a False just falls back to the monolithic post-forward transfer (still correct). + page_indices = the request's full new-token logical page ids (req_to_token).""" + mgr = getattr(self, "per_layer_transfer_manager", None) + if mgr is None or not self.server_args.enable_nsa_prefill_cp_shared_kv: + return False + infos = self.transfer_infos.get(room) + if not infos: + return False + from sglang.srt.disaggregation.utils import filter_kv_pages_for_cp_shared_kv + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + layout = CpSharedKVLayout( + page_size=self.kv_args.page_size, + cp_size=self.attn_cp_size, + cp_rank=self.attn_cp_rank, + ) + pages = np.asarray(page_indices, dtype=np.int32) + for info in infos.values(): + if getattr(info, "is_dummy", False): + continue + owned_pages, positions = filter_kv_pages_for_cp_shared_kv( + layout=layout, logical_pages=pages, chunk_page_start=0 + ) + dst_indices = np.asarray(info.dst_kv_indices, dtype=np.int32)[positions] + ctx = self.build_per_layer_context( + info.mooncake_session_id, owned_pages, dst_indices + ) + if ctx is not None: + mgr.register(room, ctx) + logger.info( + "[CP_PER_LAYER_TRANSFER] registered room=%s owned_pages=%d", + room, + len(owned_pages), + ) + return True + return False + return False def _send_kvcache_generic( self, @@ -1090,6 +1135,11 @@ class MooncakeKVManager(CommonKVManager): "Skipping chunk for room %s because it has already failed or been aborted", kv_chunk.room, ) + # Lever A: drain + drop any per-layer context for this room so an + # outstanding RDMA completes before its KV pages can be reclaimed. + per_layer_mgr = getattr(self, "per_layer_transfer_manager", None) + if per_layer_mgr is not None: + per_layer_mgr.drop(kv_chunk.room) continue reqs_to_be_processed = ( self.transfer_infos[kv_chunk.room].values() @@ -1171,7 +1221,22 @@ class MooncakeKVManager(CommonKVManager): target_rank_registration_info: KVArgsRegisterInfo = ( self.decode_kv_args_table[req.mooncake_session_id] ) - if self.is_mla_backend or ( + per_layer_mgr = getattr( + self, "per_layer_transfer_manager", None + ) + if per_layer_mgr is not None and per_layer_mgr.has_room( + kv_chunk.room + ): + # Lever A: the main KV was transferred per-layer, overlapped + # with the forward; wait those transfers here instead of the + # monolithic send (no double-send). aux/state below unchanged. + ret = per_layer_mgr.finish(kv_chunk.room) + logger.info( + "[CP_PER_LAYER_TRANSFER] finished room=%s ret=%s", + kv_chunk.room, + ret, + ) + elif self.is_mla_backend or ( self.attn_tp_size == target_rank_registration_info.dst_attn_tp_size ): diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 62eed12c1..585034780 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -647,6 +647,31 @@ class SchedulerDisaggregationPrefillMixin: return batch + def _register_per_layer_transfers(self: Scheduler, batch) -> None: + """Lever A: before run_batch, register a per-layer transfer context for each + eligible request so the per-layer notifier overlaps its main-KV transfer with + the forward. Scoped (first impl) to single-forward, no-cached-prefix requests: + the per-layer notifier transfers FORWARD-written pages, so chunked or cached- + 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: + return + page_size = self.token_to_kv_pool_allocator.page_size + for req in batch.reqs: + if getattr(req, "disagg_kv_sender", None) is None: + continue + if getattr(req, "is_chunked", 0) != 0 or req.start_send_idx != 0: + continue # chunked / partially-sent: this forward isn't the full range + if len(getattr(req, "prefix_indices", []) or []) != 0: + continue # cached prefix is HiCache-loaded, not forward-written (lever B) + end_idx = min(len(req.fill_ids), len(req.origin_input_ids)) + page_indices = _kv_locs_to_page_indices_cpu( + 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) + @torch.no_grad() def event_loop_normal_disagg_prefill(self: Scheduler) -> None: """A normal scheduler loop for prefill worker in disaggregation mode.""" @@ -718,6 +743,7 @@ class SchedulerDisaggregationPrefillMixin: _seq_summary(getattr(batch, "prefix_lens", None)), len(self.disagg_prefill_inflight_queue), ) + self._register_per_layer_transfers(batch) result = self.run_batch(batch) _cp_shared_kv_bs_gt1_prefill_timing( "event_loop_run_batch_done", @@ -783,6 +809,7 @@ class SchedulerDisaggregationPrefillMixin: # Launch the current batch if batch: + self._register_per_layer_transfers(batch) batch_result = self.run_batch(batch) self.result_queue.append((batch.copy(), batch_result)) else: 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 622eab85a..5f02cc148 100644 --- a/test/registered/unit/disaggregation/test_cp_per_layer_transfer.py +++ b/test/registered/unit/disaggregation/test_cp_per_layer_transfer.py @@ -49,8 +49,8 @@ def _blocks(layer_id): return ([base], [base + 5000], [64]) -def _ctx(engine, blocks=_blocks): - return PerLayerTransferContext(engine, "sess", blocks) +def _ctx(engine, blocks=_blocks, num_layers=8): + return PerLayerTransferContext(engine, "sess", blocks, num_layers=num_layers) class TestPerLayerTransferContext(unittest.TestCase): @@ -81,7 +81,7 @@ class TestPerLayerTransferContext(unittest.TestCase): def test_finish_waits_all_and_returns_zero(self): eng = _FakeEngine() - ctx = _ctx(eng) + ctx = _ctx(eng, num_layers=3) for L in range(3): ctx.submit_layer(L, _FakeEvent()) self.assertEqual(ctx.finish(), 0) @@ -89,19 +89,38 @@ class TestPerLayerTransferContext(unittest.TestCase): def test_submit_failure_marks_failed_and_stops(self): eng = _FakeEngine(submit_rets=[101, -1, 103]) - ctx = _ctx(eng) + ctx = _ctx(eng, num_layers=3) for L in range(3): ctx.submit_layer(L, None) self.assertTrue(ctx.failed) - self.assertEqual(len(eng.submits), 2) # stopped after the -1 + self.assertEqual(len(eng.submits), 2) # stopped submitting after the -1 self.assertEqual(ctx.finish(), -1) def test_finish_propagates_wait_failure(self): eng = _FakeEngine(wait_ret=-9) - ctx = _ctx(eng) + ctx = _ctx(eng, num_layers=1) 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 + eng = _FakeEngine() + ctx = _ctx(eng, num_layers=3) + ctx.submit_layer(0, None) + self.assertEqual(ctx.finish(timeout=0.05), -1) # incomplete -> failure + + 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) + for L in range(3): + ctx.submit_layer(L, None) + self.assertEqual(ctx.finish(timeout=0.5), -1) # completes (no hang), reports failure + class TestBuildLayerBlocks(unittest.TestCase): def test_addresses_and_lengths(self):