minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity
Files
c74a4037fb0fb7afc6384c7b70b1f9773119548f
sglang/python/sglang/srt/disaggregation/mooncake/conn.py
T
minke.yu c74a4037fb mooncake: probe pointers via cuda.bindings instead of ctypes libcudart 
The ctypes-based probing loaded a second libcudart copy into the process,
which corrupted CUDA/runtime state and segfaulted the scheduler processes
shortly after registration (reproducible on prefill ranks, intermittent on
decode). cuda.bindings.runtime.cudaPointerGetAttributes is a properly
typed binding and boots cleanly on both roles.
2026-09-21 12:46:25 +08:00

3272 lines
139 KiB
Python
Raw Blame History

from __future__ import annotations
import concurrent.futures
import dataclasses
import logging
import os
import struct
import threading
import time
from collections import defaultdict
from typing import List, Optional, Set, Tuple, Union
import numpy as np
import numpy.typing as npt
import zmq
from prometheus_client import Counter
from sglang.srt.disaggregation.base.conn import KVArgs, KVPoll, StateType
from sglang.srt.disaggregation.common.conn import (
CommonKVBootstrapServer,
CommonKVManager,
CommonKVReceiver,
CommonKVSender,
KVTransferError,
)
from sglang.srt.disaggregation.common.staging_handler import (
STAGING_WATERMARK_WAIT_S,
DecodeStagingContext,
PrefillStagingContext,
StagingManagerMixin,
StagingRegisterInfo,
StagingTransferInfo,
handle_staging_rsp,
handle_watermark_msg,
)
from sglang.srt.disaggregation.common.utils import (
AuxDataCodec,
FastQueue,
TransferKVChunk,
build_dcp_token_transfer_plan,
group_concurrent_contiguous,
pack_int_lists,
unpack_int_lists,
)
from sglang.srt.disaggregation.mooncake.utils import (
check_mooncake_custom_mem_pool_enabled,
)
from sglang.srt.disaggregation.utils import (
DisaggregationMode,
build_dsa_tail_transfer_blocks,
build_transfer_entry_pairs,
compute_mamba_state_slice_byte_blocks,
resolve_dcp_dst_entry_indices,
should_send_replicated_state,
slice_dsa_tail_dst_ptrs_for_pp,
)
from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine
from sglang.srt.environ import envs
from sglang.srt.observability.mooncake_trace import (
MooncakeRequestStage,
mooncake_trace_func,
mooncake_trace_slice,
)
from sglang.srt.observability.trace import (
TraceNullContext,
TraceReqContext,
trace_set_thread_info,
)
from sglang.srt.runtime_context import (
get_memory,
get_observability,
get_schedule,
)
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils.network import NetworkAddress
logger = logging.getLogger(__name__)
FAILED_SESSION_RECOVERIES = Counter(
"sglang:failed_session_recoveries_total",
"Number of mooncake_session_ids un-blacklisted via probe.",
)
# ---------------------------------------------------------------------------
# Intra-node NVLink transport helpers.
#
# Mooncake's IntraNodeNvlinkTransport can only register and reach *device*
# memory (it IPC-opens the remote cudaMalloc segments). Host-resident regions
# (aux buffers, some state components) cannot be registered: one host region
# makes the whole registerLocalMemoryBatch fail, and the engine then rolls
# back *every* region, leaving the segment descriptor empty and all KV
# transfers failing with "Requested address ... not found". When the
# intra-node NVLink transport is active we therefore
# 1. register only device-memory regions, and
# 2. route blocks whose source is host memory over the ordered zmq channel
# (same ordering guarantee the aux TCP path relies on) instead of the
# transfer engine.
# ---------------------------------------------------------------------------
import ctypes as _ctypes
_CUDA_MEMORY_TYPE_DEVICE = 2
try:
from cuda.bindings import runtime as _cudart
except ImportError: # pragma: no cover - cuda-python is always present in images
_cudart = None
def _is_device_pointer(ptr: int) -> bool:
"""Probe a *local* pointer with cudaPointerGetAttributes.
Only valid for pointers owned by this process (never probe remote
segment addresses). Returns False on any error so the caller falls back
to the safe host path.
"""
if _cudart is None:
# Cannot tell; assume device so behavior stays unchanged.
return True
err, attr = _cudart.cudaPointerGetAttributes(int(ptr))
if int(err) != 0:
# Clear the error so subsequent CUDA calls are not poisoned.
_cudart.cudaGetLastError()
return False
return int(attr.type) == _CUDA_MEMORY_TYPE_DEVICE
def _read_bytes_from_address(addr: int, length: int) -> Optional[bytes]:
if length <= 0:
return b""
if _is_device_pointer(addr):
buf = bytearray(length)
# cudaMemcpyDeviceToHost = 2; synchronous default-stream copy.
err, = _cudart.cudaMemcpy(
_ctypes.addressof((_ctypes.c_char * length).from_buffer(buf)),
int(addr),
length,
2,
)
if int(err) != 0:
logger.error(
f"cudaMemcpy D2H failed (err={err}) for addr {hex(addr)} len {length}"
)
return None
return bytes(buf)
return _ctypes.string_at(int(addr), length)
def _write_bytes_to_address(addr: int, data: bytes) -> bool:
if not data:
return True
if _is_device_pointer(addr):
buf = _ctypes.create_string_buffer(data, len(data))
# cudaMemcpyHostToDevice = 1; synchronous default-stream copy.
err, = _cudart.cudaMemcpy(
int(addr), _ctypes.addressof(buf), len(data), 1
)
if int(err) != 0:
logger.error(
f"cudaMemcpy H2D failed (err={err}) for addr {hex(addr)} "
f"len {len(data)}"
)
return False
return True
_ctypes.memmove(int(addr), data, len(data))
return True
_NVLINK_INTRA_ACTIVE = None
def _nvlink_intra_transport_active() -> bool:
"""Whether mooncake installed the intra-node NVLink transport.
Mirrors the env probing in mooncake's transfer_engine_impl.cpp: the
transport is installed iff MC_INTRANODE_NVLINK is set (any value), or an
equivalent protocol selection was made.
"""
global _NVLINK_INTRA_ACTIVE
if _NVLINK_INTRA_ACTIVE is None:
active = bool(
os.environ.get("MC_INTRANODE_NVLINK")
or os.environ.get("MC_INTRA_NVLINK")
)
if not active:
proto = (os.environ.get("MOONCAKE_PROTOCOL") or "").strip().lower()
active = proto in ("nvlink_intra", "nvlink-intra", "intra_nvlink")
_NVLINK_INTRA_ACTIVE = active
return _NVLINK_INTRA_ACTIVE
# decode
@dataclasses.dataclass
class TransferInfo:
room: int
endpoint: str
dst_port: int
mooncake_session_id: str
dst_kv_indices: npt.NDArray[np.int32]
dst_aux_index: int
dst_state_indices: List[List[int]] # parallel to receiver's state_types
required_dst_info_num: int
is_dummy: bool
decode_prefix_len: Optional[int] = None
dst_device_kv_indices: Optional[npt.NDArray[np.int32]] = None
# Note: always put the optional staging field at the final (it will be set through 'STAGING_RSP' pkg when needed)
staging: Optional[StagingTransferInfo] = None
@classmethod
def from_zmq(cls, msg: List[bytes]):
if msg[4] == b"" and msg[5] == b"":
is_dummy = True
dst_kv_indices = np.array([], dtype=np.int32)
dst_aux_index = None
dst_state_indices = []
else:
dst_kv_indices = np.frombuffer(msg[4], dtype=np.int32)
dst_aux_index = int(msg[5].decode("ascii"))
dst_state_indices = unpack_int_lists(msg[6], "i")
is_dummy = False
return cls(
room=int(msg[0].decode("ascii")),
endpoint=msg[1].decode("ascii"),
dst_port=int(msg[2].decode("ascii")),
mooncake_session_id=msg[3].decode("ascii"),
dst_kv_indices=dst_kv_indices,
dst_aux_index=dst_aux_index,
dst_state_indices=dst_state_indices,
required_dst_info_num=int(msg[7].decode("ascii")),
is_dummy=is_dummy,
decode_prefix_len=(
int(msg[8].decode("ascii")) if len(msg) > 8 and msg[8] != b"" else None
),
dst_device_kv_indices=(
np.frombuffer(msg[9], dtype=np.int32)
if len(msg) > 9 and msg[9] != b""
else None
),
)
# decode
@dataclasses.dataclass
class KVArgsRegisterInfo:
room: str
endpoint: str
dst_port: int
mooncake_session_id: str
dst_kv_ptrs: list[int]
dst_aux_ptrs: list[int]
dst_state_data_ptrs: List[List[int]] # parallel to state_types (same below)
dst_tp_rank: int
dst_attn_tp_size: int
dst_kv_item_len: int
# for mamba state different tp slice transfer
dst_state_item_lens: List[List[int]]
dst_state_dim_per_tensor: List[List[int]]
dst_kv_layer_ids: List[int]
dst_state_layer_ids: List[List[int]]
dst_dcp_size: int = 1
dst_dcp_rank: int = 0
requires_dcp_relayout: bool = False
dcp_token_item_lens: Optional[List[int]] = None
dst_kv_item_lens: List[int] = dataclasses.field(default_factory=list)
staging_base_ptr: int = 0
staging_total_size: int = 0
staging: Optional[StagingRegisterInfo] = None
@classmethod
def from_zmq(cls, msg: List[bytes]):
return cls(
room=str(msg[0].decode("ascii")),
endpoint=msg[1].decode("ascii"),
dst_port=int(msg[2].decode("ascii")),
mooncake_session_id=msg[3].decode("ascii"),
dst_kv_ptrs=list(struct.unpack(f"{len(msg[4]) // 8}Q", msg[4])),
dst_aux_ptrs=list(struct.unpack(f"{len(msg[5]) // 8}Q", msg[5])),
dst_state_data_ptrs=unpack_int_lists(msg[6], "Q"),
dst_tp_rank=int(msg[7].decode("ascii")),
dst_attn_tp_size=int(msg[8].decode("ascii")),
dst_kv_item_len=int(msg[9].decode("ascii")),
dst_state_item_lens=(
unpack_int_lists(msg[10], "I") if len(msg) > 10 else []
),
dst_state_dim_per_tensor=(
unpack_int_lists(msg[11], "I") if len(msg) > 11 else []
),
dst_kv_layer_ids=(
list(struct.unpack(f"{len(msg[12]) // 4}I", msg[12]))
if len(msg) > 12 and msg[12] != b""
else []
),
dst_state_layer_ids=(
unpack_int_lists(msg[13], "I")
if len(msg) > 13 and msg[13] != b""
else []
),
staging_base_ptr=(
struct.unpack("Q", msg[14])[0]
if len(msg) > 14 and len(msg[14]) == 8
else 0
),
staging_total_size=(
int(msg[15].decode("ascii")) if len(msg) > 15 and msg[15] != b"" else 0
),
dst_dcp_size=(
int(msg[16].decode("ascii")) if len(msg) > 16 and msg[16] != b"" else 1
),
dst_dcp_rank=(
int(msg[17].decode("ascii")) if len(msg) > 17 and msg[17] != b"" else 0
),
dst_kv_item_lens=(
list(struct.unpack(f"{len(msg[19]) // 8}Q", msg[19]))
if len(msg) > 19 and msg[19]
else []
),
# Note: always put the staging field at the final
staging=StagingRegisterInfo.from_zmq_fields(msg, 14, slot_ids_index=18),
)
class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
AUX_DATA_HEADER = b"AUX_DATA"
STATE_DATA_HEADER = b"STATE_DATA"
# Implements teardown() below, so runtime PD role switching is supported.
supports_role_switch = True
def __init__(
self,
args: KVArgs,
disaggregation_mode: DisaggregationMode,
server_args: ServerArgs,
is_mla_backend: Optional[bool] = False,
):
super().__init__(args, disaggregation_mode, server_args, is_mla_backend)
self.init_engine()
self.register_buffer_to_engine()
# session_id -> (endpoint, dst_port, room), used to route host-memory
# transfer blocks over zmq when the intra-node NVLink transport is
# active (it cannot reach host memory). Populated on bootstrap.
self._session_endpoint_map = {}
self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get()
self.max_transfer_batch_indices = (
envs.SGLANG_MOONCAKE_MAX_TRANSFER_BATCH_INDICES.get()
)
self.enable_trace = get_observability().enable_trace
# Set by teardown() to make worker threads exit (P<->D role switch).
self._stopped = False
self._worker_threads: List[threading.Thread] = []
if self.disaggregation_mode == DisaggregationMode.PREFILL:
self.session_failures = defaultdict(int)
self.failed_sessions = set()
self.session_lock = threading.Lock()
self.start_prefill_thread()
# Per-room count of chunks not yet transferred; teardown waits for
# zero so a deferred chunk is not dropped by an early conclude.
self._staging_outstanding = defaultdict(int)
# Determine the number of threads to use for kv sender
cpu_count = os.cpu_count()
transfer_thread_pool_size = (
envs.SGLANG_DISAGGREGATION_THREAD_POOL_SIZE.get()
)
if transfer_thread_pool_size is None:
transfer_thread_pool_size = min(max(4, int(0.5 * cpu_count) // 8), 12)
transfer_queue_size = envs.SGLANG_DISAGGREGATION_QUEUE_SIZE.get()
self.transfer_queues: List[FastQueue] = [
FastQueue() for _ in range(transfer_queue_size)
]
assert transfer_thread_pool_size >= transfer_queue_size, (
f"The environment variable SGLANG_DISAGGREGATION_THREAD_POOL_SIZE={transfer_thread_pool_size} must be "
f"greater than or equal to SGLANG_DISAGGREGATION_QUEUE_SIZE={transfer_queue_size}."
)
self.executors = [
concurrent.futures.ThreadPoolExecutor(
transfer_thread_pool_size // transfer_queue_size
)
for _ in range(transfer_queue_size)
]
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)
):
# Track the thread so teardown() can join it: otherwise every
# P->D->P flip that re-enters PREFILL leaks threads
# (each parked forever in FastQueue.get()).
t = 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
),
i,
),
daemon=True,
)
t.start()
self._worker_threads.append(t)
self.enable_failed_session_probe = (
envs.SGLANG_ENABLE_FAILED_SESSION_PROBE.get()
)
if self.enable_failed_session_probe:
self.failed_session_probe_interval = (
envs.SGLANG_FAILED_SESSION_PROBE_INTERVAL_S.get()
)
self._failed_session_probe_shutdown = threading.Event()
t = threading.Thread(
target=self._failed_session_probe_loop,
name="MooncakeFailedSessionProbe",
daemon=True,
)
t.start()
self._worker_threads.append(t)
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.start_decode_thread()
def init_engine(self):
self.engine = get_mooncake_transfer_engine()
def _registerable_regions(self) -> List[Tuple[int, int]]:
"""(ptr, len) regions to (de)register, exact duplicates removed.
Deduped because the unified memory pool reports one raw buffer as both
its KV and its mamba state component, and double registration fails in
the engine.
When the intra-node NVLink transport is active, host-memory regions
(aux buffers, some state components) are skipped: the transport only
accepts device memory, and a single host region fails the whole batch
and triggers a full engine-side rollback that would unregister the KV
pools too. Host-resident payloads are instead exchanged over the
ordered zmq channel (see _transfer_data / send_aux).
"""
regions: List[Tuple[int, int]] = []
seen: Set[Tuple[int, int]] = set()
def add(ptrs: List[int], lens: List[int]) -> None:
for ptr, length in zip(ptrs or [], lens or []):
if (ptr, length) not in seen:
seen.add((ptr, length))
regions.append((ptr, length))
add(self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens)
add(self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens)
for ptrs, lens in zip(
self.kv_args.state_data_ptrs, self.kv_args.state_data_lens
):
add(ptrs, lens)
if _nvlink_intra_transport_active():
device_regions = []
skipped = []
for ptr, length in regions:
if _is_device_pointer(ptr):
device_regions.append((ptr, length))
else:
skipped.append((ptr, length))
if skipped:
logger.info(
"Intra-node NVLink transport: skipping %d host-memory "
"regions from engine registration (they will be exchanged "
"over the zmq channel instead): %s",
len(skipped),
[(hex(p), l) for p, l in skipped[:8]],
)
regions = device_regions
return regions
def register_buffer_to_engine(self):
regions = self._registerable_regions()
if regions:
ptrs, lens = zip(*regions)
self.engine.batch_register(list(ptrs), list(lens))
def deregister_buffer_to_engine(self):
regions = self._registerable_regions()
if regions:
ptrs, _ = zip(*regions)
self.engine.batch_deregister(list(ptrs))
if hasattr(self, "connection_pool"):
with self.connection_lock:
self.connection_pool.clear()
def teardown(self) -> None:
"""Stop worker threads and release transport resources so this
KVManager can be discarded during a P<->D role switch.
The KV cache pool memory is owned by the scheduler and is NOT freed
here; only mooncake-side registrations / sockets are released.
"""
self._stopped = True
# Stop the failed-session probe loop (PREFILL role) if running.
probe_shutdown = getattr(self, "_failed_session_probe_shutdown", None)
if probe_shutdown is not None:
probe_shutdown.set()
# Stop the heartbeat checker thread (DECODE role) if running.
heartbeat_shutdown = getattr(self, "_heartbeat_shutdown", None)
if heartbeat_shutdown is not None:
heartbeat_shutdown.set()
# Transfer workers (PREFILL role) park in FastQueue.get(), which has no
# timeout; push a None sentinel per shard to wake and stop them so the
# join below returns instead of leaking the thread.
for queue in getattr(self, "transfer_queues", []):
try:
queue.put(None)
except Exception:
logger.exception(
"Failed to signal mooncake transfer worker on teardown"
)
# Shutdown thread pool executors (PREFILL role).
for executor in getattr(self, "executors", []):
try:
executor.shutdown(wait=False, cancel_futures=True)
except TypeError:
# Python < 3.9 does not support cancel_futures
executor.shutdown(wait=False)
except Exception:
logger.exception("Failed to shutdown executor on teardown")
self.executors = []
# Join workers before touching their sockets: ZMQ sockets aren't
# thread-safe, so don't close server_socket while a worker may poll it.
for t in self._worker_threads:
t.join(timeout=3.0)
self._worker_threads = []
# Drop the queues so their buffered tasks/senders are released too.
self.transfer_queues = []
# Close cached PUSH sockets (used by _connect for status sync).
with self._socket_lock:
for sock in self._socket_cache.values():
try:
sock.close(linger=0)
except Exception:
pass
for monitor in self._monitor_cache.values():
try:
monitor.close()
except Exception:
pass
self._socket_cache.clear()
self._monitor_cache.clear()
try:
self.server_socket.close(linger=0)
except Exception:
logger.exception("Failed to close mooncake server_socket during teardown")
# destroy() force-closes every socket in the context; plain term()
# would block waiting on them.
try:
self._zmq_ctx.destroy(linger=0)
except Exception:
logger.exception("Failed to destroy mooncake zmq context during teardown")
# Deregister memory from the transfer engine.
try:
self.deregister_buffer_to_engine()
except Exception:
logger.exception("Failed to deregister buffers during teardown")
logger.info(
"MooncakeKVManager torn down (was role=%s)",
self.disaggregation_mode.value,
)
# ------------------------------------------------------------------
# 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,
slot_layer_ids: Optional[List[int]] = None,
):
# slot_layer_ids follows the staging slot order (every k_buffer, then
# every v_buffer), which is not kv_args.kv_layer_ids once a draft exists.
self.kv_buffer_tensors = {
"k_buffers": k_buffers,
"v_buffers": v_buffers,
"page_size": page_size,
"slot_layer_ids": list(slot_layer_ids or []),
}
def _register_staging_memory(self, ptr: int, size: int) -> None:
self.engine.batch_register([ptr], [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._register_staging_memory,
self.kv_args,
count,
get_schedule().chunked_prefill_size,
)
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._register_staging_memory,
self.kv_args,
)
self.kv_buffer_tensors = None
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 staging chunk RDMA is complete (every chunk;
scatter is arrival-driven)."""
na = NetworkAddress(req.endpoint, req.dst_port)
self._send_multipart_locked(
na.to_tcp(),
[
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"),
],
is_ipv6=na.is_ipv6,
)
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, and CHUNK_READY notification; a chunk
that cannot fit returns -1 (the caller fails only this room) instead of
falling back to the slice path, which would leak the decode-side
allocation. deferred=True means caller should re-enqueue and break.
"""
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:
# Fail this room, not the worker thread: the same prefill still
# serves other (same-TP, non-staging) decode instances.
logger.warning_once(
"[Staging] a chunk exceeds the staging ring; failing affected "
"requests. Increase SGLANG_DISAGG_STAGING_POOL_SIZE_MB or "
"reduce chunked_prefill_size."
)
return (-1, False)
# Not ready yet: wait (bounded) for a watermark advance, then
# re-enqueue to retry. A plain block-until-ready would head-of-line
# block other rooms on this single worker thread.
with self._staging_ctx.watermark_cv:
self._staging_ctx.watermark_cv.wait(STAGING_WATERMARK_WAIT_S)
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:
# Doesn't fit the ring: fail this room (caller's ret != 0 path), do
# not fall back to the slice path (leaks the decode-side allocation).
logger.warning_once(
"[Staging] a chunk does not fit the staging ring; failing affected "
"requests. Increase SGLANG_DISAGG_STAGING_POOL_SIZE_MB or "
"reduce chunked_prefill_size."
)
return (-1, False)
if ret == 0:
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,
get_schedule().chunked_prefill_size,
self._staging_ctx.prefetch_requested,
self._staging_ctx.prefetch_sockets,
requester_pp_rank=self.pp_rank,
)
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,
dst_layer_ids: List[int],
staging_buffer=None,
dst_slot_layer_ids: Optional[List[int]] = 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
local_bytes = per_layer_bytes * num_layers * 2
if self.pp_size > 1:
# Pair staging slots in [all K, all V] order; draft buffers make
# kv_data_ptrs diverge from this layout.
src_slot_ids = (
self.kv_buffer_tensors.get("slot_layer_ids")
or self.kv_args.kv_layer_ids
)
dst_slot_ids = dst_slot_layer_ids or dst_layer_ids
pairs = build_transfer_entry_pairs(
src_slot_ids,
dst_slot_ids,
num_layers * 2,
len(dst_slot_ids),
)
dst_num_layers = len(dst_slot_ids) // 2
else:
pairs = None
dst_num_layers = num_layers
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,
dst_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(local_bytes):
logger.warning(
f"Prefill staging too small for {local_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"for {dst_num_layers} layers, 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,
)
if pairs is None:
transfer_blocks = [
(
staging_buffer.get_ptr(),
dst_staging_ptr + rank_offset,
local_bytes,
)
]
else:
transfer_blocks = [
(
staging_buffer.get_ptr() + src_idx * per_layer_bytes,
dst_staging_ptr + rank_offset + dst_idx * per_layer_bytes,
per_layer_bytes,
)
for src_idx, dst_idx in pairs
]
ret = self._transfer_data(mooncake_session_id, transfer_blocks)
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_staging_ptr + rank_offset:x}, size={local_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
if not _nvlink_intra_transport_active():
src_addrs, dst_addrs, lengths = zip(*transfer_blocks)
return self.engine.batch_transfer_sync(
mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths)
)
# Intra-node NVLink transport can only move device memory. Partition
# blocks by the *local source* pointer (probing a local pointer is
# safe; the remote dst is never probed): device-sourced blocks go
# through the engine as usual, host-sourced blocks are shipped over
# the ordered zmq channel and written into the peer's buffer by the
# receiver (see _handle_state_data). This mirrors the aux TCP path.
device_blocks = []
host_blocks = []
for src, dst, length in transfer_blocks:
if _is_device_pointer(src):
device_blocks.append((src, dst, length))
else:
host_blocks.append((src, dst, length))
rc = 0
if device_blocks:
src_addrs, dst_addrs, lengths = zip(*device_blocks)
rc = self.engine.batch_transfer_sync(
mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths)
)
if rc == 0 and host_blocks:
rc = self._send_host_blocks_tcp(mooncake_session_id, host_blocks)
return rc
def _send_host_blocks_tcp(self, mooncake_session_id, host_blocks):
target = self._session_endpoint_map.get(mooncake_session_id)
if target is None:
logger.error(
f"No zmq endpoint known for mooncake session "
f"{mooncake_session_id}; cannot deliver {len(host_blocks)} "
"host-memory transfer blocks"
)
return -1
endpoint, dst_port, room = target
na = NetworkAddress(endpoint, dst_port)
for src, dst, length in host_blocks:
data = _read_bytes_from_address(src, length)
if data is None:
return -1
self._send_multipart_locked(
na.to_tcp(),
[
MooncakeKVManager.STATE_DATA_HEADER,
str(room).encode("ascii"),
str(int(dst)).encode("ascii"),
struct.pack(">I", len(data)),
data,
],
is_ipv6=na.is_ipv6,
)
return 0
def _send_kvcache_generic(
self,
mooncake_session_id: str,
src_data_ptrs: list[int],
dst_data_ptrs: list[int],
item_lens: list[int],
prefill_data_indices: npt.NDArray[np.int32],
dst_data_indices: npt.NDArray[np.int32],
executor: concurrent.futures.ThreadPoolExecutor,
state_type: Optional[StateType] = None,
force_flat: bool = False,
src_layer_ids: Optional[List[int]] = None,
dst_layer_ids: Optional[List[int]] = None,
dst_device_data_indices: Optional[npt.NDArray[np.int32]] = None,
dst_device_data_ptrs: Optional[set[int]] = None,
) -> int:
"""
Generic KV cache transfer supporting both MHA and MLA architectures.
This method is used by both send_kvcache (full pool) and maybe_send_extra.
``force_flat`` uses the MLA-style flat (single-buffer-per-layer) layout
even on a non-MLA backend, for K-only state buffers (e.g. MiniMax sparse
index) whose per-layer list must not be half-split into K/V.
"""
# Host and device buffers may use different destination page spaces.
# Build both transfer plans once, then select per destination buffer.
prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous(
prefill_data_indices, dst_data_indices
)
device_prefill_kv_blocks = device_dst_kv_blocks = None
if dst_device_data_indices is not None:
device_prefill_kv_blocks, device_dst_kv_blocks = (
group_concurrent_contiguous(
prefill_data_indices, dst_device_data_indices
)
)
layers_params = None
# Decode pp size should be equal to prefill pp size or 1
# Published layer IDs give exact pairing; plain-MHA peers publish none
# and keep positional slicing.
has_layer_ids = bool(src_layer_ids or dst_layer_ids)
# Unified SWA publishes one page-envelope region even on an MHA backend.
is_single_region_swa = (
state_type == StateType.SWA
and len(src_data_ptrs) == 1
and len(dst_data_ptrs) == 1
)
if (
self.is_mla_backend
or self.is_hybrid_mla_backend
or force_flat
or has_layer_ids
or is_single_region_swa
):
# Layer IDs map PP-local buffers to global decode entries.
# Registrations without them retain the existing PP mapping.
if has_layer_ids:
pairs = build_transfer_entry_pairs(
src_layer_ids,
dst_layer_ids,
len(src_data_ptrs),
len(dst_data_ptrs),
allow_positional_fallback=self.pp_size == 1,
)
layers_params = [
(src_data_ptrs[i], dst_data_ptrs[j], item_lens[i]) for i, j in pairs
]
else:
src_kv_ptrs, dst_kv_ptrs, layers_current_pp_stage = (
self.get_mla_kv_ptrs_with_pp(
src_data_ptrs, dst_data_ptrs, state_type
)
)
layers_params = [
(
src_kv_ptrs[layer_id],
dst_kv_ptrs[layer_id],
item_lens[layer_id],
)
for layer_id in range(layers_current_pp_stage)
]
else:
src_k_ptrs, src_v_ptrs, dst_k_ptrs, dst_v_ptrs, layers_current_pp_stage = (
self.get_mha_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs)
)
# item_lens structure: [k_layer0, k_layer1, ..., k_layerN, v_layer0, v_layer1, ..., v_layerN]
# Use correct item lengths for K and V separately
if layers_current_pp_stage > len(dst_k_ptrs):
logger.error(
"Prefill transfer kvcache error, layers_current_pp_stage is out of range: "
f"layers_current_pp_stage={layers_current_pp_stage}, len(dst_k_ptrs)={len(dst_k_ptrs)}"
)
return -1
layers_params = [
(
src_k_ptrs[layer_id],
dst_k_ptrs[layer_id],
item_lens[layer_id], # K item length
)
for layer_id in range(layers_current_pp_stage)
] + [
(
src_v_ptrs[layer_id],
dst_v_ptrs[layer_id],
item_lens[layers_current_pp_stage + layer_id], # V item length
)
for layer_id in range(layers_current_pp_stage)
]
assert layers_params is not None
def set_transfer_blocks(
src_ptr: int, dst_ptr: int, item_len: int
) -> List[Tuple[int, int, int]]:
transfer_blocks = []
if dst_device_data_ptrs and int(dst_ptr) in dst_device_data_ptrs:
assert (
device_prefill_kv_blocks is not None
and device_dst_kv_blocks is not None
)
src_blocks, dst_blocks = (
device_prefill_kv_blocks,
device_dst_kv_blocks,
)
else:
src_blocks, dst_blocks = prefill_kv_blocks, dst_kv_blocks
for prefill_index, decode_index in zip(src_blocks, dst_blocks):
src_addr = src_ptr + int(prefill_index[0]) * item_len
dst_addr = dst_ptr + int(decode_index[0]) * item_len
length = item_len * len(prefill_index)
transfer_blocks.append((src_addr, dst_addr, length))
return transfer_blocks
# Worker function for processing a single layer
def process_layer(src_ptr: int, dst_ptr: int, item_len: int) -> int:
transfer_blocks = set_transfer_blocks(src_ptr, dst_ptr, item_len)
return self._transfer_data(mooncake_session_id, transfer_blocks)
# Worker function for processing all layers in a batch
def process_layers(layers_params: List[Tuple[int, int, int]]) -> int:
transfer_blocks = []
for src_ptr, dst_ptr, item_len in layers_params:
transfer_blocks.extend(set_transfer_blocks(src_ptr, dst_ptr, item_len))
return self._transfer_data(mooncake_session_id, transfer_blocks)
if (
self.enable_custom_mem_pool
and self.custom_mem_pool_type != "INTRA_NODE_NVLINK"
):
futures = [
executor.submit(
process_layer,
src_ptr,
dst_ptr,
item_len,
)
for (src_ptr, dst_ptr, item_len) in layers_params
]
return self._await_transfer_futures(futures)
else:
# Combining all layers' params in one batch transfer is more efficient
# compared to using multiple threads. Preserve this legacy path unless
# users explicitly opt in to bounded index batches.
max_batch_indices = self.max_transfer_batch_indices
if max_batch_indices <= 0 or prefill_data_indices.size <= max_batch_indices:
return process_layers(layers_params)
def process_index_batch(
prefill_blocks,
dst_blocks,
device_prefill_blocks=None,
device_dst_blocks=None,
) -> int:
transfer_blocks = []
for src_ptr, dst_ptr, item_len in layers_params:
if dst_device_data_ptrs and int(dst_ptr) in dst_device_data_ptrs:
assert (
device_prefill_blocks is not None
and device_dst_blocks is not None
)
src_blocks, target_blocks = (
device_prefill_blocks,
device_dst_blocks,
)
else:
src_blocks, target_blocks = prefill_blocks, dst_blocks
for prefill_index, decode_index in zip(src_blocks, target_blocks):
src_addr = src_ptr + int(prefill_index[0]) * item_len
dst_addr = dst_ptr + int(decode_index[0]) * item_len
length = item_len * len(prefill_index)
transfer_blocks.append((src_addr, dst_addr, length))
return self._transfer_data(mooncake_session_id, transfer_blocks)
for start in range(
0,
prefill_data_indices.size,
max_batch_indices,
):
batch_prefill_blocks, batch_dst_blocks = group_concurrent_contiguous(
prefill_data_indices[start : start + max_batch_indices],
dst_data_indices[start : start + max_batch_indices],
)
batch_device_prefill_blocks = batch_device_dst_blocks = None
if dst_device_data_indices is not None:
batch_device_prefill_blocks, batch_device_dst_blocks = (
group_concurrent_contiguous(
prefill_data_indices[start : start + max_batch_indices],
dst_device_data_indices[start : start + max_batch_indices],
)
)
ret = process_index_batch(
batch_prefill_blocks,
batch_dst_blocks,
batch_device_prefill_blocks,
batch_device_dst_blocks,
)
if ret != 0:
return ret
return 0
def _validate_envelope_kv_layout(
self,
dst_kv_ptrs: list[int],
dst_kv_item_len: Optional[int],
dst_attn_tp_size: Optional[int] = None,
) -> None:
"""Reject a peer whose KV registration shape differs from ours.
The unified memory pool registers ONE whole-envelope region and
addresses the destination as ``dst_ptr + page_id * item_len`` using OUR
``item_len``, so a peer on a different page size / spec, or without
unified memory, would take envelope-sized blocks at the wrong offsets.
Must run before the first RDMA write.
Scoped to unified memory by config, not by region count: a non-unified
PP stage owning a single full-attention layer also registers one region,
and `_send_kvcache_generic` pairs that with the peer by layer id.
"""
if not get_memory().enable_unified_memory:
return
if dst_attn_tp_size is not None and self.attn_tp_size != dst_attn_tp_size:
# The unified mamba state ships as one whole-slot envelope with no
# per-tensor dims, so `_send_mamba_state_slice` cannot reslice it and
# silently falls back to an unsliced copy. Reject here, before any KV
# is written, rather than in `maybe_send_extra` afterwards.
raise RuntimeError(
"--enable-unified-memory does not support different prefill / "
f"decode attention TP sizes (prefill={self.attn_tp_size}, "
f"decode={dst_attn_tp_size}): the whole-envelope state cannot "
"be TP-resliced."
)
src_item_lens = self.kv_args.kv_item_lens
if (
len(src_item_lens) != 1
or len(dst_kv_ptrs) != 1
or dst_kv_item_len is None
or src_item_lens[0] != dst_kv_item_len
):
raise RuntimeError(
"PD KV layout mismatch on the whole-envelope path: prefill has "
f"{len(src_item_lens)} KV region(s) with item_lens="
f"{src_item_lens}, decode has {len(dst_kv_ptrs)} with item_len="
f"{dst_kv_item_len}. With --enable-unified-memory both sides "
"must enable it and use the same page size and model spec."
)
def _await_transfer_futures(self, futures) -> int:
"""Await a chunk's per-layer RDMA writes; return the first non-zero status.
cancel() is a no-op for a running future, so with deferred release on we
still drain the running ones before returning (no write may outlive this
call, which the drain-ack relies on). Off: original early-return."""
ret = 0
for future in concurrent.futures.as_completed(futures):
try:
status = future.result()
except concurrent.futures.CancelledError:
continue
if status != 0 and ret == 0:
ret = status
for f in futures:
f.cancel()
if not self.enable_deferred_decode_kv_release:
return ret
return ret
def send_kvcache(
self,
mooncake_session_id: str,
prefill_kv_indices: npt.NDArray[np.int32],
dst_kv_ptrs: list[int],
dst_kv_indices: npt.NDArray[np.int32],
executor: concurrent.futures.ThreadPoolExecutor,
dst_layer_ids: Optional[List[int]] = None,
dst_device_kv_indices: Optional[npt.NDArray[np.int32]] = None,
dst_kv_item_len: Optional[int] = None,
dst_attn_tp_size: Optional[int] = None,
):
self._validate_envelope_kv_layout(
dst_kv_ptrs, dst_kv_item_len, dst_attn_tp_size
)
dst_device_kv_ptrs = None
if dst_device_kv_indices is not None:
compression_ratios = self.kv_args.mla_compression_ratios
assert compression_ratios is not None
if len(dst_kv_ptrs) == len(self.kv_args.kv_data_ptrs):
start = self.kv_args.prefill_start_layer
end = self.kv_args.prefill_end_layer
assert end is not None
compression_ratios = compression_ratios[start:end]
c4_layer_num = sum(ratio == 4 for ratio in compression_ratios)
dst_device_kv_ptrs = set(dst_kv_ptrs[c4_layer_num:])
return self._send_kvcache_generic(
mooncake_session_id=mooncake_session_id,
src_data_ptrs=self.kv_args.kv_data_ptrs,
dst_data_ptrs=dst_kv_ptrs,
item_lens=self.kv_args.kv_item_lens,
prefill_data_indices=prefill_kv_indices,
dst_data_indices=dst_kv_indices,
executor=executor,
force_flat=get_memory().enable_unified_memory,
src_layer_ids=self.kv_args.kv_layer_ids,
dst_layer_ids=dst_layer_ids,
dst_device_data_indices=dst_device_kv_indices,
dst_device_data_ptrs=dst_device_kv_ptrs,
)
def send_kvcache_dcp(
self,
mooncake_session_id: str,
prefill_kv_indices: npt.NDArray[np.int32],
dst_kv_ptrs: list[int],
dst_kv_indices: npt.NDArray[np.int32],
*,
dcp_token_item_lens: List[int],
dst_dcp_size: int,
dst_dcp_rank: int,
src_page_offset: int,
decode_prefix_len: int,
num_kv_tokens: int,
executor: concurrent.futures.ThreadPoolExecutor,
dst_layer_ids: List[int],
pack_buffer=None,
dst_kv_item_lens: Optional[List[int]] = None,
dst_tp_rank: int = 0,
dst_attn_tp_size: Optional[int] = None,
) -> int:
if num_kv_tokens is None:
raise ValueError("PD DCP transfer requires num_kv_tokens")
physical_page_size = self.kv_args.page_size
if dst_kv_item_lens and len(dst_kv_item_lens) != len(dst_kv_ptrs):
raise ValueError("PD DCP destination KV lengths must match its buffers")
src_layer_ids = self.kv_args.kv_layer_ids
if src_layer_ids or dst_layer_ids:
dst_indices = resolve_dcp_dst_entry_indices(
src_layer_ids,
dst_layer_ids,
len(self.kv_args.kv_data_ptrs),
len(dst_kv_ptrs),
)
src_kv_ptrs = self.kv_args.kv_data_ptrs
dst_kv_ptrs = [dst_kv_ptrs[j] for j in dst_indices]
if dst_kv_item_lens:
dst_kv_item_lens = [dst_kv_item_lens[j] for j in dst_indices]
else:
src_kv_ptrs, dst_kv_ptrs, _ = self.get_mla_kv_ptrs_with_pp(
self.kv_args.kv_data_ptrs,
dst_kv_ptrs,
)
if dst_kv_item_lens:
_, dst_kv_item_lens, _ = self.get_mla_kv_ptrs_with_pp(
self.kv_args.kv_item_lens, dst_kv_item_lens
)
num_draft = self.kv_args.num_draft_entries
num_target = len(src_kv_ptrs) - num_draft
plan = build_dcp_token_transfer_plan(
prefill_kv_indices,
dst_kv_indices,
physical_page_size=physical_page_size,
dcp_size=dst_dcp_size,
dcp_rank=dst_dcp_rank,
src_page_offset=src_page_offset,
decode_prefix_len=decode_prefix_len,
num_kv_tokens=num_kv_tokens,
)
if plan.empty():
return 0
target_src_kv_ptrs = src_kv_ptrs[:num_target]
src_token_indices = plan.target_src_token_indices
if pack_buffer is not None and src_token_indices.size:
from sglang.srt.disaggregation.common.dcp_pack import try_pack_dcp_src
packed = try_pack_dcp_src(
pack_buffer=pack_buffer,
kv_data_ptrs=target_src_kv_ptrs,
src_token_indices=src_token_indices,
token_item_lens=dcp_token_item_lens[:num_target],
)
if packed is not None:
target_src_kv_ptrs, src_token_indices = packed
layers_params = []
if src_token_indices.size:
target_groups = group_concurrent_contiguous(
src_token_indices,
plan.target_dst_token_indices,
)
layers_params += [
(
target_src_kv_ptrs[entry],
dst_kv_ptrs[entry],
dcp_token_item_lens[entry],
target_groups,
)
for entry in range(num_target)
]
sliced_draft_params = []
if num_draft > 0 and plan.draft_src_token_indices.size:
if not dst_kv_item_lens and dst_attn_tp_size not in (
None,
self.attn_tp_size,
):
raise ValueError(
"PD DCP with different draft TP sizes requires destination KV lengths"
)
draft_groups = group_concurrent_contiguous(
plan.draft_src_token_indices,
plan.draft_dst_token_indices,
)
for entry in range(num_target, num_target + num_draft):
src_width = dcp_token_item_lens[entry]
dst_width = src_width
if dst_kv_item_lens:
dst_width, remainder = divmod(
dst_kv_item_lens[entry], physical_page_size * dst_dcp_size
)
if remainder or dst_width <= 0:
raise ValueError("Invalid PD DCP draft destination token width")
if src_width == dst_width:
layers_params.append(
(
src_kv_ptrs[entry],
dst_kv_ptrs[entry],
src_width,
draft_groups,
)
)
continue
if self.is_mla_backend:
raise ValueError(
"PD DCP draft head slicing is unsupported for pure MLA: "
"dummy prefill senders may omit draft head shards"
)
copy_width = min(src_width, dst_width)
if max(src_width, dst_width) % copy_width:
raise ValueError("PD DCP draft KV head shards must divide evenly")
if dst_attn_tp_size is None:
raise ValueError(
"PD DCP draft head slicing requires destination TP size"
)
src_span = src_width * self.attn_tp_size
dst_span = dst_width * dst_attn_tp_size
src_rank = (self.kv_args.engine_rank % self.attn_tp_size) // max(
1, src_span // dst_span
)
dst_rank = dst_tp_rank // max(1, dst_span // src_span)
src_offset = (dst_rank * dst_width) % src_width
dst_offset = (src_rank * src_width) % dst_width
sliced_draft_params.append(
(
src_kv_ptrs[entry] + src_offset,
dst_kv_ptrs[entry] + dst_offset,
src_width,
dst_width,
copy_width,
)
)
def process_sliced_draft(params) -> int:
batch_size = self.max_transfer_batch_indices
if batch_size <= 0:
batch_size = 4096
for start in range(0, plan.draft_src_token_indices.size, batch_size):
src_indices = plan.draft_src_token_indices[start : start + batch_size]
dst_indices = plan.draft_dst_token_indices[start : start + batch_size]
blocks = []
for src_ptr, dst_ptr, src_width, dst_width, copy_width in params:
src_addrs = src_ptr + src_indices * src_width
dst_addrs = dst_ptr + dst_indices * dst_width
blocks.extend(
(int(src), int(dst), copy_width)
for src, dst in zip(src_addrs, dst_addrs)
)
ret = self._transfer_data(mooncake_session_id, blocks)
if ret != 0:
return ret
return 0
def set_transfer_blocks(
src_ptr: int, dst_ptr: int, token_item_len: int, groups
) -> List[Tuple[int, int, int]]:
src_groups, dst_groups = groups
return [
(
src_ptr + int(src_group[0]) * token_item_len,
dst_ptr + int(dst_group[0]) * token_item_len,
len(src_group) * token_item_len,
)
for src_group, dst_group in zip(src_groups, dst_groups)
]
def process_layer(
src_ptr: int, dst_ptr: int, token_item_len: int, groups
) -> int:
return self._transfer_data(
mooncake_session_id,
set_transfer_blocks(src_ptr, dst_ptr, token_item_len, groups),
)
if self.enable_custom_mem_pool:
futures = [
executor.submit(process_layer, *layer_params)
for layer_params in layers_params
]
futures.extend(
executor.submit(process_sliced_draft, [params])
for params in sliced_draft_params
)
return self._await_transfer_futures(futures)
transfer_blocks = []
for layer_params in layers_params:
transfer_blocks.extend(set_transfer_blocks(*layer_params))
ret = self._transfer_data(mooncake_session_id, transfer_blocks)
if ret != 0 or not sliced_draft_params:
return ret
return process_sliced_draft(sliced_draft_params)
def send_kvcache_slice(
self,
mooncake_session_id: str,
prefill_kv_indices: npt.NDArray[np.int32],
dst_kv_ptrs: list[int],
dst_kv_indices: npt.NDArray[np.int32],
dst_tp_rank: int,
dst_attn_tp_size: int,
dst_kv_item_len: int,
executor: concurrent.futures.ThreadPoolExecutor,
dst_layer_ids: Optional[List[int]] = None,
):
"""
Sends KV cache slices from this Prefill rank to a target Decode rank,
supporting generic M-to-N TP size configurations.
NOTE: This implementation calls the transfer engine for each token slot within
each page to ensure correctness for any page_size and head-slicing configuration.
This may introduce performance overhead (increased TTFT) for long sequences.
"""
# Extract configuration
local_tp_rank_in_group = self.kv_args.engine_rank % self.attn_tp_size
src_kv_item_len = self.kv_args.kv_item_lens[0]
dst_tp_rank_in_group = dst_tp_rank % dst_attn_tp_size
page_size = self.kv_args.page_size
# Use total KV head count (not per-rank) for correct head distribution.
# Per-rank kv_head_num is max(1, total//tp) which loses info when total < tp.
total_kv_heads = getattr(self.kv_args, "total_kv_head_num", 0)
if total_kv_heads <= 0:
total_kv_heads = self.kv_args.kv_head_num * self.attn_tp_size
src_heads_per_rank = max(1, total_kv_heads // self.attn_tp_size)
dst_heads_per_rank = max(1, total_kv_heads // dst_attn_tp_size)
bytes_per_head_slice_to_send = (
dst_kv_item_len // page_size // dst_heads_per_rank
)
# GQA replication: how many prefill ranks share the same KV head
src_replication = max(1, self.attn_tp_size // total_kv_heads)
# Determine slicing parameters based on TP configuration
if self.attn_tp_size > dst_attn_tp_size:
# Send KVCache from multiple prefill instances to 1 decode instance
src_head_start_offset = 0
num_heads_to_send = src_heads_per_rank
unique_head_idx = local_tp_rank_in_group // src_replication
dst_head_start_offset = (
unique_head_idx * src_heads_per_rank
) % dst_heads_per_rank
else:
# Send KVCache from 1 prefill instance to multiple decode instances
# GQA replication (total_kv_heads < dst_attn_tp_size): consecutive decode
# ranks share one KV head (QKVParallelLinear: tp_rank // num_kv_head_replicas),
# so map by integer division NOT modulo or ranks 1..r-1 fetch the wrong head.
dst_replication = max(1, dst_attn_tp_size // total_kv_heads)
unique_dst_head_idx = dst_tp_rank_in_group // dst_replication
src_head_start_offset = (
unique_dst_head_idx * dst_heads_per_rank
) % src_heads_per_rank
num_heads_to_send = dst_heads_per_rank
dst_head_start_offset = 0
src_data_ptrs = self.kv_args.kv_data_ptrs
src_layer_ids = self.kv_args.kv_layer_ids
if src_layer_ids or dst_layer_ids:
# Draft buffers break the flat [K block, V block] layout, so pair by
# layer ID instead of the half-split used by get_mha_kv_ptrs_with_pp.
if any(l != src_kv_item_len for l in self.kv_args.kv_item_lens):
logger.error(
f"[{mooncake_session_id}] head-sliced transfer assumes one item "
f"length for every KV entry, got {set(self.kv_args.kv_item_lens)}"
)
return -1
layer_ptr_pairs = [
(src_data_ptrs[i], dst_kv_ptrs[j])
for i, j in build_transfer_entry_pairs(
src_layer_ids,
dst_layer_ids or [],
len(src_data_ptrs),
len(dst_kv_ptrs),
allow_positional_fallback=self.pp_size == 1,
)
]
else:
src_k_ptrs, src_v_ptrs, dst_k_ptrs, dst_v_ptrs, layers_current_pp_stage = (
self.get_mha_kv_ptrs_with_pp(src_data_ptrs, dst_kv_ptrs)
)
layer_ptr_pairs = [
(src_k_ptrs[i], dst_k_ptrs[i]) for i in range(layers_current_pp_stage)
] + [(src_v_ptrs[i], dst_v_ptrs[i]) for i in range(layers_current_pp_stage)]
# Calculate precise byte offset and length for the sub-slice within the token
src_head_slice_offset = src_head_start_offset * bytes_per_head_slice_to_send
dst_head_slice_offset = dst_head_start_offset * bytes_per_head_slice_to_send
heads_bytes_per_token_to_send = num_heads_to_send * bytes_per_head_slice_to_send
# Sanity check: The data sub-slice to be sent should fit into the dst buffer.
# This means heads_bytes_per_token_to_send <= (dst_kv_item_len // page_size)
if heads_bytes_per_token_to_send > (dst_kv_item_len // page_size):
logger.error(
f"[{mooncake_session_id}] slice size ({heads_bytes_per_token_to_send}) exceeds "
f"target token slot size ({dst_kv_item_len // page_size})"
)
return -1
prefill_page_indices = prefill_kv_indices.reshape(-1, 1).astype(np.int64)
decode_page_indices = dst_kv_indices.reshape(-1, 1).astype(np.int64)
tokens_per_page = np.arange(page_size, dtype=np.int64).reshape(1, -1)
bytes_per_token_on_prefill = src_kv_item_len // page_size
bytes_per_token_on_decode = dst_kv_item_len // page_size
src_token_slot_offsets = (
tokens_per_page * bytes_per_token_on_prefill + src_head_slice_offset
)
dst_token_slot_offsets = (
tokens_per_page * bytes_per_token_on_decode + dst_head_slice_offset
)
def process_layer_tp_aware(src_layer_ptr, dst_layer_ptr):
src_page_base_addrs = src_layer_ptr + prefill_page_indices * src_kv_item_len
dst_page_base_addrs = dst_layer_ptr + decode_page_indices * dst_kv_item_len
src_slice_addrs = src_page_base_addrs + src_token_slot_offsets
dst_slice_addrs = dst_page_base_addrs + dst_token_slot_offsets
src_addr_list = src_slice_addrs.reshape(-1).tolist()
if not src_addr_list:
# Nothing to transfer for this layer.
return 0
dst_addr_list = dst_slice_addrs.reshape(-1).tolist()
total_slices = len(src_addr_list)
length_list = [heads_bytes_per_token_to_send] * total_slices
return self.engine.batch_transfer_sync(
mooncake_session_id, src_addr_list, dst_addr_list, length_list
)
futures = [
executor.submit(process_layer_tp_aware, src_layer_ptr, dst_layer_ptr)
for src_layer_ptr, dst_layer_ptr in layer_ptr_pairs
]
return self._await_transfer_futures(futures)
def send_aux(
self,
req: TransferInfo,
prefill_aux_index: int,
dst_aux_ptrs: list[int],
):
# TODO(shangming): Fix me when nvlink_transport of Mooncake is bug-free
if (
(self.enable_custom_mem_pool and self.custom_mem_pool_type == "NVLINK")
or envs.SGLANG_MOONCAKE_SEND_AUX_TCP.get()
or _nvlink_intra_transport_active()
):
return self.send_aux_tcp(req, prefill_aux_index, dst_aux_ptrs)
transfer_blocks = []
prefill_aux_ptrs = self.kv_args.aux_data_ptrs
prefill_aux_item_lens = self.kv_args.aux_item_lens
for i, dst_aux_ptr in enumerate(dst_aux_ptrs):
length = prefill_aux_item_lens[i]
src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index
dst_addr = dst_aux_ptrs[i] + length * req.dst_aux_index
transfer_blocks.append((src_addr, dst_addr, length))
return self._transfer_data(req.mooncake_session_id, transfer_blocks)
def send_aux_tcp(
self,
req: TransferInfo,
prefill_aux_index: int,
dst_aux_ptrs: list[int],
):
prefill_aux_ptrs = self.kv_args.aux_data_ptrs
prefill_aux_item_lens = self.kv_args.aux_item_lens
for i in range(len(prefill_aux_ptrs)):
length = prefill_aux_item_lens[i]
src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index
data = AuxDataCodec.serialize_data_from_buffer(src_addr, length)
self.send_aux_data_to_endpoint(
remote=req.endpoint,
dst_port=req.dst_port,
room=req.room,
buffer_index=i,
aux_index=req.dst_aux_index,
data=data,
)
return 0
def send_aux_data_to_endpoint(
self,
remote: str,
dst_port: int,
room: int,
buffer_index: int,
aux_index: int,
data: bytes,
):
na = NetworkAddress(remote, dst_port)
self._send_multipart_locked(
na.to_tcp(),
[
MooncakeKVManager.AUX_DATA_HEADER,
str(room).encode("ascii"),
str(buffer_index).encode("ascii"),
str(aux_index).encode("ascii"),
struct.pack(">I", len(data)),
data,
],
is_ipv6=na.is_ipv6,
)
def _handle_aux_data(self, msg: List[bytes]):
"""Handle AUX_DATA messages received by the decode thread."""
room = int(msg[1].decode("ascii"))
buffer_index = int(msg[2].decode("ascii"))
aux_index = int(msg[3].decode("ascii"))
data_length = struct.unpack(">I", msg[4])[0]
data = msg[5]
if len(data) != data_length:
logger.error(f"AUX_DATA length mismatch for bootstrap_room {room}")
return
AuxDataCodec.deserialize_data_to_buffer(
self.kv_args, buffer_index, aux_index, data
)
logger.debug(
f"Received AUX_DATA for bootstrap_room {room} with length:{len(data)}"
)
def _host_transfer_regions(self):
"""Address ranges this process published as transfer targets.
Used to validate STATE_DATA writes. Built lazily because kv_args is
fully populated only after registration.
"""
regions = getattr(self, "_host_transfer_regions_cache", None)
if regions is None:
regions = []
for ptr, length in zip(
self.kv_args.kv_data_ptrs or [], self.kv_args.kv_data_lens or []
):
regions.append((int(ptr), int(ptr) + int(length)))
for ptr, length in zip(
self.kv_args.aux_data_ptrs or [], self.kv_args.aux_data_lens or []
):
regions.append((int(ptr), int(ptr) + int(length)))
for ptrs, lens in zip(
self.kv_args.state_data_ptrs or [], self.kv_args.state_data_lens or []
):
for ptr, length in zip(ptrs or [], lens or []):
regions.append((int(ptr), int(ptr) + int(length)))
self._host_transfer_regions_cache = regions
return regions
def _handle_state_data(self, msg: List[bytes]):
"""Handle STATE_DATA messages received by the decode thread.
Carries one host-memory transfer block that could not go through the
intra-node NVLink transport. Written directly into the local buffer at
the destination address; ordering against the final status message is
guaranteed by the shared per-endpoint zmq socket.
"""
room = int(msg[1].decode("ascii"))
dst_addr = int(msg[2].decode("ascii"))
data_length = struct.unpack(">I", msg[3])[0]
data = msg[4]
if len(data) != data_length:
logger.error(f"STATE_DATA length mismatch for bootstrap_room {room}")
return
in_region = any(
start <= dst_addr and dst_addr + len(data) <= end
for start, end in self._host_transfer_regions()
)
if not in_region:
logger.error(
f"STATE_DATA for bootstrap_room {room} targets unknown region "
f"{hex(dst_addr)}..{hex(dst_addr + len(data))}; dropping"
)
return
if not _write_bytes_to_address(dst_addr, data):
logger.error(
f"STATE_DATA write failed for bootstrap_room {room} at "
f"{hex(dst_addr)} len {len(data)}"
)
return
logger.debug(
f"Received STATE_DATA for bootstrap_room {room} at {hex(dst_addr)} "
f"with length:{len(data)}"
)
def _get_dsa_cache_transfer_skip_flags(
self, info: Optional[KVArgsRegisterInfo]
) -> Tuple[bool, bool]:
skip_kv = False
skip_state = False
# Must be checked before the non-hybrid early return below, or every CP
# rank re-sends the same state and we transfer it cp_size times over.
# Prefill CP all-gathers before writing the pool, so every CP rank holds
# the full state regardless of whether the pool is hybrid. We assume no
# structure about the state rows, so we don't split them across CP ranks
# -- just let rank 0 send the whole thing (unless layer split already
# shards it per rank).
if self._should_skip_cp_replicated_state_transfer():
skip_state = True
if not self.is_hybrid_mla_backend:
return skip_kv, skip_state
if info is not None and self.attn_tp_size > info.dst_attn_tp_size:
sub_rank = (self.kv_args.engine_rank % self.attn_tp_size) % (
self.attn_tp_size // info.dst_attn_tp_size
)
if sub_rank != 0:
skip_kv = True
# Hybrid-MLA KV is replicated across these source ranks, but
# TP-sharded state needs every rank for the aggregation path.
return skip_kv, skip_state
def _is_generic_kvcache_state_type(self, st: StateType) -> bool:
"""State types sent via the page-indexed ``_send_kvcache_generic`` path
(not the mamba-state path); subclasses extend for hardware components."""
return st in (
StateType.SWA,
StateType.DSA,
StateType.QSA_PENDING,
StateType.QSA_COMPRESSED,
StateType.SWA_RING,
StateType.DSV4_REQUEST_STATE,
StateType.BLOCK_SCALE,
StateType.BLOCK_SCALE_SWA,
)
def _requires_exact_state_index_match(self, st: StateType) -> bool:
"""State types whose page lists are positional and must not be truncated."""
return st in (
StateType.QSA_PENDING,
StateType.QSA_COMPRESSED,
StateType.SWA_RING,
StateType.DSV4_REQUEST_STATE,
)
def maybe_send_extra(
self,
req: TransferInfo,
prefill_state_indices: List,
executor: concurrent.futures.ThreadPoolExecutor,
target_rank_registration_info: Optional[KVArgsRegisterInfo] = None,
):
rc = 0
state_types = getattr(self.kv_args, "state_types", [])
for i, st in enumerate(state_types):
indices = (
prefill_state_indices[i] if i < len(prefill_state_indices) else None
)
if indices is None:
continue
src_data_ptrs = self.kv_args.state_data_ptrs[i]
src_item_lens = self.kv_args.state_item_lens[i]
src_dim_per_tensor = (
self.kv_args.state_dim_per_tensor[i]
if i < len(self.kv_args.state_dim_per_tensor)
else []
)
src_conv_shard_groups = getattr(self.kv_args, "state_conv_shard_groups", [])
src_conv_shard_groups = (
src_conv_shard_groups[i] if i < len(src_conv_shard_groups) else []
)
src_slice_outer_counts = getattr(
self.kv_args, "state_slice_outer_counts", []
)
src_slice_outer_counts = (
src_slice_outer_counts[i] if i < len(src_slice_outer_counts) else []
)
src_state_layer_ids = self.kv_args.state_layer_ids
src_state_layer_ids = (
src_state_layer_ids[i] if i < len(src_state_layer_ids) else []
)
if target_rank_registration_info is not None:
dst_data_ptrs = (
target_rank_registration_info.dst_state_data_ptrs[i]
if i < len(target_rank_registration_info.dst_state_data_ptrs)
else []
)
dst_item_lens = (
target_rank_registration_info.dst_state_item_lens[i]
if i < len(target_rank_registration_info.dst_state_item_lens)
else []
)
dst_dim_per_tensor = (
target_rank_registration_info.dst_state_dim_per_tensor[i]
if i < len(target_rank_registration_info.dst_state_dim_per_tensor)
else []
)
dst_state_layer_ids = (
target_rank_registration_info.dst_state_layer_ids[i]
if i < len(target_rank_registration_info.dst_state_layer_ids)
else []
)
else:
dst_data_ptrs, dst_item_lens, dst_dim_per_tensor = [], [], []
dst_state_layer_ids = []
dst_indices = (
req.dst_state_indices[i] if i < len(req.dst_state_indices) else []
)
if st == StateType.MAMBA:
if (not src_dim_per_tensor or not dst_dim_per_tensor) and list(
src_item_lens
) != list(dst_item_lens):
raise RuntimeError(
"Mamba state layouts differ between prefill and decode "
f"(src item_lens={src_item_lens}, dst item_lens="
f"{dst_item_lens}) and no per-tensor dim metadata is "
"available to reslice. With --enable-unified-memory, "
"prefill and decode must both enable it and use equal "
"attention TP sizes."
)
if (
target_rank_registration_info is not None
and self.attn_tp_size
!= target_rank_registration_info.dst_attn_tp_size
):
rc = (
self._send_mamba_state_slice(
req,
indices,
src_data_ptrs,
src_item_lens,
src_dim_per_tensor,
dst_data_ptrs,
dst_indices,
dst_item_lens,
dst_dim_per_tensor,
target_rank_registration_info.dst_tp_rank,
target_rank_registration_info.dst_attn_tp_size,
src_conv_shard_groups,
src_slice_outer_counts,
src_state_layer_ids,
dst_state_layer_ids,
)
or rc
)
else:
rc = (
self._send_mamba_state(
req,
indices,
src_data_ptrs,
src_item_lens,
dst_data_ptrs,
dst_indices,
src_state_layer_ids,
dst_state_layer_ids,
dst_item_lens,
)
or rc
)
elif st == StateType.DSA_TAIL:
rc = (
self._send_slot_state(
req,
src_data_ptrs,
src_item_lens,
dst_data_ptrs,
dst_item_lens,
list(indices),
list(dst_indices),
st.value,
)
or rc
)
elif self._is_generic_kvcache_state_type(st):
is_qwen4_qsa_state = st in (
StateType.QSA_PENDING,
StateType.QSA_COMPRESSED,
)
has_heterogeneous_attn_tp = (
target_rank_registration_info is not None
and self.attn_tp_size
!= target_rank_registration_info.dst_attn_tp_size
)
if (
has_heterogeneous_attn_tp
and not self.is_mla_backend
and not self.is_hybrid_mla_backend
and not is_qwen4_qsa_state
):
raise RuntimeError(
f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {st.upper()} hybrid models yet."
)
if has_heterogeneous_attn_tp and is_qwen4_qsa_state:
if len(dst_item_lens) != len(dst_data_ptrs):
raise RuntimeError(
f"Replicated {st.upper()} destination pointer/item-length "
"metadata is inconsistent: "
f"dst ptrs={len(dst_data_ptrs)} lens={len(dst_item_lens)}"
)
qsa_entry_pairs = build_transfer_entry_pairs(
src_state_layer_ids,
dst_state_layer_ids,
len(src_data_ptrs),
len(dst_data_ptrs),
allow_positional_fallback=self.pp_size == 1,
)
layout_mismatches = [
(i, j, src_item_lens[i], dst_item_lens[j])
for i, j in qsa_entry_pairs
if src_item_lens[i] != dst_item_lens[j]
]
if layout_mismatches:
raise RuntimeError(
f"Replicated {st.upper()} layout differs between mapped "
"prefill and decode entries: "
f"{layout_mismatches}"
)
local_tp_rank_in_group = (
self.kv_args.engine_rank % self.attn_tp_size
)
if not should_send_replicated_state(
src_attn_tp_size=self.attn_tp_size,
dst_attn_tp_size=(
target_rank_registration_info.dst_attn_tp_size
),
local_tp_rank_in_group=local_tp_rank_in_group,
):
continue
src_indices = list(indices)
dst_indices_local = list(dst_indices)
if (
st == StateType.DSV4_REQUEST_STATE
and len(src_indices) == 0
and len(dst_indices_local) == 0
):
continue
if len(src_indices) != len(dst_indices_local):
# These components are position- or request-indexed:
# truncating silently misaligns rows and corrupts KV.
# Paged SWA/DSA tolerate a 1-page drift -> keep the
# lenient truncation below.
if self._requires_exact_state_index_match(st):
raise RuntimeError(
f"{st.upper()} state index length mismatch: "
f"prefill={len(src_indices)}, dst={len(dst_indices_local)}"
)
logger.warning(
f"len(prefill_state_indices) = {len(src_indices)}, len(dst_state_indices) = {len(dst_indices_local)}"
)
if len(src_indices) > len(dst_indices_local):
src_indices = src_indices[: len(dst_indices_local)]
else:
dst_indices_local = dst_indices_local[: len(src_indices)]
rc = (
self._send_kvcache_generic(
mooncake_session_id=req.mooncake_session_id,
src_data_ptrs=src_data_ptrs,
dst_data_ptrs=dst_data_ptrs,
item_lens=src_item_lens,
prefill_data_indices=np.array(src_indices, dtype=np.int32),
dst_data_indices=np.array(dst_indices_local, dtype=np.int32),
executor=executor,
state_type=st,
# Two independent reasons to keep the flat layout.
# QSA's per-layer list must not be half-split into K/V;
# neither must a unified sub-pool's single region, which
# holds every layer's K and V per slot envelope -- the
# MHA branch would compute zero layers and ship nothing
# (same reason as in `send_kvcache`).
force_flat=(
st in (StateType.QSA_PENDING, StateType.QSA_COMPRESSED)
or get_memory().enable_unified_memory
),
src_layer_ids=src_state_layer_ids,
dst_layer_ids=dst_state_layer_ids,
)
or rc
)
elif st == StateType.MINIMAX_INDEX_K:
# Equal-TP / PP=1 only. Sub-pools are compacted sparse-layer
# lists, so PP>1 mis-slices and heterogeneous TP is unsupported.
if self.pp_size is not None and self.pp_size > 1:
raise RuntimeError(
"PD disagg: PP>1 not supported for MiniMax sparse index yet."
)
if (
target_rank_registration_info is not None
and self.attn_tp_size
!= target_rank_registration_info.dst_attn_tp_size
):
raise RuntimeError(
"PD disagg: heterogeneous TP not supported for MiniMax "
"sparse index yet."
)
src_indices = list(indices)
dst_indices_local = list(dst_indices)
if len(src_indices) > len(dst_indices_local):
src_indices = src_indices[: len(dst_indices_local)]
elif len(src_indices) < len(dst_indices_local):
dst_indices_local = dst_indices_local[: len(src_indices)]
rc = (
self._send_kvcache_generic(
mooncake_session_id=req.mooncake_session_id,
src_data_ptrs=src_data_ptrs,
dst_data_ptrs=dst_data_ptrs,
item_lens=src_item_lens,
prefill_data_indices=np.array(src_indices, dtype=np.int32),
dst_data_indices=np.array(dst_indices_local, dtype=np.int32),
executor=executor,
force_flat=True,
)
or rc
)
return rc
def _send_slot_state(
self,
req: TransferInfo,
src_ptrs: list[int],
src_item_lens: list[int],
dst_ptrs: list[int],
dst_item_lens: list[int],
src_indices: list[int],
dst_indices: list[int],
label: str,
) -> int:
try:
dst_ptrs = slice_dsa_tail_dst_ptrs_for_pp(
src_ptrs,
dst_ptrs,
self.kv_args.prefill_start_layer,
self.kv_args.prefill_end_layer,
)
dst_item_lens = slice_dsa_tail_dst_ptrs_for_pp(
src_ptrs,
dst_item_lens,
self.kv_args.prefill_start_layer,
self.kv_args.prefill_end_layer,
)
transfer_blocks = build_dsa_tail_transfer_blocks(
src_ptrs,
src_item_lens,
dst_ptrs,
src_indices,
dst_indices,
dst_item_lens,
)
except ValueError as exc:
logger.error("%s: %s", label, exc)
return -1
return self._transfer_data(req.mooncake_session_id, transfer_blocks)
def _send_mamba_state(
self,
req: TransferInfo,
prefill_mamba_index: list,
src_state_data_ptrs: list[int],
src_state_item_lens: list[int],
dst_state_data_ptrs: list[int],
dst_mamba_index: list,
src_layer_ids: Optional[List[int]] = None,
dst_layer_ids: Optional[List[int]] = None,
dst_state_item_lens: Optional[list[int]] = None,
):
assert len(prefill_mamba_index) == 1, "Mamba should have single state index"
transfer_blocks = []
pairs = build_transfer_entry_pairs(
src_layer_ids or [],
dst_layer_ids or [],
len(src_state_data_ptrs),
len(dst_state_data_ptrs),
allow_positional_fallback=self.pp_size == 1,
)
for i, j in pairs:
dst_state_ptr = dst_state_data_ptrs[j]
length = src_state_item_lens[i]
if dst_state_item_lens and length != dst_state_item_lens[j]:
raise RuntimeError(
"Prefill/Decode Mamba slot size mismatch "
f"(src={length}, dst={dst_state_item_lens[j]}). "
"Configure matching persistent state layouts on both peers."
)
src_addr = src_state_data_ptrs[i] + length * int(prefill_mamba_index[0])
dst_addr = dst_state_ptr + length * int(dst_mamba_index[0])
transfer_blocks.append((src_addr, dst_addr, length))
return self._transfer_data(req.mooncake_session_id, transfer_blocks)
def _send_mamba_state_slice(
self,
req: TransferInfo,
prefill_mamba_index: list,
src_state_data_ptrs: list[int],
src_state_item_lens: list[int],
src_state_dim_per_tensor: list[int],
dst_state_data_ptrs: list[int],
dst_mamba_index: list,
dst_state_item_lens: list[int],
dst_state_dim_per_tensor: list[int],
dst_tp_rank: int,
dst_attn_tp_size: int,
src_state_conv_shard_groups: list = None,
src_state_slice_outer_counts: list[int] = None,
src_layer_ids: Optional[List[int]] = None,
dst_layer_ids: Optional[List[int]] = None,
):
"""Transfer Mamba states with TP slice support.
Mamba state layout:
- conv_state: [num_layers, size+1, conv_dim/tp, conv_kernel-1]
- temporal_state: [num_layers, size+1, num_heads/tp, head_dim, state_size]
The 3rd dimension is sliced by TP. When prefill and decode have different
attn_tp_size, we slice the state accordingly. GDN conv_state is the
concatenation [query | key | value] with each sub-block head-sharded
independently, so on the scatter path it is sliced per sub-block via
``src_state_conv_shard_groups`` (see
compute_mamba_state_slice_byte_blocks).
"""
logger.warning_once(
"Using Mamba state slice transfer for different runtime attention TP "
f"sizes: prefill={self.attn_tp_size}, decode={dst_attn_tp_size}. "
"Performance may be affected."
)
assert len(prefill_mamba_index) == 1, "Mamba should have single state index"
# If no dimension info available, fall back to regular transfer
if not src_state_dim_per_tensor or not dst_state_dim_per_tensor:
return self._send_mamba_state(
req,
prefill_mamba_index,
src_state_data_ptrs,
src_state_item_lens,
dst_state_data_ptrs,
dst_mamba_index,
src_layer_ids,
dst_layer_ids,
dst_state_item_lens,
)
local_tp_rank_in_group = self.kv_args.engine_rank % self.attn_tp_size
dst_tp_rank_in_group = dst_tp_rank % dst_attn_tp_size
transfer_blocks = []
pairs = build_transfer_entry_pairs(
src_layer_ids or [],
dst_layer_ids or [],
len(src_state_data_ptrs),
len(dst_state_data_ptrs),
allow_positional_fallback=self.pp_size == 1,
)
for i, j in pairs:
dst_state_ptr = dst_state_data_ptrs[j]
src_item_len = src_state_item_lens[i]
dst_item_len = dst_state_item_lens[j]
src_dim = src_state_dim_per_tensor[i]
dst_dim = dst_state_dim_per_tensor[j]
conv_shard_groups = (
src_state_conv_shard_groups[i]
if src_state_conv_shard_groups and i < len(src_state_conv_shard_groups)
else None
)
outer_count = (
src_state_slice_outer_counts[i]
if src_state_slice_outer_counts
and i < len(src_state_slice_outer_counts)
else 1
)
for (
src_offset,
dst_offset,
bytes_to_send,
) in compute_mamba_state_slice_byte_blocks(
src_item_len=src_item_len,
dst_item_len=dst_item_len,
src_dim=src_dim,
dst_dim=dst_dim,
outer_count=outer_count,
src_attn_tp_size=self.attn_tp_size,
dst_attn_tp_size=dst_attn_tp_size,
dst_tp_rank_in_group=dst_tp_rank_in_group,
local_tp_rank_in_group=local_tp_rank_in_group,
conv_shard_groups=conv_shard_groups,
):
src_addr = (
src_state_data_ptrs[i]
+ src_item_len * int(prefill_mamba_index[0])
+ src_offset
)
dst_addr = (
dst_state_ptr + dst_item_len * int(dst_mamba_index[0]) + dst_offset
)
transfer_blocks.append((src_addr, dst_addr, bytes_to_send))
return self._transfer_data(req.mooncake_session_id, transfer_blocks)
def transfer_worker(
self,
queue: FastQueue,
executor: concurrent.futures.ThreadPoolExecutor,
staging_buffer=None,
worker_index=0,
):
staging_strategy = None
if self.enable_trace:
trace_set_thread_info(
f"mooncake transfer worker {worker_index}",
tp_rank=self.attn_tp_rank,
dp_rank=self.attn_dp_rank,
)
while True:
try:
kv_chunk: TransferKVChunk = queue.get()
# teardown() pushes a None sentinel to unblock get() and stop
# the worker: FastQueue.get() blocks indefinitely, so checking
# _stopped alone can never wake a parked worker during a role switch.
if kv_chunk is None:
break
if self.enable_trace:
kv_chunk.trace_ctx.rebuild_thread_context()
kv_chunk.trace_ctx.trace_slice_start(
MooncakeRequestStage.MOONCAKE_WORKER_SEND.stage_name,
MooncakeRequestStage.MOONCAKE_WORKER_SEND.level,
)
# Counted at dequeue, before the status check, so
# `outstanding == 0` means nothing is dequeued or in flight --
# the predicate the abort ack relies on. The flag survives
# re-enqueue on defer.
if not kv_chunk.staging_counted:
self._staging_outstanding[kv_chunk.room] += 1
kv_chunk.staging_counted = True
if (
kv_chunk.room not in self.request_status
or self.check_status(kv_chunk.room) == KVPoll.Failed
):
logger.debug(
f"Skipping chunk for room {kv_chunk.room} because it has already failed or been aborted"
)
if self.enable_trace:
kv_chunk.trace_ctx.trace_slice_end(
MooncakeRequestStage.MOONCAKE_WORKER_SEND.stage_name,
MooncakeRequestStage.MOONCAKE_WORKER_SEND.level,
thread_finish_flag=True,
)
self._staging_outstanding.pop(kv_chunk.room, None)
if self.enable_deferred_decode_kv_release:
# Skipped => nothing written for this aborted room; ack.
self._maybe_ack_drained_abort(kv_chunk.room)
continue
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
else []
)
polls = []
dst_ranks_infos = []
# Unique id per prefill sender so decode's response set size matches expected_response_num.
prefill_unique_rank = self._prefill_unique_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:
start_ts = time.perf_counter()
if not req.is_dummy:
# Early exit if the request has failed
with self.session_lock:
if req.mooncake_session_id in self.failed_sessions:
self.conclude_failure(
bootstrap_room=kv_chunk.room,
failure_reason=(
"Decode instance could be dead, remote "
f"mooncake session {req.mooncake_session_id} "
"is not alive"
),
)
break
target_rank_registration_info: KVArgsRegisterInfo = (
self.decode_kv_args_table[req.mooncake_session_id]
)
is_dcp_transfer = (
target_rank_registration_info.requires_dcp_relayout
)
chunked_dst_device_kv_indice = None
if is_dcp_transfer:
if req.dst_device_kv_indices is not None:
raise RuntimeError(
"HiSparse destination device indices are not "
"supported by PD DCP relayout"
)
chunked_dst_kv_indice = req.dst_kv_indices
else:
chunked_dst_kv_indice = req.dst_kv_indices[
kv_chunk.index_slice
]
if req.dst_device_kv_indices is not None:
chunked_dst_device_kv_indice = (
req.dst_device_kv_indices[kv_chunk.index_slice]
)
# NOTE: This is temporarily a workaround to deal with the case where the prefill_kv_indices
# is mismatched with the dst_kv_indices when page size > 1, this should never happen.
if len(chunked_dst_kv_indice) < len(
kv_chunk.prefill_kv_indices
):
logger.warning(
f"len(chunked_dst_kv_indice) = {len(chunked_dst_kv_indice)}, len(kv_chunk.prefill_kv_indices) = {len(kv_chunk.prefill_kv_indices)}"
)
kv_chunk.prefill_kv_indices = (
kv_chunk.prefill_kv_indices[
: len(chunked_dst_kv_indice)
]
)
if chunked_dst_device_kv_indice is not None:
chunked_dst_device_kv_indice = (
chunked_dst_device_kv_indice[
: len(kv_chunk.prefill_kv_indices)
]
)
skip_kv, skip_state = self._get_dsa_cache_transfer_skip_flags(
target_rank_registration_info
)
if (
len(kv_chunk.prefill_kv_indices) == 0
or not self.kv_args.kv_data_ptrs
or skip_kv
):
ret = 0
elif is_dcp_transfer:
dcp_token_item_lens = (
target_rank_registration_info.dcp_token_item_lens
)
assert dcp_token_item_lens is not None
pack_buffer = (
self._dcp_pack_buffers[worker_index]
if self._dcp_pack_buffers
else None
)
ret = self.send_kvcache_dcp(
req.mooncake_session_id,
kv_chunk.prefill_kv_indices,
target_rank_registration_info.dst_kv_ptrs,
chunked_dst_kv_indice,
dcp_token_item_lens=dcp_token_item_lens,
dst_dcp_size=target_rank_registration_info.dst_dcp_size,
dst_dcp_rank=target_rank_registration_info.dst_dcp_rank,
src_page_offset=kv_chunk.index_slice.start or 0,
decode_prefix_len=req.decode_prefix_len or 0,
num_kv_tokens=kv_chunk.num_kv_tokens,
executor=executor,
dst_layer_ids=(
target_rank_registration_info.dst_kv_layer_ids
),
pack_buffer=pack_buffer,
dst_kv_item_lens=target_rank_registration_info.dst_kv_item_lens,
dst_tp_rank=target_rank_registration_info.dst_tp_rank,
dst_attn_tp_size=target_rank_registration_info.dst_attn_tp_size,
)
elif (
self.is_mla_backend
or self.is_hybrid_mla_backend
or self.attn_tp_size
== target_rank_registration_info.dst_attn_tp_size
):
ret = self.send_kvcache(
req.mooncake_session_id,
kv_chunk.prefill_kv_indices,
target_rank_registration_info.dst_kv_ptrs,
chunked_dst_kv_indice,
executor,
dst_layer_ids=target_rank_registration_info.dst_kv_layer_ids,
dst_device_kv_indices=chunked_dst_device_kv_indice,
dst_kv_item_len=target_rank_registration_info.dst_kv_item_len,
dst_attn_tp_size=target_rank_registration_info.dst_attn_tp_size,
)
elif (
self.enable_staging
and staging_strategy is not None
and (
target_rank_registration_info.staging_base_ptr != 0
or target_rank_registration_info.staging_total_size != 0
)
):
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,
kv_chunk.prefill_kv_indices,
target_rank_registration_info.dst_kv_ptrs,
chunked_dst_kv_indice,
target_rank_registration_info.dst_tp_rank,
target_rank_registration_info.dst_attn_tp_size,
target_rank_registration_info.dst_kv_item_len,
executor,
target_rank_registration_info.dst_kv_layer_ids,
)
if ret != 0:
with self.session_lock:
self.session_failures[req.mooncake_session_id] += 1
# Failures should never happen if the session is not dead, if the session fails once, mark it as failed
if self.session_failures[req.mooncake_session_id] >= 1:
self.failed_sessions.add(req.mooncake_session_id)
logger.error(
f"Session {req.mooncake_session_id} failed."
)
self.conclude_failure(
bootstrap_room=kv_chunk.room,
failure_reason=(
f"Failed to send kv chunk of {kv_chunk.room} to "
f"{NetworkAddress(req.endpoint, req.dst_port).to_host_port_str()}"
),
)
break
if kv_chunk.is_last_chunk:
if kv_chunk.state_indices and not skip_state:
state_rc = self.maybe_send_extra(
req,
kv_chunk.state_indices,
executor,
target_rank_registration_info,
)
if state_rc != 0:
with self.session_lock:
self.session_failures[
req.mooncake_session_id
] += 1
self.failed_sessions.add(
req.mooncake_session_id
)
self.conclude_failure(
bootstrap_room=kv_chunk.room,
failure_reason=(
"Failed to send state components of "
f"{kv_chunk.room} to "
f"{NetworkAddress(req.endpoint, req.dst_port).to_host_port_str()}"
),
)
break
# Only the last chunk we need to send the aux data
ret = self.send_aux(
req,
kv_chunk.prefill_aux_index,
target_rank_registration_info.dst_aux_ptrs,
)
polls.append(True if ret == 0 else False)
dst_ranks_infos.append((req.endpoint, req.dst_port))
# Only sync status when all the dst ranks have received the kvcache
if len(polls) == req.required_dst_info_num:
status = KVPoll.Success if all(polls) else KVPoll.Failed
self.conclude_transfer(
bootstrap_room=req.room,
status=status,
targets=dst_ranks_infos,
failure_reason=(
None
if status == KVPoll.Success
else f"Failed to send aux data of {req.room}"
),
)
else:
# Dummy request means the decode instance is not used, so its status can be marked as success directly
# Dummy request does not need to sync status to decode endpoint
if kv_chunk.is_last_chunk and req.room in self.request_status:
self.update_status(req.room, KVPoll.Success)
if self.enable_trace:
mooncake_trace_slice(
kv_chunk.trace_ctx,
MooncakeRequestStage.MOONCAKE_WORKER_SEND_SESSION,
start_ts,
)
if self.enable_trace:
kv_chunk.trace_ctx.trace_slice_end(
MooncakeRequestStage.MOONCAKE_WORKER_SEND.stage_name,
MooncakeRequestStage.MOONCAKE_WORKER_SEND.level,
thread_finish_flag=True,
)
if staging_deferred:
continue
self._staging_outstanding[kv_chunk.room] -= 1
if self.enable_deferred_decode_kv_release:
# In-flight write finished; if aborted and nothing outstanding,
# the pages are idle -> release the held ack.
self._maybe_ack_drained_abort(kv_chunk.room)
# Tear down only when no chunk is still outstanding and the room
# has concluded: already cleared, Success, or a Failed *last*
# chunk. A non-last Failed chunk keeps the room (more chunks may
# follow), not on the last chunk alone since an earlier deferred
# chunk may still need to transfer.
if self._staging_outstanding.get(kv_chunk.room, 0) <= 0 and (
kv_chunk.room not in self.request_status
or self.check_status(kv_chunk.room) == KVPoll.Success
or (
kv_chunk.is_last_chunk
and self.check_status(kv_chunk.room) == KVPoll.Failed
)
):
self._staging_outstanding.pop(kv_chunk.room, None)
if kv_chunk.room in self.transfer_infos:
for sid in self.transfer_infos[kv_chunk.room]:
self._session_endpoint_map.pop(sid, None)
self.transfer_infos.pop(kv_chunk.room)
self.req_to_decode_prefix_len.pop(kv_chunk.room, None)
if self.enable_staging:
# Purge prefetch bookkeeping for the finished room.
# Snapshot first: the scheduler thread adds concurrently.
for key in list(self._staging_ctx.prefetch_requested):
if key[0] == kv_chunk.room:
self._staging_ctx.prefetch_requested.discard(key)
self._staging_ctx.prefetched_rooms.discard(kv_chunk.room)
except Exception as e:
# NOTE(shangming): Remove this when we make sure the transfer thread is bug-free
raise RuntimeError(
f"Transfer thread failed because of {e}. Prefill instance with bootstrap_port={self.bootstrap_port} is dead."
)
def start_prefill_thread(self):
recv = self._make_worker_recv(self.server_socket)
def bootstrap_thread():
"""This thread recvs pre-alloc notification from the decode engine"""
# KVPoll.Bootstrapping -> KVPoll.WaitingForInput
while not self._stopped:
waiting_req_bytes = recv()
if waiting_req_bytes is None:
continue
room = waiting_req_bytes[0].decode("ascii")
# Staging: decode reports consumption watermark back to prefill
if room == "WATERMARK":
handle_watermark_msg(self._staging_ctx, waiting_req_bytes)
continue
# Staging: decode replies with allocated staging offset
if room == "STAGING_RSP":
handle_staging_rsp(waiting_req_bytes, self.transfer_infos)
continue
# Decode-side abort notification: mark room as failed and ACK
if room == "ABORT":
room_to_be_aborted = int(waiting_req_bytes[1].decode("ascii"))
decode_ip = waiting_req_bytes[2].decode("ascii")
decode_port = int(waiting_req_bytes[3].decode("ascii"))
room_active = (
room_to_be_aborted in self.request_status
and self.check_status(room_to_be_aborted) != KVPoll.Success
)
if self.enable_deferred_decode_kv_release:
# Mark Failed FIRST (stops add_transfer_request enqueuing
# new chunks), THEN register the ack target: registering
# first would let the worker drain+ack while the room is
# not yet Failed, so a newly enqueued chunk could still
# write to the freed pages. The worker (not this thread)
# acks once its in-flight write drains; if nothing is in
# flight, decode falls back to the release timeout.
if room_active:
self.update_status(room_to_be_aborted, KVPoll.Failed)
self.register_deferred_ack_target(
room_to_be_aborted, decode_ip, decode_port
)
# Try once: the room may already be quiescent and
# never revisited by the worker.
self._maybe_ack_drained_abort(room_to_be_aborted)
logger.debug(
f"Received abort notification for room {room_to_be_aborted}, "
f"marked as Failed; ACK deferred until transfer drains"
)
elif self._staging_outstanding.get(room_to_be_aborted, 0) == 0:
# Concluded/unknown AND quiescent: ack now. A cleared
# room is not automatically quiescent -- clear() can
# drop a room whose chunk is still transferring.
self._send_abort_ack(
decode_ip, decode_port, room_to_be_aborted
)
continue
# No need to abort the room if it has already succeeded
if room_active:
self.update_status(room_to_be_aborted, KVPoll.Failed)
logger.debug(
f"Received abort notification for room {room_to_be_aborted}, "
f"marked as Failed"
)
else:
logger.debug(
f"Received abort notification for room {room_to_be_aborted}, "
f"ignoring (already completed or unknown)"
)
# Send ACK back to decode endpoint
try:
na = NetworkAddress(decode_ip, decode_port)
self._send_multipart_locked(
na.to_tcp(),
[
b"ABORT_ACK",
str(room_to_be_aborted).encode("ascii"),
],
is_ipv6=na.is_ipv6,
)
logger.debug(
f"Sent ABORT_ACK for room {room_to_be_aborted} to "
f"{decode_ip}:{decode_port}"
)
except Exception as e:
logger.debug(
f"Failed to send ABORT_ACK for room {room_to_be_aborted}: {e}"
)
continue
mooncake_session_id = waiting_req_bytes[3].decode("ascii")
if room == "None":
decode_kv_args = KVArgsRegisterInfo.from_zmq(waiting_req_bytes)
decode_kv_args.requires_dcp_relayout = self.requires_dcp_relayout(
decode_kv_args.dst_dcp_size,
decode_kv_args.dst_dcp_rank,
)
if decode_kv_args.requires_dcp_relayout:
num_entries = len(self.kv_args.kv_item_lens)
num_draft = self.kv_args.num_draft_entries
dst_item_lens: List[Optional[int]] = [
decode_kv_args.dst_kv_item_len
] * (num_entries - num_draft) + [None] * num_draft
decode_kv_args.dcp_token_item_lens = (
self.prepare_dcp_token_item_lens(
dst_item_lens,
decode_kv_args.dst_dcp_size,
)
)
self._init_dcp_pack_buffers_once(decode_kv_args.dst_dcp_size)
self.decode_kv_args_table[mooncake_session_id] = decode_kv_args
with self.session_lock:
if mooncake_session_id in self.failed_sessions:
self.failed_sessions.remove(mooncake_session_id)
if mooncake_session_id in self.session_failures:
del self.session_failures[mooncake_session_id]
logger.debug(
f"Register KVArgs from {mooncake_session_id} successfully"
)
continue
else:
required_dst_info_num = int(waiting_req_bytes[7].decode("ascii"))
room = int(room)
if room not in self.transfer_infos:
self.transfer_infos[room] = {}
self.transfer_infos[room][mooncake_session_id] = (
TransferInfo.from_zmq(waiting_req_bytes)
)
self._session_endpoint_map[mooncake_session_id] = (
self.transfer_infos[room][mooncake_session_id].endpoint,
self.transfer_infos[room][mooncake_session_id].dst_port,
room,
)
# NOTE: after bootstrapping we can mark the req as waiting for input
if len(self.transfer_infos[room]) == required_dst_info_num:
self.resolve_kv_replica_factor(self.transfer_infos[room])
self.req_to_decode_prefix_len[room] = next(
(
info.decode_prefix_len
for info in self.transfer_infos[room].values()
if info.decode_prefix_len is not None
),
0,
)
self.update_status(room, KVPoll.WaitingForInput)
t = threading.Thread(target=bootstrap_thread, daemon=True)
t.start()
self._worker_threads.append(t)
def start_decode_thread(self):
recv = self._make_worker_recv(self.server_socket)
def decode_thread():
while not self._stopped:
msg = recv()
if msg is None:
continue
if msg[0] == MooncakeKVManager.AUX_DATA_HEADER:
self._handle_aux_data(msg)
continue
if msg[0] == MooncakeKVManager.STATE_DATA_HEADER:
self._handle_state_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")
handler = self._staging_handler
assert handler is not None, (
"CHUNK_READY received before staging handler initialized"
)
handler.handle_chunk_arrived(
room,
chunk_idx,
page_start,
num_pages,
session_id,
)
continue
# Staging: prefill pre-requests staging allocation before forward
if msg[0] == b"STAGING_REQ":
self._handle_staging_req(msg)
continue
# Prefill acknowledges abort notification
if msg[0] == b"ABORT_ACK":
ack_aborted_room = int(msg[1].decode("ascii"))
logger.debug(f"Received ABORT_ACK for room {ack_aborted_room}")
# Deferred release: the 3-frame ack carries the prefill rank
# and means its transfer drained; aggregate for is_abort_release_safe.
if self.enable_deferred_decode_kv_release and len(msg) >= 3:
self.note_abort_ack(
ack_aborted_room, int(msg[2].decode("ascii"))
)
continue
parsed = self.parse_kv_status_message(msg)
if parsed is None:
continue
room, status, prefill_rank, reason = parsed
self.apply_prefill_status(
bootstrap_room=room,
status=status,
prefill_rank=prefill_rank,
failure_reason=reason,
)
t = threading.Thread(target=decode_thread, daemon=True)
t.start()
self._worker_threads.append(t)
t = self._start_heartbeat_checker_thread()
self._worker_threads.append(t)
def add_transfer_request(
self,
bootstrap_room: int,
kv_indices: npt.NDArray[np.int32],
index_slice: slice,
is_last_chunk: bool,
aux_index: Optional[int] = None,
state_indices: Optional[List] = None,
num_kv_tokens: Optional[int] = None,
trace_ctx: Optional[Union[TraceReqContext, TraceNullContext]] = None,
):
assert self.disaggregation_mode == DisaggregationMode.PREFILL
assert not is_last_chunk or (is_last_chunk and aux_index is not None)
if (
bootstrap_room not in self.request_status
or self.check_status(bootstrap_room) == KVPoll.Failed
):
logger.debug(
"Request with bootstrap_room=%s already failed", bootstrap_room
)
return
if bootstrap_room not in self.transfer_infos:
# This means that the current rank is a dummy rank for this request,
# and it has already been marked as success, so there is no need to
# add further chunks into the transfer queue.
return
# NOTE(shangming): sharding according to the dst_infos to make sure
# requests with the same dst_sessions will be added into the same
# queue, which enables early abort with failed sessions.
dst_infos = self.transfer_infos[bootstrap_room].keys()
session_port_sum = sum(int(session.rsplit(":", 1)[1]) for session in dst_infos)
shard_idx = session_port_sum % len(self.transfer_queues)
if trace_ctx is None:
trace_ctx = TraceNullContext()
self.transfer_queues[shard_idx].put(
TransferKVChunk(
room=bootstrap_room,
prefill_kv_indices=kv_indices,
index_slice=index_slice,
is_last_chunk=is_last_chunk,
prefill_aux_index=aux_index,
state_indices=state_indices,
num_kv_tokens=num_kv_tokens,
trace_ctx=trace_ctx,
)
)
def get_session_id(self):
return self.engine.get_session_id()
def _run_one_probe_pass(self) -> None:
with self.session_lock:
snapshot = list(self.failed_sessions)
for session_id in snapshot:
send_probe = getattr(self.engine, "send_probe", None)
if send_probe is None:
rc = -1
else:
try:
rc = send_probe(session_id)
except Exception as e:
logger.warning("send_probe(%s) raised: %s", session_id, e)
continue
if rc == 0:
with self.session_lock:
was_blacklisted = session_id in self.failed_sessions
self.failed_sessions.discard(session_id)
self.session_failures.pop(session_id, None)
if was_blacklisted:
logger.info(
"Session %s recovered via probe; un-blacklisted",
session_id,
)
FAILED_SESSION_RECOVERIES.inc()
else:
logger.debug("Probe still failing for %s (rc=%d)", session_id, rc)
def _failed_session_probe_loop(self) -> None:
logger.info(
"Starting failed-session probe loop (interval=%.1fs)",
self.failed_session_probe_interval,
)
while not self._failed_session_probe_shutdown.wait(
self.failed_session_probe_interval
):
self._run_one_probe_pass()
class MooncakeFailureExceptionMixin:
"""Shared `failure_exception` for the Mooncake sender and receiver.
Both sides conclude a failed room identically: latch Failed, clear local
state, then raise with the recorded reason -- or, when no reason was
recorded locally, report it as propagated from another rank. Expects the
concrete class to provide ``conclude_state``, ``clear()``,
``bootstrap_room`` and ``kv_mgr``.
"""
def failure_exception(self):
# A room with no locally recorded reason failed on another rank.
if self.conclude_state is None:
self.conclude_state = KVPoll.Failed
self.clear()
with self.kv_mgr.failure_lock:
failure_reason = self.kv_mgr.failure_records.pop(self.bootstrap_room, None)
is_propagated = failure_reason is None
if is_propagated:
failure_reason = "Failed due to an unknown reason from another rank"
raise KVTransferError(
self.bootstrap_room, failure_reason, is_from_another_rank=is_propagated
)
class MooncakeKVSender(MooncakeFailureExceptionMixin, CommonKVSender):
def __init__(
self,
mgr: MooncakeKVManager,
bootstrap_addr: str,
bootstrap_room: int,
dest_tp_ranks: List[int],
pp_rank: int,
req_has_disagg_prefill_dp_rank: bool = False,
):
super().__init__(
mgr,
bootstrap_addr,
bootstrap_room,
dest_tp_ranks,
pp_rank,
req_has_disagg_prefill_dp_rank,
)
self.conclude_state = None
self.init_time = time.time()
self._init_trace_ctx()
@mooncake_trace_func(MooncakeRequestStage.MOONCAKE_SEND)
def send(
self,
kv_indices: npt.NDArray[np.int32],
state_indices: Optional[List] = None,
num_kv_tokens: Optional[int] = None,
):
kv_indices, index_slice, is_last_chunk, should_skip = (
self._prepare_send_indices(kv_indices, state_indices)
)
if should_skip:
return
if not is_last_chunk:
self.kv_mgr.add_transfer_request(
self.bootstrap_room,
kv_indices,
index_slice,
False,
num_kv_tokens=num_kv_tokens,
trace_ctx=self.trace_ctx.copy_for_thread(),
)
else:
self.kv_mgr.add_transfer_request(
self.bootstrap_room,
kv_indices,
index_slice,
True,
aux_index=self.aux_index,
state_indices=state_indices,
num_kv_tokens=num_kv_tokens,
trace_ctx=self.trace_ctx.copy_for_thread(),
)
self._record_transfer_indices(kv_indices, state_indices)
def poll(self) -> KVPoll:
if self.conclude_state is None:
status = self.kv_mgr.check_status(self.bootstrap_room)
# Hold Success until all staging chunks transferred: a deferred
# chunk can still be pending, and concluding now would drop it.
if (
status == KVPoll.Success
and self.kv_mgr._staging_outstanding.get(self.bootstrap_room, 0) > 0
):
return KVPoll.Transferring
if status in (KVPoll.Success, KVPoll.Failed):
self.conclude_state = status
self.trace_ctx.trace_req_finish()
elif status == KVPoll.Bootstrapping:
timeout_result = self._check_bootstrap_timeout()
if timeout_result is not None:
return timeout_result
return status
else:
return self.conclude_state
def _init_trace_ctx(self):
if self.kv_mgr.enable_trace:
self.trace_ctx = TraceReqContext(
rid=str(hex(self.bootstrap_room)),
bootstrap_room=self.bootstrap_room,
role="Sender",
module_name="mooncake",
)
if not self.trace_ctx.tracing_enable:
self.trace_ctx = TraceNullContext()
else:
self.trace_ctx = TraceNullContext()
self.trace_ctx.trace_req_start()
def abort(self):
super().abort()
self.trace_ctx.abort(abort_info={"reason": "Aborted"})
self.trace_ctx.trace_req_finish()
class MooncakeKVReceiver(MooncakeFailureExceptionMixin, CommonKVReceiver):
def __init__(
self,
mgr: MooncakeKVManager,
bootstrap_addr: str,
bootstrap_room: Optional[int] = None,
):
self.session_id = mgr.get_session_id()
self.init_time = None
super().__init__(mgr, bootstrap_addr, bootstrap_room)
def _register_kv_args(self) -> bool:
for bootstrap_info in self.bootstrap_infos:
packed_kv_data_ptrs = b"".join(
struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.kv_data_ptrs
)
packed_aux_data_ptrs = b"".join(
struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.aux_data_ptrs
)
packed_state_data_ptrs = pack_int_lists(
self.kv_mgr.kv_args.state_data_ptrs, "Q"
)
packed_state_item_lens = pack_int_lists(
self.kv_mgr.kv_args.state_item_lens, "I"
)
packed_state_dim_per_tensor = pack_int_lists(
getattr(self.kv_mgr.kv_args, "state_dim_per_tensor", []) or [], "I"
)
packed_state_layer_ids = pack_int_lists(
self.kv_mgr.kv_args.state_layer_ids, "I"
)
packed_kv_layer_ids = b"".join(
struct.pack("I", layer_id)
for layer_id in self.kv_mgr.kv_args.kv_layer_ids
)
# Note(shangming): No need to add pp rank here since decode pp size should be equal to prefill pp size or 1
tp_rank = self.kv_mgr.kv_args.engine_rank
# Some pools have no full-token contiguous KV (kv_item_lens empty)
# and ship per-pool instead, so report 0.
kv_item_len = (
self.kv_mgr.kv_args.kv_item_lens[0]
if self.kv_mgr.kv_args.kv_item_lens
else 0
)
dst_tp_rank = str(tp_rank).encode("ascii")
dst_attn_tp_size = str(self.kv_mgr.attn_tp_size).encode("ascii")
dst_kv_item_len = str(kv_item_len).encode("ascii")
dst_dcp_size = str(self.kv_mgr.dcp_size).encode("ascii")
dst_dcp_rank = str(self.kv_mgr.dcp_rank).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""
staging_slots = getattr(self.kv_mgr, "kv_buffer_tensors", None) or {}
packed_staging_slot_layer_ids = b"".join(
struct.pack("Q", layer_id)
for layer_id in (staging_slots.get("slot_layer_ids") or [])
)
try:
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
with lock:
sock.send_multipart(
[
"None".encode("ascii"),
self.kv_mgr.local_ip.encode("ascii"),
str(self.kv_mgr.rank_port).encode("ascii"),
self.session_id.encode("ascii"),
packed_kv_data_ptrs,
packed_aux_data_ptrs,
packed_state_data_ptrs,
dst_tp_rank,
dst_attn_tp_size,
dst_kv_item_len,
packed_state_item_lens,
packed_state_dim_per_tensor,
packed_kv_layer_ids,
packed_state_layer_ids,
packed_staging_base_ptr,
staging_total_size_str,
dst_dcp_size,
dst_dcp_rank,
packed_staging_slot_layer_ids,
struct.pack(
f"{len(self.kv_mgr.kv_args.kv_item_lens)}Q",
*self.kv_mgr.kv_args.kv_item_lens,
),
]
)
except zmq.ZMQError:
self.kv_mgr.record_failure(
self.bootstrap_room,
f"_register_kv_args to prefill {bootstrap_info.get('rank_ip')}:{bootstrap_info.get('rank_port')} failed",
)
self.conclude_state = KVPoll.Failed
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
return False
return True
def send_metadata(
self,
kv_indices: npt.NDArray[np.int32],
aux_index: Optional[int] = None,
state_indices: Optional[List] = None,
decode_prefix_len: Optional[int] = None,
device_kv_indices: Optional[npt.NDArray[np.int32]] = None,
):
if self.bootstrap_infos is None:
self.kv_mgr.record_failure(
self.bootstrap_room,
f"Could not fetch prefill parallel info from bootstrap_addr: {self.bootstrap_addr}",
)
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
return
self.chunk_staging_infos = []
if (
self.kv_mgr.enable_staging
and self.kv_mgr._staging_ctx.allocator is not None
):
self.kv_mgr.register_staging_room_bootstrap(
self.bootstrap_room, self.bootstrap_infos, self
)
for bootstrap_info in self.bootstrap_infos:
is_dummy = bootstrap_info["is_dummy"]
try:
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
with lock:
sock.send_multipart(
[
str(self.bootstrap_room).encode("ascii"),
self.kv_mgr.local_ip.encode("ascii"),
str(self.kv_mgr.rank_port).encode("ascii"),
self.session_id.encode("ascii"),
kv_indices.tobytes() if not is_dummy else b"",
str(aux_index).encode("ascii") if not is_dummy else b"",
(
pack_int_lists(state_indices, "i")
if not is_dummy and state_indices
else b""
),
str(self.required_dst_info_num).encode("ascii"),
str(decode_prefix_len or 0).encode("ascii"),
(
np.asarray(device_kv_indices, dtype=np.int32).tobytes()
if not is_dummy and device_kv_indices is not None
else b""
),
]
)
except zmq.ZMQError:
self.invalidate_cached_bootstrap_infos()
self.kv_mgr.record_failure(
self.bootstrap_room,
f"send_metadata to prefill {bootstrap_info.get('rank_ip')}:{bootstrap_info.get('rank_port')} failed",
)
self.conclude_state = KVPoll.Failed
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
return
self.init_time = time.time()
def poll(self) -> KVPoll:
if self.conclude_state is not None:
return self.conclude_state
status = self.kv_mgr.check_status(self.bootstrap_room)
if status in (KVPoll.Success, KVPoll.Failed):
self.conclude_state = status
elif status == KVPoll.WaitingForInput:
timeout_result = self._check_waiting_timeout()
if timeout_result is not None:
return timeout_result
return status
class MooncakeKVBootstrapServer(CommonKVBootstrapServer):
pass
Reference in New Issue View Git Blame Copy Permalink