diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 7ef85fdb6..326922645 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -148,6 +148,7 @@ class CommonKVManager(BaseKVManager): # These timeout requests should be aborted to release the tree cache. self.bootstrap_timeout = envs.SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT.get() elif self.disaggregation_mode == DisaggregationMode.DECODE: + self.enable_staging: bool = False self.connection_pool: Dict[str, Dict[str, Union[str, int]]] = {} self.connection_lock = threading.Lock() self.required_prefill_response_num_table: Dict[int, int] = {} @@ -501,6 +502,7 @@ class CommonKVReceiver(BaseKVReceiver): self.bootstrap_addr = bootstrap_addr self.kv_mgr = mgr self.conclude_state: Optional[KVPoll] = None + self.require_staging: bool = False self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room) self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping) @@ -529,6 +531,12 @@ class CommonKVReceiver(BaseKVReceiver): self.required_prefill_response_num ) + if self.kv_mgr.enable_staging: + self.require_staging = ( + self.prefill_info.attn_tp_size != 0 + and self.prefill_info.attn_tp_size != self.kv_mgr.attn_tp_size + ) + self.prefill_dp_rank = prefill_dp_rank self._setup_bootstrap_infos() self.kv_mgr.update_status(self.bootstrap_room, KVPoll.WaitingForInput) diff --git a/python/sglang/srt/disaggregation/common/staging_buffer.py b/python/sglang/srt/disaggregation/common/staging_buffer.py new file mode 100644 index 000000000..4380c3ce7 --- /dev/null +++ b/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 (NVLink incompatible!)" + 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." + ) diff --git a/python/sglang/srt/disaggregation/common/staging_handler.py b/python/sglang/srt/disaggregation/common/staging_handler.py new file mode 100644 index 000000000..b353f4609 --- /dev/null +++ b/python/sglang/srt/disaggregation/common/staging_handler.py @@ -0,0 +1,732 @@ +""" +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 + ) + prefetch_requested: 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 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) + + +# ====================================================================== +# Mooncake-specific staging protocol and 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, + ) -> Tuple[bool, int, int, int, int]: + """Check if staging offset and watermark are ready for this chunk. + + 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 not self.kv_manager._is_watermark_ready( + req.mooncake_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 init_staging_buffers(engine, kv_args, count: int) -> list: + """Create prefill-side staging buffers and register them with the engine. + + Returns list of StagingBuffer instances. + """ + from sglang.srt.disaggregation.common.staging_buffer import StagingBuffer + from sglang.srt.disaggregation.mooncake.utils import ( + init_mooncake_custom_mem_pool, + ) + 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, pool_type = init_mooncake_custom_mem_pool(device) + if custom_mem_pool is None: + logger.warning( + "No mooncake custom mem pool available for staging buffer. " + "NVLink transport will NOT work. Set SGLANG_MOONCAKE_CUSTOM_MEM_POOL." + ) + + buffers = [] + for _ in range(count): + buf = StagingBuffer(size_bytes, device, gpu_id, custom_mem_pool=custom_mem_pool) + engine.batch_register([buf.get_ptr()], [buf.get_size()]) + buffers.append(buf) + return buffers + + +def init_staging_allocator(engine, kv_args): + """Create decode-side staging ring-buffer allocator and register with engine. + + Returns a StagingAllocator instance. + """ + from sglang.srt.disaggregation.common.staging_buffer import StagingAllocator + from sglang.srt.disaggregation.mooncake.utils import ( + init_mooncake_custom_mem_pool, + ) + 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, _ = init_mooncake_custom_mem_pool(device) + allocator = StagingAllocator(pool_size_bytes, device, gpu_id, custom_mem_pool) + engine.batch_register([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(): + if tinfo.is_dummy: + 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) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index d99061fd3..92bb34bf1 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -45,6 +45,7 @@ from sglang.srt.disaggregation.utils import ( is_mla_backend, kv_to_page_indices, poll_and_all_reduce, + poll_and_all_reduce_with_staging, prepare_abort, ) from sglang.srt.environ import envs @@ -290,7 +291,10 @@ class DecodePreallocQueue: self._max_ensure_retries: int = 15 # scheduling cycles self._ensure_last_attempt_time: Dict[str, float] = {} self._ensure_retry_interval: float = 1.0 # seconds + self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get() self.kv_manager = self._init_kv_manager() + if self.enable_staging: + self.transfer_queue._init_staging_handler(self.kv_manager) if self.scheduler.tp_worker.is_hybrid_swa: # FIXME: current SWA allocation allocate full kv cache size in prefill @@ -366,6 +370,21 @@ class DecodePreallocQueue: self.scheduler.server_args, self.is_mla_backend, ) + # Staging buffer setup (only when heterogeneous TP staging is enabled) + if self.enable_staging and not self.is_mla_backend: + kv_pool_for_heads = self.token_to_kv_pool + if hasattr(kv_pool_for_heads, "full_kv_pool"): + kv_pool_for_heads = kv_pool_for_heads.full_kv_pool + per_rank_kv_heads = getattr(kv_pool_for_heads, "head_num", 0) + if per_rank_kv_heads > 0: + kv_args.kv_head_num = per_rank_kv_heads + kv_args.total_kv_head_num = per_rank_kv_heads * attn_tp_size + if hasattr(kv_manager, "set_kv_buffer_tensors"): + kv_pool = kv_pool_for_heads + if hasattr(kv_pool, "k_buffer") and hasattr(kv_pool, "v_buffer"): + kv_manager.set_kv_buffer_tensors( + kv_pool.k_buffer, kv_pool.v_buffer, kv_pool.page_size + ) return kv_manager def add(self, req: Req, is_retracted: bool = False) -> None: @@ -737,6 +756,13 @@ class DecodePreallocQueue: decode_req.kv_receiver.send_metadata( page_indices, decode_req.metadata_buffer_index, state_indices ) + if ( + self.transfer_queue.enable_staging + and decode_req.kv_receiver.require_staging + ): + self.transfer_queue.staging_handler.register_decode_req( + decode_req.req.bootstrap_room, decode_req + ) preallocated_reqs.append(decode_req) indices_to_remove.add(i) decode_req.req.time_stats.set_decode_transfer_queue_entry_time() @@ -863,12 +889,18 @@ class DecodeTransferQueue: self.scheduler = scheduler self.tree_cache = tree_cache self.spec_algorithm = scheduler.spec_algorithm + self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get() + self.staging_handler = None def add(self, decode_req: DecodeRequest) -> None: self.queue.append(decode_req) def extend(self, decode_reqs: List[DecodeRequest]) -> None: self.queue.extend(decode_reqs) + if self.enable_staging: + for dr in decode_reqs: + if dr.kv_receiver.require_staging: + self.staging_handler.register_decode_req(dr.req.bootstrap_room, dr) def _commit_transfer_to_req(self, decode_req: DecodeRequest) -> bool: """ @@ -951,18 +983,39 @@ class DecodeTransferQueue: decode_req.req.time_stats.set_wait_queue_entry_time() return True + def _poll_with_staging(self) -> list: + return poll_and_all_reduce_with_staging( + self.queue, self.staging_handler, self.gloo_group + ) + + def _init_staging_handler(self, kv_manager): + """Create staging handler from kv_manager. Must be called exactly once.""" + from sglang.srt.disaggregation.common.staging_handler import ( + DecodeStagingHandler, + ) + + self.staging_handler = DecodeStagingHandler.create( + kv_manager, self.scheduler, self.tp_rank + ) + kv_manager._staging_handler = self.staging_handler + def pop_transferred(self, rids_to_check: Optional[List[str]] = None) -> List[Req]: if not self.queue: return [] - polls = poll_and_all_reduce( - [decode_req.kv_receiver for decode_req in self.queue], self.gloo_group - ) + + if self.enable_staging: + polls = self._poll_with_staging() + else: + polls = poll_and_all_reduce( + [dr.kv_receiver for dr in self.queue], self.gloo_group + ) transferred_reqs = [] indices_to_remove = set() for i, (decode_req, poll) in enumerate(zip(self.queue, polls)): if rids_to_check is not None and decode_req.req.rid not in rids_to_check: continue + if poll == KVPoll.Failed: error_message = f"Decode transfer failed for request rank={self.tp_rank} {decode_req.req.rid=} {decode_req.req.bootstrap_room=}" try: @@ -1010,6 +1063,12 @@ class DecodeTransferQueue: raise ValueError(f"Unexpected poll case: {poll}") for i in indices_to_remove: + if self.enable_staging and self.staging_handler.is_staging_room( + self.queue[i].req.bootstrap_room + ): + self.staging_handler.unregister_decode_req( + self.queue[i].req.bootstrap_room + ) idx = self.queue[i].metadata_buffer_index assert idx != -1 self.req_to_metadata_buffer_idx_allocator.free(idx) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index e4cdcca2b..e4b23ed81 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -61,6 +61,14 @@ class TransferKVChunk: state_indices: Optional[List[int]] +from sglang.srt.disaggregation.common.staging_handler import ( + DecodeStagingContext, + PrefillStagingContext, + StagingRegisterInfo, + StagingTransferInfo, +) + + # decode @dataclasses.dataclass class TransferInfo: @@ -73,6 +81,7 @@ class TransferInfo: dst_state_indices: List[int] required_dst_info_num: int is_dummy: bool + staging: Optional[StagingTransferInfo] = None @classmethod def from_zmq(cls, msg: List[bytes]): @@ -118,6 +127,7 @@ class KVArgsRegisterInfo: # for mamba state different tp slice transfer dst_state_item_lens: list[int] dst_state_dim_per_tensor: list[int] + staging: Optional[StagingRegisterInfo] = None @classmethod def from_zmq(cls, msg: List[bytes]): @@ -142,6 +152,7 @@ class KVArgsRegisterInfo: if len(msg) > 11 and len(msg[11]) > 0 else [] ), + staging=StagingRegisterInfo.from_zmq_fields(msg, 12), ) @@ -178,6 +189,7 @@ class MooncakeKVManager(CommonKVManager): super().__init__(args, disaggregation_mode, server_args, is_mla_backend) self.init_engine() self.register_buffer_to_engine() + self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get() if self.disaggregation_mode == DisaggregationMode.PREFILL: self.start_prefill_thread() self.session_failures = defaultdict(int) @@ -204,14 +216,34 @@ class MooncakeKVManager(CommonKVManager): ) for _ in range(transfer_queue_size) ] - for queue, executor in zip(self.transfer_queues, self.executors): - threading.Thread( - target=self.transfer_worker, args=(queue, executor), daemon=True - ).start() self.enable_custom_mem_pool, self.custom_mem_pool_type = ( check_mooncake_custom_mem_pool_enabled() ) + self._staging_ctx = PrefillStagingContext() if self.enable_staging else None + if self.enable_staging: + self._init_staging_buffers(len(self.transfer_queues)) + for i, (queue, executor) in enumerate( + zip(self.transfer_queues, self.executors) + ): + threading.Thread( + target=self.transfer_worker, + args=( + queue, + executor, + ( + self._staging_ctx.buffers[i] + if self.enable_staging and self._staging_ctx.buffers + else None + ), + ), + daemon=True, + ).start() elif self.disaggregation_mode == DisaggregationMode.DECODE: + self._staging_ctx = DecodeStagingContext() if self.enable_staging else None + if self.enable_staging: + self._init_staging_allocator() + self._staging_handler = None + self._chunk_writer_counts: dict = defaultdict(lambda: defaultdict(list)) self.start_decode_thread() def init_engine(self): @@ -236,6 +268,297 @@ class MooncakeKVManager(CommonKVManager): self.kv_args.state_data_ptrs, self.kv_args.state_data_lens ) + # ------------------------------------------------------------------ + # Staging buffer methods (all delegate to staging_handler.py) + # ------------------------------------------------------------------ + + def register_staging_room_bootstrap(self, room, bootstrap_infos, receiver): + self._staging_ctx.room_bootstrap[room] = bootstrap_infos + self._staging_ctx.room_receivers[room] = receiver + + def set_kv_buffer_tensors(self, k_buffers: list, v_buffers: list, page_size: int): + self.kv_buffer_tensors = { + "k_buffers": k_buffers, + "v_buffers": v_buffers, + "page_size": page_size, + } + + def _init_staging_buffers(self, count: int): + from sglang.srt.disaggregation.common.staging_handler import ( + init_staging_buffers, + ) + + self._staging_ctx.buffers = init_staging_buffers( + self.engine, self.kv_args, count + ) + self.kv_buffer_tensors = None + + def _init_staging_allocator(self): + from sglang.srt.disaggregation.common.staging_handler import ( + init_staging_allocator, + ) + + self._staging_ctx.allocator = init_staging_allocator(self.engine, self.kv_args) + self.kv_buffer_tensors = None + + def _handle_staging_req(self, msg): + from sglang.srt.disaggregation.common.staging_handler import ( + handle_staging_req, + ) + + room = int(msg[1].decode("ascii")) + session_id = msg[4].decode("ascii") + handler = self._staging_handler + assert ( + handler is not None + ), "STAGING_REQ received before staging handler initialized" + decode_req = handler._room_to_decode_req.get(room) + if decode_req is None: + logger.warning( + "STAGING_REQ received for unregistered room=%s, skipping", + room, + ) + return + prefill_tp = decode_req.kv_receiver.prefill_info.attn_tp_size + handle_staging_req( + msg, + self._staging_ctx.allocator, + self.kv_args, + self.attn_tp_size, + prefill_tp, + getattr(self, "kv_buffer_tensors", None), + self._staging_ctx.room_receivers, + self._staging_ctx.room_bootstrap, + ) + + receiver = self._staging_ctx.room_receivers.get(room) + if receiver is not None: + handler.register_wm_subscriber(receiver, session_id) + + def _is_watermark_ready( + self, session_id: str, alloc_round: int, alloc_end: int + ) -> bool: + from sglang.srt.disaggregation.common.staging_handler import ( + is_watermark_ready, + ) + + return is_watermark_ready(self._staging_ctx, session_id, alloc_round, alloc_end) + + def _try_create_staging_strategy(self, staging_buffer): + if not self.enable_staging or self.kv_buffer_tensors is None: + return None + from sglang.srt.disaggregation.common.staging_handler import ( + PrefillStagingStrategy, + ) + + return PrefillStagingStrategy(self, staging_buffer) + + def _send_chunk_ready(self, req, chunk_idx, kv_chunk, prefill_unique_rank): + """Notify decode that a non-last staging chunk RDMA is complete.""" + try: + na = NetworkAddress(req.endpoint, req.dst_port) + self._connect( + na.to_tcp(), + is_ipv6=na.is_ipv6, + ).send_multipart( + [ + b"CHUNK_READY", + str(req.room).encode("ascii"), + str(chunk_idx).encode("ascii"), + str(kv_chunk.index_slice.start).encode("ascii"), + str(len(kv_chunk.prefill_kv_indices)).encode("ascii"), + req.mooncake_session_id.encode("ascii"), + str(prefill_unique_rank).encode("ascii"), + ] + ) + except Exception: + pass + + def _do_staging_transfer( + self, + staging_strategy, + kv_chunk, + req, + target_info, + chunked_dst_kv_indice, + executor, + queue, + prefill_unique_rank, + ): + """Execute staging transfer for one chunk. Returns (ret, deferred). + + Handles readiness check, transfer, fallback, and CHUNK_READY notification. + deferred=True means caller should re-enqueue and break. + """ + _tp = self.attn_tp_rank + ready, chunk_idx, c_offset, _, _ = staging_strategy.check_ready( + req, + kv_chunk.index_slice.start, + len(kv_chunk.prefill_kv_indices), + ) + if not ready: + from sglang.srt.disaggregation.common.staging_buffer import StagingAllocator + + if c_offset == StagingAllocator.ALLOC_OVERSIZED: + raise RuntimeError( + f"[Staging] Chunk staging allocation permanently failed: " + f"chunk exceeds ring buffer total size (room={kv_chunk.room}). " + f"Increase SGLANG_DISAGG_STAGING_POOL_SIZE_MB." + ) + queue.put(kv_chunk) + return (-1, True) + + ret = staging_strategy.transfer( + req.mooncake_session_id, + kv_chunk.prefill_kv_indices, + target_info.staging.base_ptr + c_offset, + target_info.staging.total_size - c_offset, + target_info, + ) + if ret == -1: + logger.warning( + f"[Staging][tp{_tp}] Falling back to per-token slice path " + f"(room={kv_chunk.room})" + ) + ret = self.send_kvcache_slice( + req.mooncake_session_id, + kv_chunk.prefill_kv_indices, + target_info.dst_kv_ptrs, + chunked_dst_kv_indice, + target_info.dst_tp_rank, + target_info.dst_attn_tp_size, + target_info.dst_kv_item_len, + executor, + ) + elif ret == 0 and not kv_chunk.is_last_chunk: + self._send_chunk_ready(req, chunk_idx, kv_chunk, prefill_unique_rank) + return (ret, False) + + def _prefetch_staging_reqs(self, room: int): + if not self.enable_staging or self.kv_buffer_tensors is None: + return + + room_infos = self.transfer_infos.get(room, {}) + needs_staging = any( + not tinfo.is_dummy + and self.decode_kv_args_table.get(tinfo.mooncake_session_id) is not None + and self.decode_kv_args_table[tinfo.mooncake_session_id].dst_attn_tp_size + != self.attn_tp_size + for tinfo in room_infos.values() + ) + if not needs_staging: + return + + from sglang.srt.disaggregation.common.staging_handler import ( + prefetch_staging_reqs, + ) + + prefetch_staging_reqs( + room, + self.transfer_infos, + self.kv_buffer_tensors, + self.server_args.chunked_prefill_size, + self._staging_ctx.prefetch_requested, + self._staging_ctx.prefetch_sockets, + ) + + def send_kvcache_staged( + self, + mooncake_session_id: str, + prefill_kv_indices: npt.NDArray[np.int32], + dst_staging_ptr: int, + dst_staging_size: int, + dst_tp_rank: int, + dst_attn_tp_size: int, + dst_kv_item_len: int, + staging_buffer=None, + ) -> int: + """Transfer KV cache via staging buffers (gather -> bulk RDMA -> scatter on decode).""" + from sglang.srt.disaggregation.common.staging_buffer import ( + compute_head_slice_params, + compute_staging_layout, + resolve_total_kv_heads, + ) + + if self.kv_buffer_tensors is None or staging_buffer is None: + return -1 + + k_buffers = self.kv_buffer_tensors["k_buffers"] + v_buffers = self.kv_buffer_tensors["v_buffers"] + page_size = self.kv_buffer_tensors["page_size"] + num_layers = len(k_buffers) + head_dim = k_buffers[0].shape[-1] + dtype_size = k_buffers[0].element_size() + + total_kv_heads = resolve_total_kv_heads(self.kv_args, self.attn_tp_size) + + local_tp_rank = self.kv_args.engine_rank % self.attn_tp_size + src_head_start, num_heads_to_send, _, _ = compute_head_slice_params( + self.attn_tp_size, + dst_attn_tp_size, + local_tp_rank, + dst_tp_rank, + total_kv_heads, + ) + + num_tokens = len(prefill_kv_indices) * page_size + per_layer_bytes = num_tokens * num_heads_to_send * head_dim * dtype_size + per_rank_bytes = per_layer_bytes * num_layers * 2 + + num_writers, writer_rank_bytes, total_staging_needed = compute_staging_layout( + self.attn_tp_size, + dst_attn_tp_size, + dst_tp_rank, + total_kv_heads, + num_tokens, + head_dim * dtype_size, + num_layers, + ) + writer_idx = local_tp_rank % num_writers if num_writers > 1 else 0 + rank_offset = sum(writer_rank_bytes[:writer_idx]) + + if not staging_buffer.fits(per_rank_bytes): + logger.warning( + f"Prefill staging too small for {per_rank_bytes} bytes, falling back" + ) + return -1 + if dst_staging_size < total_staging_needed: + logger.warning( + f"Decode staging too small: need {total_staging_needed} bytes " + f"({num_writers if self.attn_tp_size > dst_attn_tp_size else 1} writers " + f"x {per_rank_bytes} bytes/rank), have {dst_staging_size}, falling back" + ) + return -1 + + from sglang.srt.disaggregation.common.staging_buffer import ( + gather_all_layers_to_staging, + ) + + gather_all_layers_to_staging( + k_buffers, + v_buffers, + prefill_kv_indices, + staging_buffer, + src_head_start, + num_heads_to_send, + page_size, + self.kv_args.gpu_id, + ) + + dst_write_ptr = dst_staging_ptr + rank_offset + ret = self._transfer_data( + mooncake_session_id, + [(staging_buffer.get_ptr(), dst_write_ptr, per_rank_bytes)], + ) + if ret != 0: + raise RuntimeError( + f"[Staging] Bulk RDMA transfer failed with ret={ret}. " + f"src_ptr=0x{staging_buffer.get_ptr():x}, " + f"dst_ptr=0x{dst_write_ptr:x}, size={per_rank_bytes}. " + f"The decode staging buffer may not be properly registered." + ) + return ret + def _transfer_data(self, mooncake_session_id, transfer_blocks): if not transfer_blocks: return 0 @@ -770,11 +1093,22 @@ class MooncakeKVManager(CommonKVManager): ) def transfer_worker( - self, queue: FastQueue, executor: concurrent.futures.ThreadPoolExecutor + self, + queue: FastQueue, + executor: concurrent.futures.ThreadPoolExecutor, + staging_buffer=None, ): + staging_strategy = None + while True: try: kv_chunk: TransferKVChunk = queue.get() + if ( + self.enable_staging + and staging_strategy is None + and staging_buffer is not None + ): + staging_strategy = self._try_create_staging_strategy(staging_buffer) reqs_to_be_processed = ( self.transfer_infos[kv_chunk.room].values() if kv_chunk.room in self.transfer_infos @@ -788,6 +1122,9 @@ class MooncakeKVManager(CommonKVManager): + self.pp_rank * self.attn_cp_size + self.attn_cp_rank ) + # When staging transfer is not yet ready (watermark/allocation pending), + # the chunk is re-enqueued and we break out of the req loop to retry later. + staging_deferred = False for req in reqs_to_be_processed: if not req.is_dummy: # Early exit if the request has failed @@ -835,6 +1172,25 @@ class MooncakeKVManager(CommonKVManager): chunked_dst_kv_indice, executor, ) + elif ( + self.enable_staging + and staging_strategy is not None + and target_rank_registration_info.staging is not None + ): + ret, deferred = self._do_staging_transfer( + staging_strategy, + kv_chunk, + req, + target_rank_registration_info, + chunked_dst_kv_indice, + executor, + queue, + prefill_unique_rank, + ) + if deferred: + staging_deferred = True + # Chunk re-enqueued; stop processing remaining reqs for this chunk + break else: ret = self.send_kvcache_slice( req.mooncake_session_id, @@ -909,6 +1265,9 @@ class MooncakeKVManager(CommonKVManager): if kv_chunk.is_last_chunk and req.room in self.request_status: self.update_status(req.room, KVPoll.Success) + if staging_deferred: + continue + if ( kv_chunk.room not in self.request_status or self.check_status(kv_chunk.room) == KVPoll.Success @@ -929,6 +1288,50 @@ class MooncakeKVManager(CommonKVManager): while True: waiting_req_bytes = self.server_socket.recv_multipart() room = waiting_req_bytes[0].decode("ascii") + # Staging: decode reports consumption watermark back to prefill + if room == "WATERMARK": + wm_round = int(waiting_req_bytes[1].decode("ascii")) + wm_tail = int(waiting_req_bytes[2].decode("ascii")) + wm_session = ( + waiting_req_bytes[3].decode("ascii") + if len(waiting_req_bytes) > 3 + else "" + ) + with self._staging_ctx.watermark_cv: + prev = self._staging_ctx.remote_watermarks.get( + wm_session, (0, 0) + ) + if (wm_round, wm_tail) > prev: + self._staging_ctx.remote_watermarks[wm_session] = ( + wm_round, + wm_tail, + ) + self._staging_ctx.watermark_cv.notify_all() + continue + # Staging: decode replies with allocated staging offset + if room == "STAGING_RSP": + stg_room = int(waiting_req_bytes[1].decode("ascii")) + stg_chunk_idx = int(waiting_req_bytes[2].decode("ascii")) + stg_offset = int(waiting_req_bytes[3].decode("ascii")) + stg_round = int(waiting_req_bytes[4].decode("ascii")) + stg_end = int(waiting_req_bytes[5].decode("ascii")) + stg_session = waiting_req_bytes[6].decode("ascii") + room_infos = self.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, + ) + continue mooncake_session_id = waiting_req_bytes[3].decode("ascii") if room == "None": self.decode_kv_args_table[mooncake_session_id] = ( @@ -966,6 +1369,42 @@ class MooncakeKVManager(CommonKVManager): self._handle_aux_data(msg) continue + # Staging: prefill notifies a chunk written to staging buffer + if msg[0] == b"CHUNK_READY": + room = int(msg[1].decode("ascii")) + chunk_idx = int(msg[2].decode("ascii")) + page_start = int(msg[3].decode("ascii")) + num_pages = int(msg[4].decode("ascii")) + session_id = msg[5].decode("ascii") + self._chunk_writer_counts[room][chunk_idx].append( + (page_start, num_pages, session_id) + ) + handler = self._staging_handler + assert ( + handler is not None + ), "CHUNK_READY received before staging handler initialized" + writers_arrived = len(self._chunk_writer_counts[room][chunk_idx]) + decode_req = handler._room_to_decode_req.get(room) + if decode_req is None: + logger.warning( + "CHUNK_READY received for unregistered room=%s chunk=%d, skipping", + room, + chunk_idx, + ) + continue + num_writers = handler.num_writers_for(decode_req) + if writers_arrived >= num_writers: + handler.submit_chunk_scatter( + room, chunk_idx, page_start, num_pages + ) + del self._chunk_writer_counts[room][chunk_idx] + continue + + # Staging: prefill pre-requests staging allocation before forward + if msg[0] == b"STAGING_REQ": + self._handle_staging_req(msg) + continue + bootstrap_room, status, prefill_rank = msg status = int(status.decode("ascii")) bootstrap_room = int(bootstrap_room.decode("ascii")) @@ -981,6 +1420,11 @@ class MooncakeKVManager(CommonKVManager): self.prefill_response_tracker[bootstrap_room] ) if arrived_response_num == expected_response_num: + if self.enable_staging: + handler = self._staging_handler + if handler.is_staging_room(bootstrap_room): + handler.submit_last_scatter_async(bootstrap_room) + self._chunk_writer_counts.pop(bootstrap_room, None) self.update_status(bootstrap_room, KVPoll.Success) elif status == KVPoll.Failed: self.record_failure( @@ -1272,6 +1716,17 @@ class MooncakeKVReceiver(CommonKVReceiver): dst_attn_tp_size = str(self.kv_mgr.attn_tp_size).encode("ascii") dst_kv_item_len = str(kv_item_len).encode("ascii") + if ( + self.kv_mgr.enable_staging + and self.kv_mgr._staging_ctx.allocator is not None + ): + _alloc = self.kv_mgr._staging_ctx.allocator + packed_staging_base_ptr = struct.pack("Q", _alloc.get_base_ptr()) + staging_total_size_str = str(_alloc.get_total_size()).encode("ascii") + else: + packed_staging_base_ptr = b"" + staging_total_size_str = b"" + sock, lock = self._connect_to_bootstrap_server(bootstrap_info) with lock: sock.send_multipart( @@ -1288,6 +1743,8 @@ class MooncakeKVReceiver(CommonKVReceiver): dst_kv_item_len, packed_state_item_lens, packed_state_dim_per_tensor, + packed_staging_base_ptr, + staging_total_size_str, ] ) @@ -1311,6 +1768,15 @@ class MooncakeKVReceiver(CommonKVReceiver): self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) return + if ( + self.kv_mgr.enable_staging + and self.kv_mgr._staging_ctx.allocator is not None + ): + self.chunk_staging_infos = [] + self.kv_mgr.register_staging_room_bootstrap( + self.bootstrap_room, self.bootstrap_infos, self + ) + for bootstrap_info in self.bootstrap_infos: sock, lock = self._connect_to_bootstrap_server(bootstrap_info) is_dummy = bootstrap_info["is_dummy"] diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 7d97f123b..8eadf8195 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -42,6 +42,7 @@ from sglang.srt.disaggregation.utils import ( poll_and_all_reduce_attn_cp_tp_group, prepare_abort, ) +from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import ( FINISH_ABORT, FINISH_LENGTH, @@ -201,6 +202,21 @@ class PrefillBootstrapQueue: self.scheduler.server_args, self.is_mla_backend, ) + # Pass KV pool tensor refs to the manager for GPU gather (staging mode) + if ( + envs.SGLANG_DISAGG_STAGING_BUFFER.get() + and hasattr(kv_manager, "set_kv_buffer_tensors") + and not self.is_mla_backend + ): + kv_pool = self.token_to_kv_pool + if hasattr(kv_pool, "full_kv_pool"): + kv_pool = kv_pool.full_kv_pool + if hasattr(kv_pool, "k_buffer") and hasattr(kv_pool, "v_buffer"): + kv_manager.set_kv_buffer_tensors( + kv_pool.k_buffer, + kv_pool.v_buffer, + kv_pool.page_size, + ) return kv_manager def add(self, req: Req, num_kv_heads: int) -> None: @@ -336,6 +352,17 @@ class SchedulerDisaggregationPrefillMixin: Mixin for Scheduler to handle disaggregation prefill """ + def maybe_prefetch_staging_for_batch(self: Scheduler, batch: ScheduleBatch) -> None: + """Pre-send STAGING_REQ so decode allocates staging during GPU forward.""" + kv_mgr = self.disagg_prefill_bootstrap_queue.kv_manager + prefetch = getattr(kv_mgr, "_prefetch_staging_reqs", None) + if prefetch is None: + return + for req in batch.reqs: + room = getattr(req, "bootstrap_room", None) + if room is not None and room in kv_mgr.transfer_infos: + prefetch(room) + def get_next_disagg_prefill_batch_to_run( self: Scheduler, ) -> Optional[ScheduleBatch]: @@ -356,6 +383,7 @@ class SchedulerDisaggregationPrefillMixin: @torch.no_grad() def event_loop_normal_disagg_prefill(self: Scheduler) -> None: """A normal scheduler loop for prefill worker in disaggregation mode.""" + self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get() while True: # Receive requests @@ -371,6 +399,8 @@ class SchedulerDisaggregationPrefillMixin: # Launch the current batch if batch: + if self.enable_staging: + self.maybe_prefetch_staging_for_batch(batch) result = self.run_batch(batch) self.process_batch_result(batch, result) else: @@ -384,6 +414,7 @@ class SchedulerDisaggregationPrefillMixin: @torch.no_grad() def event_loop_overlap_disagg_prefill(self: Scheduler) -> None: self.result_queue = deque() + self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get() while True: # Receive requests @@ -399,6 +430,8 @@ class SchedulerDisaggregationPrefillMixin: # Launch the current batch if batch: + if self.enable_staging: + self.maybe_prefetch_staging_for_batch(batch) batch_result = self.run_batch(batch) self.result_queue.append((batch.copy(), batch_result)) else: diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index b7b3b0238..636b86986 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -80,6 +80,30 @@ def poll_and_all_reduce_attn_cp_tp_group( return tensor_to_reduce.tolist() +def poll_and_all_reduce_with_staging( + decode_reqs, staging_handler, gloo_group: dist.ProcessGroup +): + """Staging-aware polling: advance scatter, demote incomplete transfers, all_reduce.""" + from sglang.srt.disaggregation.base import KVPoll + + 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) + + raw_polls = [int(dr.kv_receiver.poll()) for dr in decode_reqs] + 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) + 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 ######################### diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index adc21301b..139b70b32 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -287,6 +287,14 @@ class Envs: # Max fraction of cache (by token count) that can be pinned; 0 = disable pinning. SGLANG_HICACHE_MAX_PINNED_RATIO = EnvFloat(0.0) + # Staging buffer for heterogeneous TP KV transfer + SGLANG_DISAGG_STAGING_BUFFER = EnvBool(False) + SGLANG_DISAGG_STAGING_BUFFER_SIZE_MB = EnvInt(64) + SGLANG_DISAGG_STAGING_POOL_SIZE_MB = EnvInt(4096) + # TODO(yangminl): remove SGLANG_STAGING_USE_TORCH and the torch fallback in + # staging_buffer.py once Triton kernels are fully validated in production. + SGLANG_STAGING_USE_TORCH = EnvBool(False) + # Mooncake KV Transfer SGLANG_MOONCAKE_CUSTOM_MEM_POOL = EnvStr(None) ENABLE_ASCEND_TRANSFER_WITH_MOONCAKE = EnvBool(False) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index acc426962..4d860505c 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -3238,6 +3238,17 @@ class ServerArgs: "Cuda graph is disabled for prefill server when piecewise cuda graph is not enabled." ) + if self.disaggregation_mode in ("prefill", "decode"): + if ( + envs.SGLANG_DISAGG_STAGING_BUFFER.get() + and self.disaggregation_transfer_backend != "mooncake" + ): + raise ValueError( + f"SGLANG_DISAGG_STAGING_BUFFER requires " + f"disaggregation_transfer_backend='mooncake', " + f"got '{self.disaggregation_transfer_backend}'." + ) + def _handle_encoder_disaggregation(self): if self.enable_prefix_mm_cache and not self.encoder_only: raise ValueError(