From 86412bb1746c1caa2876065fadfcfee2bd678647 Mon Sep 17 00:00:00 2001 From: Claude Code Agent Date: Wed, 13 May 2026 08:12:04 +0800 Subject: [PATCH] =?UTF-8?q?feat(sglang):=20D=E2=86=92P=20snapshot=20link?= =?UTF-8?q?=20integration=20=E2=80=94=20controller=20+=20RPC=20handlers?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Phase 2 of the D→P sync feature (Phase 1 in dc4867c verified the underlying RDMA link in isolation). This commit wires that link into each SGLang worker's scheduler so D and P can exchange session KV without going through the PD prefill pipeline. New module: third_party/sglang/python/sglang/srt/disaggregation/snapshot/ controller.py — SnapshotLinkController owns one mooncake transfer engine per worker, pre-registers all kv_pool layer buffers, and exposes prepare_receive() and push_session_kv() APIs. Receive bookkeeping via a session_id → SnapshotIngestRecord side-table. Three RPC types added to io_struct.py and full plumbing wired through: SnapshotPrepareReceiveReqInput/Output P-side alloc + return layout SnapshotDumpReqInput/Output D-side read kv_pool + RDMA push SnapshotFinalizeIngestReqInput/Output P-side radix tree insert Files touched: managers/io_struct.py 3 new ReqInput/ReqOutput pairs managers/tokenizer_communicator_mixin.py 3 communicators, 3 awaitables managers/scheduler.py init controller + 3 handlers entrypoints/http_server.py 3 HTTP endpoints under /_snapshot Activation: set SGLANG_SNAPSHOT_LINK_ENABLE=1 (and SGLANG_SNAPSHOT_LINK_HOST / _PORT / _IB_DEVICE) per worker. Controller init is opt-in and defaults off, so production PD pipeline is untouched. Subsequent work (Phase 3): agentic-pd-hybrid orchestration in _invoke_kvcache_seeded_router to call prepare_receive on P, dump on D-old, finalize_ingest on P, then trigger the existing P→D' transfer which will now hit P's radix cache (skipping re-prefill). --- .../srt/disaggregation/snapshot/__init__.py | 27 ++ .../srt/disaggregation/snapshot/controller.py | 342 ++++++++++++++++++ .../sglang/srt/entrypoints/http_server.py | 18 + .../python/sglang/srt/managers/io_struct.py | 91 +++++ .../python/sglang/srt/managers/scheduler.py | 228 ++++++++++++ .../managers/tokenizer_communicator_mixin.py | 57 +++ 6 files changed, 763 insertions(+) create mode 100644 third_party/sglang/python/sglang/srt/disaggregation/snapshot/__init__.py create mode 100644 third_party/sglang/python/sglang/srt/disaggregation/snapshot/controller.py diff --git a/third_party/sglang/python/sglang/srt/disaggregation/snapshot/__init__.py b/third_party/sglang/python/sglang/srt/disaggregation/snapshot/__init__.py new file mode 100644 index 0000000..d6e7992 --- /dev/null +++ b/third_party/sglang/python/sglang/srt/disaggregation/snapshot/__init__.py @@ -0,0 +1,27 @@ +"""D→P RDMA snapshot push subsystem. + +A minimal, role-symmetric mooncake transport that runs alongside SGLang's +existing PD pipeline. Both D and P workers can both send and receive +snapshots — direction is determined by which kv_pool we read from / +write into. + +See ``docs/D_TO_P_SYNC_DESIGN_ZH.md`` for the full design. +""" + +from sglang.srt.disaggregation.snapshot.controller import ( + SnapshotLinkController, + SnapshotIngestRecord, + SNAPSHOT_LINK_ENABLE_ENV, + SNAPSHOT_LINK_HOST_ENV, + SNAPSHOT_LINK_PORT_ENV, + SNAPSHOT_LINK_IB_DEVICE_ENV, +) + +__all__ = [ + "SnapshotLinkController", + "SnapshotIngestRecord", + "SNAPSHOT_LINK_ENABLE_ENV", + "SNAPSHOT_LINK_HOST_ENV", + "SNAPSHOT_LINK_PORT_ENV", + "SNAPSHOT_LINK_IB_DEVICE_ENV", +] diff --git a/third_party/sglang/python/sglang/srt/disaggregation/snapshot/controller.py b/third_party/sglang/python/sglang/srt/disaggregation/snapshot/controller.py new file mode 100644 index 0000000..6a93bb7 --- /dev/null +++ b/third_party/sglang/python/sglang/srt/disaggregation/snapshot/controller.py @@ -0,0 +1,342 @@ +"""SnapshotLinkController — drives D→P RDMA snapshot pushes. + +This class owns: + * A dedicated ``mooncake.engine.TransferEngine`` (independent of the PD + pipeline engine). + * Memory-region registrations covering the worker's KV pool layer buffers + (registered once at startup for zero per-snapshot overhead). + * A side-table mapping ``session_id → SnapshotIngestRecord`` for P-side + receivers tracking outstanding ingests until ``snapshot_finalize_ingest`` + is called. + +Lifecycle: + SGLang scheduler instantiates one of these per worker if + ``SGLANG_SNAPSHOT_LINK_ENABLE=1``. The scheduler is responsible for + feeding kv_pool buffer descriptors and the kv_pool_allocator so the + controller can pre-register memory and (on P) allocate slots. + +Direction symmetry: + Both D and P workers run identical controllers. Who sends and who + receives is determined by who calls ``push_session_kv`` vs + ``prepare_receive`` / ``finalize_ingest`` — there's no PREFILL/DECODE + role baked into the snapshot side. + +This is the **vendored** copy living alongside SGLang internals so the +scheduler can import it directly. ``src/agentic_pd_hybrid/snapshot_link.py`` +contains the same primitives for stand-alone smoke testing. +""" + +from __future__ import annotations + +import ctypes +import logging +import os +import threading +import time +from dataclasses import dataclass, field +from typing import List, Optional, Tuple + +logger = logging.getLogger(__name__) + +# Env-var names (also exported from package __init__) +SNAPSHOT_LINK_ENABLE_ENV = "SGLANG_SNAPSHOT_LINK_ENABLE" +SNAPSHOT_LINK_HOST_ENV = "SGLANG_SNAPSHOT_LINK_HOST" +SNAPSHOT_LINK_PORT_ENV = "SGLANG_SNAPSHOT_LINK_PORT" +SNAPSHOT_LINK_IB_DEVICE_ENV = "SGLANG_SNAPSHOT_LINK_IB_DEVICE" + + +@dataclass +class _LayerBufferDesc: + """Per-layer KV buffer descriptor on this worker.""" + base_ptr: int # data pointer of the layer's full buffer tensor + bytes_per_token: int # head_num * head_dim * dtype.itemsize + capacity_bytes: int # full buffer size in bytes + is_k: bool # True for K-buffer, False for V + + +@dataclass +class SnapshotIngestRecord: + """P-side: bookkeeping for an outstanding incoming snapshot.""" + session_id: str + slot_indices: List[int] # kv_pool slots reserved for this ingest + num_tokens: int + created_at: float = field(default_factory=time.time) + + +def _import_transfer_engine(): + try: + from mooncake.engine import TransferEngine + except ImportError as e: + raise ImportError( + "mooncake.engine.TransferEngine is required for the snapshot " + "link. Install mooncake-transfer-engine in the venv." + ) from e + return TransferEngine + + +class SnapshotLinkController: + """Mooncake engine + per-layer mem registrations + ingest bookkeeping.""" + + def __init__( + self, + host: str, + port: int, + ib_device: Optional[str], + kv_pool_layer_buffers: List[Tuple[int, int, int, bool]], + token_to_kv_pool_allocator, + protocol: Optional[str] = None, + ): + """ + Parameters + ---------- + host, port : where this worker binds its snapshot engine. + ib_device : preferred IB HCA (e.g. "mlx5_60"). + kv_pool_layer_buffers : list of ``(base_ptr, bytes_per_token, + capacity_bytes, is_k)`` tuples. Order should be K-layers + first then V-layers (matches MHATokenToKVPool layout). + token_to_kv_pool_allocator : the worker's allocator. We use + ``.alloc(N)`` on the P side to reserve receive slots. + """ + TransferEngine = _import_transfer_engine() + self.host = host + self.port = port + self.ib_device = ib_device + self.token_to_kv_pool_allocator = token_to_kv_pool_allocator + self.layer_buffers: List[_LayerBufferDesc] = [ + _LayerBufferDesc( + base_ptr=base, bytes_per_token=btok, + capacity_bytes=cap, is_k=is_k, + ) + for (base, btok, cap, is_k) in kv_pool_layer_buffers + ] + + self.engine = TransferEngine() + proto = protocol or os.environ.get("MOONCAKE_PROTOCOL", "rdma") + listen = f"{host}:{port}" + ret = self.engine.initialize(listen, "P2PHANDSHAKE", proto, ib_device or "") + if ret != 0: + raise RuntimeError( + f"SnapshotLinkController.initialize({listen}, {proto}, " + f"ib={ib_device}) returned {ret}" + ) + self._session_id = f"{host}:{self.engine.get_rpc_port()}" + + # Register all layer buffers up-front (one-shot) + ptrs = [d.base_ptr for d in self.layer_buffers] + lens = [d.capacity_bytes for d in self.layer_buffers] + try: + reg_ret = self.engine.batch_register_memory(ptrs, lens) + except Exception: + reg_ret = -1 + # Fall back to individual register_memory calls + reg_ret = 0 + for ptr, length in zip(ptrs, lens): + r = self.engine.register_memory(ptr, length) + if r != 0: + logger.warning( + "SnapshotLinkController register_memory(%s, %d) returned %d", + hex(ptr), length, r, + ) + reg_ret = r + if reg_ret != 0: + logger.warning( + "SnapshotLinkController batch_register_memory returned %d " + "(continuing — individual registrations may have succeeded)", + reg_ret, + ) + + # Receive-side bookkeeping + self._ingest_records: dict[str, SnapshotIngestRecord] = {} + self._lock = threading.Lock() + + logger.info( + "SnapshotLinkController up at %s (snapshot_session_id=%s, " + "%d layer buffers registered, ib=%s)", + listen, self._session_id, len(self.layer_buffers), ib_device, + ) + + # ----- accessors ---------------------------------------------------- + + @property + def snapshot_session_id(self) -> str: + return self._session_id + + @property + def layer_num(self) -> int: + """Number of K layers (= number of V layers = total / 2).""" + return len(self.layer_buffers) // 2 + + def get_k_base_ptrs(self) -> List[int]: + return [d.base_ptr for d in self.layer_buffers if d.is_k] + + def get_v_base_ptrs(self) -> List[int]: + return [d.base_ptr for d in self.layer_buffers if not d.is_k] + + def get_stride_k_bytes(self) -> int: + for d in self.layer_buffers: + if d.is_k: + return d.bytes_per_token + return 0 + + def get_stride_v_bytes(self) -> int: + for d in self.layer_buffers: + if not d.is_k: + return d.bytes_per_token + return 0 + + # ----- P-side: prepare to receive ---------------------------------- + + def prepare_receive(self, session_id: str, num_tokens: int) -> Optional[SnapshotIngestRecord]: + """Allocate ``num_tokens`` slots in kv_pool for an incoming snapshot. + + Returns a record with the slot indices, or ``None`` if no capacity. + Mooncake registration is already in place (whole-buffer registration + at startup), so no per-snapshot register is needed. + """ + try: + indices_tensor = self.token_to_kv_pool_allocator.alloc(num_tokens) + except Exception as e: + logger.exception("SnapshotLinkController.prepare_receive alloc failed: %s", e) + return None + if indices_tensor is None: + return None + try: + slot_indices = [int(x) for x in indices_tensor.tolist()] + except Exception: + # If allocator returns a python list directly + slot_indices = list(map(int, indices_tensor)) + record = SnapshotIngestRecord( + session_id=session_id, + slot_indices=slot_indices, + num_tokens=num_tokens, + ) + with self._lock: + old = self._ingest_records.pop(session_id, None) + if old is not None: + # Best-effort: free old slots if we're overwriting + try: + self._free_slots(old.slot_indices) + except Exception: + pass + self._ingest_records[session_id] = record + return record + + def take_record(self, session_id: str) -> Optional[SnapshotIngestRecord]: + with self._lock: + return self._ingest_records.pop(session_id, None) + + def discard_record(self, session_id: str) -> None: + """Drop a pending ingest (e.g. on timeout or D-side failure).""" + with self._lock: + rec = self._ingest_records.pop(session_id, None) + if rec is not None: + try: + self._free_slots(rec.slot_indices) + except Exception: + pass + + def _free_slots(self, slot_indices: List[int]) -> None: + import torch + t = torch.tensor(slot_indices, dtype=torch.int64, + device=self._allocator_device()) + try: + self.token_to_kv_pool_allocator.free(t) + except Exception as e: + logger.warning("SnapshotLinkController._free_slots failed: %s", e) + + def _allocator_device(self): + # Best-effort: pull device from one of the buffer tensors via the allocator + try: + return self.token_to_kv_pool_allocator.device + except AttributeError: + return "cuda" + + # ----- D-side: push session KV -------------------------------------- + + def push_session_kv( + self, + *, + target_snapshot_session_id: str, + src_slot_indices: List[int], + target_k_base_ptrs: List[int], + target_v_base_ptrs: List[int], + target_slot_indices: List[int], + target_stride_k_bytes: int, + target_stride_v_bytes: int, + ) -> Tuple[int, int]: + """Push the KV bytes at src_slot_indices to the remote slots. + + Returns ``(mooncake_return_code, bytes_pushed)``. Caller is + responsible for any post-push handshake (e.g. finalize_ingest RPC). + """ + layer_num = self.layer_num + k_src_bases = self.get_k_base_ptrs() + v_src_bases = self.get_v_base_ptrs() + stride_k = self.get_stride_k_bytes() + stride_v = self.get_stride_v_bytes() + if len(target_k_base_ptrs) != layer_num or len(target_v_base_ptrs) != layer_num: + raise ValueError( + f"target K/V base ptr count {len(target_k_base_ptrs)}/{len(target_v_base_ptrs)} " + f"!= local layer_num {layer_num}" + ) + if stride_k != target_stride_k_bytes or stride_v != target_stride_v_bytes: + raise ValueError( + f"stride mismatch: local k={stride_k}, v={stride_v}; " + f"target k={target_stride_k_bytes}, v={target_stride_v_bytes}" + ) + if len(src_slot_indices) != len(target_slot_indices): + raise ValueError( + f"slot count mismatch: src={len(src_slot_indices)}, " + f"target={len(target_slot_indices)}" + ) + + local_addrs: List[int] = [] + remote_addrs: List[int] = [] + lengths: List[int] = [] + + # Group contiguous runs on the target side to coalesce ops. + # Simple approach: per (layer, K/V) pair, walk src/target index + # tuples; merge runs where both src and target are sequential. + for layer_id in range(layer_num): + for kv_bases, stride, kv_label in ( + (k_src_bases[layer_id], stride_k, "K"), + (v_src_bases[layer_id], stride_v, "V"), + ): + src_base = kv_bases + if kv_label == "K": + tgt_base = target_k_base_ptrs[layer_id] + else: + tgt_base = target_v_base_ptrs[layer_id] + run_src_start = run_tgt_start = run_len = None + for s, t in zip(src_slot_indices, target_slot_indices): + if run_src_start is None: + run_src_start, run_tgt_start, run_len = s, t, 1 + elif s == run_src_start + run_len and t == run_tgt_start + run_len: + run_len += 1 + else: + local_addrs.append(src_base + run_src_start * stride) + remote_addrs.append(tgt_base + run_tgt_start * stride) + lengths.append(run_len * stride) + run_src_start, run_tgt_start, run_len = s, t, 1 + if run_src_start is not None: + local_addrs.append(src_base + run_src_start * stride) + remote_addrs.append(tgt_base + run_tgt_start * stride) + lengths.append(run_len * stride) + + t0 = time.perf_counter() + try: + ret = self.engine.batch_transfer_sync_write( + target_snapshot_session_id, local_addrs, remote_addrs, lengths + ) + except Exception as e: + logger.exception("SnapshotLinkController.push_session_kv threw: %s", e) + return -1, 0 + t1 = time.perf_counter() + bytes_pushed = sum(lengths) + logger.info( + "SnapshotLinkController.push_session_kv → %s: %d ops, %d B, " + "ret=%d, %.2f ms", + target_snapshot_session_id, len(lengths), bytes_pushed, ret, + (t1 - t0) * 1000.0, + ) + return ret, bytes_pushed diff --git a/third_party/sglang/python/sglang/srt/entrypoints/http_server.py b/third_party/sglang/python/sglang/srt/entrypoints/http_server.py index 07f5521..3362224 100644 --- a/third_party/sglang/python/sglang/srt/entrypoints/http_server.py +++ b/third_party/sglang/python/sglang/srt/entrypoints/http_server.py @@ -125,6 +125,9 @@ from sglang.srt.managers.io_struct import ( LoadLoRAAdapterFromTensorsReqInput, LoadLoRAAdapterReqInput, DirectAppendAdmissionReqInput, + SnapshotDumpReqInput, + SnapshotFinalizeIngestReqInput, + SnapshotPrepareReceiveReqInput, OpenSessionReqInput, ParseFunctionCallReq, PauseGenerationReqInput, @@ -1295,6 +1298,21 @@ async def admit_direct_append(obj: DirectAppendAdmissionReqInput): return await _global_state.tokenizer_manager.admit_direct_append(obj) +@app.post("/_snapshot/prepare_receive") +async def snapshot_prepare_receive(obj: SnapshotPrepareReceiveReqInput): + return await _global_state.tokenizer_manager.snapshot_prepare_receive(obj) + + +@app.post("/_snapshot/dump") +async def snapshot_dump(obj: SnapshotDumpReqInput): + return await _global_state.tokenizer_manager.snapshot_dump(obj) + + +@app.post("/_snapshot/finalize_ingest") +async def snapshot_finalize_ingest(obj: SnapshotFinalizeIngestReqInput): + return await _global_state.tokenizer_manager.snapshot_finalize_ingest(obj) + + @app.api_route("/configure_logging", methods=["GET", "POST"]) @auth_level(AuthLevel.ADMIN_OPTIONAL) async def configure_logging(obj: ConfigureLoggingReq, request: Request): diff --git a/third_party/sglang/python/sglang/srt/managers/io_struct.py b/third_party/sglang/python/sglang/srt/managers/io_struct.py index 3591976..5b2c6c0 100644 --- a/third_party/sglang/python/sglang/srt/managers/io_struct.py +++ b/third_party/sglang/python/sglang/srt/managers/io_struct.py @@ -1632,6 +1632,97 @@ class HealthCheckOutput(BaseReq): pass +# --------------------------------------------------------------------------- +# D→P snapshot ingest (Phase 2 of D→P sync feature; see +# docs/D_TO_P_SYNC_DESIGN_ZH.md). +# +# Three-step protocol orchestrated by agentic-pd-hybrid: +# 1. PrepareReceive → P allocates kv_pool slots + returns destination +# addresses for D's RDMA writes. +# 2. (out-of-band) → D uses snapshot_link to RDMA-push KV bytes +# directly to P's slot addresses. +# 3. FinalizeIngest → P inserts (token_ids, kv_indices) into its radix +# tree so subsequent prefill requests for this +# session see a cache hit. +# +# Each step is its own ReqInput/ReqOutput pair so the scheduler handlers can +# be written stateless and the orchestrator can retry / abort cleanly. +# --------------------------------------------------------------------------- + + +@dataclass +class SnapshotPrepareReceiveReqInput(BaseReq): + """P-side: allocate slots + register them with mooncake for D to push into.""" + + session_id: str + num_tokens: int # P will alloc this many contiguous slots + expected_bytes_per_layer_k: int = 0 # per-token K bytes × num_tokens (sanity) + expected_bytes_per_layer_v: int = 0 # per-token V bytes × num_tokens (sanity) + + +@dataclass +class SnapshotPrepareReceiveReqOutput(BaseReq): + ok: bool + reason: Optional[str] = None + # Layout the D side needs to address P's kv_pool slots: + # k_base_ptrs[L] = base device address of layer L's K buffer + # v_base_ptrs[L] = base device address of layer L's V buffer + # slot_indices = the contiguous range P allocated (list[int], length=num_tokens) + # stride_k_bytes = bytes per token K = head_num * head_dim * dtype.itemsize + # stride_v_bytes = bytes per token V (often equals stride_k_bytes) + # P also registers these slot regions with mooncake before returning. + k_base_ptrs: List[int] = field(default_factory=list) + v_base_ptrs: List[int] = field(default_factory=list) + slot_indices: List[int] = field(default_factory=list) + stride_k_bytes: int = 0 + stride_v_bytes: int = 0 + layer_num: int = 0 + # P's mooncake snapshot session id (host:rpc_port) for D's batch write target + snapshot_session_id: str = "" + available_tokens: int = 0 + + +@dataclass +class SnapshotDumpReqInput(BaseReq): + """D-side: dump session KV via snapshot_link to a target P endpoint.""" + + session_id: str + target_snapshot_session_id: str # P's mooncake snapshot session id + target_k_base_ptrs: List[int] = field(default_factory=list) + target_v_base_ptrs: List[int] = field(default_factory=list) + target_slot_indices: List[int] = field(default_factory=list) + target_stride_k_bytes: int = 0 + target_stride_v_bytes: int = 0 + ib_device: Optional[str] = None # for the D-side SnapshotPeer initialization + + +@dataclass +class SnapshotDumpReqOutput(BaseReq): + ok: bool + reason: Optional[str] = None + bytes_pushed: int = 0 + transfer_duration_ms: float = 0.0 + kv_committed_len: int = 0 # the actual number of tokens D had for this session + # The token_ids that go with the KV (so P can call radix_cache.insert) + token_ids: List[int] = field(default_factory=list) + + +@dataclass +class SnapshotFinalizeIngestReqInput(BaseReq): + """P-side: insert (token_ids, kv_indices) into radix tree after D's push.""" + + session_id: str + token_ids: List[int] + slot_indices: List[int] + + +@dataclass +class SnapshotFinalizeIngestReqOutput(BaseReq): + ok: bool + reason: Optional[str] = None + inserted_prefix_len: int = 0 + + class ExpertDistributionReqType(Enum): START_RECORD = 1 STOP_RECORD = 2 diff --git a/third_party/sglang/python/sglang/srt/managers/scheduler.py b/third_party/sglang/python/sglang/srt/managers/scheduler.py index 15d7558..b180a1c 100644 --- a/third_party/sglang/python/sglang/srt/managers/scheduler.py +++ b/third_party/sglang/python/sglang/srt/managers/scheduler.py @@ -96,6 +96,12 @@ from sglang.srt.managers.io_struct import ( ContinueGenerationReqInput, DirectAppendAdmissionReqInput, DirectAppendAdmissionReqOutput, + SnapshotDumpReqInput, + SnapshotDumpReqOutput, + SnapshotFinalizeIngestReqInput, + SnapshotFinalizeIngestReqOutput, + SnapshotPrepareReceiveReqInput, + SnapshotPrepareReceiveReqOutput, DestroyWeightsUpdateGroupReqInput, DetachHiCacheStorageReqInput, DetachHiCacheStorageReqOutput, @@ -844,6 +850,69 @@ class Scheduler( embedding_cache_size = envs.SGLANG_VLM_CACHE_SIZE_MB.get() init_mm_embedding_cache(embedding_cache_size * 1024 * 1024) + # ---- D→P snapshot link (Phase 2 of D→P sync feature) ------------ + # Enabled per-worker via SGLANG_SNAPSHOT_LINK_ENABLE=1. Each worker + # binds an independent mooncake transfer engine on + # SGLANG_SNAPSHOT_LINK_HOST:SGLANG_SNAPSHOT_LINK_PORT and pre- + # registers the kv_pool layer buffers for one-shot RDMA pushes / + # receives. See docs/D_TO_P_SYNC_DESIGN_ZH.md. + self.snapshot_link_controller = None + from sglang.srt.disaggregation.snapshot import ( + SnapshotLinkController as _SnapLinkCtrl, + SNAPSHOT_LINK_ENABLE_ENV, + SNAPSHOT_LINK_HOST_ENV, + SNAPSHOT_LINK_PORT_ENV, + SNAPSHOT_LINK_IB_DEVICE_ENV, + ) + if os.environ.get(SNAPSHOT_LINK_ENABLE_ENV, "0") == "1": + host = os.environ.get(SNAPSHOT_LINK_HOST_ENV, server_args.host) + port = int(os.environ.get(SNAPSHOT_LINK_PORT_ENV, + str(server_args.disaggregation_bootstrap_port + 1000))) + ib = os.environ.get(SNAPSHOT_LINK_IB_DEVICE_ENV, server_args.disaggregation_ib_device) + try: + kv_pool = self.token_to_kv_pool_allocator.get_kvcache() + except AttributeError: + # Some allocators expose the pool directly + kv_pool = getattr(self.token_to_kv_pool_allocator, "kvcache", None) + if kv_pool is None: + logger.warning("SNAPSHOT_LINK_ENABLE=1 but kv_pool unavailable; skipping init") + else: + try: + kv_data_ptrs, kv_data_lens, kv_item_lens = kv_pool.get_contiguous_buf_infos() + layer_n = len(kv_data_ptrs) // 2 + layer_buffers = [] + # K layers first, then V layers (matches MHATokenToKVPool.get_contiguous_buf_infos) + for i in range(layer_n): + layer_buffers.append(( + kv_data_ptrs[i], + kv_item_lens[i] // max(1, kv_pool.page_size), + kv_data_lens[i], + True, # is_k + )) + for i in range(layer_n): + layer_buffers.append(( + kv_data_ptrs[layer_n + i], + kv_item_lens[layer_n + i] // max(1, kv_pool.page_size), + kv_data_lens[layer_n + i], + False, # is_k=False (V) + )) + self.snapshot_link_controller = _SnapLinkCtrl( + host=host, + port=port, + ib_device=ib, + kv_pool_layer_buffers=layer_buffers, + token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, + ) + logger.info( + "Snapshot link controller initialized: %s, sid=%s, %d layer bufs", + f"{host}:{port}", + self.snapshot_link_controller.snapshot_session_id, + len(layer_buffers), + ) + except Exception as e: + logger.warning("Snapshot link init failed: %s; continuing without it", e) + self.snapshot_link_controller = None + def init_running_status(self): self.waiting_queue: List[Req] = [] self.decode_direct_waiting_queue: List[Req] = [] @@ -1219,6 +1288,9 @@ class Scheduler( (OpenSessionReqInput, self.open_session), (CloseSessionReqInput, self.close_session), (DirectAppendAdmissionReqInput, self.admit_direct_append), + (SnapshotPrepareReceiveReqInput, self.snapshot_prepare_receive), + (SnapshotDumpReqInput, self.snapshot_dump), + (SnapshotFinalizeIngestReqInput, self.snapshot_finalize_ingest), (UpdateWeightFromDiskReqInput, self.update_weights_from_disk), (InitWeightsUpdateGroupReqInput, self.init_weights_update_group), (DestroyWeightsUpdateGroupReqInput, self.destroy_weights_update_group), @@ -3673,6 +3745,162 @@ class Scheduler( ), ) + # ----- D→P snapshot link handlers (Phase 2/3) --------------------- + + def snapshot_prepare_receive( + self, recv_req: SnapshotPrepareReceiveReqInput + ) -> SnapshotPrepareReceiveReqOutput: + """P-side: alloc kv_pool slots, return slot/buffer layout for D's batch_push.""" + ctrl = self.snapshot_link_controller + if ctrl is None: + return SnapshotPrepareReceiveReqOutput( + ok=False, reason="snapshot-link-disabled", + ) + try: + available = int(self.token_to_kv_pool_allocator.available_size()) + except Exception: + available = -1 + if recv_req.num_tokens <= 0: + return SnapshotPrepareReceiveReqOutput(ok=False, reason="zero-tokens") + record = ctrl.prepare_receive(recv_req.session_id, recv_req.num_tokens) + if record is None: + return SnapshotPrepareReceiveReqOutput( + ok=False, reason="alloc-failed", + available_tokens=available, + ) + return SnapshotPrepareReceiveReqOutput( + ok=True, + k_base_ptrs=ctrl.get_k_base_ptrs(), + v_base_ptrs=ctrl.get_v_base_ptrs(), + slot_indices=record.slot_indices, + stride_k_bytes=ctrl.get_stride_k_bytes(), + stride_v_bytes=ctrl.get_stride_v_bytes(), + layer_num=ctrl.layer_num, + snapshot_session_id=ctrl.snapshot_session_id, + available_tokens=available, + ) + + def snapshot_dump( + self, recv_req: SnapshotDumpReqInput + ) -> SnapshotDumpReqOutput: + """D-side: gather session KV from kv_pool, RDMA-write to remote slots.""" + ctrl = self.snapshot_link_controller + if ctrl is None: + return SnapshotDumpReqOutput(ok=False, reason="snapshot-link-disabled") + if not isinstance(self.tree_cache, SessionAwareCache): + return SnapshotDumpReqOutput(ok=False, reason="tree-cache-not-session-aware") + slot = self.tree_cache.slots.get(recv_req.session_id) + if slot is None or slot.req_pool_idx is None: + return SnapshotDumpReqOutput(ok=False, reason="session-not-resident") + kv_committed_len = int(slot.kv_committed_len) + if kv_committed_len == 0: + return SnapshotDumpReqOutput(ok=False, reason="zero-committed-len") + # Read kv_indices for the session's prefix + try: + kv_idx_tensor = self.req_to_token_pool.req_to_token[ + slot.req_pool_idx, :kv_committed_len + ] + src_slot_indices = [int(x) for x in kv_idx_tensor.tolist()] + except Exception as e: + logger.exception("snapshot_dump: failed to read kv_indices: %s", e) + return SnapshotDumpReqOutput(ok=False, reason=f"read-indices-failed: {e!r}") + + # Truncate to the count P prepared for (must match) + target_n = len(recv_req.target_slot_indices) + if target_n > kv_committed_len: + return SnapshotDumpReqOutput( + ok=False, + reason=f"target-larger-than-source({target_n}>{kv_committed_len})", + ) + src_slot_indices = src_slot_indices[:target_n] + + try: + ret, bytes_pushed = ctrl.push_session_kv( + target_snapshot_session_id=recv_req.target_snapshot_session_id, + src_slot_indices=src_slot_indices, + target_k_base_ptrs=recv_req.target_k_base_ptrs, + target_v_base_ptrs=recv_req.target_v_base_ptrs, + target_slot_indices=recv_req.target_slot_indices[:target_n], + target_stride_k_bytes=recv_req.target_stride_k_bytes, + target_stride_v_bytes=recv_req.target_stride_v_bytes, + ) + except Exception as e: + logger.exception("snapshot_dump: push_session_kv threw: %s", e) + return SnapshotDumpReqOutput(ok=False, reason=f"push-failed: {e!r}") + + if ret != 0: + return SnapshotDumpReqOutput( + ok=False, + reason=f"mooncake-batch-write-ret={ret}", + bytes_pushed=int(bytes_pushed), + kv_committed_len=int(kv_committed_len), + ) + return SnapshotDumpReqOutput( + ok=True, bytes_pushed=int(bytes_pushed), + kv_committed_len=int(kv_committed_len), + token_ids=[], # caller already has token_ids + ) + + def snapshot_finalize_ingest( + self, recv_req: SnapshotFinalizeIngestReqInput + ) -> SnapshotFinalizeIngestReqOutput: + """P-side: insert (token_ids, slot_indices) into radix tree.""" + ctrl = self.snapshot_link_controller + if ctrl is None: + return SnapshotFinalizeIngestReqOutput( + ok=False, reason="snapshot-link-disabled", + ) + record = ctrl.take_record(recv_req.session_id) + if record is None: + return SnapshotFinalizeIngestReqOutput( + ok=False, reason="no-pending-ingest", + ) + # Sanity: the slot indices we're about to insert should match the ones we reserved. + if list(recv_req.slot_indices) != record.slot_indices: + # The caller passed back the slot indices we returned in prepare; if they + # don't match, something's gone wrong. Free reserved slots and bail. + try: + ctrl._free_slots(record.slot_indices) + except Exception: + pass + return SnapshotFinalizeIngestReqOutput( + ok=False, + reason="slot-indices-mismatch", + ) + n_tokens = min(len(recv_req.token_ids), len(record.slot_indices)) + if n_tokens == 0: + ctrl._free_slots(record.slot_indices) + return SnapshotFinalizeIngestReqOutput(ok=False, reason="empty-token-ids") + try: + import torch + from sglang.srt.mem_cache.base_prefix_cache import InsertParams + from sglang.srt.mem_cache.radix_cache import RadixKey + kv_indices = torch.tensor( + record.slot_indices[:n_tokens], + dtype=torch.int64, + device=self.tree_cache.token_to_kv_pool_allocator.device, + ) + radix_key = RadixKey(recv_req.token_ids[:n_tokens], None) + inner = ( + self.tree_cache.inner + if isinstance(self.tree_cache, SessionAwareCache) + else self.tree_cache + ) + result = inner.insert(InsertParams(key=radix_key, value=kv_indices)) + inserted = int(result.prefix_len) + except Exception as e: + logger.exception("snapshot_finalize_ingest: radix insert failed: %s", e) + try: + ctrl._free_slots(record.slot_indices) + except Exception: + pass + return SnapshotFinalizeIngestReqOutput( + ok=False, reason=f"radix-insert-failed: {e!r}", + ) + return SnapshotFinalizeIngestReqOutput( + ok=True, inserted_prefix_len=inserted, + ) + def _compute_backpressure_pause_hint( self, *, diff --git a/third_party/sglang/python/sglang/srt/managers/tokenizer_communicator_mixin.py b/third_party/sglang/python/sglang/srt/managers/tokenizer_communicator_mixin.py index a1a3c1b..ccb2f5f 100644 --- a/third_party/sglang/python/sglang/srt/managers/tokenizer_communicator_mixin.py +++ b/third_party/sglang/python/sglang/srt/managers/tokenizer_communicator_mixin.py @@ -74,6 +74,12 @@ from sglang.srt.managers.io_struct import ( SetInternalStateReqOutput, SlowDownReqInput, SlowDownReqOutput, + SnapshotDumpReqInput, + SnapshotDumpReqOutput, + SnapshotFinalizeIngestReqInput, + SnapshotFinalizeIngestReqOutput, + SnapshotPrepareReceiveReqInput, + SnapshotPrepareReceiveReqOutput, UnloadLoRAAdapterReqInput, UnloadLoRAAdapterReqOutput, UpdateWeightsFromDistributedReqInput, @@ -225,6 +231,15 @@ class TokenizerCommunicatorMixin: self.direct_append_admission_communicator = _Communicator( self.send_to_scheduler, server_args.dp_size ) + self.snapshot_prepare_receive_communicator = _Communicator( + self.send_to_scheduler, server_args.dp_size + ) + self.snapshot_dump_communicator = _Communicator( + self.send_to_scheduler, server_args.dp_size + ) + self.snapshot_finalize_ingest_communicator = _Communicator( + self.send_to_scheduler, server_args.dp_size + ) self.set_internal_state_communicator = _Communicator( self.send_to_scheduler, server_args.dp_size ) @@ -325,6 +340,18 @@ class TokenizerCommunicatorMixin: DirectAppendAdmissionReqOutput, self.direct_append_admission_communicator.handle_recv, ), + ( + SnapshotPrepareReceiveReqOutput, + self.snapshot_prepare_receive_communicator.handle_recv, + ), + ( + SnapshotDumpReqOutput, + self.snapshot_dump_communicator.handle_recv, + ), + ( + SnapshotFinalizeIngestReqOutput, + self.snapshot_finalize_ingest_communicator.handle_recv, + ), ( SetInternalStateReqOutput, self.set_internal_state_communicator.handle_recv, @@ -890,6 +917,36 @@ class TokenizerCommunicatorMixin: ) return responses[0] + async def snapshot_prepare_receive( + self: TokenizerManager, + obj: SnapshotPrepareReceiveReqInput, + ) -> SnapshotPrepareReceiveReqOutput: + self.auto_create_handle_loop() + responses: List[SnapshotPrepareReceiveReqOutput] = ( + await self.snapshot_prepare_receive_communicator(obj) + ) + return responses[0] + + async def snapshot_dump( + self: TokenizerManager, + obj: SnapshotDumpReqInput, + ) -> SnapshotDumpReqOutput: + self.auto_create_handle_loop() + responses: List[SnapshotDumpReqOutput] = ( + await self.snapshot_dump_communicator(obj) + ) + return responses[0] + + async def snapshot_finalize_ingest( + self: TokenizerManager, + obj: SnapshotFinalizeIngestReqInput, + ) -> SnapshotFinalizeIngestReqOutput: + self.auto_create_handle_loop() + responses: List[SnapshotFinalizeIngestReqOutput] = ( + await self.snapshot_finalize_ingest_communicator(obj) + ) + return responses[0] + async def set_internal_state( self: TokenizerManager, obj: SetInternalStateReq ) -> List[bool]: