chore: import upstream snapshot with attribution

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