From 74e22a07214a926351fad77557f2bd0738f72307 Mon Sep 17 00:00:00 2001 From: leavelet Date: Sun, 7 Jun 2026 00:39:29 +0000 Subject: [PATCH] feat(disagg): register per-layer transfer manager notifier (lever A, A3-step1) Wire PerLayerTransferManager into the prefill bootstrap queue: when SGLANG_CP_SHARED_KV_PER_LAYER_TRANSFER is on, create the manager (real torch.cuda.Event / current_stream) and register its on_layer_end on the KV pool's layer_backup_notifiers, exposing it as kv_manager.per_layer_transfer_manager. The forward's per-layer end hook now reaches the manager. Additive and a no-op until a request registers a transfer context (A3-step2): on_layer_end returns early when no contexts are active (verified). Remaining A3 (get_blocks + setup-before-run_batch + finish reconciled with send_kv_chunk) builds on this. Co-Authored-By: Claude Opus 4.8 (1M context) --- python/sglang/srt/disaggregation/prefill.py | 33 +++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 95e20199c..62eed12c1 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -265,6 +265,7 @@ class PrefillBootstrapQueue: self.scheduler = scheduler self.transfer_backend = transfer_backend self.kv_manager = self._init_kv_manager() + self._maybe_init_per_layer_transfer_manager() if self.scheduler.tp_worker.is_hybrid_swa: # FIXME: current SWA allocation allocate full kv cache size in prefill @@ -273,6 +274,38 @@ class PrefillBootstrapQueue: self.scheduler.tp_worker.model_runner.swa_max_total_num_tokens, ) + def _maybe_init_per_layer_transfer_manager(self) -> None: + # Lever A: create the per-layer overlapped-transfer manager and register its + # per-layer notifier on the KV pool. The forward fires + # notify_layer_end_for_backup -> layer_backup_notifiers(local_layer_id), which + # calls manager.on_layer_end. Additive/no-op until a request registers a + # transfer context (A3 wiring); gated by SGLANG_CP_SHARED_KV_PER_LAYER_TRANSFER. + self.kv_manager.per_layer_transfer_manager = None + if not envs.SGLANG_CP_SHARED_KV_PER_LAYER_TRANSFER.get(): + return + import torch + + from sglang.srt.disaggregation.cp_per_layer_transfer import ( + PerLayerTransferManager, + ) + + manager = PerLayerTransferManager( + event_factory=torch.cuda.Event, + current_stream=torch.cuda.current_stream, + ) + pool = self.token_to_kv_pool + if hasattr(pool, "register_layer_backup_notifier"): + pool.register_layer_backup_notifier(manager.on_layer_end) + self.kv_manager.per_layer_transfer_manager = manager + logger.info( + "[CP_PER_LAYER_TRANSFER] registered per-layer transfer manager notifier" + ) + else: + logger.warning( + "[CP_PER_LAYER_TRANSFER] kv pool lacks register_layer_backup_notifier; " + "per-layer overlap disabled" + ) + def _init_kv_manager(self) -> CommonKVManager: kv_args_class = get_kv_class(self.transfer_backend, KVClassType.KVARGS) kv_args = kv_args_class()