leaveletandClaude Fable 5
58f2350738
E2e caught one request wedged FOREVER in disagg_prefill_inflight_queue (probes 21s apart with zero traffic both showed #inflight-req: 1, no reap/timeout warnings ever logged). Mechanism, established by reading the full state machine: prefill Success is set locally by the transfer worker on the LAST chunk; if the decode peer is torn down between the handshake and the prefill's final send(), add_transfer_request silently drops the chunk (no transfer destinations) — Success becomes unreachable. The only external rescue, the decode ABORT notification, is best-effort (silently swallowed on send error, no-op if it races the room registration), there is no prefill-side heartbeat of decode sessions, and the sender's only timeout covers Bootstrapping — the inflight queue itself has no liveness bound. The orphan pins the request's KV pages and rides every poll collective. Two fixes, both reaped through the existing Failed branch via the CP/TP MIN-reduce poll consensus (Failed=0 wins, so one rank concluding flips every rank together — rank-uniform by construction): - add_transfer_request: a room with no transfer destinations that is NOT already Success (the dummy-rank handshake marking) now concludes Failed loudly instead of dropping the chunk silently. - Inflight residency timeout: entries stuck in a non-terminal poll state past SGLANG_DISAGGREGATION_INFLIGHT_TIMEOUT (default 300s, matching the sibling BOOTSTRAP/WAITING timeouts) get sender.abort() and reap on the next poll. Covers what the hardening cannot: lost ABORT datagrams, decode crashes. Known sibling gaps left for follow-up: the decode transfer queue has no Transferring liveness bound, and an abort that matches no queue is still a silent no-op (much narrower race than first thought — work requests are ordered before control requests within a tick). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2086 lines
92 KiB
Python
2086 lines
92 KiB
Python
from __future__ import annotations
|
|
|
|
import concurrent.futures
|
|
import ctypes
|
|
import dataclasses
|
|
import logging
|
|
import os
|
|
import struct
|
|
import threading
|
|
import time
|
|
from collections import defaultdict
|
|
from typing import List, Optional, Tuple
|
|
|
|
import numpy as np
|
|
import numpy.typing as npt
|
|
|
|
from sglang.srt.disaggregation.base.conn import KVArgs, KVPoll
|
|
from sglang.srt.disaggregation.common.conn import (
|
|
CommonKVBootstrapServer,
|
|
CommonKVManager,
|
|
CommonKVReceiver,
|
|
CommonKVSender,
|
|
)
|
|
from sglang.srt.disaggregation.common.utils import (
|
|
FastQueue,
|
|
contiguous_group_stats,
|
|
group_concurrent_contiguous,
|
|
)
|
|
from sglang.srt.disaggregation.mooncake.utils import (
|
|
check_mooncake_custom_mem_pool_enabled,
|
|
)
|
|
from sglang.srt.disaggregation.utils import (
|
|
DisaggregationMode,
|
|
filter_kv_indices_for_cp_rank,
|
|
validate_transfer_page_count_or_raise,
|
|
)
|
|
from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.server_args import ServerArgs
|
|
from sglang.srt.utils.network import NetworkAddress
|
|
|
|
logger = logging.getLogger(__name__)
|
|
_CP_SHARED_DEBUG_COUNTS: dict[str, int] = {}
|
|
|
|
|
|
def _cp_shared_debug_log(key: str, message: str, *args, limit: int = 64) -> None:
|
|
if not envs.SGLANG_DEBUG_CP_SHARED_KV.get():
|
|
return
|
|
count = _CP_SHARED_DEBUG_COUNTS.get(key, 0)
|
|
if count >= limit:
|
|
return
|
|
_CP_SHARED_DEBUG_COUNTS[key] = count + 1
|
|
logger.info("[CP_SHARED_KV_DEBUG] " + message, *args)
|
|
|
|
|
|
def _cp_draft_shared_kv_debug(message: str, *args, limit: int = 64) -> None:
|
|
if not envs.SGLANG_CP_DRAFT_SHARED_KV_DEBUG.get():
|
|
return
|
|
key = "draft:" + message.split(" ", 1)[0]
|
|
count = _CP_SHARED_DEBUG_COUNTS.get(key, 0)
|
|
if count >= limit:
|
|
return
|
|
_CP_SHARED_DEBUG_COUNTS[key] = count + 1
|
|
logger.info("[CP_DRAFT_SHARED_KV] " + message, *args)
|
|
|
|
|
|
def _mooncake_transfer_stats_enabled() -> bool:
|
|
return envs.SGLANG_DISAGGREGATION_TRANSFER_STATS.get()
|
|
|
|
|
|
def _mooncake_transfer_stats_log(message: str, *args) -> None:
|
|
if not _mooncake_transfer_stats_enabled():
|
|
return
|
|
limit = envs.SGLANG_DISAGGREGATION_TRANSFER_STATS_LIMIT.get()
|
|
if limit is not None and limit <= 0:
|
|
return
|
|
key = "mooncake_transfer_stats"
|
|
count = _CP_SHARED_DEBUG_COUNTS.get(key, 0)
|
|
if limit is not None and count >= limit:
|
|
return
|
|
_CP_SHARED_DEBUG_COUNTS[key] = count + 1
|
|
logger.info("[Mooncake-transfer-stats] " + message, *args)
|
|
|
|
|
|
def _np_summary(arr) -> str:
|
|
if arr is None:
|
|
return "None"
|
|
arr = np.asarray(arr)
|
|
if arr.size == 0:
|
|
return f"shape={arr.shape} dtype={arr.dtype} size=0"
|
|
head = arr.reshape(-1)[: min(8, arr.size)].tolist()
|
|
return (
|
|
f"shape={arr.shape} dtype={arr.dtype} size={arr.size} "
|
|
f"min={int(arr.min())} max={int(arr.max())} head={head}"
|
|
)
|
|
|
|
|
|
class KVTransferError(Exception):
|
|
def __init__(self, bootstrap_room: int, failure_reason: str):
|
|
super().__init__(failure_reason)
|
|
self.bootstrap_room = bootstrap_room
|
|
self.failure_reason = failure_reason
|
|
|
|
def __str__(self):
|
|
return f"KVTransferError(bootstrap_room={self.bootstrap_room}): {self.failure_reason}"
|
|
|
|
|
|
# prefill
|
|
@dataclasses.dataclass
|
|
class TransferKVChunk:
|
|
room: int
|
|
prefill_kv_indices: npt.NDArray[np.int32]
|
|
index_slice: Optional[slice]
|
|
logical_page_positions: Optional[npt.NDArray[np.int32]]
|
|
state_logical_page_positions: Optional[npt.NDArray[np.int32]]
|
|
is_last_chunk: bool
|
|
prefill_aux_index: Optional[int]
|
|
state_indices: Optional[List[int]]
|
|
|
|
|
|
# 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[int]
|
|
required_dst_info_num: int
|
|
is_dummy: bool
|
|
|
|
@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"))
|
|
if msg[6] == b"":
|
|
dst_state_indices = []
|
|
else:
|
|
dst_state_indices = list(np.frombuffer(msg[6], dtype=np.int32))
|
|
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
|
|
@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[int]
|
|
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[int]
|
|
dst_state_dim_per_tensor: list[int]
|
|
dst_state_layer_ids: list[int]
|
|
|
|
@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=list(struct.unpack(f"{len(msg[6])//8}Q", msg[6])),
|
|
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=(
|
|
list(struct.unpack(f"{len(msg[10])//4}I", msg[10]))
|
|
if len(msg) > 10 and len(msg[10]) > 0
|
|
else []
|
|
),
|
|
dst_state_dim_per_tensor=(
|
|
list(struct.unpack(f"{len(msg[11])//4}I", msg[11]))
|
|
if len(msg) > 11 and len(msg[11]) > 0
|
|
else []
|
|
),
|
|
dst_state_layer_ids=(
|
|
list(struct.unpack(f"{len(msg[12])//4}i", msg[12]))
|
|
if len(msg) > 12 and len(msg[12]) > 0
|
|
else []
|
|
),
|
|
)
|
|
|
|
|
|
class AuxDataCodec:
|
|
"""Handles serialization and deserialization of auxiliary data buffers"""
|
|
|
|
@staticmethod
|
|
def serialize_data_from_buffer(src_addr, data_length):
|
|
"""Serialize data from memory buffer to bytes"""
|
|
buffer = (ctypes.c_byte * data_length).from_address(src_addr)
|
|
return bytes(buffer)
|
|
|
|
@staticmethod
|
|
def deserialize_data_to_buffer(kv_args, buffer_index, aux_index, data):
|
|
"""Deserialize bytes into target memory buffer"""
|
|
dst_aux_ptr = kv_args.aux_data_ptrs[buffer_index]
|
|
item_len = kv_args.aux_item_lens[buffer_index]
|
|
dst_addr = dst_aux_ptr + item_len * aux_index
|
|
buffer = (ctypes.c_byte * len(data)).from_address(dst_addr)
|
|
buffer[:] = data
|
|
return
|
|
|
|
|
|
class MooncakeKVManager(CommonKVManager):
|
|
AUX_DATA_HEADER = b"AUX_DATA"
|
|
|
|
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()
|
|
if self.disaggregation_mode == DisaggregationMode.PREFILL:
|
|
self.start_prefill_thread()
|
|
self.session_failures = defaultdict(int)
|
|
self.failed_sessions = set()
|
|
self.session_lock = threading.Lock()
|
|
# 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)
|
|
]
|
|
for queue, executor in zip(self.transfer_queues, self.executors):
|
|
threading.Thread(
|
|
target=self.transfer_worker, args=(queue, executor), daemon=True
|
|
).start()
|
|
self.enable_custom_mem_pool, self.custom_mem_pool_type = (
|
|
check_mooncake_custom_mem_pool_enabled()
|
|
)
|
|
self.enable_per_layer_async_transfer = (
|
|
envs.SGLANG_CP_SHARED_KV_PER_LAYER_TRANSFER.get()
|
|
)
|
|
elif self.disaggregation_mode == DisaggregationMode.DECODE:
|
|
self.start_decode_thread()
|
|
|
|
def init_engine(self):
|
|
self.engine = get_mooncake_transfer_engine()
|
|
|
|
def register_buffer_to_engine(self):
|
|
# Batch register KV data buffers
|
|
if self.kv_args.kv_data_ptrs and self.kv_args.kv_data_lens:
|
|
_cp_draft_shared_kv_debug(
|
|
"register_buffers mode=%s cp_rank=%s total_kv_bufs=%s "
|
|
"draft_start=%s draft_count=%s kv_lens=%s kv_item_lens=%s "
|
|
"state_type=%s state_bufs=%s state_lens=%s state_item_lens=%s "
|
|
"draft_state_type=%s draft_state_bufs=%s",
|
|
self.disaggregation_mode,
|
|
self.attn_cp_rank,
|
|
len(self.kv_args.kv_data_ptrs),
|
|
getattr(self.kv_args, "draft_kv_buffer_start", None),
|
|
getattr(self.kv_args, "draft_kv_buffer_count", None),
|
|
_np_summary(self.kv_args.kv_data_lens),
|
|
_np_summary(self.kv_args.kv_item_lens),
|
|
getattr(self.kv_args, "state_type", None),
|
|
len(getattr(self.kv_args, "state_data_ptrs", []) or []),
|
|
_np_summary(getattr(self.kv_args, "state_data_lens", [])),
|
|
_np_summary(getattr(self.kv_args, "state_item_lens", [])),
|
|
getattr(self.kv_args, "draft_state_type", None),
|
|
getattr(self.kv_args, "draft_state_buffer_count", None),
|
|
)
|
|
self.engine.batch_register(
|
|
self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens
|
|
)
|
|
|
|
# Batch register auxiliary data buffers
|
|
if self.kv_args.aux_data_ptrs and self.kv_args.aux_data_lens:
|
|
self.engine.batch_register(
|
|
self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens
|
|
)
|
|
|
|
# Batch register state/extra pool data buffers
|
|
if self.kv_args.state_data_ptrs and self.kv_args.state_data_lens:
|
|
self.engine.batch_register(
|
|
self.kv_args.state_data_ptrs, self.kv_args.state_data_lens
|
|
)
|
|
|
|
def _transfer_data(self, mooncake_session_id, transfer_blocks):
|
|
if not transfer_blocks:
|
|
return 0
|
|
|
|
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)
|
|
)
|
|
|
|
def _transfer_layers_async(
|
|
self, mooncake_session_id, layers_params, set_transfer_blocks
|
|
):
|
|
# Submit each layer's transfer non-blocking (pipelined in the RDMA engine),
|
|
# then wait for all to complete once. Removes the per-layer blocking-sync tax
|
|
# (B1a) and is the transfer mechanism for per-layer overlap (lever A). Uses the
|
|
# safe async API (G1: batch_transfer_async_submit + wait_batch_transfers),
|
|
# never the OnCuda busy-wait/_exit path.
|
|
batch_ids = []
|
|
for src_ptr, dst_ptr, item_len in layers_params:
|
|
transfer_blocks = set_transfer_blocks(src_ptr, dst_ptr, item_len)
|
|
if not transfer_blocks:
|
|
continue
|
|
src_addrs, dst_addrs, lengths = zip(*transfer_blocks)
|
|
batch_id = self.engine.batch_transfer_async_submit(
|
|
mooncake_session_id, list(src_addrs), list(dst_addrs), list(lengths)
|
|
)
|
|
if batch_id < 0:
|
|
# Submit failed: drain whatever was already submitted, then fail.
|
|
if batch_ids:
|
|
self.engine.wait_batch_transfers(batch_ids)
|
|
return -1
|
|
batch_ids.append(batch_id)
|
|
return self.engine.wait_batch_transfers(batch_ids)
|
|
|
|
def build_per_layer_context(
|
|
self,
|
|
mooncake_session_id: str,
|
|
prefill_kv_indices: npt.NDArray[np.int32],
|
|
dst_kv_indices: npt.NDArray[np.int32],
|
|
):
|
|
"""Build a PerLayerTransferContext (lever A) for one request's main-KV pages,
|
|
from the SAME CP-filtered (prefill_kv_indices, dst_kv_indices) the post-forward
|
|
transfer would use — so the bytes moved are identical to the monolithic path.
|
|
The caller supplies the indices (reusing the send() CP filter, so this does NOT
|
|
re-derive the CP owner mapping — the #1 correctness risk). Mirrors the MLA
|
|
branch of _send_kvcache_generic exactly. Returns None if not applicable
|
|
(MHA / no decode registration yet / empty owned set)."""
|
|
if not self.is_mla_backend:
|
|
return None # first lever-A impl targets MLA (the production GLM/NSA path)
|
|
reg = self.decode_kv_args_table.get(mooncake_session_id)
|
|
if reg is None or len(prefill_kv_indices) == 0:
|
|
return None
|
|
from sglang.srt.disaggregation.cp_per_layer_transfer import (
|
|
PerLayerTransferContext,
|
|
build_layer_blocks,
|
|
)
|
|
|
|
prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous(
|
|
prefill_kv_indices, dst_kv_indices
|
|
)
|
|
src_kv_ptrs, dst_kv_ptrs, n_layers = self.get_mla_kv_ptrs_with_pp(
|
|
self.kv_args.kv_data_ptrs, reg.dst_kv_ptrs
|
|
)
|
|
item_lens = self.kv_args.kv_item_lens
|
|
|
|
def get_blocks(layer_id):
|
|
if layer_id >= n_layers:
|
|
return None
|
|
return build_layer_blocks(
|
|
src_kv_ptrs[layer_id],
|
|
dst_kv_ptrs[layer_id],
|
|
item_lens[layer_id],
|
|
prefill_kv_blocks,
|
|
dst_kv_blocks,
|
|
)
|
|
|
|
return PerLayerTransferContext(
|
|
self.engine, mooncake_session_id, get_blocks, num_layers=n_layers
|
|
)
|
|
|
|
def register_per_layer_transfer(self, room, page_indices, chunk_key=0) -> bool:
|
|
"""Lever A: before the forward, build + register a per-layer transfer context
|
|
for `room` so the per-layer notifier overlaps its main-KV transfer with the
|
|
forward. Reuses send()'s CP filter exactly (no re-derivation). Scoped to
|
|
CP-shared-KV (the target scenario). Returns True iff a context was registered;
|
|
a False just falls back to the monolithic post-forward transfer (still correct).
|
|
page_indices = the request's full new-token logical page ids (req_to_token)."""
|
|
mgr = getattr(self, "per_layer_transfer_manager", None)
|
|
if mgr is None or not self.server_args.enable_nsa_prefill_cp_shared_kv:
|
|
return False
|
|
infos = self.transfer_infos.get(room)
|
|
if not infos:
|
|
return False
|
|
# The per-layer path registers exactly ONE context per room/chunk, but the
|
|
# transfer worker iterates every non-dummy decode info for the room and calls
|
|
# finish() per info (conn.py reqs_to_be_processed loop). That is only sound
|
|
# when there is exactly one non-dummy info (required_dst_info_num == 1). For
|
|
# decode attn_tp < prefill attn_tp a single prefill rank holds >1 non-dummy
|
|
# infos; finishing once-per-info would over-pop chunk contexts and the single
|
|
# ctx only carries one info's dst_kv_indices. Fall back to the monolithic
|
|
# post-forward transfer (which fans out to all infos) in that case.
|
|
non_dummy = [
|
|
info for info in infos.values() if not getattr(info, "is_dummy", False)
|
|
]
|
|
if len(non_dummy) != 1:
|
|
return False
|
|
info = non_dummy[0]
|
|
from sglang.srt.disaggregation.utils import filter_kv_pages_for_cp_shared_kv
|
|
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
|
|
|
layout = CpSharedKVLayout(
|
|
page_size=self.kv_args.page_size,
|
|
cp_size=self.attn_cp_size,
|
|
cp_rank=self.attn_cp_rank,
|
|
)
|
|
pages = np.asarray(page_indices, dtype=np.int32)
|
|
# chunk_key is this chunk's start_send_idx in TOKENS (page-aligned for any
|
|
# non-first chunk). The CP filter's `positions` (second return) are absolute
|
|
# full-sequence page positions built from chunk_page_start, and the transfer
|
|
# indexes the FULL-request dst_kv_indices by those absolute positions (mirrors
|
|
# send(): chunk_page_start=index_slice.start, worker: dst_kv_indices[logical_
|
|
# page_positions]). Must offset by the chunk's absolute page start, else chunk
|
|
# N>0 writes its KV onto chunk 0's decode pages. page_size divides chunk_key.
|
|
chunk_page_start = int(chunk_key) // self.kv_args.page_size
|
|
owned_pages, positions = filter_kv_pages_for_cp_shared_kv(
|
|
layout=layout, logical_pages=pages, chunk_page_start=chunk_page_start
|
|
)
|
|
dst_indices = np.asarray(info.dst_kv_indices, dtype=np.int32)[positions]
|
|
ctx = self.build_per_layer_context(
|
|
info.mooncake_session_id, owned_pages, dst_indices
|
|
)
|
|
if ctx is None:
|
|
return False
|
|
mgr.register(room, ctx, chunk_key=chunk_key)
|
|
logger.debug(
|
|
"[CP_PER_LAYER_TRANSFER] registered room=%s chunk=%s owned_pages=%d",
|
|
room,
|
|
chunk_key,
|
|
len(owned_pages),
|
|
)
|
|
return True
|
|
|
|
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,
|
|
debug_room: Optional[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.
|
|
"""
|
|
# Group by indices for optimization
|
|
prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous(
|
|
prefill_data_indices, dst_data_indices
|
|
)
|
|
transfer_stats_enabled = _mooncake_transfer_stats_enabled()
|
|
grouping_stats = None
|
|
if transfer_stats_enabled:
|
|
grouping_stats = contiguous_group_stats(
|
|
prefill_data_indices,
|
|
dst_data_indices,
|
|
prefill_kv_blocks,
|
|
dst_kv_blocks,
|
|
)
|
|
|
|
layers_params = None
|
|
|
|
# Decode pp size should be equal to prefill pp size or 1
|
|
if self.is_mla_backend:
|
|
src_kv_ptrs, dst_kv_ptrs, layers_current_pp_stage = (
|
|
self.get_mla_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs)
|
|
)
|
|
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 = []
|
|
for prefill_index, decode_index in zip(prefill_kv_blocks, dst_kv_blocks):
|
|
src_addr = src_ptr + int(prefill_index[0]) * item_len
|
|
dst_addr = dst_ptr + int(decode_index[0]) * item_len
|
|
length = item_len * len(prefill_index)
|
|
transfer_blocks.append((src_addr, dst_addr, length))
|
|
return transfer_blocks
|
|
|
|
# Worker function for processing a single layer
|
|
def process_layer(src_ptr: int, dst_ptr: int, item_len: int) -> int:
|
|
transfer_blocks = set_transfer_blocks(src_ptr, dst_ptr, item_len)
|
|
return self._transfer_data(mooncake_session_id, transfer_blocks)
|
|
|
|
# Worker function for processing all layers in a batch
|
|
def process_layers(layers_params: List[Tuple[int, int, int]]) -> int:
|
|
transfer_blocks = []
|
|
for src_ptr, dst_ptr, item_len in layers_params:
|
|
transfer_blocks.extend(set_transfer_blocks(src_ptr, dst_ptr, item_len))
|
|
return self._transfer_data(mooncake_session_id, transfer_blocks)
|
|
|
|
start_time = time.perf_counter() if transfer_stats_enabled else 0.0
|
|
if self.enable_per_layer_async_transfer:
|
|
status = self._transfer_layers_async(
|
|
mooncake_session_id, layers_params, set_transfer_blocks
|
|
)
|
|
elif self.enable_custom_mem_pool:
|
|
futures = [
|
|
executor.submit(
|
|
process_layer,
|
|
src_ptr,
|
|
dst_ptr,
|
|
item_len,
|
|
)
|
|
for (src_ptr, dst_ptr, item_len) in layers_params
|
|
]
|
|
for future in concurrent.futures.as_completed(futures):
|
|
status = future.result()
|
|
if status != 0:
|
|
for f in futures:
|
|
f.cancel()
|
|
if transfer_stats_enabled:
|
|
self._log_kvcache_transfer_stats(
|
|
mooncake_session_id=mooncake_session_id,
|
|
debug_room=debug_room,
|
|
grouping_stats=grouping_stats,
|
|
layers_params=layers_params,
|
|
elapsed_ms=(time.perf_counter() - start_time) * 1000,
|
|
status=status,
|
|
custom_mem_pool=True,
|
|
)
|
|
return status
|
|
status = 0
|
|
else:
|
|
# Combining all layers' params in one batch transfer is more efficient
|
|
# compared to using multiple threads
|
|
status = process_layers(layers_params)
|
|
|
|
if transfer_stats_enabled:
|
|
self._log_kvcache_transfer_stats(
|
|
mooncake_session_id=mooncake_session_id,
|
|
debug_room=debug_room,
|
|
grouping_stats=grouping_stats,
|
|
layers_params=layers_params,
|
|
elapsed_ms=(time.perf_counter() - start_time) * 1000,
|
|
status=status,
|
|
custom_mem_pool=self.enable_custom_mem_pool,
|
|
)
|
|
return status
|
|
|
|
def _log_kvcache_transfer_stats(
|
|
self,
|
|
mooncake_session_id: str,
|
|
debug_room: Optional[int],
|
|
grouping_stats: Optional[dict[str, object]],
|
|
layers_params: List[Tuple[int, int, int]],
|
|
elapsed_ms: float,
|
|
status: int,
|
|
custom_mem_pool: bool,
|
|
) -> None:
|
|
if grouping_stats is None:
|
|
return
|
|
page_count = int(grouping_stats["pages"])
|
|
group_count = int(grouping_stats["groups"])
|
|
layer_count = len(layers_params)
|
|
item_bytes_per_page = sum(int(item_len) for _, _, item_len in layers_params)
|
|
total_bytes = page_count * item_bytes_per_page
|
|
transfer_blocks = group_count * layer_count
|
|
avg_block_bytes = (
|
|
float(total_bytes / transfer_blocks) if transfer_blocks else 0.0
|
|
)
|
|
bandwidth_gbps = (
|
|
float(total_bytes / elapsed_ms / 1e6) if elapsed_ms > 0 else 0.0
|
|
)
|
|
|
|
_mooncake_transfer_stats_log(
|
|
"cp_rank=%s room=%s session=%s status=%s custom_mem_pool=%s "
|
|
"pages=%s groups=%s avg_group_pages=%.2f max_group_pages=%s "
|
|
"layers=%s transfer_blocks=%s total_bytes=%.3fGiB "
|
|
"avg_block_bytes=%.1fKiB elapsed_ms=%.3f bandwidth=%.2fGB/s "
|
|
"src_diff_head=%s dst_diff_head=%s",
|
|
self.attn_cp_rank,
|
|
debug_room,
|
|
mooncake_session_id,
|
|
status,
|
|
custom_mem_pool,
|
|
page_count,
|
|
group_count,
|
|
float(grouping_stats["avg_group_pages"]),
|
|
grouping_stats["max_group_pages"],
|
|
layer_count,
|
|
transfer_blocks,
|
|
total_bytes / (1024**3),
|
|
avg_block_bytes / 1024,
|
|
elapsed_ms,
|
|
bandwidth_gbps,
|
|
grouping_stats["src_diff_head"],
|
|
grouping_stats["dst_diff_head"],
|
|
)
|
|
|
|
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,
|
|
debug_room: Optional[int] = None,
|
|
):
|
|
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,
|
|
debug_room=debug_room,
|
|
)
|
|
|
|
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,
|
|
):
|
|
"""
|
|
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
|
|
src_head_start_offset = (
|
|
dst_tp_rank_in_group * dst_heads_per_rank
|
|
) % src_heads_per_rank
|
|
num_heads_to_send = dst_heads_per_rank
|
|
dst_head_start_offset = 0
|
|
|
|
src_k_ptrs, src_v_ptrs, dst_k_ptrs, dst_v_ptrs, layers_current_pp_stage = (
|
|
self.get_mha_kv_ptrs_with_pp(self.kv_args.kv_data_ptrs, dst_kv_ptrs)
|
|
)
|
|
|
|
# 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 = []
|
|
for i in range(layers_current_pp_stage):
|
|
futures.append(
|
|
executor.submit(process_layer_tp_aware, src_k_ptrs[i], dst_k_ptrs[i])
|
|
)
|
|
for i in range(layers_current_pp_stage):
|
|
futures.append(
|
|
executor.submit(process_layer_tp_aware, src_v_ptrs[i], dst_v_ptrs[i])
|
|
)
|
|
|
|
for future in concurrent.futures.as_completed(futures):
|
|
status = future.result()
|
|
if status != 0:
|
|
for f in futures:
|
|
f.cancel()
|
|
return status
|
|
|
|
return 0
|
|
|
|
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():
|
|
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)
|
|
socket = self._connect(na.to_tcp(), is_ipv6=na.is_ipv6)
|
|
|
|
socket.send_multipart(
|
|
[
|
|
MooncakeKVManager.AUX_DATA_HEADER,
|
|
str(room).encode("ascii"),
|
|
str(buffer_index).encode("ascii"),
|
|
str(aux_index).encode("ascii"),
|
|
struct.pack(">I", len(data)),
|
|
data,
|
|
]
|
|
)
|
|
|
|
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 maybe_send_extra(
|
|
self,
|
|
req: TransferInfo,
|
|
prefill_state_indices: list[int],
|
|
dst_state_data_ptrs: list[int],
|
|
executor: concurrent.futures.ThreadPoolExecutor,
|
|
target_rank_registration_info: Optional[KVArgsRegisterInfo] = None,
|
|
dst_state_indices: Optional[npt.NDArray[np.int32]] = None,
|
|
):
|
|
"""Send state or extra pool data with type-specific handling."""
|
|
state_type = getattr(self.kv_args, "state_type", "none")
|
|
_cp_draft_shared_kv_debug(
|
|
"maybe_send_extra_state cp_rank=%s room=%s session=%s state_type=%s "
|
|
"src_state_bufs=%s dst_state_bufs=%s src_state_lens=%s src_state_item_lens=%s "
|
|
"prefill_state_indices=%s dst_state_indices=%s draft_state_type=%s "
|
|
"draft_state_bufs=%s target_registration_state_bufs=%s",
|
|
self.attn_cp_rank,
|
|
req.room,
|
|
req.mooncake_session_id,
|
|
state_type,
|
|
len(getattr(self.kv_args, "state_data_ptrs", []) or []),
|
|
len(dst_state_data_ptrs or []),
|
|
_np_summary(getattr(self.kv_args, "state_data_lens", [])),
|
|
_np_summary(getattr(self.kv_args, "state_item_lens", [])),
|
|
_np_summary(prefill_state_indices),
|
|
_np_summary(
|
|
dst_state_indices
|
|
if dst_state_indices is not None
|
|
else getattr(req, "dst_state_indices", [])
|
|
),
|
|
getattr(self.kv_args, "draft_state_type", None),
|
|
getattr(self.kv_args, "draft_state_buffer_count", None),
|
|
(
|
|
len(target_rank_registration_info.dst_state_data_ptrs)
|
|
if target_rank_registration_info is not None
|
|
else None
|
|
),
|
|
)
|
|
|
|
if state_type == "mamba":
|
|
# Check if we need slice transfer for different TP sizes
|
|
if (
|
|
target_rank_registration_info is not None
|
|
and self.attn_tp_size != target_rank_registration_info.dst_attn_tp_size
|
|
):
|
|
return self._send_mamba_state_slice(
|
|
req,
|
|
prefill_state_indices,
|
|
dst_state_data_ptrs,
|
|
target_rank_registration_info.dst_state_item_lens,
|
|
target_rank_registration_info.dst_state_dim_per_tensor,
|
|
target_rank_registration_info.dst_tp_rank,
|
|
target_rank_registration_info.dst_attn_tp_size,
|
|
)
|
|
else:
|
|
return self._send_mamba_state(
|
|
req,
|
|
prefill_state_indices,
|
|
dst_state_data_ptrs,
|
|
)
|
|
elif state_type in ["swa", "nsa"]:
|
|
# SWA and NSA hybrid models do not support different TP sizes yet
|
|
if (
|
|
target_rank_registration_info is not None
|
|
and not self.is_mla_backend
|
|
and self.attn_tp_size != target_rank_registration_info.dst_attn_tp_size
|
|
):
|
|
raise RuntimeError(
|
|
f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {state_type.upper()} hybrid models yet."
|
|
)
|
|
src_state_layer_ids = list(
|
|
getattr(self.kv_args, "state_layer_ids", []) or []
|
|
)
|
|
dst_state_layer_ids = (
|
|
list(getattr(target_rank_registration_info, "dst_state_layer_ids", []) or [])
|
|
if target_rank_registration_info is not None
|
|
else []
|
|
)
|
|
if src_state_layer_ids or dst_state_layer_ids:
|
|
if src_state_layer_ids != dst_state_layer_ids:
|
|
raise RuntimeError(
|
|
"[CP_SHARED_KV_FAIL_FAST][state_layer_ids] "
|
|
f"prefill={src_state_layer_ids} decode={dst_state_layer_ids} "
|
|
f"state_type={state_type} room={req.room} "
|
|
f"session={req.mooncake_session_id}"
|
|
)
|
|
effective_dst_state_indices = (
|
|
np.asarray(dst_state_indices, dtype=np.int32)
|
|
if dst_state_indices is not None
|
|
else np.asarray(req.dst_state_indices, dtype=np.int32)
|
|
)
|
|
# Clip src/dst state indices to the shorter side so the transfer never
|
|
# reads past prefill_state_indices or writes past dst_state_indices.
|
|
# Ported from upstream sgl-project/sglang #23323 (the prior logic only
|
|
# handled prefill<dst and clipped the wrong array, a no-op).
|
|
if len(prefill_state_indices) > len(effective_dst_state_indices):
|
|
logger.warning(
|
|
f"len(prefill_state_indices) = {len(prefill_state_indices)}, len(dst_state_indices) = {len(effective_dst_state_indices)}"
|
|
)
|
|
prefill_state_indices = prefill_state_indices[
|
|
: len(effective_dst_state_indices)
|
|
]
|
|
elif len(prefill_state_indices) < len(effective_dst_state_indices):
|
|
logger.warning(
|
|
f"len(prefill_state_indices) = {len(prefill_state_indices)}, len(dst_state_indices) = {len(effective_dst_state_indices)}"
|
|
)
|
|
effective_dst_state_indices = effective_dst_state_indices[
|
|
: len(prefill_state_indices)
|
|
]
|
|
src_state_data_ptrs = self.kv_args.state_data_ptrs
|
|
dst_state_ptrs = dst_state_data_ptrs
|
|
state_item_lens = self.kv_args.state_item_lens
|
|
if len(src_state_data_ptrs) != len(dst_state_ptrs):
|
|
if src_state_layer_ids or dst_state_layer_ids:
|
|
raise RuntimeError(
|
|
"[CP_SHARED_KV_FAIL_FAST][state_buffer_count] "
|
|
f"src={len(src_state_data_ptrs)} dst={len(dst_state_ptrs)} "
|
|
f"src_layers={src_state_layer_ids} "
|
|
f"dst_layers={dst_state_layer_ids} state_type={state_type} "
|
|
f"room={req.room} session={req.mooncake_session_id}"
|
|
)
|
|
transfer_buf_count = min(len(src_state_data_ptrs), len(dst_state_ptrs))
|
|
logger.warning(
|
|
"State buffer count mismatch during PD transfer: src=%s dst=%s "
|
|
"state_type=%s draft_state_type=%s draft_state_bufs=%s room=%s "
|
|
"session=%s; transferring first %s buffers only",
|
|
len(src_state_data_ptrs),
|
|
len(dst_state_ptrs),
|
|
state_type,
|
|
getattr(self.kv_args, "draft_state_type", None),
|
|
getattr(self.kv_args, "draft_state_buffer_count", None),
|
|
req.room,
|
|
req.mooncake_session_id,
|
|
transfer_buf_count,
|
|
)
|
|
src_state_data_ptrs = src_state_data_ptrs[:transfer_buf_count]
|
|
dst_state_ptrs = dst_state_ptrs[:transfer_buf_count]
|
|
state_item_lens = state_item_lens[:transfer_buf_count]
|
|
# Reuse _send_kvcache_generic interface to send extra pool data
|
|
prefill_state_indices = np.array(prefill_state_indices, dtype=np.int32)
|
|
return self._send_kvcache_generic(
|
|
mooncake_session_id=req.mooncake_session_id,
|
|
src_data_ptrs=src_state_data_ptrs,
|
|
dst_data_ptrs=dst_state_ptrs,
|
|
item_lens=state_item_lens,
|
|
prefill_data_indices=prefill_state_indices,
|
|
dst_data_indices=effective_dst_state_indices,
|
|
executor=executor,
|
|
)
|
|
else:
|
|
return 0
|
|
|
|
def _send_mamba_state(
|
|
self,
|
|
req: TransferInfo,
|
|
prefill_mamba_index: list[int],
|
|
dst_state_data_ptrs: list[int],
|
|
):
|
|
"""Transfer Mamba states."""
|
|
assert len(prefill_mamba_index) == 1, "Mamba should have single state index"
|
|
|
|
transfer_blocks = []
|
|
prefill_state_data_ptrs = self.kv_args.state_data_ptrs
|
|
prefill_state_item_lens = self.kv_args.state_item_lens
|
|
|
|
for i, dst_state_ptr in enumerate(dst_state_data_ptrs):
|
|
length = prefill_state_item_lens[i]
|
|
src_addr = prefill_state_data_ptrs[i] + length * int(prefill_mamba_index[0])
|
|
dst_addr = dst_state_ptr + length * int(req.dst_state_indices[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[int],
|
|
dst_state_data_ptrs: list[int],
|
|
dst_state_item_lens: list[int],
|
|
dst_state_dim_per_tensor: list[int],
|
|
dst_tp_rank: int,
|
|
dst_attn_tp_size: int,
|
|
):
|
|
"""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 need to slice the state accordingly.
|
|
"""
|
|
logger.warning_once(
|
|
"Using Mamba state slice transfer for different TP sizes between prefill and decode. "
|
|
f"Prefill attn_tp_size={self.attn_tp_size}, Decode attn_tp_size={dst_attn_tp_size}. "
|
|
"Performance may be affected."
|
|
)
|
|
assert len(prefill_mamba_index) == 1, "Mamba should have single state index"
|
|
|
|
transfer_blocks = []
|
|
prefill_state_data_ptrs = self.kv_args.state_data_ptrs
|
|
prefill_state_item_lens = self.kv_args.state_item_lens
|
|
src_state_dim_per_tensor = getattr(self.kv_args, "state_dim_per_tensor", [])
|
|
|
|
# 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, dst_state_data_ptrs)
|
|
|
|
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
|
|
|
|
for i, dst_state_ptr in enumerate(dst_state_data_ptrs):
|
|
src_item_len = prefill_state_item_lens[i]
|
|
dst_item_len = dst_state_item_lens[i]
|
|
src_dim = src_state_dim_per_tensor[i]
|
|
dst_dim = dst_state_dim_per_tensor[i]
|
|
|
|
# Calculate bytes per dimension slice
|
|
# item_len = dim * trailing_dims_size, so trailing_dims_size = item_len / dim
|
|
src_bytes_per_dim = src_item_len // src_dim
|
|
dst_bytes_per_dim = dst_item_len // dst_dim
|
|
|
|
# Determine slicing parameters based on TP configuration
|
|
if self.attn_tp_size > dst_attn_tp_size:
|
|
# Multiple prefill ranks send to 1 decode rank
|
|
# Each prefill sends all its dims to the appropriate offset in decode
|
|
src_dim_start = 0
|
|
num_dims_to_send = src_dim
|
|
writers_per_decode = self.attn_tp_size // dst_attn_tp_size
|
|
local_writer_idx = local_tp_rank_in_group % writers_per_decode
|
|
dst_dim_start = local_writer_idx * src_dim
|
|
else:
|
|
# 1 prefill rank sends to multiple decode ranks
|
|
# Prefill sends a slice of its dims to each decode rank
|
|
src_dim_start = (dst_tp_rank_in_group * dst_dim) % src_dim
|
|
num_dims_to_send = dst_dim
|
|
dst_dim_start = 0
|
|
|
|
# Calculate byte offsets
|
|
src_dim_offset = src_dim_start * src_bytes_per_dim
|
|
dst_dim_offset = dst_dim_start * dst_bytes_per_dim
|
|
bytes_to_send = num_dims_to_send * src_bytes_per_dim
|
|
|
|
# Calculate addresses for this state tensor
|
|
src_addr = (
|
|
prefill_state_data_ptrs[i]
|
|
+ src_item_len * int(prefill_mamba_index[0])
|
|
+ src_dim_offset
|
|
)
|
|
dst_addr = (
|
|
dst_state_ptr
|
|
+ dst_item_len * int(req.dst_state_indices[0])
|
|
+ dst_dim_offset
|
|
)
|
|
|
|
transfer_blocks.append((src_addr, dst_addr, bytes_to_send))
|
|
|
|
return self._transfer_data(req.mooncake_session_id, transfer_blocks)
|
|
|
|
def sync_status_to_decode_endpoint(
|
|
self, remote: str, dst_port: int, room: int, status: int, prefill_rank: int
|
|
):
|
|
na = NetworkAddress(remote, dst_port)
|
|
self._connect(na.to_tcp(), is_ipv6=na.is_ipv6).send_multipart(
|
|
[
|
|
str(room).encode("ascii"),
|
|
str(status).encode("ascii"),
|
|
str(prefill_rank).encode("ascii"),
|
|
]
|
|
)
|
|
|
|
def transfer_worker(
|
|
self, queue: FastQueue, executor: concurrent.futures.ThreadPoolExecutor
|
|
):
|
|
while True:
|
|
try:
|
|
kv_chunk: TransferKVChunk = queue.get()
|
|
# Skip a chunk whose room has already failed or been aborted, so we never
|
|
# transfer into KV pages that may have been reclaimed. Port of upstream
|
|
# sgl-project/sglang #27372 (abort KV-cache corruption guard). The
|
|
# decode->prefill ABORT/ABORT_ACK notification half of that PR is deferred
|
|
# until the PD test harness exists (it interleaves with staging/tracing
|
|
# infra this branch does not carry).
|
|
if (
|
|
kv_chunk.room not in self.request_status
|
|
or self.check_status(kv_chunk.room) == KVPoll.Failed
|
|
):
|
|
logger.debug(
|
|
"Skipping chunk for room %s because it has already failed or been aborted",
|
|
kv_chunk.room,
|
|
)
|
|
# Lever A: drain + drop any per-layer context for this room so an
|
|
# outstanding RDMA completes before its KV pages can be reclaimed.
|
|
per_layer_mgr = getattr(self, "per_layer_transfer_manager", None)
|
|
if per_layer_mgr is not None:
|
|
per_layer_mgr.drop(kv_chunk.room)
|
|
continue
|
|
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.attn_tp_rank * (self.pp_size * self.attn_cp_size)
|
|
+ self.pp_rank * self.attn_cp_size
|
|
+ self.attn_cp_rank
|
|
)
|
|
for req in reqs_to_be_processed:
|
|
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.record_failure(
|
|
kv_chunk.room,
|
|
f"Decode instance could be dead, remote mooncake session {req.mooncake_session_id} is not alive",
|
|
)
|
|
self.update_status(kv_chunk.room, KVPoll.Failed)
|
|
self.sync_status_to_decode_endpoint(
|
|
req.endpoint,
|
|
req.dst_port,
|
|
req.room,
|
|
KVPoll.Failed,
|
|
prefill_unique_rank,
|
|
)
|
|
break
|
|
|
|
if kv_chunk.logical_page_positions is not None:
|
|
chunked_dst_kv_indice = req.dst_kv_indices[
|
|
kv_chunk.logical_page_positions
|
|
]
|
|
else:
|
|
assert kv_chunk.index_slice is not None
|
|
chunked_dst_kv_indice = req.dst_kv_indices[
|
|
kv_chunk.index_slice
|
|
]
|
|
_cp_draft_shared_kv_debug(
|
|
"transfer_pages cp_rank=%s room=%s prefill_pages=%s "
|
|
"logical_positions=%s dst_pages=%s is_last=%s",
|
|
self.attn_cp_rank,
|
|
kv_chunk.room,
|
|
_np_summary(kv_chunk.prefill_kv_indices),
|
|
_np_summary(kv_chunk.logical_page_positions),
|
|
_np_summary(chunked_dst_kv_indice),
|
|
kv_chunk.is_last_chunk,
|
|
)
|
|
if envs.SGLANG_DEBUG_CP_SHARED_KV.get():
|
|
_cp_shared_debug_log(
|
|
"transfer_worker_kv",
|
|
"transfer worker cp_rank=%s room=%s prefill_pages=%s "
|
|
"logical_positions=%s dst_pages=%s is_last=%s",
|
|
self.attn_cp_rank,
|
|
kv_chunk.room,
|
|
_np_summary(kv_chunk.prefill_kv_indices),
|
|
_np_summary(kv_chunk.logical_page_positions),
|
|
_np_summary(chunked_dst_kv_indice),
|
|
kv_chunk.is_last_chunk,
|
|
)
|
|
|
|
validate_transfer_page_count_or_raise(
|
|
prefill_indices=kv_chunk.prefill_kv_indices,
|
|
dst_indices=chunked_dst_kv_indice,
|
|
room=kv_chunk.room,
|
|
cp_rank=self.attn_cp_rank,
|
|
logical_page_positions=kv_chunk.logical_page_positions,
|
|
index_slice=kv_chunk.index_slice,
|
|
is_cp_shared_kv=(
|
|
kv_chunk.logical_page_positions is not None
|
|
),
|
|
path="mooncake_kv",
|
|
)
|
|
|
|
target_rank_registration_info: KVArgsRegisterInfo = (
|
|
self.decode_kv_args_table[req.mooncake_session_id]
|
|
)
|
|
per_layer_mgr = getattr(
|
|
self, "per_layer_transfer_manager", None
|
|
)
|
|
if per_layer_mgr is not None and per_layer_mgr.has_room(
|
|
kv_chunk.room
|
|
):
|
|
# Lever A: the main KV was transferred per-layer, overlapped
|
|
# with the forward; wait those transfers here instead of the
|
|
# monolithic send (no double-send). aux/state below unchanged.
|
|
ret = per_layer_mgr.finish(kv_chunk.room)
|
|
if ret != 0:
|
|
logger.warning(
|
|
"[CP_PER_LAYER_TRANSFER] finished room=%s ret=%s",
|
|
kv_chunk.room,
|
|
ret,
|
|
)
|
|
else:
|
|
logger.debug(
|
|
"[CP_PER_LAYER_TRANSFER] finished room=%s ret=%s",
|
|
kv_chunk.room,
|
|
ret,
|
|
)
|
|
elif self.is_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,
|
|
debug_room=kv_chunk.room,
|
|
)
|
|
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,
|
|
)
|
|
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.record_failure(
|
|
kv_chunk.room,
|
|
f"Failed to send kv chunk of {kv_chunk.room} to {req.endpoint}:{req.dst_port}",
|
|
)
|
|
self.update_status(kv_chunk.room, KVPoll.Failed)
|
|
self.sync_status_to_decode_endpoint(
|
|
req.endpoint,
|
|
req.dst_port,
|
|
req.room,
|
|
KVPoll.Failed,
|
|
prefill_unique_rank,
|
|
)
|
|
break
|
|
|
|
if kv_chunk.is_last_chunk:
|
|
if kv_chunk.state_indices is not None:
|
|
dst_state_indices = None
|
|
if kv_chunk.state_logical_page_positions is not None:
|
|
from sglang.srt.disaggregation.utils import (
|
|
select_pages_by_request_positions,
|
|
)
|
|
|
|
dst_state_indices = select_pages_by_request_positions(
|
|
req.dst_state_indices,
|
|
kv_chunk.state_logical_page_positions,
|
|
)
|
|
if envs.SGLANG_DEBUG_CP_SHARED_KV.get():
|
|
_cp_shared_debug_log(
|
|
"transfer_worker_state",
|
|
"transfer worker state cp_rank=%s room=%s prefill_state_pages=%s "
|
|
"state_positions=%s dst_state_pages=%s",
|
|
self.attn_cp_rank,
|
|
kv_chunk.room,
|
|
_np_summary(kv_chunk.state_indices),
|
|
_np_summary(
|
|
kv_chunk.state_logical_page_positions
|
|
),
|
|
_np_summary(dst_state_indices),
|
|
)
|
|
self.maybe_send_extra(
|
|
req,
|
|
kv_chunk.state_indices,
|
|
target_rank_registration_info.dst_state_data_ptrs,
|
|
executor,
|
|
target_rank_registration_info,
|
|
dst_state_indices,
|
|
)
|
|
|
|
# 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, req.room)
|
|
)
|
|
|
|
# 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.update_status(req.room, status)
|
|
for endpoint, dst_port, room in dst_ranks_infos:
|
|
self.sync_status_to_decode_endpoint(
|
|
endpoint,
|
|
dst_port,
|
|
room,
|
|
status,
|
|
prefill_unique_rank,
|
|
)
|
|
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 (
|
|
kv_chunk.room not in self.request_status
|
|
or self.check_status(kv_chunk.room) == KVPoll.Success
|
|
):
|
|
if kv_chunk.room in self.transfer_infos:
|
|
self.transfer_infos.pop(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):
|
|
def bootstrap_thread():
|
|
"""This thread recvs pre-alloc notification from the decode engine"""
|
|
# KVPoll.Bootstrapping -> KVPoll.WaitingForInput
|
|
while True:
|
|
waiting_req_bytes = self.server_socket.recv_multipart()
|
|
room = waiting_req_bytes[0].decode("ascii")
|
|
# Decode-side abort notification: mark the room Failed and ACK, so the
|
|
# transfer worker (which skips Failed rooms) stops RDMA-writing into KV
|
|
# pages the decode has freed/reused. Port of upstream sgl-project/sglang
|
|
# #27372 (decode->prefill half). MUST be handled here, before the
|
|
# UNCONDITIONAL waiting_req_bytes[3] decode below: a 4-field ABORT
|
|
# message [ABORT, room, ip, port] would otherwise fall through into the
|
|
# else branch and crash on waiting_req_bytes[7].
|
|
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"))
|
|
if (
|
|
room_to_be_aborted in self.request_status
|
|
and self.check_status(room_to_be_aborted) != KVPoll.Success
|
|
):
|
|
self.update_status(room_to_be_aborted, KVPoll.Failed)
|
|
try:
|
|
na = NetworkAddress(decode_ip, decode_port)
|
|
self._connect(
|
|
na.to_tcp(), is_ipv6=na.is_ipv6
|
|
).send_multipart(
|
|
[b"ABORT_ACK", str(room_to_be_aborted).encode("ascii")]
|
|
)
|
|
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":
|
|
self.decode_kv_args_table[mooncake_session_id] = (
|
|
KVArgsRegisterInfo.from_zmq(waiting_req_bytes)
|
|
)
|
|
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)
|
|
)
|
|
# NOTE: after bootstrapping we can mark the req as waiting for input
|
|
if len(self.transfer_infos[room]) == required_dst_info_num:
|
|
self.update_status(room, KVPoll.WaitingForInput)
|
|
|
|
threading.Thread(target=bootstrap_thread).start()
|
|
|
|
def start_decode_thread(self):
|
|
def decode_thread():
|
|
while True:
|
|
msg = self.server_socket.recv_multipart()
|
|
if msg[0] == MooncakeKVManager.AUX_DATA_HEADER:
|
|
self._handle_aux_data(msg)
|
|
continue
|
|
|
|
if msg[0] == b"ABORT_ACK":
|
|
# Prefill acknowledged our abort; the room is already Failed locally,
|
|
# so nothing more to do. Port of upstream #27372. Must precede the
|
|
# 3-tuple unpack below (a 2-field ACK would otherwise raise).
|
|
continue
|
|
|
|
bootstrap_room, status, prefill_rank = msg
|
|
status = int(status.decode("ascii"))
|
|
bootstrap_room = int(bootstrap_room.decode("ascii"))
|
|
prefill_rank = int(prefill_rank.decode("ascii"))
|
|
|
|
if status == KVPoll.Success:
|
|
if bootstrap_room in self.request_status:
|
|
self.prefill_response_tracker[bootstrap_room].add(prefill_rank)
|
|
expected_response_num = (
|
|
self.required_prefill_response_num_table[bootstrap_room]
|
|
)
|
|
arrived_response_num = len(
|
|
self.prefill_response_tracker[bootstrap_room]
|
|
)
|
|
if arrived_response_num == expected_response_num:
|
|
self.update_status(bootstrap_room, KVPoll.Success)
|
|
elif status == KVPoll.Failed:
|
|
self.record_failure(
|
|
bootstrap_room,
|
|
"Failed to get kvcache from prefill instance, it might be dead",
|
|
)
|
|
self.update_status(bootstrap_room, status)
|
|
|
|
def heartbeat_checker():
|
|
while True:
|
|
time.sleep(self.heartbeat_interval)
|
|
with self.connection_lock:
|
|
addresses = list(self.prefill_info_table.keys())
|
|
|
|
for bootstrap_addr in addresses:
|
|
session = None
|
|
try:
|
|
with self.session_pool_lock:
|
|
session = self.session_pool[bootstrap_addr]
|
|
response = session.get(
|
|
f"http://{bootstrap_addr}/health",
|
|
timeout=(2, 3),
|
|
headers={"Connection": "keep-alive"},
|
|
)
|
|
if response.status_code == 200:
|
|
self.heartbeat_failures[bootstrap_addr] = 0
|
|
|
|
current_rooms = self.addr_to_rooms_tracker[
|
|
bootstrap_addr
|
|
].copy()
|
|
|
|
for bootstrap_room in current_rooms:
|
|
# Remove KVPoll.Success requests from the tracker
|
|
if bootstrap_room not in self.request_status:
|
|
self.addr_to_rooms_tracker[bootstrap_addr].discard(
|
|
bootstrap_room
|
|
)
|
|
else:
|
|
logger.info(
|
|
f"Attempting to reconnect to {bootstrap_addr}..."
|
|
)
|
|
self.heartbeat_failures[bootstrap_addr] = (
|
|
self.heartbeat_failures.get(bootstrap_addr, 0) + 1
|
|
)
|
|
with self.session_pool_lock:
|
|
if bootstrap_addr in self.session_pool:
|
|
del self.session_pool[bootstrap_addr]
|
|
except Exception:
|
|
logger.info(f"Attempting to reconnect to {bootstrap_addr}...")
|
|
self.heartbeat_failures[bootstrap_addr] = (
|
|
self.heartbeat_failures.get(bootstrap_addr, 0) + 1
|
|
)
|
|
|
|
if (
|
|
self.heartbeat_failures.get(bootstrap_addr, 0)
|
|
>= self.max_failures
|
|
):
|
|
self._handle_node_failure(bootstrap_addr)
|
|
with self.session_pool_lock:
|
|
if bootstrap_addr in self.session_pool:
|
|
del self.session_pool[bootstrap_addr]
|
|
|
|
threading.Thread(target=decode_thread).start()
|
|
threading.Thread(target=heartbeat_checker).start()
|
|
|
|
def add_transfer_request(
|
|
self,
|
|
bootstrap_room: int,
|
|
kv_indices: npt.NDArray[np.int32],
|
|
index_slice: Optional[slice],
|
|
is_last_chunk: bool,
|
|
aux_index: Optional[int] = None,
|
|
state_indices: Optional[List[int]] = None,
|
|
logical_page_positions: Optional[npt.NDArray[np.int32]] = None,
|
|
state_logical_page_positions: Optional[npt.NDArray[np.int32]] = 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:
|
|
if self.check_status(bootstrap_room) == KVPoll.Success:
|
|
# Dummy rank for this request: it was already marked Success
|
|
# at handshake time, so there is nothing left to transfer.
|
|
return
|
|
# Non-Success room with no transfer destinations is an anomaly
|
|
# (decode peer torn down between handshake and this send, abort
|
|
# notification lost). Silently dropping the chunk used to wedge
|
|
# the request in WaitingForInput FOREVER — Success can never be
|
|
# set without the last chunk, and nothing else fails the room.
|
|
# Conclude it loudly instead; the scheduler's poll consensus
|
|
# (MIN-reduce, Failed=0) reaps it on every rank.
|
|
logger.warning(
|
|
"KV chunk for bootstrap_room=%s has no transfer destinations "
|
|
"(status=%s); failing the request instead of dropping the "
|
|
"chunk silently.",
|
|
bootstrap_room,
|
|
self.check_status(bootstrap_room),
|
|
)
|
|
self.record_failure(
|
|
bootstrap_room,
|
|
f"Request {bootstrap_room} lost its transfer destinations "
|
|
"before the KV chunk was submitted (decode peer gone?)",
|
|
)
|
|
self.update_status(bootstrap_room, KVPoll.Failed)
|
|
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)
|
|
|
|
self.transfer_queues[shard_idx].put(
|
|
TransferKVChunk(
|
|
room=bootstrap_room,
|
|
prefill_kv_indices=kv_indices,
|
|
index_slice=index_slice,
|
|
logical_page_positions=logical_page_positions,
|
|
state_logical_page_positions=state_logical_page_positions,
|
|
is_last_chunk=is_last_chunk,
|
|
prefill_aux_index=aux_index,
|
|
state_indices=state_indices,
|
|
)
|
|
)
|
|
|
|
def get_session_id(self):
|
|
return self.engine.get_session_id()
|
|
|
|
def _handle_node_failure(self, failed_bootstrap_addr):
|
|
with self.connection_lock:
|
|
keys_to_remove = [
|
|
k for k in self.connection_pool if k.startswith(failed_bootstrap_addr)
|
|
]
|
|
for k in keys_to_remove:
|
|
del self.connection_pool[k]
|
|
|
|
possible_affected_rooms = self.addr_to_rooms_tracker.get(
|
|
failed_bootstrap_addr, []
|
|
)
|
|
self.prefill_info_table.pop(failed_bootstrap_addr, None)
|
|
self.addr_to_rooms_tracker.pop(failed_bootstrap_addr, None)
|
|
|
|
# Report the requests associated with the failed bootstrap addr and mark their status as KVPoll.Failed
|
|
affected_rooms = []
|
|
for room in possible_affected_rooms:
|
|
if (
|
|
room in self.request_status
|
|
and self.check_status(room) != KVPoll.Success
|
|
):
|
|
self.record_failure(
|
|
room,
|
|
f"Losing connection with prefill instance (bootstrap_addr: {failed_bootstrap_addr})",
|
|
)
|
|
self.update_status(room, KVPoll.Failed)
|
|
affected_rooms.append(room)
|
|
logger.error(
|
|
f"Losing connection with prefill instance (bootstrap_addr: {failed_bootstrap_addr}), {len(affected_rooms)} requests affected"
|
|
)
|
|
|
|
|
|
class MooncakeKVSender(CommonKVSender):
|
|
|
|
def __init__(
|
|
self,
|
|
mgr: MooncakeKVManager,
|
|
bootstrap_addr: str,
|
|
bootstrap_room: int,
|
|
dest_tp_ranks: List[int],
|
|
pp_rank: int,
|
|
):
|
|
super().__init__(mgr, bootstrap_addr, bootstrap_room, dest_tp_ranks, pp_rank)
|
|
self.conclude_state = None
|
|
self.init_time = time.time()
|
|
|
|
def send(
|
|
self,
|
|
kv_indices: npt.NDArray[np.int32],
|
|
state_indices: Optional[List[int]] = None,
|
|
):
|
|
index_slice = slice(self.curr_idx, self.curr_idx + len(kv_indices))
|
|
self.curr_idx += len(kv_indices)
|
|
is_last_chunk = self.curr_idx == self.num_kv_indices
|
|
logical_page_positions = None
|
|
state_logical_page_positions = None
|
|
orig_kv_indices = None
|
|
orig_state_indices = None
|
|
|
|
if self.kv_mgr.server_args.enable_nsa_prefill_cp_shared_kv:
|
|
from sglang.srt.disaggregation.utils import filter_kv_pages_for_cp_shared_kv
|
|
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
|
|
|
chunk_page_start = index_slice.start
|
|
orig_kv_indices = np.asarray(kv_indices, dtype=np.int32).copy()
|
|
if state_indices is not None:
|
|
orig_state_indices = np.asarray(state_indices, dtype=np.int32).copy()
|
|
layout = CpSharedKVLayout(
|
|
page_size=self.kv_mgr.kv_args.page_size,
|
|
cp_size=self.kv_mgr.attn_cp_size,
|
|
cp_rank=self.kv_mgr.attn_cp_rank,
|
|
)
|
|
kv_indices, logical_page_positions = filter_kv_pages_for_cp_shared_kv(
|
|
layout=layout,
|
|
logical_pages=kv_indices,
|
|
chunk_page_start=index_slice.start,
|
|
)
|
|
if state_indices is not None:
|
|
state_indices, state_logical_page_positions = (
|
|
filter_kv_pages_for_cp_shared_kv(
|
|
layout=layout,
|
|
logical_pages=state_indices,
|
|
chunk_page_start=0,
|
|
)
|
|
)
|
|
index_slice = None
|
|
if envs.SGLANG_DEBUG_CP_SHARED_KV.get():
|
|
_cp_shared_debug_log(
|
|
"sender_filter",
|
|
"sender filter cp_rank=%s room=%s page_start=%s orig_kv_pages=%s "
|
|
"filtered_kv_pages=%s kv_positions=%s orig_state_pages=%s "
|
|
"filtered_state_pages=%s state_positions=%s is_last=%s",
|
|
self.kv_mgr.attn_cp_rank,
|
|
self.bootstrap_room,
|
|
chunk_page_start,
|
|
_np_summary(orig_kv_indices),
|
|
_np_summary(kv_indices),
|
|
_np_summary(logical_page_positions),
|
|
_np_summary(orig_state_indices),
|
|
_np_summary(state_indices),
|
|
_np_summary(state_logical_page_positions),
|
|
is_last_chunk,
|
|
)
|
|
_cp_draft_shared_kv_debug(
|
|
"sender_filter cp_rank=%s room=%s page_start=%s orig_kv_pages=%s "
|
|
"filtered_kv_pages=%s kv_positions=%s orig_state_pages=%s "
|
|
"filtered_state_pages=%s state_positions=%s is_last=%s draft_bufs=%s",
|
|
self.kv_mgr.attn_cp_rank,
|
|
self.bootstrap_room,
|
|
chunk_page_start,
|
|
_np_summary(orig_kv_indices),
|
|
_np_summary(kv_indices),
|
|
_np_summary(logical_page_positions),
|
|
_np_summary(orig_state_indices),
|
|
_np_summary(state_indices),
|
|
_np_summary(state_logical_page_positions),
|
|
is_last_chunk,
|
|
getattr(self.kv_mgr.kv_args, "draft_kv_buffer_count", None),
|
|
)
|
|
# Special handling for cp
|
|
elif self.kv_mgr.enable_all_cp_ranks_for_transfer:
|
|
kv_indices, index_slice = filter_kv_indices_for_cp_rank(
|
|
self.kv_mgr,
|
|
kv_indices,
|
|
index_slice,
|
|
)
|
|
elif self.kv_mgr.is_dummy_cp_rank:
|
|
if not is_last_chunk:
|
|
return
|
|
else:
|
|
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Success)
|
|
return
|
|
|
|
if not is_last_chunk:
|
|
self.kv_mgr.add_transfer_request(
|
|
self.bootstrap_room,
|
|
kv_indices,
|
|
index_slice,
|
|
False,
|
|
logical_page_positions=logical_page_positions,
|
|
state_logical_page_positions=state_logical_page_positions,
|
|
)
|
|
else:
|
|
self.kv_mgr.add_transfer_request(
|
|
self.bootstrap_room,
|
|
kv_indices,
|
|
index_slice,
|
|
True,
|
|
aux_index=self.aux_index,
|
|
state_indices=state_indices,
|
|
logical_page_positions=logical_page_positions,
|
|
state_logical_page_positions=state_logical_page_positions,
|
|
)
|
|
# Record the actually-sent (post CP shared-KV filter) indices for accurate
|
|
# transfer metrics. Ported from upstream sgl-project/sglang #24416.
|
|
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)
|
|
if status in (KVPoll.Success, KVPoll.Failed):
|
|
self.conclude_state = status
|
|
elif status == KVPoll.Bootstrapping:
|
|
if self.init_time is not None:
|
|
now = time.time()
|
|
elapsed = now - self.init_time
|
|
if elapsed >= self.kv_mgr.bootstrap_timeout:
|
|
logger.warning_once(
|
|
"Some requests timed out when bootstrapping, "
|
|
"which means prefill instances fail to receive the KV indices from the decode instance of this request. "
|
|
"If a greater mean TTFT is acceptable, you can 'export SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT=600' (10 minutes) to relax the timeout condition. "
|
|
)
|
|
self.kv_mgr.record_failure(
|
|
self.bootstrap_room,
|
|
f"Request {self.bootstrap_room} timed out after {elapsed:.1f}s in KVPoll.Bootstrapping",
|
|
)
|
|
self.conclude_state = KVPoll.Failed
|
|
return KVPoll.Failed
|
|
|
|
return status
|
|
else:
|
|
return self.conclude_state
|
|
|
|
def clear(self) -> None:
|
|
if self.bootstrap_room in self.kv_mgr.request_status:
|
|
self.kv_mgr.request_status.pop(self.bootstrap_room)
|
|
|
|
def failure_exception(self):
|
|
# Explicitly set the status to failure since this request has failed in 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, "Failed due to an unknown reason from another rank"
|
|
)
|
|
raise KVTransferError(self.bootstrap_room, failure_reason)
|
|
|
|
def abort(self):
|
|
self.kv_mgr.record_failure(
|
|
self.bootstrap_room,
|
|
"Aborted by AbortReq.",
|
|
)
|
|
# Mark the manager status Failed (not only the local conclude_state) so the
|
|
# transfer worker and other pollers observe the abort. Port of upstream
|
|
# sgl-project/sglang #24522.
|
|
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
|
|
# Explicitly set the status to failure since this request has been aborted
|
|
self.conclude_state = KVPoll.Failed
|
|
|
|
|
|
class MooncakeKVReceiver(CommonKVReceiver):
|
|
def __init__(
|
|
self,
|
|
mgr: MooncakeKVManager,
|
|
bootstrap_addr: str,
|
|
bootstrap_room: Optional[int] = None,
|
|
prefill_dp_rank: Optional[int] = None,
|
|
):
|
|
self.session_id = mgr.get_session_id()
|
|
self.conclude_state = None
|
|
self.init_time = None
|
|
self.abort_notified = False
|
|
super().__init__(mgr, bootstrap_addr, bootstrap_room, prefill_dp_rank)
|
|
|
|
self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room)
|
|
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.WaitingForInput)
|
|
|
|
def _register_kv_args(self):
|
|
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 = b"".join(
|
|
struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.state_data_ptrs
|
|
)
|
|
# Pack state_item_lens and state_dim_per_tensor for mamba state slice transfer
|
|
packed_state_item_lens = b"".join(
|
|
struct.pack("I", item_len)
|
|
for item_len in self.kv_mgr.kv_args.state_item_lens
|
|
)
|
|
state_dim_per_tensor = getattr(
|
|
self.kv_mgr.kv_args, "state_dim_per_tensor", []
|
|
)
|
|
packed_state_dim_per_tensor = b"".join(
|
|
struct.pack("I", dim) for dim in state_dim_per_tensor
|
|
)
|
|
packed_state_layer_ids = b"".join(
|
|
struct.pack("i", int(layer_id))
|
|
for layer_id in getattr(self.kv_mgr.kv_args, "state_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
|
|
kv_item_len = self.kv_mgr.kv_args.kv_item_lens[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")
|
|
_cp_draft_shared_kv_debug(
|
|
"decode_register_kv_args cp_rank=%s room=%s session=%s "
|
|
"kv_bufs=%s aux_bufs=%s state_type=%s state_bufs=%s "
|
|
"state_item_lens=%s draft_start=%s draft_count=%s "
|
|
"draft_state_type=%s draft_state_bufs=%s",
|
|
self.kv_mgr.attn_cp_rank,
|
|
self.bootstrap_room,
|
|
self.session_id,
|
|
len(self.kv_mgr.kv_args.kv_data_ptrs),
|
|
len(self.kv_mgr.kv_args.aux_data_ptrs),
|
|
getattr(self.kv_mgr.kv_args, "state_type", None),
|
|
len(self.kv_mgr.kv_args.state_data_ptrs),
|
|
_np_summary(self.kv_mgr.kv_args.state_item_lens),
|
|
getattr(self.kv_mgr.kv_args, "draft_kv_buffer_start", None),
|
|
getattr(self.kv_mgr.kv_args, "draft_kv_buffer_count", None),
|
|
getattr(self.kv_mgr.kv_args, "draft_state_type", None),
|
|
getattr(self.kv_mgr.kv_args, "draft_state_buffer_count", None),
|
|
)
|
|
|
|
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_state_layer_ids,
|
|
]
|
|
)
|
|
|
|
def init(
|
|
self,
|
|
kv_indices: npt.NDArray[np.int32],
|
|
aux_index: Optional[int] = None,
|
|
state_indices: Optional[List[int]] = 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
|
|
|
|
for bootstrap_info in self.bootstrap_infos:
|
|
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
|
|
is_dummy = bootstrap_info["is_dummy"]
|
|
|
|
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"",
|
|
(
|
|
np.array(
|
|
state_indices,
|
|
dtype=np.int32,
|
|
).tobytes()
|
|
if not is_dummy and state_indices is not None
|
|
else b""
|
|
),
|
|
str(self.required_dst_info_num).encode("ascii"),
|
|
]
|
|
)
|
|
self.init_time = time.time()
|
|
|
|
def poll(self) -> KVPoll:
|
|
if self.conclude_state is None:
|
|
status = self.kv_mgr.check_status(self.bootstrap_room)
|
|
if status in (KVPoll.Success, KVPoll.Failed):
|
|
self.conclude_state = status
|
|
elif status == KVPoll.WaitingForInput:
|
|
if self.init_time is not None:
|
|
now = time.time()
|
|
elapsed = now - self.init_time
|
|
if elapsed >= self.kv_mgr.waiting_timeout:
|
|
logger.warning_once(
|
|
"Some requests fail to receive KV Cache transfer done signal after bootstrapping. "
|
|
"If a greater mean TTFT is acceptable, you can 'export SGLANG_DISAGGREGATION_WAITING_TIMEOUT=600' (10 minutes) to relax the timeout condition. "
|
|
)
|
|
self.kv_mgr.record_failure(
|
|
self.bootstrap_room,
|
|
f"Request {self.bootstrap_room} timed out after {elapsed:.1f}s in KVPoll.WaitingForInput",
|
|
)
|
|
# Tell prefill peers to stop transferring into our pages.
|
|
self._send_abort_notification()
|
|
self.conclude_state = KVPoll.Failed
|
|
return KVPoll.Failed
|
|
|
|
return status
|
|
|
|
else:
|
|
return self.conclude_state
|
|
|
|
def clear(self) -> None:
|
|
if self.bootstrap_room in self.kv_mgr.request_status:
|
|
self.kv_mgr.request_status.pop(self.bootstrap_room)
|
|
|
|
if self.bootstrap_room in self.kv_mgr.required_prefill_response_num_table:
|
|
self.kv_mgr.required_prefill_response_num_table.pop(self.bootstrap_room)
|
|
|
|
if self.bootstrap_room in self.kv_mgr.prefill_response_tracker:
|
|
self.kv_mgr.prefill_response_tracker.pop(self.bootstrap_room)
|
|
|
|
def failure_exception(self):
|
|
# Explicitly set the status to failure since this request has failed in 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, "Failed due to an unknown reason from another rank"
|
|
)
|
|
raise KVTransferError(self.bootstrap_room, failure_reason)
|
|
|
|
def _send_abort_notification(self):
|
|
# Notify every prefill peer that this room is aborting, so the prefill stops
|
|
# RDMA-writing into KV pages this decode is about to free/reuse. Without this,
|
|
# the prefill never learns of a decode-initiated abort (its request_status is
|
|
# per-process) and the worker's skip-Failed guard never fires. Port of upstream
|
|
# sgl-project/sglang #27372 (decode->prefill half). Sent at most once.
|
|
if self.abort_notified:
|
|
return
|
|
bootstrap_infos = getattr(self, "bootstrap_infos", None)
|
|
if not bootstrap_infos:
|
|
return
|
|
self.abort_notified = True
|
|
for bootstrap_info in bootstrap_infos:
|
|
try:
|
|
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
|
|
with lock:
|
|
sock.send_multipart(
|
|
[
|
|
b"ABORT",
|
|
str(self.bootstrap_room).encode("ascii"),
|
|
self.kv_mgr.local_ip.encode("ascii"),
|
|
str(self.kv_mgr.rank_port).encode("ascii"),
|
|
]
|
|
)
|
|
except Exception as e:
|
|
logger.debug(
|
|
f"Failed to send ABORT for room {self.bootstrap_room}: {e}"
|
|
)
|
|
|
|
def abort(self):
|
|
self.kv_mgr.record_failure(
|
|
self.bootstrap_room,
|
|
"Aborted by AbortReq.",
|
|
)
|
|
# Mark the manager status Failed (not only the local conclude_state) so the
|
|
# transfer worker and other pollers observe the abort. Port of upstream
|
|
# sgl-project/sglang #24522.
|
|
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
|
|
# Tell the prefill peers to stop transferring into our (soon-freed) pages.
|
|
self._send_abort_notification()
|
|
# Explicitly set the status to failure since this request has been aborted
|
|
self.conclude_state = KVPoll.Failed
|
|
|
|
|
|
class MooncakeKVBootstrapServer(CommonKVBootstrapServer):
|
|
pass
|