1120 lines
41 KiB
Python
1120 lines
41 KiB
Python
from __future__ import annotations
|
|
|
|
import ctypes
|
|
import dataclasses
|
|
import logging
|
|
import os
|
|
import struct
|
|
import threading
|
|
import time
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
import msgspec
|
|
import numpy as np
|
|
import numpy.typing as npt
|
|
from mori.cpp import TransferStatus
|
|
from mori.io import (
|
|
BackendType,
|
|
EngineDesc,
|
|
IOEngine,
|
|
IOEngineConfig,
|
|
MemoryDesc,
|
|
MemoryLocationType,
|
|
PollCqMode,
|
|
RdmaBackendConfig,
|
|
)
|
|
|
|
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 group_concurrent_contiguous
|
|
from sglang.srt.disaggregation.utils import (
|
|
DisaggregationMode,
|
|
filter_kv_indices_for_cp_rank,
|
|
)
|
|
from sglang.srt.server_args import ServerArgs
|
|
from sglang.srt.utils.common import (
|
|
get_int_env_var,
|
|
get_local_ip_auto,
|
|
)
|
|
from sglang.srt.utils.network import NetworkAddress
|
|
|
|
logger = logging.getLogger(__name__)
|
|
MORI_GUARD = b"MoriMsgGuard"
|
|
|
|
|
|
def _pack_mem_desc_list(mems: List[MemoryDesc]) -> bytes:
|
|
if not mems:
|
|
return b""
|
|
packed_descs = [mem.pack() for mem in mems]
|
|
return msgspec.msgpack.encode(packed_descs)
|
|
|
|
|
|
def _unpack_mem_desc_list(blob: bytes) -> List[MemoryDesc]:
|
|
if not blob:
|
|
return []
|
|
desc_blobs = msgspec.msgpack.decode(blob)
|
|
return [MemoryDesc.unpack(b) for b in desc_blobs]
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class TransferInfo:
|
|
room: int
|
|
endpoint: str
|
|
dst_port: int
|
|
engine_key: str
|
|
dst_kv_indices: npt.NDArray[np.int32]
|
|
dst_aux_index: int
|
|
required_dst_info_num: int
|
|
is_dummy: bool
|
|
|
|
@classmethod
|
|
def from_zmq(cls, payload: List[bytes]) -> TransferInfo:
|
|
room = int(payload[0].decode("ascii"))
|
|
endpoint = payload[1].decode("ascii")
|
|
dst_port = int(payload[2].decode("ascii"))
|
|
engine_key = payload[3].decode("ascii")
|
|
|
|
if payload[4]:
|
|
dst_kv_indices = np.frombuffer(payload[4], dtype=np.int32)
|
|
else:
|
|
dst_kv_indices = np.array([], dtype=np.int32)
|
|
|
|
if payload[5]:
|
|
dst_aux_index = int(payload[5].decode("ascii"))
|
|
else:
|
|
dst_aux_index = -1
|
|
|
|
required_dst_info_num = (
|
|
int(payload[7].decode("ascii")) if len(payload) > 7 else 1
|
|
)
|
|
is_dummy = dst_kv_indices.size == 0 and dst_aux_index < 0
|
|
return cls(
|
|
room=room,
|
|
endpoint=endpoint,
|
|
dst_port=dst_port,
|
|
engine_key=engine_key,
|
|
dst_kv_indices=dst_kv_indices,
|
|
dst_aux_index=dst_aux_index,
|
|
required_dst_info_num=required_dst_info_num,
|
|
is_dummy=is_dummy,
|
|
)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class KVArgsRegisterInfo:
|
|
endpoint: str
|
|
dst_port: int
|
|
engine_desc: EngineDesc
|
|
dst_kv_mem_descs: List[MemoryDesc]
|
|
dst_aux_mem_descs: List[MemoryDesc]
|
|
dst_state_mem_descs: List[MemoryDesc]
|
|
gpu_id: int
|
|
decode_tp_size: int
|
|
decode_tp_rank: int
|
|
dst_kv_item_len: int
|
|
|
|
@property
|
|
def engine_key(self) -> str:
|
|
return self.engine_desc.key
|
|
|
|
@classmethod
|
|
def from_zmq(cls, payload: List[bytes]) -> KVArgsRegisterInfo:
|
|
endpoint = payload[1].decode("ascii")
|
|
dst_port = int(payload[2].decode("ascii"))
|
|
engine_desc = EngineDesc.unpack(payload[3])
|
|
dst_kv_mem_descs = _unpack_mem_desc_list(payload[4])
|
|
dst_aux_mem_descs = _unpack_mem_desc_list(payload[5])
|
|
dst_state_mem_descs = _unpack_mem_desc_list(payload[6])
|
|
gpu_id = int(payload[7].decode("ascii"))
|
|
decode_tp_size = int(payload[8].decode("ascii"))
|
|
decode_tp_rank = int(payload[9].decode("ascii"))
|
|
dst_kv_item_len = int(payload[10].decode("ascii"))
|
|
return cls(
|
|
endpoint=endpoint,
|
|
dst_port=dst_port,
|
|
engine_desc=engine_desc,
|
|
dst_kv_mem_descs=dst_kv_mem_descs,
|
|
dst_aux_mem_descs=dst_aux_mem_descs,
|
|
dst_state_mem_descs=dst_state_mem_descs,
|
|
gpu_id=gpu_id,
|
|
decode_tp_size=decode_tp_size,
|
|
decode_tp_rank=decode_tp_rank,
|
|
dst_kv_item_len=dst_kv_item_len,
|
|
)
|
|
|
|
|
|
class AuxDataCodec:
|
|
@staticmethod
|
|
def serialize_data_from_buffer(src_addr, data_length):
|
|
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):
|
|
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
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class TPSliceConfig:
|
|
page_size: int
|
|
src_item_len: int
|
|
dst_item_len: int
|
|
bytes_per_token_src: int
|
|
bytes_per_token_dst: int
|
|
src_head_slice_offset: int
|
|
dst_head_slice_offset: int
|
|
heads_bytes_per_token_to_send: int
|
|
|
|
|
|
class MoriKVManager(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.engine = self._init_engine()
|
|
self.engine_desc = self.engine.get_engine_desc()
|
|
self.kv_mem_descs: List[MemoryDesc] = []
|
|
self.aux_mem_descs: List[MemoryDesc] = []
|
|
self.state_mem_descs: List[MemoryDesc] = []
|
|
self.transfer_lock = threading.Lock()
|
|
self._register_local_buffers()
|
|
if self.disaggregation_mode == DisaggregationMode.PREFILL:
|
|
self._start_bootstrap_thread()
|
|
elif self.disaggregation_mode == DisaggregationMode.DECODE:
|
|
self.room_to_bootstrap_addr: Dict[int, str] = {}
|
|
self._start_decode_thread()
|
|
|
|
def _init_engine(self) -> IOEngine:
|
|
if self.kv_args.ib_device:
|
|
os.environ["MORI_RDMA_DEVICES"] = self.kv_args.ib_device
|
|
|
|
self.local_ip = get_local_ip_auto()
|
|
config = IOEngineConfig(host=self.local_ip, port=0)
|
|
|
|
engine_key = (
|
|
f"io-{self.disaggregation_mode.value}-"
|
|
f"dp{self.system_dp_rank}-tp{self.attn_tp_rank}-"
|
|
f"pid{os.getpid()}-{self.local_ip}"
|
|
)
|
|
|
|
engine = IOEngine(engine_key, config)
|
|
poll_mode = PollCqMode.POLLING
|
|
|
|
# Number of RDMA Queue Pairs (QPs) used per transfer operation.
|
|
# Higher values can increase parallelism and bandwidth utilization.
|
|
# Default: 1
|
|
qp_per_transfer = get_int_env_var("SGLANG_MORI_QP_PER_TRANSFER", 1)
|
|
|
|
# Number of RDMA work requests posted in a single batch to each QP.
|
|
# Larger batch sizes reduce per-operation overhead and improve throughput
|
|
# at the cost of higher latency. Use -1 for automatic sizing based on
|
|
# the number of merged work requests and available endpoints.
|
|
# Default: -1 (automatic)
|
|
post_batch_size = get_int_env_var("SGLANG_MORI_POST_BATCH_SIZE", -1)
|
|
|
|
# Number of worker threads in the RDMA executor thread pool.
|
|
# Each worker handles RDMA operations on a separate CPU core (with affinity).
|
|
# More workers can improve parallelism for large batch transfers across
|
|
# multiple QPs, but excessive threads may cause contention.
|
|
# Default: 1
|
|
num_worker_threads = get_int_env_var("SGLANG_MORI_NUM_WORKERS", 1)
|
|
|
|
rdma_cfg = RdmaBackendConfig(
|
|
qp_per_transfer,
|
|
post_batch_size,
|
|
num_worker_threads,
|
|
poll_mode,
|
|
False,
|
|
)
|
|
engine.create_backend(BackendType.RDMA, rdma_cfg)
|
|
actual_port = engine.get_engine_desc().port
|
|
assert actual_port > 0, f"Failed to bind port for engine {engine_key}"
|
|
logger.debug(
|
|
"Initialized Mori IOEngine %s at %s:%s (qp_per_transfer=%s, workers=%s, poll_mode=%s)",
|
|
engine_key,
|
|
self.local_ip,
|
|
actual_port,
|
|
qp_per_transfer,
|
|
num_worker_threads,
|
|
poll_mode.name,
|
|
)
|
|
return engine
|
|
|
|
def _register_local_buffers(self) -> None:
|
|
for ptr, length in zip(self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens):
|
|
mem_desc = self.engine.register_memory(
|
|
ptr,
|
|
length,
|
|
self.kv_args.gpu_id,
|
|
MemoryLocationType.GPU,
|
|
)
|
|
self.kv_mem_descs.append(mem_desc)
|
|
for ptr, length in zip(self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens):
|
|
desc = self.engine.register_memory(
|
|
ptr,
|
|
length,
|
|
-1,
|
|
MemoryLocationType.CPU,
|
|
)
|
|
self.aux_mem_descs.append(desc)
|
|
for ptr, length in zip(
|
|
self.kv_args.state_data_ptrs, getattr(self.kv_args, "state_data_lens", [])
|
|
):
|
|
desc = self.engine.register_memory(
|
|
ptr,
|
|
length,
|
|
self.kv_args.gpu_id,
|
|
MemoryLocationType.GPU,
|
|
)
|
|
self.state_mem_descs.append(desc)
|
|
|
|
def _handle_register_message(self, payload: List[bytes]) -> None:
|
|
try:
|
|
register_info = KVArgsRegisterInfo.from_zmq(payload)
|
|
self._add_remote_peer(register_info)
|
|
except Exception:
|
|
logger.exception("Failed to register remote peer")
|
|
|
|
def _handle_transfer_message(self, payload: List[bytes]) -> None:
|
|
try:
|
|
transfer_info = TransferInfo.from_zmq(payload)
|
|
infos = self.transfer_infos.setdefault(transfer_info.room, {})
|
|
infos[transfer_info.engine_key] = transfer_info
|
|
|
|
if len(infos) >= transfer_info.required_dst_info_num:
|
|
logger.debug(
|
|
"Bootstrap room %s got enough transfer info (%s)",
|
|
transfer_info.room,
|
|
len(infos),
|
|
)
|
|
self.update_status(transfer_info.room, KVPoll.WaitingForInput)
|
|
except Exception:
|
|
logger.exception("Failed to parse transfer info message")
|
|
|
|
def _validate_message(self, msg: List[bytes]) -> Optional[List[bytes]]:
|
|
if not msg or msg[0] != MORI_GUARD:
|
|
logger.warning("Received malformed bootstrap message")
|
|
return None
|
|
payload = msg[1:]
|
|
if not payload:
|
|
return None
|
|
return payload
|
|
|
|
def _start_bootstrap_thread(self) -> None:
|
|
def bootstrap_worker():
|
|
while True:
|
|
try:
|
|
msg = self.server_socket.recv_multipart()
|
|
payload = self._validate_message(msg)
|
|
if payload is None:
|
|
continue
|
|
room = payload[0].decode("ascii")
|
|
|
|
if room == "None":
|
|
self._handle_register_message(payload)
|
|
else:
|
|
self._handle_transfer_message(payload)
|
|
except Exception:
|
|
logger.exception("Bootstrap worker failed")
|
|
|
|
threading.Thread(target=bootstrap_worker, daemon=True).start()
|
|
|
|
def _cleanup_room_tracking(self, bootstrap_room: int) -> None:
|
|
bootstrap_addr = self.room_to_bootstrap_addr.pop(bootstrap_room, None)
|
|
if bootstrap_addr is not None:
|
|
rooms = self.addr_to_rooms_tracker.get(bootstrap_addr)
|
|
if rooms is not None:
|
|
rooms.discard(bootstrap_room)
|
|
if not rooms:
|
|
self.addr_to_rooms_tracker.pop(bootstrap_addr, None)
|
|
|
|
def _start_decode_thread(self) -> None:
|
|
def decode_worker():
|
|
while True:
|
|
try:
|
|
msg = self.server_socket.recv_multipart()
|
|
if msg and msg[0] == MoriKVManager.AUX_DATA_HEADER:
|
|
self._handle_aux_data(msg)
|
|
continue
|
|
|
|
if not msg or msg[0] != MORI_GUARD:
|
|
logger.warning(
|
|
"Received malformed status message on decode worker"
|
|
)
|
|
continue
|
|
payload = msg[1:]
|
|
if len(payload) < 3:
|
|
logger.warning("Incomplete status payload received")
|
|
continue
|
|
bootstrap_room = int(payload[0].decode("ascii"))
|
|
status_code = int(payload[1].decode("ascii"))
|
|
prefill_rank = int(payload[2].decode("ascii"))
|
|
failure_reason = (
|
|
payload[3].decode("utf-8")
|
|
if len(payload) > 3 and payload[3]
|
|
else None
|
|
)
|
|
|
|
if status_code == KVPoll.Success:
|
|
tracker = self.prefill_response_tracker[bootstrap_room]
|
|
tracker.add(prefill_rank)
|
|
expected = self.required_prefill_response_num_table.get(
|
|
bootstrap_room, 1
|
|
)
|
|
if len(tracker) >= expected:
|
|
self.prefill_response_tracker.pop(bootstrap_room, None)
|
|
self.update_status(bootstrap_room, KVPoll.Success)
|
|
self._cleanup_room_tracking(bootstrap_room)
|
|
elif status_code == KVPoll.Failed:
|
|
if failure_reason:
|
|
self.record_failure(bootstrap_room, failure_reason)
|
|
self.prefill_response_tracker.pop(bootstrap_room, None)
|
|
self.update_status(bootstrap_room, KVPoll.Failed)
|
|
self._cleanup_room_tracking(bootstrap_room)
|
|
else:
|
|
logger.warning(
|
|
"Unknown status code %s received for room %s",
|
|
status_code,
|
|
bootstrap_room,
|
|
)
|
|
except Exception:
|
|
logger.exception("Decode status worker failed")
|
|
|
|
threading.Thread(target=decode_worker, daemon=True).start()
|
|
|
|
def notify_decode_status(
|
|
self,
|
|
infos: List[TransferInfo],
|
|
bootstrap_room: int,
|
|
status: KVPoll,
|
|
failure_reason: Optional[str] = None,
|
|
) -> None:
|
|
if not infos:
|
|
return
|
|
payload = [
|
|
MORI_GUARD,
|
|
str(bootstrap_room).encode("ascii"),
|
|
str(int(status)).encode("ascii"),
|
|
str(self.attn_tp_rank * self.pp_size + self.pp_rank).encode("ascii"),
|
|
failure_reason.encode("utf-8") if failure_reason else b"",
|
|
]
|
|
for info in infos:
|
|
try:
|
|
na = NetworkAddress(info.endpoint, info.dst_port)
|
|
socket = self._connect(na.to_tcp(), is_ipv6=na.is_ipv6)
|
|
socket.send_multipart(payload)
|
|
except Exception:
|
|
logger.exception(
|
|
"Failed to sync status %s to decode endpoint %s:%s for room %s",
|
|
status,
|
|
info.endpoint,
|
|
info.dst_port,
|
|
bootstrap_room,
|
|
)
|
|
|
|
def _add_remote_peer(self, register_info: KVArgsRegisterInfo) -> None:
|
|
engine_key = register_info.engine_key
|
|
if engine_key in self.decode_kv_args_table:
|
|
logger.debug("Remote peer %s already registered. Skipping.", engine_key)
|
|
return
|
|
self.engine.register_remote_engine(register_info.engine_desc)
|
|
self.decode_kv_args_table[engine_key] = register_info
|
|
logger.debug(
|
|
"Registered decode peer %s (%s:%s)",
|
|
engine_key,
|
|
register_info.endpoint,
|
|
register_info.dst_port,
|
|
)
|
|
|
|
def _get_mha_mem_desc_slices(
|
|
self, dst_mem_descs: List[MemoryDesc]
|
|
) -> tuple[
|
|
List[MemoryDesc], List[MemoryDesc], List[MemoryDesc], List[MemoryDesc], int
|
|
]:
|
|
src_descs = self.kv_mem_descs
|
|
if not src_descs:
|
|
raise RuntimeError("KV memory descriptors are empty on prefill side")
|
|
|
|
num_local_layers = len(src_descs) // 2
|
|
src_k_descs = src_descs[:num_local_layers]
|
|
src_v_descs = src_descs[num_local_layers:]
|
|
|
|
start_layer = self.kv_args.prefill_start_layer
|
|
end_layer = start_layer + num_local_layers
|
|
dst_total_layers = len(dst_mem_descs) // 2
|
|
if len(dst_mem_descs) < 2 or end_layer > dst_total_layers:
|
|
raise ValueError(
|
|
"Destination KV descriptors do not match prefill pp configuration"
|
|
)
|
|
dst_k_descs = dst_mem_descs[start_layer:end_layer]
|
|
dst_v_descs = dst_mem_descs[
|
|
dst_total_layers + start_layer : dst_total_layers + end_layer
|
|
]
|
|
return src_k_descs, src_v_descs, dst_k_descs, dst_v_descs, num_local_layers
|
|
|
|
def _get_mla_mem_desc_slices(
|
|
self, dst_mem_descs: List[MemoryDesc]
|
|
) -> tuple[List[MemoryDesc], List[MemoryDesc], int]:
|
|
src_descs = self.kv_mem_descs
|
|
num_local_layers = len(src_descs)
|
|
start_layer = self.kv_args.prefill_start_layer
|
|
end_layer = start_layer + num_local_layers
|
|
if end_layer > len(dst_mem_descs):
|
|
raise ValueError(
|
|
"Destination MLA KV descriptors do not match prefill pp configuration"
|
|
)
|
|
dst_slice = dst_mem_descs[start_layer:end_layer]
|
|
return src_descs, dst_slice, num_local_layers
|
|
|
|
def _issue_layer_transfers(
|
|
self,
|
|
src_desc: MemoryDesc,
|
|
dst_desc: MemoryDesc,
|
|
kv_item_len: int,
|
|
src_groups: List[List[int]],
|
|
dst_groups: List[List[int]],
|
|
) -> List[TransferStatus]:
|
|
if not src_groups:
|
|
return []
|
|
local_offsets = [int(src_group[0]) * kv_item_len for src_group in src_groups]
|
|
remote_offsets = [int(dst_group[0]) * kv_item_len for dst_group in dst_groups]
|
|
sizes = [len(src_group) * kv_item_len for src_group in src_groups]
|
|
|
|
transfer_uid = self.engine.allocate_transfer_uid()
|
|
|
|
statuses = self.engine.batch_write(
|
|
[src_desc],
|
|
[local_offsets],
|
|
[dst_desc],
|
|
[remote_offsets],
|
|
[sizes],
|
|
[transfer_uid],
|
|
)
|
|
return statuses
|
|
|
|
def _build_tp_slice_config(self, peer_info: KVArgsRegisterInfo) -> TPSliceConfig:
|
|
page_size = self.kv_args.page_size
|
|
|
|
src_item_len = self.kv_args.kv_item_lens[0]
|
|
dst_item_len = peer_info.dst_kv_item_len
|
|
|
|
bytes_per_token_src = src_item_len // page_size
|
|
bytes_per_token_dst = dst_item_len // page_size
|
|
|
|
prefill_tp_size = self.attn_tp_size
|
|
decode_tp_size = peer_info.decode_tp_size
|
|
|
|
num_kv_heads = self.kv_args.kv_head_num
|
|
src_heads_per_rank = num_kv_heads
|
|
dst_heads_per_rank = num_kv_heads * prefill_tp_size // decode_tp_size
|
|
if dst_heads_per_rank == 0:
|
|
raise ValueError("Destination heads per rank evaluates to zero")
|
|
|
|
bytes_per_head_slice = bytes_per_token_dst // dst_heads_per_rank
|
|
if bytes_per_head_slice == 0:
|
|
raise ValueError("Head slice size evaluates to zero")
|
|
|
|
local_tp_rank = self.kv_args.engine_rank % prefill_tp_size
|
|
dst_tp_rank = peer_info.decode_tp_rank % decode_tp_size
|
|
|
|
if prefill_tp_size > decode_tp_size:
|
|
src_head_start = 0
|
|
num_heads_to_send = src_heads_per_rank
|
|
dst_head_start = local_tp_rank * src_heads_per_rank
|
|
else:
|
|
src_head_start = (dst_tp_rank * dst_heads_per_rank) % src_heads_per_rank
|
|
num_heads_to_send = dst_heads_per_rank
|
|
dst_head_start = 0
|
|
|
|
src_head_slice_offset = src_head_start * bytes_per_head_slice
|
|
dst_head_slice_offset = dst_head_start * bytes_per_head_slice
|
|
heads_bytes_per_token = num_heads_to_send * bytes_per_head_slice
|
|
|
|
if heads_bytes_per_token > bytes_per_token_dst:
|
|
raise ValueError(
|
|
"Slice size exceeds destination token capacity for TP slice transfer"
|
|
)
|
|
|
|
return TPSliceConfig(
|
|
page_size=page_size,
|
|
src_item_len=src_item_len,
|
|
dst_item_len=dst_item_len,
|
|
bytes_per_token_src=bytes_per_token_src,
|
|
bytes_per_token_dst=bytes_per_token_dst,
|
|
src_head_slice_offset=src_head_slice_offset,
|
|
dst_head_slice_offset=dst_head_slice_offset,
|
|
heads_bytes_per_token_to_send=heads_bytes_per_token,
|
|
)
|
|
|
|
def _issue_tp_slice_transfers(
|
|
self,
|
|
src_desc: MemoryDesc,
|
|
dst_desc: MemoryDesc,
|
|
kv_indices: npt.NDArray[np.int32],
|
|
dst_indices: npt.NDArray[np.int32],
|
|
tp_cfg: TPSliceConfig,
|
|
) -> List[TransferStatus]:
|
|
if kv_indices.size == 0 or dst_indices.size == 0:
|
|
return []
|
|
|
|
limit = min(kv_indices.size, dst_indices.size)
|
|
if not limit:
|
|
return []
|
|
|
|
src_pages = kv_indices[:limit].astype(np.int64)
|
|
dst_pages = dst_indices[:limit].astype(np.int64)
|
|
token_slots = np.arange(tp_cfg.page_size, dtype=np.int64)
|
|
|
|
src_page_bases = src_pages * tp_cfg.src_item_len
|
|
dst_page_bases = dst_pages * tp_cfg.dst_item_len
|
|
|
|
src_token_offsets = token_slots * tp_cfg.bytes_per_token_src
|
|
dst_token_offsets = token_slots * tp_cfg.bytes_per_token_dst
|
|
|
|
local_offsets = (
|
|
(
|
|
src_page_bases[:, np.newaxis]
|
|
+ src_token_offsets
|
|
+ tp_cfg.src_head_slice_offset
|
|
)
|
|
.flatten()
|
|
.tolist()
|
|
)
|
|
remote_offsets = (
|
|
(
|
|
dst_page_bases[:, np.newaxis]
|
|
+ dst_token_offsets
|
|
+ tp_cfg.dst_head_slice_offset
|
|
)
|
|
.flatten()
|
|
.tolist()
|
|
)
|
|
|
|
num_transfers = limit * tp_cfg.page_size
|
|
sizes = [tp_cfg.heads_bytes_per_token_to_send] * num_transfers
|
|
|
|
if not local_offsets:
|
|
return []
|
|
|
|
transfer_uid = self.engine.allocate_transfer_uid()
|
|
statuses = self.engine.batch_write(
|
|
[src_desc],
|
|
[local_offsets],
|
|
[dst_desc],
|
|
[remote_offsets],
|
|
[sizes],
|
|
[transfer_uid],
|
|
)
|
|
return statuses
|
|
|
|
def send_kvcache(
|
|
self,
|
|
peer_info: KVArgsRegisterInfo,
|
|
prefill_kv_indices: npt.NDArray[np.int32],
|
|
dst_kv_indices: npt.NDArray[np.int32],
|
|
) -> List[TransferStatus]:
|
|
src_groups, dst_groups = group_concurrent_contiguous(
|
|
prefill_kv_indices, dst_kv_indices
|
|
)
|
|
statuses = []
|
|
kv_item_len = self.kv_args.kv_item_lens[0]
|
|
if self.is_mla_backend:
|
|
(
|
|
src_descs,
|
|
dst_descs,
|
|
layers_current_pp_stage,
|
|
) = self._get_mla_mem_desc_slices(peer_info.dst_kv_mem_descs)
|
|
for layer_id in range(layers_current_pp_stage):
|
|
statuses.extend(
|
|
self._issue_layer_transfers(
|
|
src_descs[layer_id],
|
|
dst_descs[layer_id],
|
|
kv_item_len,
|
|
src_groups,
|
|
dst_groups,
|
|
)
|
|
)
|
|
else:
|
|
tp_mismatch = peer_info.decode_tp_size != self.attn_tp_size
|
|
(
|
|
src_k_descs,
|
|
src_v_descs,
|
|
dst_k_descs,
|
|
dst_v_descs,
|
|
layers_current_pp_stage,
|
|
) = self._get_mha_mem_desc_slices(peer_info.dst_kv_mem_descs)
|
|
|
|
if tp_mismatch:
|
|
tp_cfg = self._build_tp_slice_config(peer_info)
|
|
for layer_id in range(layers_current_pp_stage):
|
|
statuses.extend(
|
|
self._issue_tp_slice_transfers(
|
|
src_k_descs[layer_id],
|
|
dst_k_descs[layer_id],
|
|
prefill_kv_indices,
|
|
dst_kv_indices,
|
|
tp_cfg,
|
|
)
|
|
)
|
|
statuses.extend(
|
|
self._issue_tp_slice_transfers(
|
|
src_v_descs[layer_id],
|
|
dst_v_descs[layer_id],
|
|
prefill_kv_indices,
|
|
dst_kv_indices,
|
|
tp_cfg,
|
|
)
|
|
)
|
|
else:
|
|
src_groups, dst_groups = group_concurrent_contiguous(
|
|
prefill_kv_indices, dst_kv_indices
|
|
)
|
|
for layer_id in range(layers_current_pp_stage):
|
|
statuses.extend(
|
|
self._issue_layer_transfers(
|
|
src_k_descs[layer_id],
|
|
dst_k_descs[layer_id],
|
|
kv_item_len,
|
|
src_groups,
|
|
dst_groups,
|
|
)
|
|
)
|
|
statuses.extend(
|
|
self._issue_layer_transfers(
|
|
src_v_descs[layer_id],
|
|
dst_v_descs[layer_id],
|
|
kv_item_len,
|
|
src_groups,
|
|
dst_groups,
|
|
)
|
|
)
|
|
|
|
return statuses
|
|
|
|
def send_aux(
|
|
self,
|
|
peer_info: KVArgsRegisterInfo,
|
|
prefill_aux_index: int,
|
|
dst_aux_index: int,
|
|
room: int,
|
|
) -> List[TransferStatus]:
|
|
return self.send_aux_tcp(peer_info, prefill_aux_index, dst_aux_index, room)
|
|
|
|
def send_aux_tcp(
|
|
self,
|
|
peer_info: KVArgsRegisterInfo,
|
|
prefill_aux_index: int,
|
|
dst_aux_index: int,
|
|
room: int,
|
|
) -> List[TransferStatus]:
|
|
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=peer_info.endpoint,
|
|
dst_port=peer_info.dst_port,
|
|
room=room,
|
|
buffer_index=i,
|
|
aux_index=dst_aux_index,
|
|
data=data,
|
|
)
|
|
|
|
return []
|
|
|
|
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(
|
|
[
|
|
MoriKVManager.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 add_transfer_request(
|
|
self,
|
|
bootstrap_room: int,
|
|
kv_indices: npt.NDArray[np.int32],
|
|
index_slice: slice,
|
|
is_last: bool,
|
|
aux_index: Optional[int] = None,
|
|
state_indices: Optional[npt.NDArray[np.int32]] = None,
|
|
) -> Tuple[List[TransferStatus], Optional[List[TransferInfo]]]:
|
|
assert self.disaggregation_mode == DisaggregationMode.PREFILL
|
|
transfer_infos = self.transfer_infos.get(bootstrap_room)
|
|
if not transfer_infos:
|
|
raise RuntimeError(
|
|
f"No transfer info found for bootstrap_room={bootstrap_room}"
|
|
)
|
|
result_statuses = []
|
|
target_infos_snapshot: Optional[List[TransferInfo]] = None
|
|
with self.transfer_lock:
|
|
self.update_status(bootstrap_room, KVPoll.Transferring)
|
|
for info in transfer_infos.values():
|
|
peer_info = self.decode_kv_args_table.get(info.engine_key)
|
|
if not peer_info:
|
|
self.record_failure(
|
|
bootstrap_room,
|
|
f"Peer info missing for engine {info.engine_key}",
|
|
)
|
|
raise RuntimeError(
|
|
f"Missing decode peer info for {info.engine_key}"
|
|
)
|
|
if not info.is_dummy:
|
|
dst_indices_chunk = info.dst_kv_indices[index_slice]
|
|
statuses = self.send_kvcache(
|
|
peer_info, kv_indices, dst_indices_chunk
|
|
)
|
|
result_statuses.extend(statuses)
|
|
if (
|
|
is_last
|
|
and aux_index is not None
|
|
and info.dst_aux_index >= 0
|
|
and self.pp_group.is_last_rank
|
|
):
|
|
result_statuses.extend(
|
|
self.send_aux(
|
|
peer_info, aux_index, info.dst_aux_index, bootstrap_room
|
|
)
|
|
)
|
|
if is_last:
|
|
self.update_status(bootstrap_room, KVPoll.Success)
|
|
target_infos_snapshot = list(transfer_infos.values())
|
|
self.transfer_infos.pop(bootstrap_room, None)
|
|
return result_statuses, target_infos_snapshot
|
|
|
|
|
|
class MoriKVSender(CommonKVSender):
|
|
def __init__(
|
|
self,
|
|
mgr: MoriKVManager,
|
|
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.transfer_statuses: List[TransferStatus] = []
|
|
self.pending_infos: Optional[List[TransferInfo]] = None
|
|
self.sent_last_chunk = False
|
|
self.conclude_state: Optional[KVPoll] = None
|
|
self.status_notified = False
|
|
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 = self.curr_idx == self.num_kv_indices
|
|
|
|
# Special handling for cp
|
|
if 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:
|
|
return
|
|
else:
|
|
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Success)
|
|
return
|
|
statuses, infos = self.kv_mgr.add_transfer_request(
|
|
self.bootstrap_room,
|
|
kv_indices,
|
|
index_slice,
|
|
is_last,
|
|
aux_index=self.aux_index if is_last else None,
|
|
)
|
|
self.transfer_statuses.extend(statuses)
|
|
if infos is not None:
|
|
self.pending_infos = infos
|
|
self.sent_last_chunk = True
|
|
|
|
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 == KVPoll.Bootstrapping:
|
|
elapsed = time.time() - self.init_time
|
|
if elapsed >= self.kv_mgr.bootstrap_timeout:
|
|
reason = (
|
|
f"Request {self.bootstrap_room} timed out after {elapsed:.1f}s "
|
|
"waiting for decode handshake"
|
|
)
|
|
self.kv_mgr.record_failure(self.bootstrap_room, reason)
|
|
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
|
|
self._finalize_failure(reason)
|
|
return KVPoll.Failed
|
|
return status
|
|
|
|
if status == KVPoll.Failed:
|
|
self._finalize_failure()
|
|
return KVPoll.Failed
|
|
|
|
transfers_done = self._all_transfers_finished()
|
|
if transfers_done:
|
|
if self._has_transfer_error():
|
|
reason = self._collect_failure_reason()
|
|
self.kv_mgr.record_failure(self.bootstrap_room, reason)
|
|
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
|
|
self._finalize_failure(reason)
|
|
return KVPoll.Failed
|
|
self._notify_decode(KVPoll.Success)
|
|
self.conclude_state = KVPoll.Success
|
|
return KVPoll.Success
|
|
return KVPoll.Transferring if status == KVPoll.Success else status
|
|
|
|
def _all_transfers_finished(self) -> bool:
|
|
if not self.sent_last_chunk:
|
|
return False
|
|
if not self.transfer_statuses:
|
|
return True
|
|
return all(not status.InProgress() for status in self.transfer_statuses)
|
|
|
|
def _has_transfer_error(self) -> bool:
|
|
return any(status.Failed() for status in self.transfer_statuses)
|
|
|
|
def _collect_failure_reason(self) -> str:
|
|
for status in self.transfer_statuses:
|
|
if status.Failed():
|
|
return f"KV transfer failed: {status.Message()}"
|
|
return "KV transfer failed due to unknown reason"
|
|
|
|
def _notify_decode(
|
|
self, status: KVPoll, failure_reason: Optional[str] = None
|
|
) -> None:
|
|
if self.status_notified:
|
|
return
|
|
if self.pending_infos:
|
|
self.kv_mgr.notify_decode_status(
|
|
self.pending_infos, self.bootstrap_room, status, failure_reason
|
|
)
|
|
self.status_notified = True
|
|
|
|
def _finalize_failure(self, failure_reason: Optional[str] = None) -> None:
|
|
if self.conclude_state == KVPoll.Failed:
|
|
return
|
|
if failure_reason is None:
|
|
failure_reason = self.kv_mgr.failure_records.get(
|
|
self.bootstrap_room, "KV transfer failed"
|
|
)
|
|
self._notify_decode(KVPoll.Failed, failure_reason)
|
|
self.conclude_state = KVPoll.Failed
|
|
|
|
def clear(self) -> None:
|
|
self.kv_mgr.request_status.pop(self.bootstrap_room, None)
|
|
|
|
def failure_exception(self):
|
|
if self.conclude_state is None:
|
|
self._finalize_failure()
|
|
self.clear()
|
|
with self.kv_mgr.failure_lock:
|
|
failure_reason = self.kv_mgr.failure_records.pop(
|
|
self.bootstrap_room, "KV transfer failed"
|
|
)
|
|
raise RuntimeError(failure_reason)
|
|
|
|
def abort(self):
|
|
reason = "Aborted by AbortReq."
|
|
self.kv_mgr.record_failure(self.bootstrap_room, reason)
|
|
self._notify_decode(KVPoll.Failed, reason)
|
|
self.conclude_state = KVPoll.Failed
|
|
|
|
|
|
class MoriKVReceiver(CommonKVReceiver):
|
|
|
|
def __init__(
|
|
self,
|
|
mgr: MoriKVManager,
|
|
bootstrap_addr: str,
|
|
bootstrap_room: Optional[int] = None,
|
|
prefill_dp_rank: Optional[int] = None,
|
|
):
|
|
super().__init__(mgr, bootstrap_addr, bootstrap_room, prefill_dp_rank)
|
|
self.conclude_state: Optional[KVPoll] = None
|
|
self.init_time: Optional[float] = None
|
|
if self.bootstrap_room is None or self.bootstrap_infos is None:
|
|
return
|
|
self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room)
|
|
self.kv_mgr.room_to_bootstrap_addr[self.bootstrap_room] = self.bootstrap_addr
|
|
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.WaitingForInput)
|
|
self._register_kv_args()
|
|
|
|
def _register_kv_args(self):
|
|
if self.bootstrap_infos is None:
|
|
return
|
|
engine_desc_blob = self.kv_mgr.engine_desc.pack()
|
|
packed_kv_descs = _pack_mem_desc_list(self.kv_mgr.kv_mem_descs)
|
|
packed_aux_descs = _pack_mem_desc_list(self.kv_mgr.aux_mem_descs)
|
|
packed_state_descs = _pack_mem_desc_list(self.kv_mgr.state_mem_descs)
|
|
gpu_id = str(self.kv_mgr.kv_args.gpu_id).encode("ascii")
|
|
decode_tp_size = str(self.kv_mgr.attn_tp_size).encode("ascii")
|
|
decode_tp_rank = str(self.kv_mgr.kv_args.engine_rank).encode("ascii")
|
|
kv_item_len = str(self.kv_mgr.kv_args.kv_item_lens[0]).encode("ascii")
|
|
|
|
for bootstrap_info in self.bootstrap_infos:
|
|
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
|
|
with lock:
|
|
sock.send_multipart(
|
|
[
|
|
MORI_GUARD,
|
|
"None".encode("ascii"),
|
|
self.kv_mgr.local_ip.encode("ascii"),
|
|
str(self.kv_mgr.rank_port).encode("ascii"),
|
|
engine_desc_blob,
|
|
packed_kv_descs,
|
|
packed_aux_descs,
|
|
packed_state_descs,
|
|
gpu_id,
|
|
decode_tp_size,
|
|
decode_tp_rank,
|
|
kv_item_len,
|
|
]
|
|
)
|
|
|
|
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 or self.bootstrap_room is None:
|
|
return
|
|
|
|
kv_indices_bytes = (
|
|
np.asarray(kv_indices, dtype=np.int32).tobytes() if kv_indices.size else b""
|
|
)
|
|
aux_bytes = str(aux_index).encode("ascii") if aux_index is not None else b""
|
|
state_bytes = b""
|
|
|
|
for bootstrap_info in self.bootstrap_infos:
|
|
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
|
|
is_dummy = bootstrap_info.get("is_dummy", False)
|
|
with lock:
|
|
sock.send_multipart(
|
|
[
|
|
MORI_GUARD,
|
|
str(self.bootstrap_room).encode("ascii"),
|
|
self.kv_mgr.local_ip.encode("ascii"),
|
|
str(self.kv_mgr.rank_port).encode("ascii"),
|
|
self.kv_mgr.engine_desc.key.encode("ascii"),
|
|
kv_indices_bytes if not is_dummy else b"",
|
|
aux_bytes if not is_dummy else b"",
|
|
state_bytes,
|
|
str(self.required_dst_info_num).encode("ascii"),
|
|
]
|
|
)
|
|
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
|
|
return status
|
|
|
|
if status == KVPoll.WaitingForInput and self.init_time is not None:
|
|
elapsed = time.time() - self.init_time
|
|
if elapsed >= self.kv_mgr.waiting_timeout:
|
|
reason = f"Request {self.bootstrap_room} timed out after {elapsed:.1f}s waiting for KV transfer"
|
|
self.kv_mgr.record_failure(self.bootstrap_room, reason)
|
|
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
|
|
self.conclude_state = KVPoll.Failed
|
|
return KVPoll.Failed
|
|
|
|
return status
|
|
|
|
def clear(self) -> None:
|
|
if self.bootstrap_room is None:
|
|
return
|
|
self.kv_mgr.request_status.pop(self.bootstrap_room, None)
|
|
self.kv_mgr.required_prefill_response_num_table.pop(self.bootstrap_room, None)
|
|
self.kv_mgr.prefill_response_tracker.pop(self.bootstrap_room, None)
|
|
self.kv_mgr._cleanup_room_tracking(self.bootstrap_room)
|
|
|
|
def failure_exception(self):
|
|
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, "KV transfer failed"
|
|
)
|
|
raise RuntimeError(failure_reason)
|
|
|
|
def abort(self):
|
|
if self.bootstrap_room is None:
|
|
return
|
|
reason = "Aborted by AbortReq."
|
|
self.kv_mgr.record_failure(self.bootstrap_room, reason)
|
|
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed)
|
|
self.conclude_state = KVPoll.Failed
|
|
self.clear()
|
|
|
|
|
|
class MoriKVBootstrapServer(CommonKVBootstrapServer):
|
|
pass
|