chore: import upstream snapshot with attribution

python/sglang/srt/disaggregation/ascend/__init__.py

@@ -0,0 +1,6 @@
+from sglang.srt.disaggregation.ascend.conn import (
+    AscendKVBootstrapServer,
+    AscendKVManager,
+    AscendKVReceiver,
+    AscendKVSender,
+)

python/sglang/srt/disaggregation/ascend/conn.py

@@ -0,0 +1,191 @@
+import concurrent.futures
+import logging
+from typing import List, Tuple
+
+import numpy as np
+import numpy.typing as npt
+
+from sglang.srt.disaggregation.ascend.transfer_engine import AscendTransferEngine
+from sglang.srt.disaggregation.common.utils import group_concurrent_contiguous
+from sglang.srt.disaggregation.mooncake.conn import (
+    MooncakeKVBootstrapServer,
+    MooncakeKVManager,
+    MooncakeKVReceiver,
+    MooncakeKVSender,
+)
+from sglang.srt.utils.network import get_local_ip_auto
+
+logger = logging.getLogger(__name__)
+
+
+class AscendKVManager(MooncakeKVManager):
+    def init_engine(self):
+        # TransferEngine initialized on ascend.
+        local_ip = get_local_ip_auto()
+        self.engine = AscendTransferEngine(
+            hostname=local_ip,
+            npu_id=self.kv_args.gpu_id,
+            disaggregation_mode=self.disaggregation_mode,
+        )
+
+    def register_buffer_to_engine(self):
+        self.engine.batch_register(self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens)
+        # The Ascend backend optimize batch registration for small memory blocks.
+        self.engine.batch_register(
+            self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens
+        )
+        # Batch register state/extra pool data buffers
+        for component_ptrs, component_lens in zip(
+            self.kv_args.state_data_ptrs or [],
+            self.kv_args.state_data_lens or [],
+        ):
+            self.engine.batch_register(component_ptrs, component_lens)
+
+    def get_mla_kv_ptrs_with_pp(
+        self, src_kv_ptrs: List[int], dst_kv_ptrs: List[int]
+    ) -> Tuple[List[int], List[int], int]:
+        # src_kv_ptrs: k_data, v_data, index_k_data(optional)
+        # dst_kv_ptrs: k_data, v_data, index_k_data(optional)
+        start_layer = self.kv_args.prefill_start_layer
+        kv_buf_groups = getattr(self.kv_args, "kv_buf_groups", 1)
+        total_kv_layers = getattr(self.kv_args, "total_kv_layers", 0)
+        src_layers = len(src_kv_ptrs) // kv_buf_groups
+        # When only speculative-algorithm is enabled for decode
+        # the KV has one more layer than prefill.
+        # The draft layer needs to be skipped.
+        dst_total_layers = (
+            min(len(dst_kv_ptrs) // kv_buf_groups, total_kv_layers)
+            if total_kv_layers
+            else len(dst_kv_ptrs) // kv_buf_groups
+        )
+        end_layer = start_layer + src_layers
+        if src_layers == dst_total_layers:
+            sliced_dst_kv_ptrs = dst_kv_ptrs
+        else:
+            sliced_dst_kv_ptrs = []
+            for i in range(kv_buf_groups):
+                layer_offset = i * dst_total_layers
+                sliced_dst_kv_ptrs.extend(
+                    dst_kv_ptrs[layer_offset + start_layer : layer_offset + end_layer]
+                )
+        layers_current_pp_stage = len(src_kv_ptrs)
+        return src_kv_ptrs, sliced_dst_kv_ptrs, layers_current_pp_stage
+
+    def send_kvcache(
+        self,
+        mooncake_session_id: str,
+        prefill_kv_indices: npt.NDArray[np.int32],
+        dst_kv_ptrs: list[int],
+        dst_kv_indices: npt.NDArray[np.int32],
+        executor: concurrent.futures.ThreadPoolExecutor,
+    ):
+        # Group by indices
+        prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous(
+            prefill_kv_indices, dst_kv_indices
+        )
+
+        if self.pp_size > 1:
+            if self.is_mla_backend:
+                src_kv_ptrs, sliced_dst_kv_ptrs, layers_current_pp_stage = (
+                    self.get_mla_kv_ptrs_with_pp(self.kv_args.kv_data_ptrs, dst_kv_ptrs)
+                )
+                layers_params = [
+                    (
+                        src_kv_ptrs[layer_id],
+                        sliced_dst_kv_ptrs[layer_id],
+                        self.kv_args.kv_item_lens[layer_id],
+                    )
+                    for layer_id in range(layers_current_pp_stage)
+                ]
+            else:
+                (
+                    src_k_ptrs,
+                    src_v_ptrs,
+                    dst_k_ptrs,
+                    dst_v_ptrs,
+                    layers_current_pp_stage,
+                ) = self.get_mha_kv_ptrs_with_pp(self.kv_args.kv_data_ptrs, dst_kv_ptrs)
+
+                layers_params = [
+                    (
+                        src_k_ptrs[layer_id],
+                        dst_k_ptrs[layer_id],
+                        self.kv_args.kv_item_lens[layer_id],
+                    )
+                    for layer_id in range(layers_current_pp_stage)
+                ] + [
+                    (
+                        src_v_ptrs[layer_id],
+                        dst_v_ptrs[layer_id],
+                        self.kv_args.kv_item_lens[layers_current_pp_stage + layer_id],
+                    )
+                    for layer_id in range(layers_current_pp_stage)
+                ]
+        else:
+            num_layers = len(self.kv_args.kv_data_ptrs)
+            layers_params = [
+                (
+                    self.kv_args.kv_data_ptrs[layer_id],
+                    dst_kv_ptrs[layer_id],
+                    self.kv_args.kv_item_lens[layer_id],
+                )
+                for layer_id in range(num_layers)
+            ]
+
+        def set_transfer_blocks(
+            src_ptr: int, dst_ptr: int, item_len: int
+        ) -> List[Tuple[int, int, int]]:
+            transfer_blocks = []
+            for prefill_index, decode_index in zip(prefill_kv_blocks, dst_kv_blocks):
+                src_addr = src_ptr + int(prefill_index[0]) * item_len
+                dst_addr = dst_ptr + int(decode_index[0]) * item_len
+                length = item_len * len(prefill_index)
+                transfer_blocks.append((src_addr, dst_addr, length))
+            return transfer_blocks
+
+        # Worker function for processing a single layer
+        def process_layer(src_ptr: int, dst_ptr: int, item_len: int) -> int:
+            transfer_blocks = set_transfer_blocks(src_ptr, dst_ptr, item_len)
+            return self._transfer_data(mooncake_session_id, transfer_blocks)
+
+        # Worker function for processing all layers in a batch
+        def process_layers(layers_params: List[Tuple[int, int, int]]) -> int:
+            transfer_blocks = []
+            for src_ptr, dst_ptr, item_len in layers_params:
+                transfer_blocks.extend(set_transfer_blocks(src_ptr, dst_ptr, item_len))
+            return self._transfer_data(mooncake_session_id, transfer_blocks)
+
+        if self.enable_custom_mem_pool:
+            futures = [
+                executor.submit(
+                    process_layer,
+                    src_ptr,
+                    dst_ptr,
+                    item_len,
+                )
+                for (src_ptr, dst_ptr, item_len) in layers_params
+            ]
+            for future in concurrent.futures.as_completed(futures):
+                status = future.result()
+                if status != 0:
+                    for f in futures:
+                        f.cancel()
+                    return status
+        else:
+            # Combining all layers' params in one batch transfer is more efficient
+            # compared to using multiple threads
+            return process_layers(layers_params)
+
+        return 0
+
+
+class AscendKVSender(MooncakeKVSender):
+    pass
+
+
+class AscendKVReceiver(MooncakeKVReceiver):
+    pass
+
+
+class AscendKVBootstrapServer(MooncakeKVBootstrapServer):
+    pass

python/sglang/srt/disaggregation/ascend/transfer_engine.py

@@ -0,0 +1,103 @@
+import logging
+import os
+from typing import List
+
+import torch
+
+from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
+    MooncakeTransferEngine,
+)
+from sglang.srt.utils.network import NetworkAddress
+
+try:
+    from memfabric_hybrid import TransferEngine
+
+    import_error = None
+except ImportError as e:
+    import_error = e
+    pass
+
+logger = logging.getLogger(__name__)
+
+
+class AscendTransferEngine(MooncakeTransferEngine):
+
+    def __init__(
+        self,
+        hostname: str,
+        npu_id: int,
+        disaggregation_mode: DisaggregationMode,
+    ):
+        if import_error is not None:
+            logger.warning(
+                "Please install memfabric_hybrid, for details, see docs/backend/pd_disaggregation.md"
+            )
+            raise import_error
+
+        self.engine = TransferEngine()
+        self.hostname = hostname
+        self.npu_id = npu_id
+
+        # Centralized storage address of the AscendTransferEngine
+        self.store_url = os.getenv("ASCEND_MF_STORE_URL")
+        if disaggregation_mode == DisaggregationMode.PREFILL:
+            self.role = "Prefill"
+        elif disaggregation_mode == DisaggregationMode.DECODE:
+            self.role = "Decode"
+        else:
+            logger.error(f"Unsupported DisaggregationMode: {disaggregation_mode}")
+            raise ValueError(f"Unsupported DisaggregationMode: {disaggregation_mode}")
+        self.session_id = NetworkAddress(
+            self.hostname, self.engine.get_rpc_port()
+        ).to_host_port_str()
+        self.initialize()
+
+    def initialize(self) -> None:
+        from sglang.srt.distributed.parallel_state import (
+            get_world_group,
+            get_world_size,
+        )
+
+        transfer_protocol = self._get_transfer_protocol()
+        if transfer_protocol is None or transfer_protocol == "sdma":
+            trans_op_type = TransferEngine.TransDataOpType.SDMA
+        else:
+            trans_op_type = TransferEngine.TransDataOpType.DEVICE_RDMA
+            """with device RDMA for PD transfer"""
+            tmp_tensor = torch.zeros(1, device="npu")
+            output_tensor_list = [
+                torch.empty_like(tmp_tensor) for _ in range(get_world_size())
+            ]
+            # Initialize hccl in advance through all_gather to avoid conflicts with rdma initialization.
+            torch.distributed.all_gather(
+                output_tensor_list, tmp_tensor, group=get_world_group().device_group
+            )
+        """Initialize the ascend transfer instance."""
+        ret_value = self.engine.initialize(
+            self.store_url, self.session_id, self.role, self.npu_id, trans_op_type
+        )
+        if ret_value != 0:
+            logger.error("Ascend Transfer Engine initialization failed.")
+            raise RuntimeError("Ascend Transfer Engine initialization failed.")
+
+    def batch_register(self, ptrs: List[int], lengths: List[int]):
+        try:
+            ret_value = self.engine.batch_register_memory(ptrs, lengths)
+        except Exception:
+            # Mark register as failed
+            ret_value = -1
+        if ret_value != 0:
+            logger.debug(f"Ascend memory registration for ptr {ptrs} failed.")
+
+    @staticmethod
+    def _get_transfer_protocol():
+        protocol = os.getenv("ASCEND_MF_TRANSFER_PROTOCOL")
+        allowed_protocols = {"device_rdma", "sdma"}
+        if protocol and protocol.lower() in allowed_protocols:
+            return protocol.lower()
+        else:
+            logger.warning(
+                "Invalid or no transfer protocol specified, using default protocol."
+            )
+            return None

python/sglang/srt/disaggregation/base/__init__.py

@@ -0,0 +1,8 @@
+from sglang.srt.disaggregation.base.conn import (
+    BaseKVBootstrapServer,
+    BaseKVManager,
+    BaseKVReceiver,
+    BaseKVSender,
+    KVArgs,
+    KVPoll,
+)

python/sglang/srt/disaggregation/base/conn.py

@@ -0,0 +1,223 @@
+from __future__ import annotations
+
+import dataclasses
+import enum
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, List, Optional
+
+import numpy as np
+import numpy.typing as npt
+
+from sglang.srt.server_args import ServerArgs
+
+if TYPE_CHECKING:
+    from sglang.srt.disaggregation.utils import DisaggregationMode
+
+
+class StateType(str, enum.Enum):
+    MAMBA = "mamba"
+    SWA = "swa"
+    DSA = "dsa"
+    MINIMAX_INDEX_K = "minimax_index_k"
+    # DeepSeek-V4 unified_kv SWA ring: addressed per-row by ring slot
+    # (req_pool_idx * ring_stride + pos % ring_stride), needs its own component.
+    SWA_RING = "swa_ring"
+    # DeepSeek-V4 online C128 request-scoped state.
+    C128_STATE = "c128_state"
+
+
+@dataclasses.dataclass
+class KVTransferMetric:
+    # Backends that cannot isolate transfer latency can leave this as None.
+    transfer_latency_s: Optional[float] = None
+    # Backends that cannot isolate allocation wait latency can leave this as None.
+    alloc_latency_s: Optional[float] = None
+    transfer_total_bytes: Optional[int] = None
+
+
+class KVArgs:
+    engine_rank: int
+    kv_data_ptrs: List[int]
+    kv_data_lens: List[int]
+    kv_item_lens: List[int]
+    aux_data_ptrs: List[int]
+    aux_data_lens: List[int]
+    aux_item_lens: List[int]
+    state_types: List[StateType]
+    state_data_ptrs: List[List[int]]
+    state_data_lens: List[List[int]]
+    state_item_lens: List[List[int]]
+    # Per-tensor TP slice dim, used when prefill/decode attn_tp_size differ.
+    state_dim_per_tensor: List[List[int]]
+    is_hybrid_mla_backend: bool
+    ib_device: str
+    ib_traffic_class: str
+    gpu_id: int
+    kv_head_num: int
+    total_kv_head_num: int
+    page_size: int
+    # for system dp
+    system_dp_rank: int
+    # for pp prefill
+    pp_rank: int
+    prefill_start_layer: int
+    # Absolute end layer (exclusive) for this prefill PP stage. Needed to
+    # reconstruct PP sub-ranges when kv_data_ptrs does not use a flat
+    # layer-indexed layout (e.g. DeepSeek V4's buffer-type-organized flat
+    # list).
+    prefill_end_layer: Optional[int]
+    # For DeepSeek V4 (and other compressed-MLA) memory pools only.
+    # Full-model compression ratio per layer (entries are 0/4/128). Used by
+    # the connection layer to slice the buffer-type-organized flat list in a
+    # PP-aware manner.
+    mla_compression_ratios: Optional[List[int]]
+    # Only used of npu, for kv buf groups
+    kv_buf_groups: int
+    # Only used of npu, for decode total kv layers
+    total_kv_layers: int
+
+
+class KVPoll:
+    Failed = 0
+    Bootstrapping = 1
+    WaitingForInput = 2
+    Transferring = 3
+    Success = 4
+
+
+class BaseKVManager(ABC):
+    """Base class for managing transfer states"""
+
+    @abstractmethod
+    def __init__(
+        self,
+        args: KVArgs,
+        disaggregation_mode: DisaggregationMode,
+        server_args: ServerArgs,
+        is_mla_backend: Optional[bool] = False,
+    ): ...
+
+    @abstractmethod
+    def register_to_bootstrap(self):
+        """Register prefill server info to the bootstrap server."""
+        ...
+
+
+class BaseKVSender(ABC):
+    @abstractmethod
+    def __init__(
+        self,
+        mgr: BaseKVManager,
+        bootstrap_addr: str,
+        bootstrap_room: int,
+        dest_tp_ranks: List[int],
+        pp_rank: int,
+        req_has_disagg_prefill_dp_rank: bool = False,
+    ): ...
+
+    @abstractmethod
+    def init(self, num_kv_indices: int, aux_index: Optional[int] = None):
+        """
+        Set req's index metadata locally or notify the decoder server about the kv indices length and aux index.
+        """
+        ...
+
+    @abstractmethod
+    def send(
+        self,
+        kv_indices: npt.NDArray[np.int32],
+        state_indices: Optional[List] = None,
+    ):
+        """
+        Send the kv cache at the given kv indices and the extra cache/state at the given indices to the decoder server.
+        """
+        ...
+
+    def pop_decode_prefix_len(self) -> int:
+        return 0
+
+    def should_send_kv_chunk(self, num_pages: int, last_chunk: bool) -> bool:
+        return num_pages > 0
+
+    @abstractmethod
+    def get_transfer_metric(self) -> KVTransferMetric:
+        """Return backend-specific transfer metrics for this sender."""
+        ...
+
+    @abstractmethod
+    def poll(self) -> KVPoll:
+        """
+        Check the status of the kv cache transfer.
+        """
+        ...
+
+    @abstractmethod
+    def failure_exception(self):
+        """
+        Raise an exception if the kv cache transfer fails.
+        """
+        ...
+
+
+class BaseKVReceiver(ABC):
+    @abstractmethod
+    def __init__(
+        self,
+        mgr: BaseKVManager,
+        bootstrap_addr: str,
+        bootstrap_room: Optional[int] = None,
+    ): ...
+
+    @abstractmethod
+    def init(
+        self,
+        prefill_dp_rank: int,
+    ):
+        """
+        Resolve bootstrap metadata and mark the receiver ready for transfer metadata.
+        """
+        ...
+
+    @abstractmethod
+    def send_metadata(
+        self,
+        kv_indices: npt.NDArray[np.int32],
+        aux_index: Optional[int] = None,
+        state_indices: Optional[List] = None,
+        decode_prefix_len: Optional[int] = None,
+    ):
+        """
+        Notify the prefill server about the kv indices, aux index, and state_indices.
+        """
+        ...
+
+    @abstractmethod
+    def poll(self) -> KVPoll:
+        """
+        Check the status of the kv cache transfer.
+        """
+        ...
+
+    @abstractmethod
+    def failure_exception(self):
+        """
+        Raise an exception if the kv cache transfer fails.
+        """
+        ...
+
+    def clear(self):
+        """
+        Clear any internal states.
+        """
+        pass
+
+    def abort(self):
+        """
+        Abort the current transfer.
+        """
+        pass
+
+
+class BaseKVBootstrapServer(ABC):
+    @abstractmethod
+    def __init__(self, host: str, port: int): ...

python/sglang/srt/disaggregation/common/__init__.py

@@ -0,0 +1,5 @@
+from sglang.srt.disaggregation.common.conn import (
+    CommonKVBootstrapServer,
+    CommonKVManager,
+    CommonKVReceiver,
+)

python/sglang/srt/disaggregation/common/staging_buffer.py

@@ -0,0 +1,768 @@
+"""
+GPU Staging Buffer for heterogeneous TP KV cache transfer.
+
+When prefill attn_tp_size != decode attn_tp_size, the per-token RDMA approach
+generates O(tokens * layers) small RDMA requests. This module provides a staging
+buffer mechanism that gathers scattered head slices into contiguous GPU memory,
+enabling bulk RDMA transfers that reduce request count to O(layers) or O(1).
+
+Usage:
+    Activated by setting SGLANG_DISAGG_STAGING_BUFFER=1.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import threading
+from typing import List, Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+logger = logging.getLogger(__name__)
+
+# TODO(yangminl): remove torch fallback implementations once the Triton kernels
+# have been validated in production across all configurations.
+_USE_TRITON_STAGING = not bool(os.environ.get("SGLANG_STAGING_USE_TORCH", ""))
+
+
+@triton.jit
+def _fused_gather_to_staging_kernel(
+    layer_ptrs,
+    page_indices,
+    staging,
+    num_tokens,
+    stride_pool_token,
+    head_offset,
+    per_layer_elems,
+    ELEMS_PER_TOKEN: tl.constexpr,
+    PAGE_SIZE: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    layer_id = tl.program_id(0)
+    block_id = tl.program_id(1)
+
+    layer_ptr = tl.load(layer_ptrs + layer_id).to(staging.dtype)
+
+    offsets = block_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < per_layer_elems
+
+    t_idx = offsets // ELEMS_PER_TOKEN
+    e_idx = offsets % ELEMS_PER_TOKEN
+
+    page_id = t_idx // PAGE_SIZE
+    intra_page = t_idx % PAGE_SIZE
+    page_val = tl.load(page_indices + page_id, mask=mask, other=0)
+    pool_token = page_val * PAGE_SIZE + intra_page
+
+    src_offsets = (
+        pool_token * stride_pool_token.to(tl.int64) + head_offset.to(tl.int64) + e_idx
+    )
+    vals = tl.load(layer_ptr + src_offsets, mask=mask)
+
+    dst_offsets = tl.program_id(0).to(tl.int64) * per_layer_elems.to(tl.int64) + offsets
+    tl.store(staging + dst_offsets, vals, mask=mask)
+
+
+@triton.jit
+def _fused_scatter_from_staging_kernel(
+    layer_ptrs,
+    page_indices,
+    staging,
+    writer_head_offsets,
+    num_tokens,
+    stride_pool_token,
+    per_layer_elems,
+    ELEMS_PER_TOKEN: tl.constexpr,
+    PAGE_SIZE: tl.constexpr,
+    NUM_LAYERS_X2: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    prog_id = tl.program_id(0)
+    block_id = tl.program_id(1)
+
+    writer_id = prog_id // NUM_LAYERS_X2
+    layer_kv_id = prog_id % NUM_LAYERS_X2
+
+    layer_ptr = tl.load(layer_ptrs + layer_kv_id).to(staging.dtype)
+    head_offset = tl.load(writer_head_offsets + writer_id)
+
+    offsets = block_id * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < per_layer_elems
+
+    t_idx = offsets // ELEMS_PER_TOKEN
+    e_idx = offsets % ELEMS_PER_TOKEN
+
+    page_id = t_idx // PAGE_SIZE
+    intra_page = t_idx % PAGE_SIZE
+    page_val = tl.load(page_indices + page_id, mask=mask, other=0)
+    pool_token = page_val * PAGE_SIZE + intra_page
+
+    per_rank_elems = per_layer_elems.to(tl.int64) * NUM_LAYERS_X2
+    src_offsets = (
+        writer_id.to(tl.int64) * per_rank_elems
+        + layer_kv_id.to(tl.int64) * per_layer_elems.to(tl.int64)
+        + offsets
+    )
+    vals = tl.load(staging + src_offsets, mask=mask)
+
+    dst_offsets = (
+        pool_token * stride_pool_token.to(tl.int64) + head_offset.to(tl.int64) + e_idx
+    )
+    tl.store(layer_ptr + dst_offsets, vals, mask=mask)
+
+
+class StagingBuffer:
+    """Pre-allocated GPU staging buffer for bulk KV transfer.
+
+    When a custom_mem_pool is provided (e.g., mooncake NVLink allocator),
+    the buffer is allocated within that pool so it's compatible with
+    NVLink/MNNVL transport (requires cuMemCreate-backed memory).
+    """
+
+    def __init__(
+        self,
+        size_bytes: int,
+        device: str,
+        gpu_id: int,
+        custom_mem_pool=None,
+    ):
+        self.size_bytes = size_bytes
+        self.device = device
+        self.gpu_id = gpu_id
+
+        torch.cuda.set_device(gpu_id)
+        if custom_mem_pool is not None:
+            with torch.cuda.use_mem_pool(custom_mem_pool):
+                self.buffer = torch.empty(size_bytes, dtype=torch.uint8, device=device)
+            alloc_method = "custom_mem_pool (cuMemCreate)"
+        else:
+            self.buffer = torch.empty(size_bytes, dtype=torch.uint8, device=device)
+            alloc_method = "cudaMalloc"
+        self.data_ptr = self.buffer.data_ptr()
+
+        logger.info(
+            f"StagingBuffer allocated: {size_bytes / (1024*1024):.1f} MB "
+            f"on {device}, method={alloc_method}, ptr=0x{self.data_ptr:x}"
+        )
+
+    def get_ptr(self) -> int:
+        return self.data_ptr
+
+    def get_size(self) -> int:
+        return self.size_bytes
+
+    def fits(self, required_bytes: int) -> bool:
+        return required_bytes <= self.size_bytes
+
+
+class StagingAllocator:
+    """Decode-side dynamic staging ring buffer allocator with overcommit.
+
+    One large pre-allocated GPU buffer used as a ring buffer. Each request
+    gets a (alloc_id, offset, round) triple based on its actual byte
+    requirement. Allocation (assign) is overcommit — it always succeeds
+    as long as the request fits in the buffer. Overlap safety is enforced
+    on the prefill side before RDMA, using a watermark that tracks the
+    oldest un-freed allocation.
+
+    The watermark (round, tail_offset) is periodically sent to prefill.
+    Prefill transfer workers wait before writing if their target region
+    overlaps with not-yet-freed data from a previous round.
+    """
+
+    # Permanent alloc failure: chunk exceeds ring buffer total size.
+    ALLOC_OVERSIZED = -2
+
+    def __init__(
+        self,
+        total_size_bytes: int,
+        device: str,
+        gpu_id: int,
+        custom_mem_pool=None,
+    ):
+        self.buffer = StagingBuffer(total_size_bytes, device, gpu_id, custom_mem_pool)
+        self.total_size = total_size_bytes
+        self.base_ptr = self.buffer.data_ptr
+        self.head = 0
+        self.round = 0
+        self.allocations: dict = {}  # alloc_id -> (offset, size, round)
+        self.alloc_order: List[int] = []
+        self.next_alloc_id = 0
+        self.watermark_round = 0
+        self.watermark_tail = 0
+        self.lock = threading.Lock()
+
+        logger.info(
+            f"StagingAllocator (ring+overcommit): "
+            f"{total_size_bytes / (1024*1024):.1f} MB "
+            f"on {device}, ptr=0x{self.base_ptr:x}"
+        )
+
+    def assign(self, required_bytes: int) -> Optional[Tuple[int, int, int]]:
+        """Allocate a region. Returns (alloc_id, offset, round) or None."""
+        with self.lock:
+            if required_bytes > self.total_size:
+                return None
+
+            space_at_end = self.total_size - self.head
+            if required_bytes <= space_at_end:
+                offset = self.head
+                self.head += required_bytes
+            else:
+                self.round += 1
+                offset = 0
+                self.head = required_bytes
+
+            alloc_id = self.next_alloc_id
+            self.next_alloc_id += 1
+            self.allocations[alloc_id] = (offset, required_bytes, self.round)
+            self.alloc_order.append(alloc_id)
+            return (alloc_id, offset, self.round)
+
+    def free(self, alloc_id: int):
+        """Free an allocation and advance watermark past consecutive freed entries."""
+        with self.lock:
+            if alloc_id not in self.allocations:
+                return
+            self.allocations.pop(alloc_id)
+
+            while self.alloc_order and self.alloc_order[0] not in self.allocations:
+                self.alloc_order.pop(0)
+
+            if not self.allocations:
+                self.watermark_round = self.round
+                self.watermark_tail = self.head
+            elif self.alloc_order:
+                off, _, rnd = self.allocations[self.alloc_order[0]]
+                self.watermark_round = rnd
+                self.watermark_tail = off
+
+    def get_watermark(self) -> Tuple[int, int]:
+        """Return (round, tail_offset). Everything before this is safe to write."""
+        with self.lock:
+            return (self.watermark_round, self.watermark_tail)
+
+    def get_ptr(self, alloc_id: int) -> int:
+        offset, _, _ = self.allocations[alloc_id]
+        return self.base_ptr + offset
+
+    def get_offset(self, alloc_id: int) -> int:
+        offset, _, _ = self.allocations[alloc_id]
+        return offset
+
+    def get_round(self, alloc_id: int) -> int:
+        _, _, rnd = self.allocations[alloc_id]
+        return rnd
+
+    def get_base_ptr(self) -> int:
+        return self.base_ptr
+
+    def get_total_size(self) -> int:
+        return self.total_size
+
+
+def gather_kv_head_slices(
+    kv_buffer_tensor: torch.Tensor,
+    gather_idx: torch.Tensor,
+    head_start: int,
+    num_heads: int,
+    staging_tensor: torch.Tensor,
+):
+    """Gather KV head slices from scattered pages into contiguous staging buffer.
+
+    Uses torch.gather(out=) to write directly into staging_tensor without
+    allocating temporary tensors (avoids CUDA caching allocator stalls).
+
+    Args:
+        kv_buffer_tensor: [pool_size, head_num, head_dim], one layer.
+        gather_idx: [num_tokens, num_heads, head_dim] int64, pre-computed
+            token indices expanded for gather on dim=0.
+        head_start: Starting head index for the slice.
+        num_heads: Number of heads to gather.
+        staging_tensor: Output tensor, shape [num_tokens, num_heads, head_dim].
+    """
+    src = kv_buffer_tensor[:, head_start : head_start + num_heads, :]
+    torch.gather(src, 0, gather_idx, out=staging_tensor)
+
+
+def scatter_kv_head_slices(
+    staging_tensor: torch.Tensor,
+    kv_buffer_tensor: torch.Tensor,
+    page_indices: torch.Tensor,
+    head_start: int,
+    num_heads: int,
+    page_size: int = 1,
+):
+    """Scatter KV head slices from contiguous staging buffer to KV cache.
+
+    Args:
+        staging_tensor: Input tensor from staging buffer (contiguous packed data).
+        kv_buffer_tensor: The KV buffer for one layer, shape [pool_size, head_num, head_dim].
+        page_indices: [num_pages] int32/int64 tensor of page indices.
+        head_start: Starting head index for the slice.
+        num_heads: Number of heads to scatter.
+        page_size: Number of tokens per page.
+    """
+    head_dim = kv_buffer_tensor.shape[-1]
+    if page_size == 1:
+        num_tokens = page_indices.shape[0]
+        data = staging_tensor.reshape(num_tokens, num_heads, head_dim)
+        kv_buffer_tensor[page_indices, head_start : head_start + num_heads, :] = data
+    else:
+        num_tokens = page_indices.shape[0] * page_size
+        offsets = torch.arange(page_size, device=page_indices.device)
+        token_indices = (page_indices.unsqueeze(1) * page_size + offsets).reshape(-1)
+        data = staging_tensor.reshape(num_tokens, num_heads, head_dim)
+        kv_buffer_tensor[token_indices, head_start : head_start + num_heads, :] = data
+
+
+def _gather_all_layers_torch(
+    k_buffers: list,
+    v_buffers: list,
+    page_indices_np,
+    staging_buffer: StagingBuffer,
+    src_head_start: int,
+    num_heads: int,
+    page_size: int,
+    gpu_id: int,
+) -> int:
+    """torch.gather path: zero per-layer allocation, one kernel per layer."""
+    import numpy as np
+
+    num_layers = len(k_buffers)
+    head_dim = k_buffers[0].shape[-1]
+    dtype_size = k_buffers[0].element_size()
+    num_tokens = len(page_indices_np) * page_size
+    per_layer_bytes = num_tokens * num_heads * head_dim * dtype_size
+
+    device = f"cuda:{gpu_id}"
+    torch.cuda.set_device(gpu_id)
+    page_idx_tensor = torch.from_numpy(page_indices_np.astype(np.int64)).to(device)
+
+    if page_size == 1:
+        token_indices = page_idx_tensor
+    else:
+        offsets = torch.arange(page_size, device=device)
+        token_indices = (page_idx_tensor.unsqueeze(1) * page_size + offsets).reshape(-1)
+
+    gather_idx = token_indices.view(-1, 1, 1).expand(num_tokens, num_heads, head_dim)
+
+    if not hasattr(staging_buffer, "_gather_stream"):
+        staging_buffer._gather_stream = torch.cuda.Stream(device=device)
+
+    staging_buffer._gather_stream.wait_stream(
+        torch.cuda.default_stream(torch.device(device))
+    )
+
+    staging_view = staging_buffer.buffer
+    offset = 0
+    with torch.cuda.stream(staging_buffer._gather_stream):
+        for layer_id in range(num_layers):
+            dst = (
+                staging_view[offset : offset + per_layer_bytes]
+                .view(k_buffers[layer_id].dtype)
+                .reshape(num_tokens, num_heads, head_dim)
+            )
+            gather_kv_head_slices(
+                k_buffers[layer_id],
+                gather_idx,
+                src_head_start,
+                num_heads,
+                dst,
+            )
+            offset += per_layer_bytes
+        for layer_id in range(num_layers):
+            dst = (
+                staging_view[offset : offset + per_layer_bytes]
+                .view(v_buffers[layer_id].dtype)
+                .reshape(num_tokens, num_heads, head_dim)
+            )
+            gather_kv_head_slices(
+                v_buffers[layer_id],
+                gather_idx,
+                src_head_start,
+                num_heads,
+                dst,
+            )
+            offset += per_layer_bytes
+
+    staging_buffer._gather_stream.synchronize()
+    return offset
+
+
+def _gather_all_layers_triton(
+    k_buffers: list,
+    v_buffers: list,
+    page_indices_np,
+    staging_buffer: StagingBuffer,
+    src_head_start: int,
+    num_heads: int,
+    page_size: int,
+    gpu_id: int,
+) -> int:
+    """Triton fused kernel path: single kernel launch for all layers."""
+    import numpy as np
+
+    num_layers = len(k_buffers)
+    head_dim = k_buffers[0].shape[-1]
+    total_heads = k_buffers[0].shape[1]
+    dtype_size = k_buffers[0].element_size()
+    num_tokens = len(page_indices_np) * page_size
+    elems_per_token = num_heads * head_dim
+    per_layer_elems = num_tokens * elems_per_token
+    per_layer_bytes = per_layer_elems * dtype_size
+    total_bytes = per_layer_bytes * num_layers * 2
+
+    device = f"cuda:{gpu_id}"
+    torch.cuda.set_device(gpu_id)
+    page_idx_tensor = torch.from_numpy(page_indices_np.astype(np.int64)).to(device)
+
+    layer_ptrs = torch.tensor(
+        [buf.data_ptr() for buf in k_buffers] + [buf.data_ptr() for buf in v_buffers],
+        dtype=torch.int64,
+        device=device,
+    )
+    # Use integer dtype matching element size for bit-preserving copy
+    int_dtype_map = {1: torch.int8, 2: torch.int16, 4: torch.int32}
+    int_dtype = int_dtype_map.get(dtype_size, torch.int16)
+    staging_typed = staging_buffer.buffer[:total_bytes].view(int_dtype)
+
+    if not hasattr(staging_buffer, "_gather_stream"):
+        staging_buffer._gather_stream = torch.cuda.Stream(device=device)
+
+    staging_buffer._gather_stream.wait_stream(
+        torch.cuda.default_stream(torch.device(device))
+    )
+
+    BLOCK_SIZE = 1024
+    grid = (2 * num_layers, triton.cdiv(per_layer_elems, BLOCK_SIZE))
+
+    with torch.cuda.stream(staging_buffer._gather_stream):
+        _fused_gather_to_staging_kernel[grid](
+            layer_ptrs,
+            page_idx_tensor,
+            staging_typed,
+            num_tokens,
+            total_heads * head_dim,
+            src_head_start * head_dim,
+            per_layer_elems,
+            elems_per_token,
+            page_size,
+            BLOCK_SIZE,
+        )
+
+    staging_buffer._gather_stream.synchronize()
+    return total_bytes
+
+
+def gather_all_layers_to_staging(
+    k_buffers: list,
+    v_buffers: list,
+    page_indices_np,
+    staging_buffer: StagingBuffer,
+    src_head_start: int,
+    num_heads: int,
+    page_size: int,
+    gpu_id: int,
+) -> int:
+    """Gather all layers' K and V head slices into a staging buffer.
+
+    Returns total bytes written.
+    Dispatches to Triton fused kernel when available, falls back to torch.gather.
+    """
+    if _USE_TRITON_STAGING:
+        return _gather_all_layers_triton(
+            k_buffers,
+            v_buffers,
+            page_indices_np,
+            staging_buffer,
+            src_head_start,
+            num_heads,
+            page_size,
+            gpu_id,
+        )
+    return _gather_all_layers_torch(
+        k_buffers,
+        v_buffers,
+        page_indices_np,
+        staging_buffer,
+        src_head_start,
+        num_heads,
+        page_size,
+        gpu_id,
+    )
+
+
+def _scatter_staging_to_kv_torch(
+    staging_buffer_view: torch.Tensor,
+    k_buffers: list,
+    v_buffers: list,
+    page_idx_tensor: torch.Tensor,
+    page_size: int,
+    prefill_attn_tp_size: int,
+    decode_attn_tp_size: int,
+    dst_tp_rank: int,
+    total_kv_heads: int,
+) -> None:
+    """torch path for scatter."""
+    num_layers = len(k_buffers)
+    head_dim = k_buffers[0].shape[-1]
+    dtype_size = k_buffers[0].element_size()
+    num_tokens = page_idx_tensor.shape[0] * page_size
+
+    if prefill_attn_tp_size > decode_attn_tp_size:
+        num_writers = prefill_attn_tp_size // max(1, decode_attn_tp_size)
+    else:
+        num_writers = 1
+
+    for writer_rank in range(num_writers):
+        _, num_heads, dst_head_start, _ = compute_head_slice_params(
+            prefill_attn_tp_size,
+            decode_attn_tp_size,
+            writer_rank,
+            dst_tp_rank,
+            total_kv_heads,
+        )
+        per_layer_bytes = num_tokens * num_heads * head_dim * dtype_size
+        per_rank_bytes = per_layer_bytes * num_layers * 2
+        rank_base = writer_rank * per_rank_bytes
+
+        offset = rank_base
+        for layer_id in range(num_layers):
+            layer_data = (
+                staging_buffer_view[offset : offset + per_layer_bytes]
+                .view(k_buffers[layer_id].dtype)
+                .reshape(num_tokens, num_heads, head_dim)
+            )
+            scatter_kv_head_slices(
+                layer_data,
+                k_buffers[layer_id],
+                page_idx_tensor,
+                dst_head_start,
+                num_heads,
+                page_size,
+            )
+            offset += per_layer_bytes
+        for layer_id in range(num_layers):
+            layer_data = (
+                staging_buffer_view[offset : offset + per_layer_bytes]
+                .view(v_buffers[layer_id].dtype)
+                .reshape(num_tokens, num_heads, head_dim)
+            )
+            scatter_kv_head_slices(
+                layer_data,
+                v_buffers[layer_id],
+                page_idx_tensor,
+                dst_head_start,
+                num_heads,
+                page_size,
+            )
+            offset += per_layer_bytes
+
+
+def _scatter_staging_to_kv_triton(
+    staging_buffer_view: torch.Tensor,
+    k_buffers: list,
+    v_buffers: list,
+    page_idx_tensor: torch.Tensor,
+    page_size: int,
+    prefill_attn_tp_size: int,
+    decode_attn_tp_size: int,
+    dst_tp_rank: int,
+    total_kv_heads: int,
+) -> None:
+    """Triton fused kernel path for scatter."""
+    num_layers = len(k_buffers)
+    head_dim = k_buffers[0].shape[-1]
+    total_heads = k_buffers[0].shape[1]
+    dtype_size = k_buffers[0].element_size()
+    num_tokens = page_idx_tensor.shape[0] * page_size
+    device = page_idx_tensor.device
+
+    if prefill_attn_tp_size > decode_attn_tp_size:
+        num_writers = prefill_attn_tp_size // max(1, decode_attn_tp_size)
+    else:
+        num_writers = 1
+
+    # All writers share the same num_heads; only dst_head_start differs
+    _, num_heads, _, _ = compute_head_slice_params(
+        prefill_attn_tp_size,
+        decode_attn_tp_size,
+        0,
+        dst_tp_rank,
+        total_kv_heads,
+    )
+    elems_per_token = num_heads * head_dim
+    per_layer_elems = num_tokens * elems_per_token
+
+    layer_ptrs = torch.tensor(
+        [buf.data_ptr() for buf in k_buffers] + [buf.data_ptr() for buf in v_buffers],
+        dtype=torch.int64,
+        device=device,
+    )
+
+    writer_head_offsets = torch.tensor(
+        [
+            compute_head_slice_params(
+                prefill_attn_tp_size,
+                decode_attn_tp_size,
+                wr,
+                dst_tp_rank,
+                total_kv_heads,
+            )[2]
+            * head_dim
+            for wr in range(num_writers)
+        ],
+        dtype=torch.int64,
+        device=device,
+    )
+
+    int_dtype_map = {1: torch.int8, 2: torch.int16, 4: torch.int32}
+    int_dtype = int_dtype_map.get(dtype_size, torch.int16)
+    total_staging_bytes = (
+        num_tokens * elems_per_token * dtype_size * num_layers * 2 * num_writers
+    )
+    staging_typed = staging_buffer_view[:total_staging_bytes].view(int_dtype)
+
+    BLOCK_SIZE = 1024
+    num_layers_x2 = 2 * num_layers
+    grid = (num_writers * num_layers_x2, triton.cdiv(per_layer_elems, BLOCK_SIZE))
+
+    _fused_scatter_from_staging_kernel[grid](
+        layer_ptrs,
+        page_idx_tensor,
+        staging_typed,
+        writer_head_offsets,
+        num_tokens,
+        total_heads * head_dim,
+        per_layer_elems,
+        elems_per_token,
+        page_size,
+        num_layers_x2,
+        BLOCK_SIZE,
+    )
+
+
+def scatter_staging_to_kv(
+    staging_buffer_view: torch.Tensor,
+    k_buffers: list,
+    v_buffers: list,
+    page_idx_tensor: torch.Tensor,
+    page_size: int,
+    prefill_attn_tp_size: int,
+    decode_attn_tp_size: int,
+    dst_tp_rank: int,
+    total_kv_heads: int,
+) -> None:
+    """Scatter data from a contiguous staging region into KV cache buffers."""
+    if _USE_TRITON_STAGING:
+        return _scatter_staging_to_kv_triton(
+            staging_buffer_view,
+            k_buffers,
+            v_buffers,
+            page_idx_tensor,
+            page_size,
+            prefill_attn_tp_size,
+            decode_attn_tp_size,
+            dst_tp_rank,
+            total_kv_heads,
+        )
+    return _scatter_staging_to_kv_torch(
+        staging_buffer_view,
+        k_buffers,
+        v_buffers,
+        page_idx_tensor,
+        page_size,
+        prefill_attn_tp_size,
+        decode_attn_tp_size,
+        dst_tp_rank,
+        total_kv_heads,
+    )
+
+
+def compute_head_slice_params(
+    src_attn_tp_size: int,
+    dst_attn_tp_size: int,
+    src_tp_rank: int,
+    dst_tp_rank: int,
+    total_kv_heads: int,
+) -> Tuple[int, int, int, int]:
+    """Compute head slicing parameters for heterogeneous TP transfer.
+
+    Returns:
+        (src_head_start, num_heads_to_send, dst_head_start, num_heads_to_send)
+    """
+    src_heads_per_rank = max(1, total_kv_heads // src_attn_tp_size)
+    dst_heads_per_rank = max(1, total_kv_heads // dst_attn_tp_size)
+
+    local_tp_rank = src_tp_rank % src_attn_tp_size
+    dst_tp_rank_in_group = dst_tp_rank % dst_attn_tp_size
+
+    if src_attn_tp_size > dst_attn_tp_size:
+        src_head_start = 0
+        num_heads_to_send = src_heads_per_rank
+        src_replication = max(1, src_attn_tp_size // total_kv_heads)
+        unique_head_idx = local_tp_rank // src_replication
+        dst_head_start = (unique_head_idx * src_heads_per_rank) % dst_heads_per_rank
+    else:
+        src_head_start = (
+            dst_tp_rank_in_group * dst_heads_per_rank
+        ) % src_heads_per_rank
+        num_heads_to_send = dst_heads_per_rank
+        dst_head_start = 0
+
+    return src_head_start, num_heads_to_send, dst_head_start, num_heads_to_send
+
+
+def compute_staging_layout(
+    src_attn_tp_size: int,
+    dst_attn_tp_size: int,
+    dst_tp_rank: int,
+    total_kv_heads: int,
+    num_tokens: int,
+    bytes_per_head_token: int,
+    num_layers: int,
+) -> Tuple[int, List[int], int]:
+    """Compute per-writer byte layout for a staging region.
+
+    Returns:
+        (num_writers, writer_bytes_list, total_bytes)
+        where writer_bytes_list[i] = bytes for writer i covering all layers (K+V).
+    """
+    if src_attn_tp_size > dst_attn_tp_size:
+        num_writers = src_attn_tp_size // max(1, dst_attn_tp_size)
+    else:
+        num_writers = 1
+
+    writer_bytes = []
+    for wr in range(num_writers):
+        _, nh, _, _ = compute_head_slice_params(
+            src_attn_tp_size,
+            dst_attn_tp_size,
+            wr,
+            dst_tp_rank,
+            total_kv_heads,
+        )
+        writer_bytes.append(num_tokens * nh * bytes_per_head_token * num_layers * 2)
+    return num_writers, writer_bytes, sum(writer_bytes)
+
+
+def resolve_total_kv_heads(
+    kv_args,
+    attn_tp_size: int,
+) -> int:
+    """Resolve the global total KV head count from kv_args metadata."""
+    total = getattr(kv_args, "total_kv_head_num", 0)
+    if total > 0:
+        return total
+    per_rank = getattr(kv_args, "kv_head_num", 0)
+    if per_rank > 0:
+        return per_rank * attn_tp_size
+    raise ValueError(
+        "Cannot resolve total_kv_heads: kv_args has neither total_kv_head_num "
+        "nor kv_head_num. "
+        "Ensure DecodePreallocQueue._init_kv_manager sets kv_args.kv_head_num."
+    )

python/sglang/srt/disaggregation/common/staging_handler.py

@@ -0,0 +1,840 @@
+"""
+Staging handler for heterogeneous TP KV cache transfer.
+
+Isolates staging scatter lifecycle from decode.py and conn.py.
+Generic (backend-agnostic) code is at the top; mooncake-specific
+protocol code is at the bottom.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import logging
+import struct
+import threading
+from typing import TYPE_CHECKING, List, Optional, Tuple
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from sglang.srt.disaggregation.decode import DecodeRequest
+
+
+# ======================================================================
+# Generic staging state and handler (backend-agnostic)
+# ======================================================================
+
+
+@dataclasses.dataclass
+class DecodeStagingContext:
+    """Staging-specific context for decode mode."""
+
+    allocator: object = None
+    room_bootstrap: dict = dataclasses.field(default_factory=dict)
+    room_receivers: dict = dataclasses.field(default_factory=dict)
+
+
+@dataclasses.dataclass
+class PrefillStagingContext:
+    """Staging-specific context for prefill mode."""
+
+    buffers: list = dataclasses.field(default_factory=list)
+    remote_watermarks: dict = dataclasses.field(default_factory=dict)
+    watermark_cv: threading.Condition = dataclasses.field(
+        default_factory=threading.Condition
+    )
+    # (room, chunk_idx, session_id) keys for chunks already requested.
+    prefetch_requested: set = dataclasses.field(default_factory=set)
+    # Rooms that have already had their full prefetch fan-out triggered. Used
+    # to short-circuit per-room prefetch entry on every chunk after the first.
+    prefetched_rooms: set = dataclasses.field(default_factory=set)
+    prefetch_sockets: dict = dataclasses.field(default_factory=dict)
+
+
+class DecodeStagingHandler:
+    """Decode-side staging scatter lifecycle manager.
+
+    Scatter submission can be called from the decode_thread (background) as
+    soon as all writers/ranks have arrived, while event checking and freeing
+    always run on the scheduler main thread.
+    """
+
+    def __init__(
+        self,
+        kv_manager,
+        staging_allocator,
+        kv_buffer_info: dict,
+        decode_tp: int,
+        total_kv_heads: int,
+        tp_rank: int,
+        scheduler,
+    ):
+        self.kv_manager = kv_manager
+        self.staging_allocator = staging_allocator
+        self.kv_buffer_info = kv_buffer_info
+        self.decode_tp = decode_tp
+        self.total_kv_heads = total_kv_heads
+        self.tp_rank = tp_rank
+        self.scheduler = scheduler
+        self._room_to_decode_req: dict = {}
+        self._wm_subscribers: dict = {}
+
+    def register_wm_subscriber(self, receiver, session_id: str) -> None:
+        """Register a prefill's bootstrap connection for watermark broadcasts."""
+        if receiver is None or not getattr(receiver, "bootstrap_infos", None):
+            return
+        key = tuple(str(bi) for bi in receiver.bootstrap_infos)
+        if key not in self._wm_subscribers:
+            self._wm_subscribers[key] = (receiver, session_id)
+
+    def num_writers_for(self, decode_req) -> int:
+        """Compute num_writers for a specific request based on its prefill TP."""
+        prefill_tp = decode_req.kv_receiver.prefill_info.attn_tp_size
+        if prefill_tp > self.decode_tp:
+            return prefill_tp // max(1, self.decode_tp)
+        return 1
+
+    @classmethod
+    def create(cls, kv_manager, scheduler, tp_rank: int) -> DecodeStagingHandler:
+        """Factory: create handler. Raises if staging infra is missing."""
+        staging_allocator = kv_manager._staging_ctx.allocator
+        if staging_allocator is None:
+            raise RuntimeError(
+                "Staging is enabled but kv_manager._staging_ctx.allocator is None. "
+                "Check that the transfer backend correctly initializes the staging allocator."
+            )
+        kv_buffer_info = kv_manager.kv_buffer_tensors
+        if kv_buffer_info is None:
+            raise RuntimeError(
+                "Staging is enabled but kv_manager.kv_buffer_tensors is None. "
+                "Check that set_kv_buffer_tensors() was called during kv_manager init."
+            )
+        decode_tp = kv_manager.attn_tp_size
+
+        from sglang.srt.disaggregation.common.staging_buffer import (
+            resolve_total_kv_heads,
+        )
+
+        total_kv_heads = resolve_total_kv_heads(kv_manager.kv_args, decode_tp)
+        return cls(
+            kv_manager=kv_manager,
+            staging_allocator=staging_allocator,
+            kv_buffer_info=kv_buffer_info,
+            decode_tp=decode_tp,
+            total_kv_heads=total_kv_heads,
+            tp_rank=tp_rank,
+            scheduler=scheduler,
+        )
+
+    # ------------------------------------------------------------------
+    # Registration: called from main thread (DecodeTransferQueue)
+    # ------------------------------------------------------------------
+
+    def register_decode_req(self, room: int, decode_req: DecodeRequest) -> None:
+        self._room_to_decode_req[room] = decode_req
+
+    def unregister_decode_req(self, room: int) -> None:
+        self._room_to_decode_req.pop(room, None)
+
+    # ------------------------------------------------------------------
+    # Scatter submission: called from decode_thread (background)
+    # ------------------------------------------------------------------
+
+    def submit_chunk_scatter(
+        self, room: int, chunk_idx: int, page_start: int, num_pages: int
+    ) -> bool:
+        """Submit scatter for an intermediate chunk whose writers all arrived.
+
+        Called from decode_thread.  Records a CUDA event on decode_req so
+        the main thread can later check completion and free the allocation.
+        """
+        decode_req = self._room_to_decode_req.get(room)
+        if decode_req is None:
+            logger.warning(
+                "[STAGING] submit_chunk_scatter: room=%s not registered, "
+                "chunk_idx=%s. This should not happen if register_decode_req "
+                "is called at kv_receiver.init() time.",
+                room,
+                chunk_idx,
+            )
+            return False
+        chunk_infos = getattr(decode_req.kv_receiver, "chunk_staging_infos", [])
+        if chunk_idx >= len(chunk_infos):
+            return False
+        alloc_id, staging_offset, _, _, _ = chunk_infos[chunk_idx]
+        if staging_offset < 0 or alloc_id < 0:
+            return False
+
+        ok = self._scatter_region(staging_offset, page_start, num_pages, decode_req)
+        if ok:
+            event = torch.cuda.Event()
+            event.record(self.staging_allocator._scatter_stream)
+            if not hasattr(decode_req, "_chunk_events"):
+                decode_req._chunk_events = []
+            decode_req._chunk_events.append((event, alloc_id))
+            chunk_infos[chunk_idx] = (-1, -1, 0, -1, 0)
+        else:
+            logger.warning(
+                "submit_chunk_scatter failed room=%s chunk_idx=%s tp_rank=%s",
+                room,
+                chunk_idx,
+                self.tp_rank,
+            )
+        return ok
+
+    def is_staging_room(self, room: int) -> bool:
+        """Check if a room is registered for staging scatter."""
+        return room in self._room_to_decode_req
+
+    def handle_chunk_arrived(
+        self,
+        room: int,
+        chunk_idx: int,
+        page_start: int,
+        num_pages: int,
+        writer_id: str,
+        chunk_writer_counts: dict,
+    ) -> bool:
+        """Process a staging chunk arrival from any transport (NIXL RDMA notif or ZMQ CHUNK_READY).
+
+        Accumulates writer arrivals in *chunk_writer_counts* and submits scatter
+        once all writers for this chunk have reported in. Returns True if scatter
+        was submitted.
+        """
+        chunk_writer_counts[room][chunk_idx].append((page_start, num_pages, writer_id))
+        decode_req = self._room_to_decode_req.get(room)
+        if decode_req is None:
+            logger.warning(
+                "Staging chunk arrived for unregistered room=%s chunk=%d, skipping",
+                room,
+                chunk_idx,
+            )
+            return False
+        writers_arrived = len(chunk_writer_counts[room][chunk_idx])
+        num_writers = self.num_writers_for(decode_req)
+        if writers_arrived >= num_writers:
+            self.submit_chunk_scatter(room, chunk_idx, page_start, num_pages)
+            del chunk_writer_counts[room][chunk_idx]
+            return True
+        return False
+
+    def submit_last_scatter_async(self, room: int) -> bool:
+        """Submit scatter for the last chunk when all ranks report Success.
+
+        Called from decode_thread.  Sets ``_scatter_event`` **before**
+        ``_staging_last_scatter_submitted`` so the main thread sees the
+        event when it checks the flag (CPython GIL guarantees ordering).
+        """
+        decode_req = self._room_to_decode_req.get(room)
+        if decode_req is None:
+            logger.warning(
+                "[STAGING] submit_last_scatter_async: room=%s not registered. "
+                "This should not happen if register_decode_req is called at "
+                "kv_receiver.init() time.",
+                room,
+            )
+            return False
+        alloc_id = self._submit_last_scatter(decode_req)
+        if alloc_id >= 0:
+            event = torch.cuda.Event()
+            event.record(self.staging_allocator._scatter_stream)
+            decode_req._scatter_event = event
+            decode_req._scatter_alloc_id = alloc_id
+            decode_req._staging_last_scatter_submitted = True
+        else:
+            decode_req._staging_scatter_done = True
+        return True
+
+    # ------------------------------------------------------------------
+    # Event check + free: called from main thread (pop_transferred)
+    # ------------------------------------------------------------------
+
+    def is_done(self, decode_req: DecodeRequest) -> bool:
+        """Return True if staging scatter is complete for this request."""
+        if not getattr(decode_req, "_staging_scatter_done", False):
+            return False
+        return not getattr(decode_req, "_chunk_events", None)
+
+    def advance_scatter(self, decode_req: DecodeRequest) -> None:
+        """Check CUDA events and free completed staging allocations.
+
+        Scatter kernels have already been submitted by the decode_thread
+        (via submit_chunk_scatter / submit_last_scatter_async).  This
+        method only polls the recorded events and releases staging memory.
+        """
+        room = decode_req.req.bootstrap_room
+        chunk_events = getattr(decode_req, "_chunk_events", None)
+        if chunk_events:
+            for i in range(len(chunk_events) - 1, -1, -1):
+                event, alloc_id = chunk_events[i]
+                if event.query():
+                    chunk_events.pop(i)
+                    self._free_and_send_watermark(alloc_id, decode_req)
+
+        if not getattr(decode_req, "_staging_last_scatter_submitted", False):
+            return
+
+        event = getattr(decode_req, "_scatter_event", None)
+        if event is not None and event.query():
+            self._free_and_send_watermark(decode_req._scatter_alloc_id, decode_req)
+            decode_req._scatter_event = None
+            decode_req._scatter_alloc_id = -1
+            decode_req._staging_scatter_done = True
+
+    # ------------------------------------------------------------------
+    # Internal methods
+    # ------------------------------------------------------------------
+
+    def _scatter_region(
+        self,
+        staging_offset: int,
+        page_start: int,
+        num_pages: int,
+        decode_req: DecodeRequest,
+    ) -> bool:
+        """Submit scatter kernels for a staging region to scatter_stream.
+
+        May be called from the decode_thread (background).  All GPU work
+        runs on scatter_stream so that the decode_thread never blocks on
+        the default stream (which carries the main-thread forward pass).
+        """
+        from sglang.srt.disaggregation.common.staging_buffer import (
+            scatter_staging_to_kv,
+        )
+
+        k_buffers = self.kv_buffer_info["k_buffers"]
+        v_buffers = self.kv_buffer_info["v_buffers"]
+        page_size = self.kv_buffer_info["page_size"]
+        dst_tp_rank = self.kv_manager.kv_args.engine_rank % self.decode_tp
+
+        device = k_buffers[0].device
+        torch.cuda.set_device(device)
+
+        if not hasattr(self.staging_allocator, "_scatter_stream"):
+            self.staging_allocator._scatter_stream = torch.cuda.Stream(device=device)
+
+        scatter_stream = self.staging_allocator._scatter_stream
+
+        staging_view = self.staging_allocator.buffer.buffer[staging_offset:]
+
+        req_pool_idx = decode_req.req.req_pool_idx
+        token_start = page_start * page_size
+        token_end = token_start + num_pages * page_size
+        prefill_tp = decode_req.kv_receiver.prefill_info.attn_tp_size
+
+        with torch.cuda.stream(scatter_stream):
+            kv_indices = self.scheduler.req_to_token_pool.req_to_token[
+                req_pool_idx, token_start:token_end
+            ]
+            if page_size > 1:
+                page_idx_tensor = kv_indices[::page_size] // page_size
+            else:
+                page_idx_tensor = kv_indices
+
+            scatter_staging_to_kv(
+                staging_view,
+                k_buffers,
+                v_buffers,
+                page_idx_tensor,
+                page_size,
+                prefill_tp,
+                self.decode_tp,
+                dst_tp_rank,
+                self.total_kv_heads,
+            )
+
+        return True
+
+    def _submit_last_scatter(self, decode_req: DecodeRequest) -> int:
+        """Submit scatter for the last chunk. Returns alloc_id >= 0, or -1."""
+        receiver = decode_req.kv_receiver
+        chunk_infos = getattr(receiver, "chunk_staging_infos", [])
+        if not chunk_infos:
+            return -1
+
+        last_info = chunk_infos[-1]
+        alloc_id, staging_offset, _, _, last_num_pages = last_info
+        if staging_offset < 0 or alloc_id < 0:
+            return -1
+
+        seq_len = len(decode_req.req.origin_input_ids)
+        ps = self.scheduler.token_to_kv_pool_allocator.page_size
+        total_pages = (seq_len + ps - 1) // ps
+        page_start = total_pages - last_num_pages
+
+        ok = self._scatter_region(
+            staging_offset, page_start, last_num_pages, decode_req
+        )
+        return alloc_id if ok else -1
+
+    def _free_and_send_watermark(
+        self, alloc_id: int, decode_req: DecodeRequest
+    ) -> None:
+        """Free a staging allocation and broadcast watermark to all prefills."""
+        self.staging_allocator.free(alloc_id)
+        post_wm = self.staging_allocator.get_watermark()
+        room = decode_req.req.bootstrap_room
+        wm_round, wm_tail = post_wm
+        wm_round_b = str(wm_round).encode("ascii")
+        wm_tail_b = str(wm_tail).encode("ascii")
+        for _key, (receiver, session_id) in list(self._wm_subscribers.items()):
+            sid_b = session_id.encode("ascii")
+            for bootstrap_info in receiver.bootstrap_infos:
+                try:
+                    sock, lock = receiver._connect_to_bootstrap_server(bootstrap_info)
+                    with lock:
+                        sock.send_multipart(
+                            [b"WATERMARK", wm_round_b, wm_tail_b, sid_b]
+                        )
+                except Exception:
+                    pass
+
+
+def is_watermark_ready(
+    staging_state, session_id: str, alloc_round: int, alloc_end: int
+) -> bool:
+    """Non-blocking check: is the staging region safe to write?"""
+    if alloc_round <= 0:
+        return True
+    prev_round = alloc_round - 1
+    wm_round, wm_tail = staging_state.remote_watermarks.get(session_id, (0, 0))
+    return prev_round < wm_round or (prev_round == wm_round and alloc_end <= wm_tail)
+
+
+def handle_watermark_msg(staging_ctx, msg_parts) -> None:
+    """Process a WATERMARK message and update remote watermark tracking."""
+    wm_round = int(msg_parts[1].decode("ascii"))
+    wm_tail = int(msg_parts[2].decode("ascii"))
+    wm_session = msg_parts[3].decode("ascii") if len(msg_parts) > 3 else ""
+    with staging_ctx.watermark_cv:
+        prev = staging_ctx.remote_watermarks.get(wm_session, (0, 0))
+        if (wm_round, wm_tail) > prev:
+            staging_ctx.remote_watermarks[wm_session] = (
+                wm_round,
+                wm_tail,
+            )
+        staging_ctx.watermark_cv.notify_all()
+
+
+def handle_staging_rsp(msg_parts, transfer_infos: dict) -> None:
+    """Process a STAGING_RSP message and update transfer info with allocation."""
+    stg_room = int(msg_parts[1].decode("ascii"))
+    stg_chunk_idx = int(msg_parts[2].decode("ascii"))
+    stg_offset = int(msg_parts[3].decode("ascii"))
+    stg_round = int(msg_parts[4].decode("ascii"))
+    stg_end = int(msg_parts[5].decode("ascii"))
+    stg_session = msg_parts[6].decode("ascii")
+    room_infos = transfer_infos.get(stg_room, {})
+    tinfo = room_infos.get(stg_session)
+    if tinfo is not None:
+        if tinfo.staging is None:
+            tinfo.staging = StagingTransferInfo()
+        tinfo.staging.set_chunk(stg_chunk_idx, stg_offset, stg_round, stg_end)
+    else:
+        logger.warning(
+            "STAGING_RSP RECV but tinfo=None room=%s chunk=%d session=%s",
+            stg_room,
+            stg_chunk_idx,
+            stg_session,
+        )
+
+
+# ======================================================================
+# Staging data structures and protocol utilities
+# ======================================================================
+
+
+@dataclasses.dataclass
+class StagingTransferInfo:
+    """Per-chunk staging allocation info attached to a TransferInfo."""
+
+    offsets: List[int] = dataclasses.field(default_factory=lambda: [-1])
+    rounds: List[int] = dataclasses.field(default_factory=lambda: [0])
+    ends: List[int] = dataclasses.field(default_factory=lambda: [-1])
+
+    def set_chunk(self, idx: int, offset: int, rnd: int, end: int):
+        while len(self.offsets) <= idx:
+            self.offsets.append(-1)
+            self.rounds.append(0)
+            self.ends.append(-1)
+        self.offsets[idx] = offset
+        self.rounds[idx] = rnd
+        self.ends[idx] = end
+
+
+@dataclasses.dataclass
+class StagingRegisterInfo:
+    """Staging buffer registration info attached to a KVArgsRegisterInfo."""
+
+    base_ptr: int = 0
+    total_size: int = 0
+
+    @classmethod
+    def from_zmq_fields(
+        cls, msg: list, msg_start_offset: int
+    ) -> Optional[StagingRegisterInfo]:
+        i = msg_start_offset
+        base_ptr = (
+            struct.unpack("Q", msg[i])[0] if len(msg) > i and len(msg[i]) == 8 else 0
+        )
+        total_size = (
+            int(msg[i + 1].decode("ascii"))
+            if len(msg) > i + 1 and len(msg[i + 1]) > 0
+            else 0
+        )
+        if base_ptr == 0 and total_size == 0:
+            return None
+        return cls(base_ptr=base_ptr, total_size=total_size)
+
+
+class PrefillStagingStrategy:
+    """Prefill-side staging transfer: readiness check + gather-RDMA execution.
+
+    Encapsulates the decision logic (chunk index calculation, staging offset
+    lookup, watermark readiness) and delegates actual RDMA to the kv_manager.
+    """
+
+    def __init__(self, kv_manager, staging_buffer):
+        self.kv_manager = kv_manager
+        self.staging_buffer = staging_buffer
+        page_size = kv_manager.kv_buffer_tensors["page_size"]
+        cps = kv_manager.server_args.chunked_prefill_size or 8192
+        self.full_chunk_pages = max(1, cps // page_size)
+
+    def check_ready(
+        self,
+        req,
+        kv_chunk_index_start: int,
+        num_chunk_pages: int,
+        session_id: Optional[str] = None,
+    ) -> Tuple[bool, int, int, int, int]:
+        """Check if staging offset and watermark are ready for this chunk.
+
+        Args:
+            req: transfer request with a ``.staging`` attribute.
+            kv_chunk_index_start: page-level start index for this chunk.
+            num_chunk_pages: number of pages in this chunk.
+            session_id: identifier used for watermark lookup. Falls back to
+                ``req.mooncake_session_id`` when *None* (mooncake compat).
+
+        Returns (ready, chunk_idx, offset, round, end).
+        offset == ALLOC_OVERSIZED means permanent failure (fall back to slice).
+        offset == -1 means allocation pending (re-enqueue).
+        """
+        from sglang.srt.disaggregation.common.staging_buffer import StagingAllocator
+
+        chunk_idx = (
+            kv_chunk_index_start // self.full_chunk_pages
+            if self.full_chunk_pages > 0
+            else 0
+        )
+
+        stg = req.staging
+        if stg is None or chunk_idx >= len(stg.offsets):
+            return (False, chunk_idx, -1, 0, -1)
+
+        c_offset = stg.offsets[chunk_idx]
+        if c_offset == StagingAllocator.ALLOC_OVERSIZED:
+            return (False, chunk_idx, StagingAllocator.ALLOC_OVERSIZED, 0, -1)
+        if c_offset < 0:
+            return (False, chunk_idx, -1, 0, -1)
+
+        c_round = stg.rounds[chunk_idx]
+        c_end = stg.ends[chunk_idx]
+
+        if session_id is None:
+            session_id = req.mooncake_session_id
+        if not self.kv_manager._is_watermark_ready(session_id, c_round, c_end):
+            return (False, chunk_idx, c_offset, c_round, c_end)
+
+        return (True, chunk_idx, c_offset, c_round, c_end)
+
+    def transfer(
+        self,
+        session_id: str,
+        prefill_kv_indices,
+        dst_staging_ptr: int,
+        dst_staging_size: int,
+        target_info,
+    ) -> int:
+        """Execute staged transfer (gather + RDMA).
+
+        Returns 0 on success, -1 to signal fallback to slice path.
+        """
+        try:
+            return self.kv_manager.send_kvcache_staged(
+                session_id,
+                prefill_kv_indices,
+                dst_staging_ptr,
+                dst_staging_size,
+                target_info.dst_tp_rank,
+                target_info.dst_attn_tp_size,
+                target_info.dst_kv_item_len,
+                staging_buffer=self.staging_buffer,
+            )
+        except Exception as e:
+            raise RuntimeError(
+                f"[Staging] KV transfer via staging buffer failed: {e}. "
+                f"session={session_id}"
+            ) from e
+
+
+def _get_custom_mem_pool(device: str):
+    """Get custom memory pool for staging buffer allocation (backend-agnostic).
+
+    Returns (custom_mem_pool, pool_type) tuple. custom_mem_pool may be None
+    if no custom pool is configured.
+    """
+    from sglang.srt.disaggregation.mooncake.utils import (
+        init_mooncake_custom_mem_pool,
+    )
+
+    _, custom_mem_pool, pool_type = init_mooncake_custom_mem_pool(device)
+    if custom_mem_pool is None:
+        logger.info(
+            "Staging buffer using cudaMalloc (no custom mem pool). "
+            "This works for all GPU architectures. "
+            "For NVLink/MNNVL transport, set SGLANG_MOONCAKE_CUSTOM_MEM_POOL."
+        )
+    return custom_mem_pool, pool_type
+
+
+def init_staging_buffers(register_fn, kv_args, count: int) -> list:
+    """Create prefill-side staging buffers and register them with the transport.
+
+    Args:
+        register_fn: callable(ptr: int, size: int) that registers a memory
+            region with the transport backend.
+        kv_args: KVArgs with gpu_id.
+        count: number of staging buffers to create.
+
+    Returns list of StagingBuffer instances.
+    """
+    from sglang.srt.disaggregation.common.staging_buffer import StagingBuffer
+    from sglang.srt.environ import envs
+
+    size_mb = envs.SGLANG_DISAGG_STAGING_BUFFER_SIZE_MB.get()
+    size_bytes = size_mb * 1024 * 1024
+    gpu_id = kv_args.gpu_id
+    device = f"cuda:{gpu_id}"
+
+    custom_mem_pool, _ = _get_custom_mem_pool(device)
+
+    buffers = []
+    for _ in range(count):
+        buf = StagingBuffer(size_bytes, device, gpu_id, custom_mem_pool=custom_mem_pool)
+        register_fn(buf.get_ptr(), buf.get_size())
+        buffers.append(buf)
+    return buffers
+
+
+def init_staging_allocator(register_fn, kv_args):
+    """Create decode-side staging ring-buffer allocator and register with transport.
+
+    Args:
+        register_fn: callable(ptr: int, size: int) that registers a memory
+            region with the transport backend.
+        kv_args: KVArgs with gpu_id.
+
+    Returns a StagingAllocator instance.
+    """
+    from sglang.srt.disaggregation.common.staging_buffer import StagingAllocator
+    from sglang.srt.environ import envs
+
+    pool_size_mb = envs.SGLANG_DISAGG_STAGING_POOL_SIZE_MB.get()
+    pool_size_bytes = pool_size_mb * 1024 * 1024
+    gpu_id = kv_args.gpu_id
+    device = f"cuda:{gpu_id}"
+
+    custom_mem_pool, _ = _get_custom_mem_pool(device)
+    allocator = StagingAllocator(pool_size_bytes, device, gpu_id, custom_mem_pool)
+    register_fn(allocator.get_base_ptr(), allocator.get_total_size())
+    return allocator
+
+
+def handle_staging_req(
+    msg,
+    staging_allocator,
+    kv_args,
+    attn_tp_size: int,
+    prefill_attn_tp_size: int,
+    kv_buffer_tensors,
+    room_receivers: dict,
+    room_bootstrap: dict,
+):
+    """Allocate staging for a chunk on-demand and send STAGING_RSP to prefill.
+
+    Deduplicates: multiple prefill TP ranks requesting the same (room, chunk_idx)
+    only allocate once.  Sends ALLOC_OVERSIZED on permanent failure.
+    """
+    from sglang.srt.disaggregation.common.staging_buffer import StagingAllocator
+
+    room = int(msg[1].decode("ascii"))
+    chunk_idx = int(msg[2].decode("ascii"))
+    chunk_num_pages = int(msg[3].decode("ascii"))
+    session_id = msg[4].decode("ascii")
+
+    if staging_allocator is None:
+        logger.warning(
+            "STAGING_REQ ignored: allocator is None room=%s chunk=%s",
+            room,
+            chunk_idx,
+        )
+        return
+
+    receiver = room_receivers.get(room)
+    if receiver is None:
+        logger.warning(
+            "STAGING_REQ dropped: no receiver for room=%s chunk=%s session=%s",
+            room,
+            chunk_idx,
+            session_id,
+        )
+        return
+    infos = getattr(receiver, "chunk_staging_infos", [])
+
+    if chunk_idx < len(infos) and infos[chunk_idx][0] >= 0:
+        _, offset, rnd, end, _ = infos[chunk_idx]
+    elif (
+        chunk_idx < len(infos)
+        and infos[chunk_idx][1] == StagingAllocator.ALLOC_OVERSIZED
+    ):
+        offset, rnd, end = StagingAllocator.ALLOC_OVERSIZED, 0, -1
+    else:
+        from sglang.srt.disaggregation.common.staging_buffer import (
+            compute_staging_layout,
+            resolve_total_kv_heads,
+        )
+
+        page_size = kv_args.page_size
+        kv_item_lens = kv_args.kv_item_lens
+        num_kv_layers = len(kv_item_lens) // 2
+        decode_bytes_per_token = kv_item_lens[0] // page_size
+        total_kv_heads = resolve_total_kv_heads(kv_args, attn_tp_size)
+        dst_heads_per_rank = max(1, total_kv_heads // max(1, attn_tp_size))
+        bytes_per_head_per_token = decode_bytes_per_token // dst_heads_per_rank
+        dst_tp_rank = kv_args.engine_rank % max(1, attn_tp_size)
+
+        chunk_tokens = chunk_num_pages * page_size
+        _, _, required = compute_staging_layout(
+            prefill_attn_tp_size,
+            attn_tp_size,
+            dst_tp_rank,
+            total_kv_heads,
+            chunk_tokens,
+            bytes_per_head_per_token,
+            num_kv_layers,
+        )
+        result = staging_allocator.assign(required)
+        if result is None:
+            logger.error(
+                "[STAGING_REQ] alloc failed room=%s chunk=%d (need %d bytes, "
+                "buffer total=%d bytes). Increase SGLANG_DISAGG_STAGING_POOL_SIZE_MB.",
+                room,
+                chunk_idx,
+                required,
+                staging_allocator.total_size,
+            )
+            offset, rnd, end = StagingAllocator.ALLOC_OVERSIZED, 0, -1
+            while len(infos) <= chunk_idx:
+                infos.append((-1, -1, 0, -1, 0))
+            infos[chunk_idx] = (
+                -1,
+                StagingAllocator.ALLOC_OVERSIZED,
+                0,
+                -1,
+                chunk_num_pages,
+            )
+        else:
+            alloc_id, offset, rnd = result
+            end = offset + required
+            while len(infos) <= chunk_idx:
+                infos.append((-1, -1, 0, -1, 0))
+            infos[chunk_idx] = (alloc_id, offset, rnd, end, chunk_num_pages)
+
+    bootstrap_infos = room_bootstrap.get(room)
+    if bootstrap_infos:
+        for bi in bootstrap_infos:
+            try:
+                sock, lock = receiver._connect_to_bootstrap_server(bi)
+                with lock:
+                    sock.send_multipart(
+                        [
+                            b"STAGING_RSP",
+                            str(room).encode("ascii"),
+                            str(chunk_idx).encode("ascii"),
+                            str(offset).encode("ascii"),
+                            str(rnd).encode("ascii"),
+                            str(end).encode("ascii"),
+                            session_id.encode("ascii"),
+                        ]
+                    )
+            except Exception:
+                pass
+
+
+def prefetch_staging_reqs(
+    room: int,
+    transfer_infos: dict,
+    kv_buffer_tensors: dict,
+    chunked_prefill_size: int,
+    staging_requested: set,
+    prefetch_sockets: dict,
+) -> None:
+    """Send STAGING_REQ for all chunks before the prefill forward starts.
+
+    Called from the scheduler right after batch formation, so that decode
+    allocates staging during the GPU forward pass.
+    """
+    import zmq
+
+    from sglang.srt.utils.network import NetworkAddress
+
+    page_size = kv_buffer_tensors["page_size"]
+    cps = chunked_prefill_size or 8192
+    full_chunk_pages = max(1, cps // page_size)
+
+    for session_id, tinfo in transfer_infos[room].items():
+        # mooncake exposes is_dummy as a dataclass bool field, NIXL exposes it
+        # as a method (it consults decode_prefix_len). Normalize via callable()
+        # so this shared helper works for either backend; treating a bound
+        # method as truthy (the previous behavior) silently dropped every
+        # STAGING_REQ on NIXL and deadlocked the prefill transfer worker.
+        is_dummy_attr = tinfo.is_dummy
+        if is_dummy_attr() if callable(is_dummy_attr) else is_dummy_attr:
+            continue
+        total_pages = len(tinfo.dst_kv_indices)
+        if total_pages == 0:
+            continue
+        num_chunks = (total_pages + full_chunk_pages - 1) // full_chunk_pages
+
+        for chunk_idx in range(num_chunks):
+            stg_key = (room, chunk_idx, session_id)
+            if stg_key in staging_requested:
+                continue
+            staging_requested.add(stg_key)
+
+            remaining = total_pages - chunk_idx * full_chunk_pages
+            chunk_pages = min(full_chunk_pages, remaining)
+            try:
+                na = NetworkAddress(tinfo.endpoint, tinfo.dst_port)
+                ep = na.to_tcp()
+                if ep not in prefetch_sockets:
+                    sock = zmq.Context().socket(zmq.PUSH)
+                    if na.is_ipv6:
+                        sock.setsockopt(zmq.IPV6, 1)
+                    sock.connect(ep)
+                    prefetch_sockets[ep] = sock
+                prefetch_sockets[ep].send_multipart(
+                    [
+                        b"STAGING_REQ",
+                        str(room).encode("ascii"),
+                        str(chunk_idx).encode("ascii"),
+                        str(chunk_pages).encode("ascii"),
+                        session_id.encode("ascii"),
+                    ]
+                )
+            except Exception:
+                staging_requested.discard(stg_key)

python/sglang/srt/disaggregation/common/utils.py

@@ -0,0 +1,129 @@
+import ctypes
+import dataclasses
+import struct
+import threading
+from collections import deque
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import numpy.typing as npt
+
+from sglang.srt.observability.trace import (
+    TraceNullContext,
+    TraceReqContext,
+)
+
+
+@dataclasses.dataclass
+class TransferKVChunk:
+    """Work unit for KV cache transfer from prefill to decode."""
+
+    room: int
+    prefill_kv_indices: npt.NDArray[np.int32]
+    index_slice: slice
+    is_last_chunk: bool
+    prefill_aux_index: Optional[int]
+    state_indices: Optional[List]
+    chunk_id: Optional[int] = None
+    trace_ctx: Union[TraceReqContext, TraceNullContext] = dataclasses.field(
+        default_factory=TraceNullContext
+    )
+
+
+def pack_list_of_buffers(buffers: List[bytes]) -> bytes:
+    if not buffers:
+        return b""
+    n = len(buffers)
+    header = struct.pack(f"<{n+1}I", n, *(len(b) for b in buffers))
+    return header + b"".join(buffers)
+
+
+def unpack_list_of_buffers(buf: bytes) -> List[bytes]:
+    if buf == b"":
+        return []
+    (n,) = struct.unpack("<I", buf[:4])
+    lens = struct.unpack(f"<{n}I", buf[4 : 4 + 4 * n])
+    out = []
+    offset = 4 + 4 * n
+    for length in lens:
+        out.append(buf[offset : offset + length])
+        offset += length
+    return out
+
+
+def pack_int_lists(lists, fmt: str) -> bytes:
+    return pack_list_of_buffers([struct.pack(f"<{len(a)}{fmt}", *a) for a in lists])
+
+
+def unpack_int_lists(buf: bytes, fmt: str) -> List[List[int]]:
+    width = struct.calcsize(fmt)
+    return [
+        list(struct.unpack(f"<{len(b)//width}{fmt}", b))
+        for b in unpack_list_of_buffers(buf)
+    ]
+
+
+class FastQueue:
+    def __init__(self):
+        self._buf = deque()
+        self._cond = threading.Condition()
+
+    def put(self, item):
+        with self._cond:
+            self._buf.append(item)
+            # wake up a thread of wait()
+            self._cond.notify()
+
+    def get(self):
+        with self._cond:
+            # if queue is empty  ,block until is notified()
+            while not self._buf:
+                self._cond.wait()
+            return self._buf.popleft()
+
+
+class AuxDataCodec:
+    """Handles serialization and deserialization of auxiliary data buffers."""
+
+    @staticmethod
+    def serialize_data_from_buffer(src_addr, data_length):
+        """Serialize data from memory buffer to bytes."""
+        buffer = (ctypes.c_byte * data_length).from_address(src_addr)
+        return bytes(buffer)
+
+    @staticmethod
+    def deserialize_data_to_buffer(kv_args, buffer_index, aux_index, data):
+        """Deserialize bytes into target memory buffer."""
+        dst_aux_ptr = kv_args.aux_data_ptrs[buffer_index]
+        item_len = kv_args.aux_item_lens[buffer_index]
+        dst_addr = dst_aux_ptr + item_len * aux_index
+        buffer = (ctypes.c_byte * len(data)).from_address(dst_addr)
+        buffer[:] = data
+        return
+
+
+def group_concurrent_contiguous(
+    src_indices: npt.NDArray[np.int32], dst_indices: npt.NDArray[np.int32]
+) -> Tuple[List[npt.NDArray[np.int32]], List[npt.NDArray[np.int32]]]:
+    """Vectorised NumPy implementation."""
+    # src/dst indices are transferred pairwise, so an empty side means there is
+    # nothing to transfer. Guarding both sides (not just src) avoids a cryptic
+    # NumPy broadcast error from np.diff() below when only one side is empty, e.g.
+    # a non-empty prefill DSA/SWA state list paired with an empty decode registration.
+    if src_indices.size == 0 or dst_indices.size == 0:
+        return [], []
+
+    if src_indices.size != dst_indices.size:
+        raise ValueError(
+            "group_concurrent_contiguous requires equal-length src/dst index arrays, "
+            f"got {src_indices.size} and {dst_indices.size}"
+        )
+
+    brk = np.where((np.diff(src_indices) != 1) | (np.diff(dst_indices) != 1))[0] + 1
+    src_groups = np.split(src_indices, brk)
+    dst_groups = np.split(dst_indices, brk)
+
+    src_groups = [g.tolist() for g in src_groups]
+    dst_groups = [g.tolist() for g in dst_groups]
+
+    return src_groups, dst_groups

python/sglang/srt/disaggregation/decode_hicache_mixin.py

@@ -0,0 +1,309 @@
+"""HiCache integration mixins for the decode side of PD disaggregation"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING, Any, List, Optional
+
+import torch
+
+from sglang.srt.disaggregation.base import KVPoll
+from sglang.srt.managers.schedule_policy import match_prefix_for_req
+from sglang.srt.mem_cache.base_prefix_cache import InitLoadBackParams
+
+if TYPE_CHECKING:
+    from sglang.srt.disaggregation.decode import DecodeRequest
+    from sglang.srt.managers.schedule_batch import Req
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class DecodePrefixMatch:
+    prefix_indices: torch.Tensor
+    l2_host_hit_length: int
+    l3_storage_hit_length: int
+    last_device_node: Any
+    last_host_node: Any = None
+    prefetch_registered: bool = False
+
+    @property
+    def l1_prefix_len(self) -> int:
+        return len(self.prefix_indices)
+
+    @property
+    def decode_prefix_len(self) -> int:
+        return self.l1_prefix_len + self.l2_host_hit_length + self.l3_storage_hit_length
+
+    @property
+    def needs_local_restore(self) -> bool:
+        return self.decode_prefix_len > self.l1_prefix_len
+
+    @property
+    def restore_token_count(self) -> int:
+        """Number of tokens that need L2/L3 load_back to device."""
+        return self.decode_prefix_len - self.l1_prefix_len
+
+
+class HiCacheRestoreResult(Enum):
+    """Outcome of one tick of the HiCache local-restore state machine."""
+
+    PENDING = "pending"
+    READY = "ready"
+    FAILED = "failed"
+
+
+class DecodeHiCachePreallocMixin:
+    """HiCache hooks for ``DecodePreallocQueue``: issue prefetch + reserve tokens."""
+
+    def _build_decode_prefix_match(self, req: Req, result: Any) -> DecodePrefixMatch:
+        """Convert a ``match_prefix_for_req`` result into ``DecodePrefixMatch``.
+
+        Performs the optional L3 storage hit length query when decode-side
+        HiCache is enabled and the last host node is backed up.
+        """
+        prefix_indices = result.device_indices
+        l1_prefix_len = len(prefix_indices)
+        l2_host_hit_length = result.host_hit_length
+
+        l3_storage_hit_length = 0
+        last_host_node = None
+        if self.scheduler.enable_decode_hicache:
+            last_host_node = result.last_host_node
+            if last_host_node.backuped or last_host_node is self.tree_cache.root_node:
+                matched_len = l1_prefix_len + l2_host_hit_length
+                suffix_tokens = req.origin_input_ids[matched_len:]
+                last_hash = last_host_node.get_last_hash_value()
+                prefix_keys = (
+                    last_host_node.get_prefix_hash_values(last_host_node.parent)
+                    if self.tree_cache.hicache_storage_pass_prefix_keys
+                    else None
+                )
+                l3_storage_hit_length = self.tree_cache.query_storage_hit_length(
+                    last_host_node,
+                    suffix_tokens,
+                    last_hash,
+                    prefix_keys,
+                )
+
+        return DecodePrefixMatch(
+            prefix_indices=prefix_indices,
+            l2_host_hit_length=l2_host_hit_length,
+            l3_storage_hit_length=l3_storage_hit_length,
+            last_device_node=result.last_device_node,
+            last_host_node=last_host_node if l3_storage_hit_length > 0 else None,
+        )
+
+    def _start_hicache_prefetch(
+        self, req: Req, prefix_match: Optional[DecodePrefixMatch]
+    ) -> None:
+        """Issue L3 storage prefetch after admission succeeds.
+
+        On failure, degrades to L2-only restore by clearing l3 fields.
+        """
+        if (
+            prefix_match is None
+            or prefix_match.l3_storage_hit_length <= 0
+            or prefix_match.last_host_node is None
+        ):
+            return
+        try:
+            node = prefix_match.last_host_node
+            matched_len = prefix_match.l1_prefix_len + prefix_match.l2_host_hit_length
+            suffix = req.origin_input_ids[
+                matched_len : matched_len + prefix_match.l3_storage_hit_length
+            ]
+            last_hash = node.get_last_hash_value()
+            prefix_keys = (
+                node.get_prefix_hash_values(node.parent)
+                if self.tree_cache.hicache_storage_pass_prefix_keys
+                else None
+            )
+            self.tree_cache.prefetch_from_storage(
+                req.rid, node, suffix, last_hash, prefix_keys
+            )
+            prefix_match.prefetch_registered = (
+                req.rid in self.tree_cache.ongoing_prefetch
+            )
+        except Exception as e:
+            logger.warning(
+                "HiCache L3 prefetch failed for rid=%s: %s; falling back to L2-only LoadingBack",
+                req.rid,
+                e,
+            )
+            prefix_match.l3_storage_hit_length = 0
+            prefix_match.prefetch_registered = False
+
+    def _hicache_pending_restore_tokens(self) -> int:
+        """Total device tokens reserved for pending HiCache L2/L3 load_back."""
+        if not self.scheduler.enable_decode_hicache:
+            return 0
+        return sum(
+            dr.prefix_match.restore_token_count
+            for dr in self.transfer_queue.queue
+            if dr.prefix_match is not None
+            and dr.hicache_restore_status == HiCacheRestoreResult.PENDING
+            and dr.hicache_restored_node is None
+        )
+
+
+class HiCacheRestoreGatedKVReceiver:
+    """Wraps a kv_receiver so KVPoll.Success is gated on HiCache restore READY."""
+
+    def __init__(self, decode_req: DecodeRequest):
+        self.decode_req = decode_req
+
+    def poll(self) -> KVPoll:
+        poll = self.decode_req.kv_receiver.poll()
+        if (
+            poll == KVPoll.Success
+            and self.decode_req.hicache_restore_status == HiCacheRestoreResult.PENDING
+        ):
+            return KVPoll.Transferring
+        return poll
+
+
+class DecodeHiCacheTransferMixin:
+    """HiCache hooks for ``DecodeTransferQueue``: drive restore state machine."""
+
+    def _clean_hicache_prefetch_resources(self, decode_req: DecodeRequest) -> None:
+        if (
+            decode_req.prefix_match is not None
+            and decode_req.prefix_match.prefetch_registered
+        ):
+            self.tree_cache.release_aborted_request(decode_req.req.rid)
+        if decode_req.hicache_restored_node is not None:
+            self.tree_cache.dec_lock_ref(decode_req.hicache_restored_node)
+            decode_req.hicache_restored_node = None
+
+    def _try_hicache_queue_load_back(self, dr: DecodeRequest) -> bool:
+        """Queue one L2->L1 load_back op for ``dr``; True iff a DMA was queued.
+
+        On success, ``dr.hicache_restored_node`` and ``hicache_restored_kv_indices``
+        are populated, and an inc_lock_ref is held until commit/abort.
+        Trivial cases (all-on-device / no needed coverage) auto-flip to READY.
+        Failback paths flip to FAILED.
+        """
+        pm = dr.prefix_match
+
+        # Wait for L3 -> L2 prefetch to drain (skip when no L3 hit).
+        if pm.l3_storage_hit_length > 0:
+            if not self.tree_cache.check_prefetch_progress(dr.req.rid):
+                return False
+            self.tree_cache.pop_prefetch_loaded_tokens(dr.req.rid)
+
+        # Re-match: req.last_node / prefix_indices updated to current device state.
+        rematch = match_prefix_for_req(
+            self.tree_cache,
+            dr.req,
+            dr.req.origin_input_ids,
+            cow_mamba=False,
+            include_req=True,
+        )
+        new_indices, restored_node = self.tree_cache.init_load_back(
+            InitLoadBackParams(
+                best_match_node=rematch.best_match_node,
+                host_hit_length=rematch.host_hit_length,
+                req=dr.req,
+            )
+        )
+        # Failback: total coverage < required prefix means device alloc likely failed.
+        if len(rematch.device_indices) + len(new_indices) < pm.decode_prefix_len:
+            logger.warning(
+                "HiCache load_back failed for rid=%s: device_indices=%d, "
+                "new_indices=%d, expected decode_prefix_len=%d (l1=%d, l2=%d, l3=%d)",
+                dr.req.rid,
+                len(rematch.device_indices),
+                len(new_indices),
+                pm.decode_prefix_len,
+                pm.l1_prefix_len,
+                pm.l2_host_hit_length,
+                pm.l3_storage_hit_length,
+            )
+            dr.hicache_restore_status = HiCacheRestoreResult.FAILED
+            return False
+
+        dr.hicache_restored_kv_indices = torch.cat(
+            [rematch.device_indices[pm.l1_prefix_len :], new_indices]
+        )
+        dr.hicache_restored_node = restored_node
+        self.tree_cache.inc_lock_ref(restored_node)
+
+        if len(new_indices) == 0:
+            # Whole prefix already on device; no DMA needed.
+            dr.hicache_restore_status = HiCacheRestoreResult.READY
+            return False
+        return True
+
+    def _process_hicache_local_restores(self, decode_reqs: List[DecodeRequest]) -> None:
+        if not hasattr(self.tree_cache, "is_load_back_event_done"):
+            return
+
+        # Filter once: keep only PENDING reqs that still need restore work;
+        # trivially-done reqs (no prefix_match / nothing to restore) flip to READY.
+        active: List[DecodeRequest] = []
+        for dr in decode_reqs:
+            if dr.hicache_restore_status != HiCacheRestoreResult.PENDING:
+                continue
+            pm = dr.prefix_match
+            if pm is None or not pm.needs_local_restore:
+                dr.hicache_restore_status = HiCacheRestoreResult.READY
+                continue
+            active.append(dr)
+
+        # Phase A: advance in-flight DMAs to READY.
+        for dr in active:
+            if (
+                dr.hicache_restored_node is not None
+                and self.tree_cache.is_load_back_event_done(
+                    dr.hicache_load_consumer_index
+                )
+            ):
+                dr.hicache_restore_status = HiCacheRestoreResult.READY
+
+        # Phase B: queue new load_back ops if the next slot is free.
+        # The (producer_index + 1) check ensures we never overwrite a still-in-flight slot:
+        # if a previous req holds that slot and isn't done, its event won't be signaled.
+        counter = self.tree_cache.cache_controller.layer_done_counter
+        if not self.tree_cache.is_load_back_event_done(
+            (counter.producer_index + 1) % counter.num_counters
+        ):
+            return
+        queued = [
+            dr
+            for dr in active
+            if dr.hicache_restored_node is None
+            and self._try_hicache_queue_load_back(dr)
+        ]
+        if not queued:
+            return
+
+        # Phase C: kick off merged DMA, bind consumer_index for Phase A polling next tick.
+        consumer_index = self.tree_cache.ready_to_load_host_cache()
+        if consumer_index < 0:
+            for dr in queued:
+                dr.hicache_restore_status = HiCacheRestoreResult.READY
+            return
+        for dr in queued:
+            dr.hicache_load_consumer_index = consumer_index
+
+    def _commit_hicache_local_restore_to_req(self, decode_req: DecodeRequest) -> None:
+        prefix_match = decode_req.prefix_match
+        if prefix_match is None or not prefix_match.needs_local_restore:
+            return
+
+        self.tree_cache.dec_lock_ref(prefix_match.last_device_node)
+
+        self.tree_cache.req_to_token_pool.write(
+            (
+                decode_req.req.req_pool_idx,
+                slice(prefix_match.l1_prefix_len, prefix_match.decode_prefix_len),
+            ),
+            decode_req.hicache_restored_kv_indices,
+        )
+        decode_req.req.prefix_indices = torch.cat(
+            [prefix_match.prefix_indices, decode_req.hicache_restored_kv_indices]
+        )
+        decode_req.req.last_node = decode_req.hicache_restored_node

python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py

@@ -0,0 +1,351 @@
+from __future__ import annotations
+
+import json
+import logging
+import threading
+import time
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.disaggregation.kv_events import OffloadedState
+from sglang.srt.environ import envs
+from sglang.srt.managers.cache_controller import HiCacheController
+from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+from sglang.srt.mem_cache.memory_pool import (
+    MHATokenToKVPool,
+    MLATokenToKVPool,
+    ReqToTokenPool,
+)
+from sglang.srt.mem_cache.memory_pool_host import MLATokenToKVPoolHost
+from sglang.srt.mem_cache.pool_host.mha import get_mha_host_pool_cls
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils.common import ceil_align
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req
+
+logger = logging.getLogger(__name__)
+
+
+class DecodeKVCacheOffloadManager:
+    """Manage decode-side KV cache offloading lifecycle and operations."""
+
+    def __init__(
+        self,
+        req_to_token_pool: ReqToTokenPool,
+        token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
+        tp_group: torch.distributed.ProcessGroup,
+        tree_cache: BasePrefixCache,
+        server_args: ServerArgs,
+    ) -> None:
+        self.req_to_token_pool = req_to_token_pool
+        self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
+        self.page_size = server_args.page_size
+        self.server_args = server_args
+        self.request_counter = 0
+        self.tree_cache = tree_cache
+        env_stride = envs.SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE.get()
+        if env_stride is None or env_stride <= 0:
+            self.offload_stride = self.page_size
+        else:
+            self.offload_stride = max(
+                self.page_size, (env_stride // self.page_size) * self.page_size
+            )
+        kv_cache = self.token_to_kv_pool_allocator.get_kvcache()
+        if isinstance(kv_cache, MHATokenToKVPool):
+            self.decode_host_mem_pool = get_mha_host_pool_cls(kv_cache)(
+                kv_cache,
+                server_args.hicache_ratio,
+                server_args.hicache_size,
+                self.page_size,
+                server_args.hicache_mem_layout,
+            )
+        elif isinstance(kv_cache, MLATokenToKVPool):
+            self.decode_host_mem_pool = MLATokenToKVPoolHost(
+                kv_cache,
+                server_args.hicache_ratio,
+                server_args.hicache_size,
+                self.page_size,
+                server_args.hicache_mem_layout,
+            )
+        else:
+            raise ValueError("Unsupported KV cache type for decode offload")
+
+        self.tp_group = tp_group
+        self.tp_world_size = torch.distributed.get_world_size(group=self.tp_group)
+
+        hicache_storage_backend_extra_config = {}
+        if server_args.hicache_storage_backend_extra_config:
+            try:
+                hicache_storage_backend_extra_config = json.loads(
+                    server_args.hicache_storage_backend_extra_config
+                )
+            except json.JSONDecodeError as e:
+                raise ValueError(
+                    f"Invalid hicache storage backend extra config JSON: {e}"
+                )
+
+        self.cache_controller = HiCacheController(
+            token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+            mem_pool_host=self.decode_host_mem_pool,
+            page_size=self.page_size,
+            tp_group=tp_group,
+            io_backend=server_args.hicache_io_backend,
+            load_cache_event=threading.Event(),
+            storage_backend=server_args.hicache_storage_backend,
+            model_name=server_args.served_model_name,
+            storage_backend_extra_config=hicache_storage_backend_extra_config,
+        )
+
+        self.ongoing_offload = {}
+        self.ongoing_backup = {}
+        self.offloaded_state = {}
+        self.offload_inflight = {}
+        logger.info("Enable offload kv cache for decode side")
+
+    def _mark_offload_started(self, rid):
+        self.offload_inflight[rid] = self.offload_inflight.get(rid, 0) + 1
+
+    def _mark_offload_finished(self, rid):
+        count = self.offload_inflight.get(rid, 0)
+        if count <= 1:
+            self.offload_inflight.pop(rid, None)
+        else:
+            self.offload_inflight[rid] = count - 1
+
+    def _has_inflight_offload(self, rid):
+        return self.offload_inflight.get(rid, 0) > 0
+
+    def offload_kv_cache(self, req) -> bool:
+        """Offload incremental KV cache for decode side."""
+
+        if self.cache_controller is None or self.decode_host_mem_pool is None:
+            return False
+
+        if req.req_pool_idx == -1 or len(req.output_ids) == 0:
+            return False
+
+        token_indices = self.req_to_token_pool.req_to_token[req.req_pool_idx]
+        if token_indices.dim() == 0 or token_indices.numel() == 0:
+            return False
+
+        # Prefill side offloads page-aligned origin_input_ids, decode side offloads the incremental part
+        all_tokens = req.origin_input_ids + req.output_ids[:-1]
+        prefill_offloaded_len = (
+            len(req.origin_input_ids) // self.page_size * self.page_size
+        )
+        state = self.offloaded_state.get(req.rid)
+        if state is None:
+            prefill_hashes = self._compute_prefix_hash(
+                req.origin_input_ids[:prefill_offloaded_len]
+            )
+            last_prefill_hash = (
+                prefill_hashes[-1] if prefill_offloaded_len > 0 else None
+            )
+            state = OffloadedState(
+                prefill_len=prefill_offloaded_len,
+                inc_len=0,
+                last_hash=last_prefill_hash,
+            )
+            self.offloaded_state[req.rid] = state
+        incremental_total = len(all_tokens) - state.prefill_len
+        incremental_new = incremental_total - state.inc_len
+        incremental_aligned_len = (
+            incremental_new // self.offload_stride * self.offload_stride
+        )
+
+        if incremental_aligned_len == 0:
+            return False
+
+        # Extract incremental tokens and indices for the newly available chunk
+        start = state.prefill_len + state.inc_len
+        end = start + incremental_aligned_len
+        incremental_tokens = all_tokens[start:end]
+        incremental_indices = token_indices[start:end]
+
+        # Prefill-aligned GPU slots are freed at request finish in
+        # _release_finished_req, NOT here. The decoding request
+        # continues to attend to those slots via req_to_token; freeing
+        # them mid-decode races with concurrent admission, which can
+        # reuse the slots and produce cross-pollinated KV reads.
+
+        # Asynchronously offload incremental KV cache from device to host
+        self.request_counter += 1
+        ack_id = self.request_counter
+        host_indices = self.cache_controller.write(
+            device_indices=incremental_indices.long(),
+            node_id=ack_id,
+        )
+        if host_indices is None:
+            logger.error(f"Not enough host memory for request {req.rid}")
+            return False
+
+        self._mark_offload_started(req.rid)
+        self.ongoing_offload[ack_id] = (
+            req,
+            host_indices,
+            incremental_tokens,
+            time.time(),
+            start,
+            end,
+        )
+        state.inc_len += incremental_aligned_len
+        return True
+
+    def check_offload_progress(self):
+        """Check the progress of offload from device to host and backup from host to storage."""
+        cc = self.cache_controller
+
+        qsizes = torch.tensor(
+            [
+                len(cc.ack_write_queue),
+                cc.ack_backup_queue.qsize(),
+            ],
+            dtype=torch.int,
+        )
+        if self.tp_world_size > 1:
+            torch.distributed.all_reduce(
+                qsizes, op=torch.distributed.ReduceOp.MIN, group=self.tp_group
+            )
+
+        n_write, n_backup = map(int, qsizes.tolist())
+        self._check_offload_progress(n_write)
+        self._check_backup_progress(n_backup)
+
+    def _check_offload_progress(self, finish_count):
+        """Check the progress of offload from device to host."""
+        while finish_count > 0:
+            _, finish_event, ack_list = self.cache_controller.ack_write_queue.pop(0)
+            finish_event.synchronize()
+            for ack_id in ack_list:
+                (
+                    req,
+                    host_indices,
+                    incremental_tokens,
+                    start_time,
+                    start,
+                    end,
+                ) = self.ongoing_offload.pop(ack_id)
+
+                self._mark_offload_finished(req.rid)
+                prior_hash = (
+                    self.offloaded_state[req.rid].last_hash
+                    if req.rid in self.offloaded_state
+                    else None
+                )
+                last_hash = self._trigger_backup(
+                    req, host_indices, incremental_tokens, start_time, prior_hash
+                )
+                if req.rid in self.offloaded_state:
+                    self.offloaded_state[req.rid].last_hash = last_hash
+
+                if req.finished() and not self._has_inflight_offload(req.rid):
+                    state = self.offloaded_state.get(req.rid)
+                    start_offset = state.prefill_len if state is not None else start
+                    self._release_finished_req(req, start_offset)
+            finish_count -= 1
+
+    def _release_finished_req(self, req: Req, start_offset: int):
+        # Defensive guard: ReqToTokenPool.free sets req_pool_idx to None,
+        # so a previously-released request must be skipped here to avoid
+        # non-idempotent side effects (e.g. tree_cache.protected_size_
+        # double-decrement, host pool double-free).
+        if req.req_pool_idx is None or req.req_pool_idx == -1:
+            return
+
+        kv_committed_len = req.pop_committed_kv_cache()
+
+        # Free the prefill-aligned slots. Previously this was done
+        # eagerly in offload_kv_cache (mid-decode), which raced with
+        # concurrent admission. Now consolidated here at request
+        # finish, where the request is guaranteed to no longer attend
+        # to those slots.
+        state = self.offloaded_state.get(req.rid)
+        if state is not None and state.prefill_len > 0:
+            prefill_indices = self.req_to_token_pool.req_to_token[
+                req.req_pool_idx, : state.prefill_len
+            ]
+            self.token_to_kv_pool_allocator.free(prefill_indices)
+        start = start_offset
+        end = kv_committed_len
+        # Free the incremental part of the request (DSA-aware)
+        kv_indices = self.req_to_token_pool.req_to_token[req.req_pool_idx, start:end]
+        self.token_to_kv_pool_allocator.free(kv_indices)
+
+        # Free over-allocated KV cache slots (e.g. from speculative decoding v2).
+        # Without spec v2, start_p == end_p so this is a no-op.
+        start_p, end_p = req.pop_overallocated_kv_cache()
+        if self.page_size > 1:
+            start_p = ceil_align(start_p, self.page_size)
+        if start_p < end_p:
+            overalloc_indices = self.req_to_token_pool.req_to_token[
+                req.req_pool_idx, start_p:end_p
+            ]
+            self.token_to_kv_pool_allocator.free(overalloc_indices)
+
+        self.req_to_token_pool.free(req)
+        self.tree_cache.protected_size_ -= len(req.prefix_indices)
+        if req.rid in self.offloaded_state:
+            del self.offloaded_state[req.rid]
+
+    def _check_backup_progress(self, finish_count):
+        """Check the progress of backup from host to storage."""
+        for _ in range(finish_count):
+            storage_operation = self.cache_controller.ack_backup_queue.get()
+            ack_id = storage_operation.id
+            req_id, host_indices, start_time = self.ongoing_backup.pop(ack_id)
+
+            # Release host memory
+            self.decode_host_mem_pool.free(host_indices)
+
+            logger.debug(
+                f"Finished backup request {req_id}, free host memory, len:{len(host_indices)}, cost time:{time.time() - start_time:.2f} seconds."
+            )
+
+    def _trigger_backup(
+        self, req, host_indices, incremental_tokens, start_time, prior_hash
+    ):
+        """Trigger async backup from host to storage."""
+        page_hashes = self._compute_prefix_hash(incremental_tokens, prior_hash)
+        ack_id = self.cache_controller.write_storage(
+            host_indices,
+            incremental_tokens,
+            hash_value=page_hashes,
+        )
+        self.ongoing_backup[ack_id] = (req.rid, host_indices, start_time)
+        return page_hashes[-1] if len(page_hashes) > 0 else prior_hash
+
+    def _compute_prefix_hash(self, tokens, prior_hash=""):
+        page_hashes = []
+        last_hash = prior_hash
+        for offset in range(0, len(tokens), self.page_size):
+            page_tokens = tokens[offset : offset + self.page_size]
+            last_hash = self.cache_controller.get_hash_str(page_tokens, last_hash)
+            page_hashes.append(last_hash)
+        return page_hashes
+
+    def finalize_release_on_finish(self, req: Req):
+        """Free any remaining tail KV that was not offloaded due to non-aligned length."""
+        # ReqToTokenPool.free sets req_pool_idx to None on release, so
+        # guard against both sentinels here.
+        if req.req_pool_idx is None or req.req_pool_idx == -1:
+            return
+        state = self.offloaded_state.get(req.rid)
+        if state is None:
+            prefill_len = len(req.origin_input_ids) // self.page_size * self.page_size
+            inc_len = 0
+        else:
+            prefill_len = state.prefill_len
+            inc_len = state.inc_len
+        # Prefill-aligned slots are freed by _release_finished_req. Make
+        # sure state exists so it can find prefill_len.
+        if state is None:
+            self.offloaded_state[req.rid] = OffloadedState(
+                prefill_len=prefill_len, inc_len=0, last_hash=None
+            )
+        if self._has_inflight_offload(req.rid):
+            return
+        start_offset = prefill_len
+        self._release_finished_req(req, start_offset)

python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+import logging
+from array import array
+from http import HTTPStatus
+from typing import TYPE_CHECKING, List
+
+import torch
+
+from sglang.srt.managers.overlap_utils import RelayPayload
+from sglang.srt.mem_cache.common import maybe_cache_unfinished_req
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.overlap_utils import FutureMap
+    from sglang.srt.managers.schedule_batch import ScheduleBatch
+    from sglang.srt.server_args import ServerArgs
+
+
+class ScheduleBatchDisaggregationDecodeMixin:
+
+    def prepare_for_prebuilt(self: ScheduleBatch):
+        """
+        Prepare a prebuilt extend by populate metadata
+        Adapted from .prepare_for_extend().
+        """
+
+        self.forward_mode = ForwardMode.PREBUILT
+        reqs = self.reqs
+        input_ids = [r.get_fill_ids()[len(r.prefix_indices) :] for r in reqs]
+        extend_num_tokens = sum(len(ids) for ids in input_ids)
+        seq_lens = []
+        pre_lens = []
+        req_pool_indices = []
+
+        # Pre-calculate total size
+        total_size = sum(req.extend_range.length for req in reqs)
+        out_cache_loc = torch.empty(total_size, dtype=torch.int64, device=self.device)
+
+        # Fill the tensor in one pass
+        offset = 0
+        for i, req in enumerate(reqs):
+            req_pool_indices.append(req.req_pool_idx)
+            pre_len = len(req.prefix_indices)
+
+            chunk = self.req_to_token_pool.req_to_token[req.req_pool_idx][
+                pre_len : pre_len + req.extend_range.length
+            ]
+            assert (
+                offset + req.extend_range.length <= total_size
+            ), f"Exceeds total size: offset={offset}, req.extend_range.length={req.extend_range.length}, total_size={total_size}"
+            out_cache_loc[offset : offset + req.extend_range.length] = chunk
+            offset += req.extend_range.length
+
+            seq_len = len(req.origin_input_ids) + max(0, len(req.output_ids) - 1)
+            seq_lens.append(seq_len)
+            if len(req.output_ids) == 0:
+                assert (
+                    seq_len - pre_len == req.extend_range.length
+                ), f"seq_len={seq_len}, pre_len={pre_len}, req.extend_range.length={req.extend_range.length}"
+
+            if not req.retracted_stain:
+                # Clamp to avoid double-counting: already_computed is seeded from
+                # the prefill-reported cached_tokens in _commit_transfer_to_req, so
+                # a decode-side prefix shorter than the prefill report must not
+                # subtract from cached_tokens.
+                delta = max(0, pre_len - req.already_computed)
+                req.cached_tokens += delta
+                req.cached_tokens_device += delta
+                req.already_computed = seq_len
+            req.is_retracted = False
+            if getattr(req, "pd_rebootstrap_in_progress", False):
+                req.pd_rebootstrap_in_progress = False
+            pre_lens.append(pre_len)
+
+        # Set fields
+        self.input_ids = torch.tensor(
+            sum(input_ids, array("q")), dtype=torch.int32, device=self.device
+        )
+        self.req_pool_indices = torch.tensor(
+            req_pool_indices, dtype=torch.int64, device=self.device
+        )
+        self.req_pool_indices_cpu = torch.tensor(req_pool_indices, dtype=torch.int64)
+        self.seq_lens = torch.tensor(seq_lens, dtype=torch.int64, device=self.device)
+        self.seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.int64)
+        self.orig_seq_lens = torch.tensor(
+            seq_lens, dtype=torch.int32, device=self.device
+        )
+        self.out_cache_loc = out_cache_loc
+        self.seq_lens_sum = sum(seq_lens)
+
+        if self.return_logprob:
+            self.top_logprobs_nums = [r.logprob.top_logprobs_num for r in reqs]
+            self.token_ids_logprobs = [r.logprob.token_ids_logprob for r in reqs]
+
+        self.extend_num_tokens = extend_num_tokens
+        self.prefix_lens = [len(r.prefix_indices) for r in reqs]
+        self.extend_lens = [r.extend_range.length for r in reqs]
+        self.extend_logprob_start_lens = None
+        self.extend_input_logprob_token_ids = None
+        self.multimodal_inputs = [r.multimodal_inputs for r in reqs]
+
+        # Build sampling info
+        self.sampling_info = SamplingBatchInfo.from_schedule_batch(
+            self,
+            self.model_config.vocab_size,
+        )
+
+    def process_prebuilt(
+        self: ScheduleBatch,
+        server_args: ServerArgs,
+        future_map: FutureMap,
+    ):
+        """Assign the buffered last input id to schedule batch"""
+        last_tokens: List[int] = []
+        for req in self.reqs:
+            last_tokens.append(req.output_ids[-1])
+            maybe_cache_unfinished_req(req, self.tree_cache)
+            if req.grammar is not None:
+                # FIXME: this try-except block is for handling unexpected xgrammar issue.
+                try:
+                    # if it is not None, then the grammar is from a retracted request, and we should not
+                    # accept the token as it's already accepted
+                    if req.grammar.current_token is None:
+                        req.grammar.accept_token(req.output_ids[-1])
+                except ValueError as e:
+                    from sglang.srt.managers.schedule_batch import FINISH_ABORT
+
+                    # Grammar accept_token can raise ValueError if the token is not in the grammar.
+                    # This can happen if the grammar is not set correctly or the token is invalid.
+                    # Use to_finish (not finished_reason) so that process_batch_result_prebuilt
+                    # handles the release via update_finish_state -> release_kv_cache in one place.
+                    error_message = f"Grammar accept_token failed for req {req.rid} with token {req.output_ids[-1]}: {e}"
+                    req.to_finish = FINISH_ABORT(
+                        error_message, HTTPStatus.INTERNAL_SERVER_ERROR
+                    )
+                req.grammar.finished = req.finished()
+        last_tokens_tensor = torch.tensor(
+            last_tokens, dtype=torch.int64, device=self.device
+        )
+
+        spec_info = self.spec_algorithm.build_disagg_draft_input(
+            self,
+            server_args,
+            last_tokens_tensor,
+            future_map,
+        )
+        if spec_info is not None:
+            self.spec_info = spec_info
+        else:
+            # Non-spec: stash last token into the relay so the first DECODE's
+            # resolve_forward_inputs gathers it like any other decode iter.
+            future_map.stash(
+                self.req_pool_indices, RelayPayload(bonus_tokens=last_tokens_tensor)
+            )
+            self.input_ids = None

python/sglang/srt/disaggregation/encode_grpc_server.py

@@ -0,0 +1,275 @@
+"""
+gRPC Encoder Server for SGLang EPD (Encode-Prefill-Decode) mode.
+
+This server provides gRPC-based encoding for multimodal inputs.
+
+Usage:
+    python -m sglang.launch_server --model-path <model> --encoder-only --grpc-mode
+"""
+
+import asyncio
+import logging
+import multiprocessing as mp
+import traceback
+from concurrent import futures
+from typing import List
+
+import grpc
+import zmq
+import zmq.asyncio
+from grpc_health.v1 import health_pb2, health_pb2_grpc
+from grpc_reflection.v1alpha import reflection
+from smg_grpc_proto import sglang_encoder_pb2, sglang_encoder_pb2_grpc
+
+from sglang.srt.disaggregation.encode_server import (
+    MMEncoder,
+    handle_scheduler_receive_url_request,
+    launch_encoder,
+)
+from sglang.srt.managers.io_struct import async_sock_send, wrap_as_pickle
+from sglang.srt.managers.schedule_batch import Modality
+from sglang.srt.server_args import PortArgs, ServerArgs
+from sglang.srt.utils import random_uuid
+from sglang.srt.utils.network import NetworkAddress, get_zmq_socket
+
+logger = logging.getLogger(__name__)
+SGLangEncoderServicer = sglang_encoder_pb2_grpc.SglangEncoderServicer
+add_SGLangEncoderServicer_to_server = (
+    sglang_encoder_pb2_grpc.add_SglangEncoderServicer_to_server
+)
+
+
+class EncoderHealthServicer(health_pb2_grpc.HealthServicer):
+    """
+    Standard gRPC health check service for encoder server.
+    Implements grpc.health.v1.Health for Kubernetes probes.
+    """
+
+    OVERALL_SERVER = ""
+    ENCODER_SERVICE = "sglang.grpc.encoder.SglangEncoder"
+
+    def __init__(self):
+        self._serving = False
+
+    def set_serving(self):
+        self._serving = True
+
+    def set_not_serving(self):
+        self._serving = False
+
+    async def Check(self, request, context) -> health_pb2.HealthCheckResponse:
+        if self._serving:
+            return health_pb2.HealthCheckResponse(
+                status=health_pb2.HealthCheckResponse.SERVING
+            )
+        return health_pb2.HealthCheckResponse(
+            status=health_pb2.HealthCheckResponse.NOT_SERVING
+        )
+
+    async def Watch(self, request, context):
+        yield await self.Check(request, context)
+
+
+class SGLangEncoderServer(SGLangEncoderServicer):
+    """
+    gRPC service implementation for SGLang encoder.
+    """
+
+    def __init__(
+        self,
+        encoder: MMEncoder,
+        send_sockets: List[zmq.Socket],
+        server_args: ServerArgs,
+    ):
+        self.encoder = encoder
+        self.send_sockets = send_sockets
+        self.server_args = server_args
+
+    async def Encode(
+        self, request: sglang_encoder_pb2.EncodeRequest, context
+    ) -> sglang_encoder_pb2.EncodeResponse:
+        try:
+            request_dict = {
+                "mm_items": list(request.mm_items),
+                "req_id": request.req_id,
+                "num_parts": request.num_parts,
+                "part_idx": request.part_idx,
+            }
+            for socket in self.send_sockets:
+                await async_sock_send(socket, wrap_as_pickle(request_dict))
+
+            # gRPC encode is image-only; encoder.encode() requires modality
+            (
+                nbytes,
+                embedding_len,
+                embedding_dim,
+                error_msg,
+                error_code,
+            ) = await self.encoder.encode(
+                mm_items=list(request.mm_items),
+                modality=Modality.IMAGE,
+                req_id=request.req_id,
+                num_parts=request.num_parts,
+                part_idx=request.part_idx,
+            )
+            if error_msg is not None:
+                context.set_code(grpc.StatusCode.INTERNAL)
+                context.set_details(error_msg)
+                return sglang_encoder_pb2.EncodeResponse()
+
+            if self.server_args.encoder_transfer_backend == "mooncake":
+                return sglang_encoder_pb2.EncodeResponse(
+                    embedding_size=nbytes,
+                    embedding_len=embedding_len,
+                    embedding_dim=embedding_dim,
+                )
+            elif self.server_args.encoder_transfer_backend == "zmq_to_scheduler":
+                embedding_ports = list(request.embedding_port)
+                logger.info(f"embedding_port = {embedding_ports}")
+                if not embedding_ports:
+                    await self.encoder.send_with_url(req_id=request.req_id)
+                else:
+                    tasks = []
+                    for embedding_port in embedding_ports:
+                        tasks.append(
+                            self.encoder.send(
+                                req_id=request.req_id,
+                                prefill_host=request.prefill_host,
+                                embedding_port=embedding_port,
+                            )
+                        )
+                    await asyncio.gather(*tasks)
+                    self.encoder.embedding_to_send.pop(request.req_id, None)
+                return sglang_encoder_pb2.EncodeResponse()
+            elif self.server_args.encoder_transfer_backend == "zmq_to_tokenizer":
+                embedding_port = (
+                    request.embedding_port[0] if request.embedding_port else 0
+                )
+                await self.encoder.send(
+                    req_id=request.req_id,
+                    prefill_host=request.prefill_host,
+                    embedding_port=embedding_port,
+                )
+                self.encoder.embedding_to_send.pop(request.req_id, None)
+                return sglang_encoder_pb2.EncodeResponse()
+
+            return sglang_encoder_pb2.EncodeResponse()
+
+        except Exception as e:
+            logger.error(f"Encode error: {e}")
+            traceback.print_exc()
+            context.set_code(grpc.StatusCode.INTERNAL)
+            context.set_details(str(e))
+            return sglang_encoder_pb2.EncodeResponse()
+
+    async def Send(
+        self, request: sglang_encoder_pb2.SendRequest, context
+    ) -> sglang_encoder_pb2.SendResponse:
+        try:
+            await self.encoder.send(
+                req_id=request.req_id,
+                prefill_host=request.prefill_host,
+                embedding_port=request.embedding_port,
+                session_id=request.session_id if request.session_id else None,
+                buffer_address=(
+                    request.buffer_address if request.buffer_address else None
+                ),
+            )
+            self.encoder.embedding_to_send.pop(request.req_id, None)
+            return sglang_encoder_pb2.SendResponse()
+
+        except Exception as e:
+            logger.error(f"Send error: {e}")
+            traceback.print_exc()
+            context.set_code(grpc.StatusCode.INTERNAL)
+            context.set_details(str(e))
+            return sglang_encoder_pb2.SendResponse()
+
+    async def SchedulerReceiveUrl(
+        self, request: sglang_encoder_pb2.SchedulerReceiveUrlRequest, context
+    ) -> sglang_encoder_pb2.SchedulerReceiveUrlResponse:
+        try:
+            await handle_scheduler_receive_url_request(
+                {
+                    "req_id": request.req_id,
+                    "receive_count": request.receive_count,
+                    "receive_url": request.receive_url,
+                }
+            )
+            return sglang_encoder_pb2.SchedulerReceiveUrlResponse()
+
+        except Exception as e:
+            logger.error(f"SchedulerReceiveUrl error: {e}")
+            traceback.print_exc()
+            context.set_code(grpc.StatusCode.INTERNAL)
+            context.set_details(str(e))
+            return sglang_encoder_pb2.SchedulerReceiveUrlResponse()
+
+
+async def serve_grpc_encoder(server_args: ServerArgs):
+    ctx = mp.get_context("spawn")
+    zmq_ctx = zmq.asyncio.Context(10)
+    ipc_path_prefix = random_uuid()
+    port_args = PortArgs.init_new(server_args)
+
+    if server_args.dist_init_addr:
+        na = NetworkAddress.parse(server_args.dist_init_addr)
+        dist_init_method = na.to_tcp()
+    else:
+        dist_init_method = NetworkAddress(
+            server_args.host or "127.0.0.1", port_args.nccl_port
+        ).to_tcp()
+
+    send_sockets: List[zmq.Socket] = []
+    for rank in range(1, server_args.tp_size):
+        schedule_path = f"ipc:///tmp/{ipc_path_prefix}_schedule_{rank}"
+        send_sockets.append(
+            get_zmq_socket(zmq_ctx, zmq.PUSH, schedule_path, bind=False)
+        )
+        ctx.Process(
+            target=launch_encoder,
+            args=(server_args, schedule_path, dist_init_method, rank),
+            daemon=True,
+        ).start()
+
+    encoder = MMEncoder(server_args, dist_init_method=dist_init_method)
+
+    server = grpc.aio.server(
+        futures.ThreadPoolExecutor(max_workers=10),
+        options=[
+            ("grpc.max_send_message_length", 1024 * 1024 * 256),
+            ("grpc.max_receive_message_length", 1024 * 1024 * 256),
+        ],
+    )
+
+    health_servicer = EncoderHealthServicer()
+    health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server)
+
+    encoder_servicer = SGLangEncoderServer(
+        encoder=encoder,
+        send_sockets=send_sockets,
+        server_args=server_args,
+    )
+    add_SGLangEncoderServicer_to_server(encoder_servicer, server)
+
+    SERVICE_NAMES = (
+        sglang_encoder_pb2.DESCRIPTOR.services_by_name["SglangEncoder"].full_name,
+        "grpc.health.v1.Health",
+        reflection.SERVICE_NAME,
+    )
+    reflection.enable_server_reflection(SERVICE_NAMES, server)
+
+    listen_addr = NetworkAddress(server_args.host, server_args.port).to_host_port_str()
+    server.add_insecure_port(listen_addr)
+
+    await server.start()
+    logger.info(f"gRPC encoder server listening on {listen_addr}")
+
+    health_servicer.set_serving()
+
+    try:
+        await server.wait_for_termination()
+    except KeyboardInterrupt:
+        logger.info("Shutting down gRPC encoder server...")
+        health_servicer.set_not_serving()
+        await server.stop(grace=5)

python/sglang/srt/disaggregation/fake/__init__.py

@@ -0,0 +1,5 @@
+from sglang.srt.disaggregation.fake.conn import (
+    FakeKVManager,
+    FakeKVReceiver,
+    FakeKVSender,
+)

python/sglang/srt/disaggregation/fake/conn.py

@@ -0,0 +1,139 @@
+import logging
+from typing import List, Optional
+
+import numpy as np
+import numpy.typing as npt
+
+from sglang.srt.disaggregation.base.conn import (
+    BaseKVManager,
+    BaseKVReceiver,
+    BaseKVSender,
+    KVArgs,
+    KVPoll,
+    KVTransferMetric,
+)
+from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.server_args import ServerArgs
+
+logger = logging.getLogger(__name__)
+
+
+# For warmup reqs, we don't kv transfer, we use the fake manager, sender and receiver
+class FakeKVManager(BaseKVManager):
+    def __init__(
+        self,
+        args: KVArgs,
+        disaggregation_mode: DisaggregationMode,
+        server_args: ServerArgs,
+        is_mla_backend: Optional[bool] = False,
+    ):
+        super().__init__(args, disaggregation_mode, server_args, is_mla_backend)
+        self.kv_args = args
+        self.req_to_decode_prefix_len = {}
+
+    def register_to_bootstrap(self):
+        pass
+
+
+class FakeKVSender(BaseKVSender):
+    def __init__(
+        self,
+        mgr: BaseKVManager,
+        bootstrap_addr: str,
+        bootstrap_room: int,
+        dest_tp_ranks: List[int],
+        pp_rank: int,
+        req_has_disagg_prefill_dp_rank: bool = False,
+    ):
+        self.kv_mgr = mgr
+        self.has_sent = False
+        self.conclude_state: Optional[KVPoll] = None
+
+    def poll(self) -> KVPoll:
+        if self.conclude_state is not None:
+            return self.conclude_state
+        if not self.has_sent:
+            # Assume handshake completed instantly
+            return KVPoll.WaitingForInput
+
+        # Assume transfer completed instantly
+        logger.debug("FakeKVSender poll success")
+        self.conclude_state = KVPoll.Success
+        return KVPoll.Success
+
+    def get_transfer_metric(self) -> KVTransferMetric:
+        return KVTransferMetric()
+
+    def init(
+        self,
+        kv_indices: list[int],
+        aux_index: Optional[int] = None,
+    ):
+        logger.debug(
+            f"FakeKVSender init with kv_indices: {kv_indices}, aux_index: {aux_index}"
+        )
+        pass
+
+    def send(
+        self,
+        kv_indices: npt.NDArray[np.int32],
+        state_indices: Optional[List] = None,
+    ):
+        self.has_sent = True
+        logger.debug(
+            f"FakeKVSender send with kv_indices: {kv_indices}, state_indices: {state_indices}"
+        )
+
+    def failure_exception(self):
+        raise Exception("Fake KVSender Exception")
+
+    def abort(self):
+        self.conclude_state = KVPoll.Failed
+
+
+class FakeKVReceiver(BaseKVReceiver):
+    def __init__(
+        self,
+        mgr: BaseKVManager,
+        bootstrap_addr: str,
+        bootstrap_room: Optional[int] = None,
+    ):
+        self.bootstrap_done = False
+        self.has_sent_metadata = False
+        self.require_staging: bool = False
+        self.conclude_state: Optional[KVPoll] = None
+
+    def poll(self) -> KVPoll:
+        if self.conclude_state is not None:
+            return self.conclude_state
+        if not self.bootstrap_done:
+            return KVPoll.Bootstrapping
+        if not self.has_sent_metadata:
+            return KVPoll.WaitingForInput
+        logger.debug("FakeKVReceiver poll success")
+        self.conclude_state = KVPoll.Success
+        return KVPoll.Success
+
+    def init(
+        self,
+        prefill_dp_rank: int,
+    ):
+        self.bootstrap_done = True
+
+    def send_metadata(
+        self,
+        kv_indices: list[int],
+        aux_index: Optional[int] = None,
+        state_indices: Optional[List] = None,
+        decode_prefix_len: Optional[int] = None,
+    ):
+        self.has_sent_metadata = True
+        logger.debug(
+            f"FakeKVReceiver send_metadata with kv_indices: {kv_indices}, aux_index: {aux_index}, state_indices: {state_indices}"
+        )
+
+    def failure_exception(self):
+        raise Exception("Fake KVReceiver Exception")
+
+    def abort(self):
+        self.conclude_state = KVPoll.Failed

python/sglang/srt/disaggregation/kv_events.py

@@ -0,0 +1,470 @@
+"""
+Copyright 2025 SGLang Team
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+"""
+KV caching events
+"""
+
+import atexit
+import enum
+import logging
+import queue
+import threading
+import time
+from abc import ABC, abstractmethod
+from collections import deque
+from itertools import count
+from queue import Queue
+from typing import Any, Callable, Optional, Union
+
+import msgspec
+import zmq
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+
+def select_kv_publisher_dp_rank(
+    attn_dp_size: int, attn_dp_rank: int, dp_rank: Optional[int]
+) -> int:
+    """Index used to offset this scheduler's KV-event publisher port.
+
+    Each independent KV cache must publish on its own port so a consumer can
+    subscribe per replica. There are always ``dp_size`` such publishers; which
+    rank distinguishes them depends on the parallelism mode:
+
+    - DP-attention (``attn_dp_size > 1``): each attention-DP rank owns a KV
+      cache shard, so distinguish by ``attn_dp_rank``.
+    - Pure DP (``attn_dp_size == 1``): every worker has ``attn_dp_rank == 0``,
+      so distinguish by ``dp_rank`` (the data-parallel replica index).
+
+    Both span ``0..dp_size-1``, matching the ``dp_size`` advertised in
+    ``/server_info`` and the per-rank ports the router subscribes to.
+    """
+    if attn_dp_size > 1:
+        return attn_dp_rank
+    return dp_rank or 0
+
+
+class EventBatch(
+    msgspec.Struct,
+    array_like=True,  # type: ignore[call-arg]
+    omit_defaults=True,  # type: ignore[call-arg]
+    gc=False,  # type: ignore[call-arg]
+):
+    ts: float
+    events: list[Any]
+    attn_dp_rank: Optional[int] = None
+
+
+class KVCacheEvent(
+    msgspec.Struct,
+    array_like=True,  # type: ignore[call-arg]
+    omit_defaults=True,  # type: ignore[call-arg]
+    gc=False,  # type: ignore[call-arg]
+    tag=True,
+):
+    """Base class for all KV cache-related events"""
+
+
+class StorageMedium(str, enum.Enum):
+    """Storage tier for KV cache events."""
+
+    GPU = "GPU"  # L1: device HBM
+    CPU = "CPU_PINNED"  # L2: host pinned memory
+    DISK = "DISK"  # L3: SSD / NVMe
+    EXTERNAL = "EXTERNAL"  # L4: shared / remote pool (e.g. Mooncake)
+
+
+class OffloadedState:
+    """
+    OffloadedState represents the state of a KV cache block offloaded to the hicache.
+
+    - prefill_len (int): The length of the prefill part of the KV cache block.
+    - inc_len (int): The length of the incremental part of the KV cache block.
+    - last_hash (Optional[str]): The hash of the last token in the KV cache block.
+    """
+
+    def __init__(
+        self, prefill_len: int, inc_len: int = 0, last_hash: Optional[str] = None
+    ):
+        self.prefill_len = prefill_len
+        self.inc_len = inc_len
+        self.last_hash = last_hash
+
+
+class BlockStored(KVCacheEvent):
+    block_hashes: list[int]
+    parent_block_hash: Optional[int]
+    token_ids: list[int]
+    block_size: int
+    lora_id: Optional[int]
+    medium: Optional[str] = None
+
+
+class BlockRemoved(KVCacheEvent):
+    block_hashes: list[int]
+    medium: Optional[str] = None
+
+
+class AllBlocksCleared(KVCacheEvent):
+    pass
+
+
+class KVEventBatch(EventBatch):
+    events: list[Union[BlockStored, BlockRemoved, AllBlocksCleared]]
+
+
+class EventPublisher(ABC):
+    """
+    Lightweight publisher for EventBatch batches with
+    support for DP attention.
+
+    In DP attention - each rank has its own Scheduler and
+    KV cache instance in order to avoid duplicate events
+    and ensure proper event attribution. In our implementation
+
+    - Each DP rank has its own EventPublisher
+    - Publishers annotate events with the dp rank
+    - This allows consumers to distinguish events from different DP ranks
+    """
+
+    @abstractmethod
+    def publish(self, events: EventBatch) -> None:
+        """Emit events in order.
+
+        Implementations should guarantee at-least-once delivery and
+        monotonic ordering (e.g., via sequence numbers).
+        """
+
+    @abstractmethod
+    def shutdown(self) -> None:
+        """Shutdown the publisher."""
+
+
+class NullEventPublisher(EventPublisher):
+    """No-op implementation (default when disabled)."""
+
+    def publish(self, events) -> None:
+        return
+
+    def shutdown(self) -> None:
+        return
+
+
+class ZmqEventPublisher(EventPublisher):
+    """Reliable PUB/ROUTER publisher with an in-memory replay buffer.
+
+    Spawns a separate thread to handle publishing from a queue.
+
+    Parameters
+    ----------
+    endpoint:
+        PUB address. Use ``tcp://*:5557`` to bind or ``tcp://host:5557`` to
+        connect.
+    replay_endpoint:
+        Optional ROUTER address for replay requests. When given, subscribers can
+        request missed batches by sending the starting sequence number as an
+        8-byte big-endian integer.
+    buffer_steps:
+        Number of past batches to keep for replay.
+    hwm:
+        ZeroMQ high-water-mark for PUB socket.
+    max_queue_size:
+        Maximum number of events to buffer in memory.
+    topic:
+        Topic to publish events to.
+    """
+
+    SHUTDOWN_TIMEOUT: float = 1.0
+    END_SEQ = (-1).to_bytes(8, "big", signed=True)
+
+    def __init__(
+        self,
+        attn_dp_rank: int,
+        endpoint: str = "tcp://*:5557",
+        replay_endpoint: Optional[str] = None,
+        buffer_steps: int = 10_000,
+        hwm: int = 100_000,
+        max_queue_size: int = 100_000,
+        topic: str = "",
+    ) -> None:
+        # Storage
+        self._event_queue = Queue[Optional[EventBatch]](maxsize=max_queue_size)
+        self._buffer = deque[tuple[int, bytes]](maxlen=buffer_steps)
+
+        # ZMQ sockets
+        self._ctx = zmq.Context.instance()
+        self._pub: Optional[zmq.Socket] = None
+        self._replay: Optional[zmq.Socket] = None
+        self._dp_rank = attn_dp_rank
+        self._endpoint = self.offset_endpoint_port(endpoint, self._dp_rank)
+        self._replay_endpoint = self.offset_endpoint_port(
+            replay_endpoint, self._dp_rank
+        )
+        self._hwm = hwm
+        self._socket_setup()
+
+        # Payload
+        self._seq_gen = count()
+        self._topic_bytes = topic.encode("utf-8")
+
+        # Thread
+        self._running = True
+        logger.info("Starting ZMQ publisher thread")
+
+        self._thread = threading.Thread(
+            target=self._publisher_thread, daemon=True, name="zmq-publisher"
+        )
+        self._thread.start()
+
+        atexit.register(self.shutdown)
+
+    def publish(self, events: EventBatch) -> None:
+        if not self._running:
+            raise RuntimeError("Publisher is closed")
+        if events.attn_dp_rank is None:
+            events.attn_dp_rank = self._dp_rank
+        self._event_queue.put(events)
+
+    def shutdown(self) -> None:
+        """Stop the publisher thread and clean up resources."""
+        self._running = False
+        self._event_queue.put_nowait(None)
+
+        start = time.time()
+        pending_items = True
+        while pending_items and (time.time() - start < self.SHUTDOWN_TIMEOUT):
+            pending_items = not self._event_queue.empty()
+            if pending_items:
+                time.sleep(0.1)
+
+        if pending_items:
+            logger.warning(
+                "Warning: Queue still has %s items after %s seconds timeout",
+                self._event_queue.qsize(),
+                self.SHUTDOWN_TIMEOUT,
+            )
+
+        if self._thread.is_alive():
+            self._thread.join(timeout=self.SHUTDOWN_TIMEOUT)
+
+        # Clean up ZMQ resources
+        try:
+            if self._pub is not None:
+                self._pub.close(linger=0)
+            if self._replay is not None:
+                self._replay.close(linger=0)
+        finally:
+            pass  # Do not terminate context; other sockets may use it
+
+    def _socket_setup(self) -> None:
+        """Initialize sockets
+        https://pyzmq.readthedocs.io/en/v19.0.0/morethanbindings.html#thread-safety
+        """
+        if self._pub is None:
+            self._pub = self._ctx.socket(zmq.PUB)
+            self._pub.set_hwm(self._hwm)
+            # Heuristic: bind if wildcard / * present, else connect.
+            # bind stable, connect volatile convention.
+            # ``0.0.0.0`` is the IPv4 bind-all wildcard alongside ``*``
+            # and ``::``; ``/server_info`` advertises it as a wildcard,
+            # so the publisher must bind it for the advertised endpoint
+            # to actually be listening.
+            if (
+                "*" in self._endpoint
+                or "::" in self._endpoint
+                or "0.0.0.0" in self._endpoint
+                or self._endpoint.startswith("ipc://")
+                or self._endpoint.startswith("inproc://")
+            ):
+                logger.debug(
+                    f"ZmqEventPublisher socket publisher_endpoint bind to {self._endpoint}"
+                )
+                self._pub.bind(self._endpoint)
+            else:
+                self._pub.connect(self._endpoint)
+
+        # Set up replay socket: use ROUTER
+        # 1) handles multiple REQ clients (identities)
+        # 2) lets us send back one request → many replies (streamed events)
+        # 3) works in our non‑blocking poll loop alongside PUB
+        if self._replay_endpoint is not None:
+            self._replay = self._ctx.socket(zmq.ROUTER)
+            logger.debug(
+                f"ZmqEventPublisher socket replay_endpoint bind to {self._replay_endpoint}"
+            )
+            self._replay.bind(self._replay_endpoint)
+
+    def _publisher_thread(self) -> None:
+        """Background thread that processes the event queue."""
+        self._pack = msgspec.msgpack.Encoder()
+
+        assert self._pub is not None  # narrows type for mypy
+
+        while self._running or self._event_queue.qsize() > 0:
+            # --- replay (non-critical) ---------------------------------
+            if self._replay is not None and self._replay.poll(0):
+                try:
+                    self._service_replay()
+                except Exception as e:
+                    logger.exception("Error in replay: %s", e)
+
+            # --- main queue (critical) ---------------------------------
+            try:
+                event = self._event_queue.get(timeout=0.1)
+                if event is None:
+                    break  # Sentinel received, exit thread
+            except queue.Empty:
+                continue
+
+            try:
+                seq = next(self._seq_gen)
+
+                payload = self._pack.encode(event)
+                seq_bytes = seq.to_bytes(8, "big")
+                self._pub.send_multipart((self._topic_bytes, seq_bytes, payload))
+
+                self._buffer.append((seq, payload))
+                self._event_queue.task_done()
+
+            except Exception as e:
+                # Publishing failed;  back-off a bit to avoid a tight error loop
+                logger.exception("Error in publisher thread: %s", e)
+                time.sleep(0.1)
+
+    def _service_replay(self) -> None:
+        """If a replay request is waiting, send buffered batches."""
+        assert self._replay is not None  # narrows type for mypy
+
+        frame = self._replay.recv_multipart()
+        if len(frame) != 3:
+            logger.warning("Invalid replay request: %s", frame)
+            return
+        client_id, _, start_seq_bytes = frame
+        start_seq = int.from_bytes(start_seq_bytes, "big")
+
+        for seq, buf in self._buffer:
+            if seq >= start_seq:
+                # [identity, empty_delim, seq_bytes, payload]
+                # (identity, empty_delim) are stripped off by the router
+                # receiving payload is (seq_bytes, payload)
+                self._replay.send_multipart(
+                    (client_id, b"", seq.to_bytes(8, "big"), buf)
+                )
+        # Send end of sequence marker
+        # receiving payload is (-1, b""")
+        self._replay.send_multipart((client_id, b"", self.END_SEQ, b""))
+
+    @staticmethod
+    def offset_endpoint_port(
+        endpoint: Optional[str], data_parallel_rank: int
+    ) -> Optional[str]:
+        """Helper function to offset the port in an endpoint by
+            the data parallel rank.
+
+        Args:
+            endpoint: The endpoint string
+                (e.g., "tcp://*:5557" or "inproc://cache")
+            data_parallel_rank: The data parallel rank to offset by
+
+        Returns:
+            The endpoint with the port offset by data_parallel_rank
+                or suffix appended
+        """
+        # Do nothing if input is None or data_parallel_rank is 0
+        if not endpoint or data_parallel_rank == 0:
+            return endpoint
+
+        if "inproc" in endpoint:
+            return f"{endpoint}_dp{data_parallel_rank}"
+        if "tcp" in endpoint:
+            if endpoint and ":" in endpoint:
+                # Get everything after the last colon (the port)
+                last_colon_idx = endpoint.rfind(":")
+                base_addr = endpoint[:last_colon_idx]
+                base_port = int(endpoint[last_colon_idx + 1 :])
+                new_port = base_port + data_parallel_rank
+                return f"{base_addr}:{new_port}"
+            return endpoint
+        raise ValueError("Invalid endpoint: must contain 'inproc' or 'tcp'")
+
+
+class KVEventsConfig(BaseModel):
+    """Configuration for KV event publishing."""
+
+    publisher: str = "null"
+    """The publisher to use for publishing kv events. Can be "null", "zmq".
+    """
+
+    endpoint: str = "tcp://*:5557"
+    """The zmq endpoint to use for publishing kv events.
+    """
+
+    replay_endpoint: Optional[str] = None
+    """The zmq endpoint to use for replaying kv events.
+    """
+
+    buffer_steps: int = 10_000
+    """The number of steps to cache for replay endpoint. Will only save
+    events from the last N steps for the replay endpoint.
+    """
+
+    hwm: int = 100_000
+    """The zmq high water mark for the event publisher. After queueing N events,
+    events will start dropping if the consumer is not keeping up.
+    """
+
+    max_queue_size: int = 100_000
+    """The maximum number of events to queue while waiting for publishing.
+    """
+
+    topic: str = ""
+    """The topic to use for the event publisher. Consumers can subscribe to
+    this topic to receive events.
+    """
+
+    @classmethod
+    def from_cli(cls, cli_value: str) -> "KVEventsConfig":
+        """Parse the CLI value for the event publisher config."""
+        return KVEventsConfig.model_validate_json(cli_value)
+
+
+class EventPublisherFactory:
+    _registry: dict[str, Callable[..., EventPublisher]] = {
+        "null": NullEventPublisher,
+        "zmq": ZmqEventPublisher,
+    }
+
+    @classmethod
+    def register_publisher(cls, name: str, ctor: Callable[..., EventPublisher]) -> None:
+        if name in cls._registry:
+            raise KeyError(f"publisher '{name}' already registered")
+        cls._registry[name] = ctor
+
+    @classmethod
+    def create(cls, config: Optional[str], attn_dp_rank: int = 0) -> EventPublisher:
+        """Create publisher from a config mapping."""
+        if not config:
+            return NullEventPublisher()
+        config = KVEventsConfig.from_cli(config)
+        config_dict = config.model_dump()
+
+        kind = config_dict.pop("publisher", "null")
+        try:
+            constructor = cls._registry[kind]
+        except KeyError as exc:
+            raise ValueError(f"Unknown event publisher '{kind}'") from exc
+        return constructor(attn_dp_rank=attn_dp_rank, **config_dict)

python/sglang/srt/disaggregation/mooncake/__init__.py

@@ -0,0 +1,6 @@
+from sglang.srt.disaggregation.mooncake.conn import (
+    MooncakeKVBootstrapServer,
+    MooncakeKVManager,
+    MooncakeKVReceiver,
+    MooncakeKVSender,
+)

python/sglang/srt/disaggregation/mooncake/utils.py

@@ -0,0 +1,112 @@
+# Copyright 2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Mooncake-specific utilities for custom memory pool management."""
+
+import logging
+from typing import Any, Optional, Tuple
+
+import torch
+
+from sglang.srt.environ import envs
+
+logger = logging.getLogger(__name__)
+
+# Global constants for custom memory pool types
+SUPPORTED_MOONCAKE_CUSTOM_MEM_POOL_TYPES = ["NVLINK", "BAREX", "INTRA_NODE_NVLINK"]
+
+
+def init_mooncake_custom_mem_pool(
+    device: str,
+) -> Tuple[bool, Optional[Any], Optional[str]]:
+    """
+    Initialize custom memory pool based on environment variable.
+
+    Args:
+        device: The device to allocate memory on
+
+    Returns:
+        Tuple of (enable_custom_mem_pool, custom_mem_pool, custom_mem_pool_type)
+    """
+    enable_custom_mem_pool, custom_mem_pool_type = (
+        check_mooncake_custom_mem_pool_enabled()
+    )
+
+    custom_mem_pool = None
+
+    if enable_custom_mem_pool:
+        try:
+            # TODO(shangming): abstract custom allocator class for more backends
+            if custom_mem_pool_type == "NVLINK":
+                from mooncake.allocator import NVLinkAllocator
+
+                allocator = NVLinkAllocator.get_allocator(device)
+            elif custom_mem_pool_type == "BAREX":
+                from mooncake.allocator import BarexAllocator
+
+                allocator = BarexAllocator.get_allocator(device)
+            elif custom_mem_pool_type == "INTRA_NODE_NVLINK":
+                return False, None, None
+            else:
+                # This should not happen due to the enable_custom_mem_pool check above
+                raise ValueError(
+                    f"Unsupported custom mem pool type: {custom_mem_pool_type}"
+                )
+
+            custom_mem_pool = torch.cuda.MemPool(allocator.allocator())
+            logger.debug(
+                f"Initialized custom memory pool: {custom_mem_pool_type} on device {device}"
+            )
+        except ImportError as e:
+            logger.warning(
+                f"Failed to import mooncake allocator for {custom_mem_pool_type}: {e}. "
+                f"Falling back to default memory pool."
+            )
+            enable_custom_mem_pool = False
+            custom_mem_pool = None
+            custom_mem_pool_type = None
+        except Exception as e:
+            logger.error(
+                f"Failed to initialize custom memory pool {custom_mem_pool_type}: {e}. "
+                f"Falling back to default memory pool."
+            )
+            enable_custom_mem_pool = False
+            custom_mem_pool = None
+            custom_mem_pool_type = None
+    else:
+        return False, None, None
+
+    return enable_custom_mem_pool, custom_mem_pool, custom_mem_pool_type
+
+
+def check_mooncake_custom_mem_pool_enabled() -> Tuple[bool, Optional[str]]:
+    """
+    Check if custom memory pool is enabled without importing allocators.
+
+    Returns:
+        Tuple of (enable_custom_mem_pool, custom_mem_pool_type)
+    """
+    custom_mem_pool_type = envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get()
+
+    if custom_mem_pool_type is not None:
+        # Handle boolean True as NVLINK
+        if custom_mem_pool_type.lower() == "true":
+            custom_mem_pool_type = "NVLINK"
+        enable_custom_mem_pool = (
+            custom_mem_pool_type in SUPPORTED_MOONCAKE_CUSTOM_MEM_POOL_TYPES
+        )
+    else:
+        enable_custom_mem_pool = False
+        custom_mem_pool_type = None
+
+    return enable_custom_mem_pool, custom_mem_pool_type

python/sglang/srt/disaggregation/mori/__init__.py

@@ -0,0 +1,6 @@
+from sglang.srt.disaggregation.mori.conn import (
+    MoriKVBootstrapServer,
+    MoriKVManager,
+    MoriKVReceiver,
+    MoriKVSender,
+)

python/sglang/srt/disaggregation/nixl/__init__.py

@@ -0,0 +1,6 @@
+from sglang.srt.disaggregation.nixl.conn import (
+    NixlKVBootstrapServer,
+    NixlKVManager,
+    NixlKVReceiver,
+    NixlKVSender,
+)

python/sglang/srt/disaggregation/utils.py

@@ -0,0 +1,878 @@
+from __future__ import annotations
+
+import os
+import random
+from collections import deque
+from contextlib import nullcontext
+from enum import Enum
+from typing import TYPE_CHECKING, List, Literal, Optional, Tuple, Type, overload
+
+import numpy as np
+import torch
+import torch.distributed as dist
+
+from sglang.srt.disaggregation.base import KVPoll
+from sglang.srt.environ import envs
+from sglang.srt.utils import is_hip, is_npu
+
+if TYPE_CHECKING:
+    from sglang.srt.disaggregation.base.conn import KVArgs, StateType
+    from sglang.srt.disaggregation.common.conn import (
+        CommonKVBootstrapServer,
+        CommonKVManager,
+        CommonKVReceiver,
+        CommonKVSender,
+    )
+    from sglang.srt.managers.schedule_batch import Req
+    from sglang.srt.server_args import ServerArgs
+
+#########################
+# Constants & Enums
+#########################
+FAKE_BOOTSTRAP_HOST = "2.2.2.2"
+_IS_HIP = is_hip()
+
+
+def is_dsv4_c128_online_enabled() -> bool:
+    """Return whether DSV4 C128 uses request-scoped online state."""
+    return not _IS_HIP and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get()
+
+
+def get_dsv4_c128_state_indices(
+    req_pool_idx: int,
+    seq_len: int,
+    *,
+    online: bool,
+    ring_size: int,
+) -> np.ndarray:
+    """Return the PD transfer row/page indices for DSV4 C128 state."""
+    if seq_len == 0 or seq_len % 128 == 0:
+        return np.empty((0,), dtype=np.int32)
+    if online:
+        return np.array([int(req_pool_idx)], dtype=np.int32)
+
+    assert ring_size % 128 == 0, f"C128 ring_size must be 128-aligned, got {ring_size}"
+    pages_per_req = ring_size // 128
+    page = int(req_pool_idx) * pages_per_req + ((seq_len - 1) % ring_size) // 128
+    return np.array([page], dtype=np.int32)
+
+
+class DisaggregationMode(Enum):
+    NULL = "null"
+    PREFILL = "prefill"
+    DECODE = "decode"
+
+    @staticmethod
+    def to_engine_type(mode: str) -> str:
+        if mode == DisaggregationMode.PREFILL.value:
+            return "prefill"
+        elif mode == DisaggregationMode.DECODE.value:
+            return "decode"
+        return "unified"
+
+
+#########################
+# Synchronization
+#########################
+
+
+def _get_failure_prob() -> float:
+    try:
+        return float(envs.SGLANG_TEST_DISAGG_FAILURE_PROB.get())
+    except Exception:
+        # fallback to legacy env var
+        return float(os.getenv("DISAGGREGATION_TEST_FAILURE_PROB", "0"))
+
+
+def _poll_with_failure_injection(pollers) -> List[int]:
+    if (failure_prob := _get_failure_prob()) > 0:
+        return [
+            int(KVPoll.Failed) if random.random() < failure_prob else int(poller.poll())
+            for poller in pollers
+        ]
+    return [int(poller.poll()) for poller in pollers]
+
+
+def _is_fake_transfer(req: Req, server_args: ServerArgs) -> bool:
+    return req.bootstrap_host == FAKE_BOOTSTRAP_HOST or (
+        req.bootstrap_host is None
+        and server_args.disaggregation_transfer_backend == "fake"
+    )
+
+
+def _apply_metadata_gate(polls, decode_reqs, metadata_buffers, server_args) -> None:
+    """Downgrade Success → Transferring for requests whose metadata hasn't landed.
+
+    Mutates `polls` in-place. Called before all-reduce so that MIN across TP
+    ranks naturally prevents any rank from committing before all ranks are ready.
+    """
+    for i, poll_val in enumerate(polls):
+        if poll_val == int(KVPoll.Success):
+            decode_req = decode_reqs[i]
+            if _is_fake_transfer(decode_req.req, server_args):
+                continue
+            actual_room = metadata_buffers.bootstrap_room[
+                decode_req.metadata_buffer_index, 0
+            ].item()
+            if actual_room == 0:
+                polls[i] = int(KVPoll.Transferring)
+
+
+def poll_and_all_reduce(
+    pollers,
+    gloo_group: dist.ProcessGroup,
+    decode_reqs=None,
+    metadata_buffers: Optional[MetadataBuffers] = None,
+    server_args: Optional[ServerArgs] = None,
+):
+    # at a certain prob, the poll is failed to simulate failure
+    polls = _poll_with_failure_injection(pollers)
+
+    # Apply metadata gate on the decode requests to downgrade Success → Transferring for requests whose metadata hasn't landed.
+    if (
+        decode_reqs is not None
+        and metadata_buffers is not None
+        and server_args is not None
+    ):
+        _apply_metadata_gate(polls, decode_reqs, metadata_buffers, server_args)
+    tensor_to_reduce = torch.tensor(polls, dtype=torch.uint8, device="cpu")
+    dist.all_reduce(tensor_to_reduce, op=dist.ReduceOp.MIN, group=gloo_group)
+    return tensor_to_reduce.tolist()
+
+
+def poll_and_all_reduce_attn_cp_tp_group(
+    pollers,
+    attn_cp_cpu_group: dist.ProcessGroup,
+    attn_tp_cpu_group: dist.ProcessGroup,
+):
+    # First sync across attn-tp ranks so all TP participants for a given (dp, cp)
+    # shard observe the same status transitions.
+    polls = poll_and_all_reduce(pollers, attn_tp_cpu_group)
+
+    # Then sync across attn-cp ranks, so all TPxCP participants in one DP shard
+    # converge to the same global status.
+    tensor_to_reduce = torch.tensor(polls, dtype=torch.uint8, device="cpu")
+    dist.all_reduce(
+        tensor_to_reduce,
+        op=dist.ReduceOp.MIN,
+        group=attn_cp_cpu_group,
+    )
+    return tensor_to_reduce.tolist()
+
+
+def poll_and_all_reduce_with_staging(
+    decode_reqs,
+    staging_handler,
+    gloo_group: dist.ProcessGroup,
+    metadata_buffers: Optional[MetadataBuffers] = None,
+    server_args: Optional[ServerArgs] = None,
+):
+    """Staging-aware polling: advance scatter, demote incomplete transfers, all_reduce."""
+    for decode_req in decode_reqs:
+        if decode_req.kv_receiver.require_staging and not staging_handler.is_done(
+            decode_req
+        ):
+            staging_handler.advance_scatter(decode_req)
+
+    # allow test injection of failure probability at runtime
+    receivers = [dr.kv_receiver for dr in decode_reqs]
+    raw_polls = _poll_with_failure_injection(receivers)
+    for i, decode_req in enumerate(decode_reqs):
+        if raw_polls[i] == int(KVPoll.Success):
+            if decode_req.kv_receiver.require_staging and not staging_handler.is_done(
+                decode_req
+            ):
+                raw_polls[i] = int(KVPoll.Transferring)
+    # Apply metadata gate on the decode requests to downgrade Success → Transferring for requests whose metadata hasn't landed.
+    if metadata_buffers is not None and server_args is not None:
+        _apply_metadata_gate(raw_polls, decode_reqs, metadata_buffers, server_args)
+    poll_tensor = torch.tensor(raw_polls, dtype=torch.uint8, device="cpu")
+    dist.all_reduce(poll_tensor, op=dist.ReduceOp.MIN, group=gloo_group)
+    return poll_tensor.tolist()
+
+
+#########################
+# Metadata Buffers
+#########################
+
+
+class ReqToMetadataIdxAllocator:
+    """A memory pool that maps a request to its first output token location."""
+
+    def __init__(
+        self,
+        size: int,
+    ):
+        self.size = size
+        self.free_slots = deque(list(range(size)))
+
+    def available_size(self):
+        return len(self.free_slots)
+
+    def alloc(self) -> Optional[int]:
+        if len(self.free_slots) == 0:
+            return None
+
+        return self.free_slots.popleft()
+
+    def free(self, free_index: int):
+        self.free_slots.append(free_index)
+
+
+class MetadataBuffers:
+    def __init__(
+        self,
+        size: int,
+        hidden_size: int,
+        hidden_states_dtype: torch.dtype,
+        max_top_logprobs_num: int = 128,
+        custom_mem_pool: torch.cuda.MemPool = None,
+    ):
+        self.custom_mem_pool = custom_mem_pool
+        bootstrap_room_dtype = torch.uint64
+        device = "cpu"
+        if is_npu():
+            # For ascend backend, output tokens are placed in the NPU and will be transferred by D2D channel.
+            device = "npu"
+            # TODO: Fix me when npu backend supports torch.uint64
+            bootstrap_room_dtype = torch.int64
+        elif self.custom_mem_pool:
+            # TODO(shangming): Fix me (use 'cuda') when nvlink_transport of Mooncake is bug-free
+            device = "cpu"
+        elif envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get() == "INTRA_NODE_NVLINK":
+            device = "cuda"
+        with (
+            torch.cuda.use_mem_pool(self.custom_mem_pool)
+            if self.custom_mem_pool
+            else nullcontext()
+        ):
+            # TODO: abort top_logprobs_num > 128 in PD
+
+            # We transfer the metadata of first output token to decode
+            # The minimal size for RDMA is 64Bytes, so we pad it to > 64Bytes
+            self.output_ids = torch.zeros((size, 16), dtype=torch.int32, device=device)
+            self.cached_tokens = torch.zeros(
+                (size, 16), dtype=torch.int32, device=device
+            )
+            self.output_token_logprobs_val = torch.zeros(
+                (size, 16), dtype=torch.float32, device=device
+            )
+            self.output_token_logprobs_idx = torch.zeros(
+                (size, 16), dtype=torch.int32, device=device
+            )
+            self.output_top_logprobs_val = torch.zeros(
+                (size, max_top_logprobs_num), dtype=torch.float32, device=device
+            )
+            self.output_top_logprobs_idx = torch.zeros(
+                (size, max_top_logprobs_num), dtype=torch.int32, device=device
+            )
+            # For PD + spec decode
+            self.output_topk_p = torch.zeros(
+                (size, 16), dtype=torch.float32, device=device
+            )
+            self.output_topk_index = torch.zeros(
+                (size, 16), dtype=torch.int64, device=device
+            )
+            self.output_hidden_states = torch.zeros(
+                (size, hidden_size), dtype=hidden_states_dtype, device=device
+            )
+            # Request validation: store bootstrap_room to detect metadata corruption
+            self.bootstrap_room = torch.zeros(
+                (size, 8), dtype=bootstrap_room_dtype, device=device
+            )
+
+    def get_buf_infos(self):
+        ptrs = [
+            self.output_ids.data_ptr(),
+            self.cached_tokens.data_ptr(),
+            self.output_token_logprobs_val.data_ptr(),
+            self.output_token_logprobs_idx.data_ptr(),
+            self.output_top_logprobs_val.data_ptr(),
+            self.output_top_logprobs_idx.data_ptr(),
+            self.output_topk_p.data_ptr(),
+            self.output_topk_index.data_ptr(),
+            self.output_hidden_states.data_ptr(),
+            self.bootstrap_room.data_ptr(),
+        ]
+        data_lens = [
+            self.output_ids.nbytes,
+            self.cached_tokens.nbytes,
+            self.output_token_logprobs_val.nbytes,
+            self.output_token_logprobs_idx.nbytes,
+            self.output_top_logprobs_val.nbytes,
+            self.output_top_logprobs_idx.nbytes,
+            self.output_topk_p.nbytes,
+            self.output_topk_index.nbytes,
+            self.output_hidden_states.nbytes,
+            self.bootstrap_room.nbytes,
+        ]
+        item_lens = [
+            self.output_ids[0].nbytes,
+            self.cached_tokens[0].nbytes,
+            self.output_token_logprobs_val[0].nbytes,
+            self.output_token_logprobs_idx[0].nbytes,
+            self.output_top_logprobs_val[0].nbytes,
+            self.output_top_logprobs_idx[0].nbytes,
+            self.output_topk_p[0].nbytes,
+            self.output_topk_index[0].nbytes,
+            self.output_hidden_states[0].nbytes,
+            self.bootstrap_room[0].nbytes,
+        ]
+        return ptrs, data_lens, item_lens
+
+    def get_buf(self, idx: int):
+        return (
+            self.output_ids[idx].clone(),
+            self.cached_tokens[idx].clone(),
+            self.output_token_logprobs_val[idx].clone(),
+            self.output_token_logprobs_idx[idx].clone(),
+            self.output_top_logprobs_val[idx].clone(),
+            self.output_top_logprobs_idx[idx].clone(),
+            self.output_topk_p[idx].clone(),
+            self.output_topk_index[idx].clone(),
+            self.output_hidden_states[idx].clone(),
+            self.bootstrap_room[idx].clone(),
+        )
+
+    def set_buf(self, req: Req):
+
+        self.output_ids[req.metadata_buffer_index][0] = req.output_ids[0]
+        # The cached_tokens buffer is (size, 16); slots 0-3 hold cached token
+        # counts and slots 4-6 are reused for multimodal prompt token counts
+        # (slots 7-15 remain spare). This avoids adding new RDMA buffers.
+        # Slot map: 0=cached 1=device 2=host 3=storage 4=image 5=audio 6=video.
+        self.cached_tokens[req.metadata_buffer_index][0] = req.cached_tokens
+        self.cached_tokens[req.metadata_buffer_index][1] = req.cached_tokens_device
+        self.cached_tokens[req.metadata_buffer_index][2] = req.cached_tokens_host
+        self.cached_tokens[req.metadata_buffer_index][3] = req.cached_tokens_storage
+
+        # Compute multimodal prompt token counts on the prefill node so decode
+        # can report them in usage.
+        if req.multimodal_inputs:
+            image_t, audio_t, video_t = req.multimodal_inputs.compute_mm_token_counts()
+        else:
+            image_t = audio_t = video_t = 0
+        self.cached_tokens[req.metadata_buffer_index][4] = image_t
+        self.cached_tokens[req.metadata_buffer_index][5] = audio_t
+        self.cached_tokens[req.metadata_buffer_index][6] = video_t
+        if req.return_logprob:
+            if req.logprob.output_token_logprobs_val:  # not none or empty list
+                self.output_token_logprobs_val[req.metadata_buffer_index][0] = (
+                    req.logprob.output_token_logprobs_val[0]
+                )
+            if req.logprob.output_token_logprobs_idx:  # not none or empty list
+                self.output_token_logprobs_idx[req.metadata_buffer_index][0] = (
+                    req.logprob.output_token_logprobs_idx[0]
+                )
+
+            if req.logprob.output_top_logprobs_val:  # not none or empty list
+                self.output_top_logprobs_val[req.metadata_buffer_index][
+                    : len(req.logprob.output_top_logprobs_val[0])
+                ] = torch.tensor(
+                    req.logprob.output_top_logprobs_val[0],
+                    dtype=torch.float32,
+                    device="cpu",
+                )
+            if req.logprob.output_top_logprobs_idx:  # not none or empty list
+                self.output_top_logprobs_idx[req.metadata_buffer_index][
+                    : len(req.logprob.output_top_logprobs_idx[0])
+                ] = torch.tensor(
+                    req.logprob.output_top_logprobs_idx[0],
+                    dtype=torch.int32,
+                    device="cpu",
+                )
+        # For PD + spec decode
+        if req.hidden_states_tensor is not None:
+            # speculative_eagle_topk should not be greater than 16 currently
+            topk = req.output_topk_p.size(0)
+
+            self.output_topk_p[req.metadata_buffer_index, :topk].copy_(
+                req.output_topk_p
+            )
+            self.output_topk_index[req.metadata_buffer_index, :topk].copy_(
+                req.output_topk_index
+            )
+            self.output_hidden_states[req.metadata_buffer_index].copy_(
+                req.hidden_states_tensor
+            )
+        # Store bootstrap_room for validation on decode side
+        self.bootstrap_room[req.metadata_buffer_index, 0] = (
+            req.bootstrap_room if req.bootstrap_room is not None else 0
+        )
+
+
+#########################
+# Transfer Backend
+#########################
+
+
+class TransferBackend(Enum):
+    MOONCAKE = "mooncake"
+    MORI = "mori"
+    NIXL = "nixl"
+    ASCEND = "ascend"
+    FAKE = "fake"
+
+
+class KVClassType(Enum):
+    KVARGS = "kvargs"
+    MANAGER = "manager"
+    SENDER = "sender"
+    RECEIVER = "receiver"
+    BOOTSTRAP_SERVER = "bootstrap_server"
+
+
+@overload
+def get_kv_class(
+    transfer_backend: TransferBackend, class_type: Literal[KVClassType.KVARGS]
+) -> Type[KVArgs]: ...
+@overload
+def get_kv_class(
+    transfer_backend: TransferBackend, class_type: Literal[KVClassType.MANAGER]
+) -> Type[CommonKVManager]: ...
+@overload
+def get_kv_class(
+    transfer_backend: TransferBackend, class_type: Literal[KVClassType.SENDER]
+) -> Type[CommonKVSender]: ...
+@overload
+def get_kv_class(
+    transfer_backend: TransferBackend, class_type: Literal[KVClassType.RECEIVER]
+) -> Type[CommonKVReceiver]: ...
+@overload
+def get_kv_class(
+    transfer_backend: TransferBackend, class_type: Literal[KVClassType.BOOTSTRAP_SERVER]
+) -> Type[CommonKVBootstrapServer]: ...
+
+
+def get_kv_class(
+    transfer_backend: TransferBackend, class_type: KVClassType
+) -> Optional[Type]:
+    from sglang.srt.disaggregation.fake import FakeKVReceiver, FakeKVSender
+
+    if transfer_backend == TransferBackend.MOONCAKE:
+        from sglang.srt.disaggregation.base import KVArgs
+        from sglang.srt.disaggregation.mooncake import (
+            MooncakeKVBootstrapServer,
+            MooncakeKVManager,
+            MooncakeKVReceiver,
+            MooncakeKVSender,
+        )
+
+        class_mapping = {
+            KVClassType.KVARGS: KVArgs,
+            KVClassType.MANAGER: MooncakeKVManager,
+            KVClassType.SENDER: MooncakeKVSender,
+            KVClassType.RECEIVER: (MooncakeKVReceiver),
+            KVClassType.BOOTSTRAP_SERVER: MooncakeKVBootstrapServer,
+        }
+        return class_mapping.get(class_type)
+    elif transfer_backend == TransferBackend.MORI:
+        from sglang.srt.disaggregation.base import KVArgs
+        from sglang.srt.disaggregation.mori import (
+            MoriKVBootstrapServer,
+            MoriKVManager,
+            MoriKVReceiver,
+            MoriKVSender,
+        )
+
+        class_mapping = {
+            KVClassType.KVARGS: KVArgs,
+            KVClassType.MANAGER: MoriKVManager,
+            KVClassType.SENDER: MoriKVSender,
+            KVClassType.RECEIVER: (MoriKVReceiver),
+            KVClassType.BOOTSTRAP_SERVER: MoriKVBootstrapServer,
+        }
+        return class_mapping.get(class_type)
+    elif transfer_backend == TransferBackend.ASCEND:
+        from sglang.srt.disaggregation.ascend import (
+            AscendKVBootstrapServer,
+            AscendKVManager,
+            AscendKVReceiver,
+            AscendKVSender,
+        )
+        from sglang.srt.disaggregation.base import KVArgs
+
+        class_mapping = {
+            KVClassType.KVARGS: KVArgs,
+            KVClassType.MANAGER: AscendKVManager,
+            KVClassType.SENDER: AscendKVSender,
+            KVClassType.RECEIVER: (AscendKVReceiver),
+            KVClassType.BOOTSTRAP_SERVER: AscendKVBootstrapServer,
+        }
+        return class_mapping.get(class_type)
+    elif transfer_backend == TransferBackend.NIXL:
+        from sglang.srt.disaggregation.base import KVArgs
+        from sglang.srt.disaggregation.nixl import (
+            NixlKVBootstrapServer,
+            NixlKVManager,
+            NixlKVReceiver,
+            NixlKVSender,
+        )
+
+        class_mapping = {
+            KVClassType.KVARGS: KVArgs,
+            KVClassType.MANAGER: NixlKVManager,
+            KVClassType.SENDER: NixlKVSender,
+            KVClassType.RECEIVER: (NixlKVReceiver),
+            KVClassType.BOOTSTRAP_SERVER: NixlKVBootstrapServer,
+        }
+        return class_mapping.get(class_type)
+    elif transfer_backend == TransferBackend.FAKE:
+        from sglang.srt.disaggregation.base import KVArgs
+        from sglang.srt.disaggregation.fake import (
+            FakeKVManager,
+            FakeKVReceiver,
+            FakeKVSender,
+        )
+
+        class_mapping = {
+            KVClassType.KVARGS: KVArgs,
+            KVClassType.MANAGER: FakeKVManager,
+            KVClassType.SENDER: FakeKVSender,
+            KVClassType.RECEIVER: (FakeKVReceiver),
+        }
+        return class_mapping.get(class_type)
+
+    raise ValueError(f"Unsupported transfer backend: {transfer_backend}")
+
+
+def _get_cp_rank_page_bounds(
+    total_pages: int, cp_rank: int, cp_size: int
+) -> Tuple[int, int]:
+    base = total_pages // cp_size
+    rem = total_pages % cp_size
+    local_start = cp_rank * base + min(cp_rank, rem)
+    n_pages = base + (1 if cp_rank < rem else 0)
+    return local_start, local_start + n_pages
+
+
+def page_indices_to_cp_rank_page_indices(
+    page_indices: np.ndarray,
+    total_pages: int,
+    cp_rank: int,
+    cp_size: int,
+) -> np.ndarray:
+    """
+    Filter page_indices (which are *global* page ids in the KV pool) to those
+    belonging to the given CP rank for this request.
+
+    For a single request, its pages occupy a contiguous global range
+    [first_page, first_page + total_pages). We first compute the local
+    split [0, total_pages) across cp_size ranks, then shift that local
+    range by first_page back into the global page id space and take
+    the intersection with page_indices.
+
+    Returns:
+        Subset of page_indices that fall in this rank's global
+        [start_page, end_page) slice for the given CP rank.
+    """
+    if cp_size <= 1:
+        return page_indices
+
+    if page_indices.size == 0:
+        return np.asarray(page_indices)
+
+    first_page = int(page_indices.min())
+    base = total_pages // cp_size
+    rem = total_pages % cp_size
+
+    if rem == 0:
+        local_start = cp_rank * base
+        local_end = local_start + base
+    else:
+        local_start = cp_rank * base + min(cp_rank, rem)
+        n_pages = base + (1 if cp_rank < rem else 0)
+        local_end = local_start + n_pages
+
+    # Map back to global page ids.
+    start_page = first_page + local_start
+    end_page = first_page + local_end
+
+    mask = (page_indices >= start_page) & (page_indices < end_page)
+    return np.asarray(page_indices)[mask]
+
+
+def filter_kv_indices_for_cp_rank(
+    kv_mgr: CommonKVManager,
+    kv_indices: np.ndarray,
+    index_slice: slice,
+    total_pages: Optional[int] = None,
+) -> Tuple[np.ndarray, slice]:
+    """Filters kv_indices and index_slice for the current CP rank."""
+    if total_pages is None:
+        total_pages = len(kv_indices)
+    cp_rank = kv_mgr.attn_cp_rank
+    cp_size = kv_mgr.attn_cp_size
+
+    if cp_size <= 1:
+        return kv_indices, index_slice
+
+    rank_start, rank_end = _get_cp_rank_page_bounds(total_pages, cp_rank, cp_size)
+    chunk_start = index_slice.start if index_slice.start is not None else 0
+    chunk_end = index_slice.stop if index_slice.stop is not None else total_pages
+    first_pos = max(rank_start, chunk_start) - chunk_start
+    last_pos = min(rank_end, chunk_end) - chunk_start
+
+    if last_pos <= first_pos:
+        new_kv_indices = kv_indices[:0]
+        new_index_slice = slice(chunk_start, chunk_start)
+    else:
+        new_kv_indices = kv_indices[first_pos:last_pos]
+        new_index_slice = slice(
+            chunk_start + first_pos,
+            chunk_start + last_pos,
+        )
+    return new_kv_indices, new_index_slice
+
+
+#########################
+# Misc
+#########################
+
+
+def is_mla_backend(target_kv_pool) -> bool:
+    from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
+    from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
+
+    return isinstance(target_kv_pool, (MLATokenToKVPool, DeepSeekV4TokenToKVPool))
+
+
+def append_state_component(
+    kv_args: KVArgs,
+    state_type: StateType,
+    data_ptrs: List[int],
+    data_lens: List[int],
+    item_lens: List[int],
+    dim_per_tensor: Optional[List[int]] = None,
+) -> None:
+    """Append one state component. Caller orders state_types consistently
+    on prefill and decode sides."""
+    kv_args.state_types.append(state_type)
+    kv_args.state_data_ptrs.append(data_ptrs)
+    kv_args.state_data_lens.append(data_lens)
+    kv_args.state_item_lens.append(item_lens)
+    kv_args.state_dim_per_tensor.append(dim_per_tensor or [])
+
+
+def setup_state_kv_args(
+    kv_args: KVArgs,
+    token_to_kv_pool,
+    draft_token_to_kv_pool=None,
+    total_kv_layers: int = None,
+    req_to_token_pool=None,
+) -> None:
+    """Populate ``kv_args`` state-buffer fields from the given pool.
+    Shared by prefill and decode bootstrap paths so the state_type dispatch
+    lives in one place.
+    """
+    from sglang.srt.disaggregation.base.conn import StateType
+    from sglang.srt.hardware_backend.npu.memory_pool_npu import NPUMLATokenToKVPool
+    from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
+    from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
+    from sglang.srt.mem_cache.memory_pool import (
+        DSATokenToKVPool,
+        HybridLinearKVPool,
+        MiniMaxSparseKVPool,
+    )
+
+    kv_args.state_types = []
+    kv_args.state_data_ptrs = []
+    kv_args.state_data_lens = []
+    kv_args.state_item_lens = []
+    kv_args.state_dim_per_tensor = []
+    kv_args.is_hybrid_mla_backend = False
+
+    if isinstance(token_to_kv_pool, MiniMaxSparseKVPool):
+        if token_to_kv_pool.index_kv_pool is not None:
+            raise NotImplementedError(
+                "PD disaggregation for MiniMax sparse layers with index value "
+                "(index_kv_pool) is not yet supported; only K-only sparse layers are."
+            )
+        if token_to_kv_pool.index_k_pool is not None:
+            dp, dl, il = token_to_kv_pool.get_index_k_state_buf_infos()
+            append_state_component(kv_args, StateType.MINIMAX_INDEX_K, dp, dl, il)
+    elif hasattr(token_to_kv_pool, "get_state_buf_infos"):
+        data_ptrs, data_lens, item_lens = token_to_kv_pool.get_state_buf_infos()
+
+        # DeepSeekV4TokenToKVPool inherits BaseSWAKVPool; its heterogeneous
+        # state list is described per-entry via get_state_buf_infos.
+        if isinstance(token_to_kv_pool, BaseSWAKVPool):
+            append_state_component(
+                kv_args, StateType.SWA, data_ptrs, data_lens, item_lens
+            )
+            # unified_kv: the SWA ring lives in the unified buffers (no separate
+            # swa_kv_pool) and is addressed per-row, so ship it as SWA_RING.
+            if getattr(token_to_kv_pool, "_unified_kv", False) and hasattr(
+                token_to_kv_pool, "get_unified_swa_ring_buf_infos"
+            ):
+                ring_ptrs, ring_lens, ring_item_lens = (
+                    token_to_kv_pool.get_unified_swa_ring_buf_infos()
+                )
+                if ring_ptrs:
+                    append_state_component(
+                        kv_args,
+                        StateType.SWA_RING,
+                        ring_ptrs,
+                        ring_lens,
+                        ring_item_lens,
+                    )
+            if hasattr(token_to_kv_pool, "get_c128_state_buf_infos"):
+                c128_ptrs, c128_lens, c128_item_lens = (
+                    token_to_kv_pool.get_c128_state_buf_infos()
+                )
+                if c128_ptrs:
+                    append_state_component(
+                        kv_args,
+                        StateType.C128_STATE,
+                        c128_ptrs,
+                        c128_lens,
+                        c128_item_lens,
+                    )
+        elif isinstance(token_to_kv_pool, HybridLinearKVPool):
+            dim = (
+                token_to_kv_pool.get_state_dim_per_tensor()
+                if hasattr(token_to_kv_pool, "get_state_dim_per_tensor")
+                else None
+            )
+            kv_args.is_hybrid_mla_backend = is_mla_backend(
+                token_to_kv_pool.full_kv_pool
+            )
+            append_state_component(
+                kv_args, StateType.MAMBA, data_ptrs, data_lens, item_lens, dim
+            )
+        elif isinstance(token_to_kv_pool, (DSATokenToKVPool, NPUMLATokenToKVPool)):
+            if draft_token_to_kv_pool is not None and isinstance(
+                draft_token_to_kv_pool, DSATokenToKVPool
+            ):
+                (
+                    draft_data_ptrs,
+                    draft_data_lens,
+                    draft_item_lens,
+                ) = draft_token_to_kv_pool.get_state_buf_infos()
+                data_ptrs = data_ptrs + draft_data_ptrs
+                data_lens = data_lens + draft_data_lens
+                item_lens = item_lens + draft_item_lens
+            if isinstance(token_to_kv_pool, NPUMLATokenToKVPool):
+                kv_args.kv_buf_groups = (
+                    len(kv_args.kv_data_ptrs) // token_to_kv_pool.layer_num
+                )
+                kv_args.total_kv_layers = total_kv_layers
+            else:
+                append_state_component(
+                    kv_args, StateType.DSA, data_ptrs, data_lens, item_lens
+                )
+
+    # DSV4 NextN shares the target allocator, so target and draft use the same
+    # local SWA indices. Keep draft buffers in a separate positional component
+    # to avoid mixing them into the target's heterogeneous state layout, while
+    # reusing the existing SWA transport dispatch. NPU has a different paged
+    # state layout and is intentionally left unchanged.
+    if (
+        not is_npu()
+        and isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool)
+        and isinstance(draft_token_to_kv_pool, DeepSeekV4TokenToKVPool)
+    ):
+        if not draft_token_to_kv_pool.compression_ratios or not all(
+            ratio == 0 for ratio in draft_token_to_kv_pool.compression_ratios
+        ):
+            raise RuntimeError(
+                "DSV4 draft state transfer expects SWA-only NextN layers"
+            )
+        if token_to_kv_pool._unified_kv != draft_token_to_kv_pool._unified_kv:
+            raise RuntimeError(
+                "DSV4 target and draft pools must use the same unified-KV mode"
+            )
+
+        if token_to_kv_pool._unified_kv:
+            target_geometry = (
+                token_to_kv_pool.unified_swa_window,
+                token_to_kv_pool.unified_swa_ring_size,
+                token_to_kv_pool.unified_swa_pages,
+            )
+            draft_geometry = (
+                draft_token_to_kv_pool.unified_swa_window,
+                draft_token_to_kv_pool.unified_swa_ring_size,
+                draft_token_to_kv_pool.unified_swa_pages,
+            )
+            if target_geometry != draft_geometry:
+                raise RuntimeError(
+                    "DSV4 target and draft pools must share SWA ring geometry: "
+                    f"target={target_geometry}, draft={draft_geometry}"
+                )
+            draft_ptrs, draft_lens, draft_item_lens = (
+                draft_token_to_kv_pool.get_unified_swa_ring_buf_infos()
+            )
+            draft_state_type = StateType.SWA_RING
+        else:
+            if (
+                token_to_kv_pool.full_to_swa_index_mapping
+                is not draft_token_to_kv_pool.full_to_swa_index_mapping
+            ):
+                raise RuntimeError(
+                    "DSV4 target and draft pools must share the SWA index mapping"
+                )
+            target_geometry = (
+                token_to_kv_pool.page_size,
+                token_to_kv_pool.sliding_window,
+            )
+            draft_geometry = (
+                draft_token_to_kv_pool.page_size,
+                draft_token_to_kv_pool.sliding_window,
+            )
+            if target_geometry != draft_geometry:
+                raise RuntimeError(
+                    "DSV4 target and draft pools must share paged SWA geometry: "
+                    f"target={target_geometry}, draft={draft_geometry}"
+                )
+            draft_ptrs, draft_lens, draft_item_lens = (
+                draft_token_to_kv_pool.get_state_buf_infos()
+            )
+            draft_state_type = StateType.SWA
+
+        if draft_ptrs:
+            append_state_component(
+                kv_args,
+                draft_state_type,
+                draft_ptrs,
+                draft_lens,
+                draft_item_lens,
+            )
+
+    if (
+        StateType.MAMBA not in kv_args.state_types
+        and req_to_token_pool is not None
+        and hasattr(req_to_token_pool, "get_state_buf_infos")
+    ):
+        data_ptrs, data_lens, item_lens = req_to_token_pool.get_state_buf_infos()
+        if data_ptrs:
+            dim = (
+                req_to_token_pool.get_state_dim_per_tensor()
+                if hasattr(req_to_token_pool, "get_state_dim_per_tensor")
+                else None
+            )
+            append_state_component(
+                kv_args, StateType.MAMBA, data_ptrs, data_lens, item_lens, dim
+            )
+
+
+def prepare_abort(req: Req, error_message: str, status_code=None):
+    from sglang.srt.managers.schedule_batch import FINISH_ABORT
+
+    # populate finish metadata and stream output
+    req.finished_reason = FINISH_ABORT(error_message, status_code)
+
+    if req.return_logprob:
+        req.logprob.input_token_logprobs_val = []
+        req.logprob.input_token_logprobs_idx = []
+        req.logprob.input_top_logprobs_val = []
+        req.logprob.input_top_logprobs_idx = []
+        req.logprob.input_token_ids_logprobs_val = []
+        req.logprob.input_token_ids_logprobs_idx = []
+
+
+def is_aborted(req: Req) -> bool:
+    from sglang.srt.managers.schedule_batch import FINISH_ABORT
+
+    return isinstance(req.to_finish, FINISH_ABORT) or isinstance(
+        req.finished_reason, FINISH_ABORT
+    )