697 lines
26 KiB
Python
697 lines
26 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import dataclasses
|
|
import logging
|
|
import queue
|
|
import socket
|
|
import struct
|
|
import threading
|
|
import uuid
|
|
from collections import defaultdict
|
|
from functools import cache
|
|
from typing import Dict, List, Optional, Set, Tuple, TypeAlias, Union
|
|
|
|
import numpy as np
|
|
import numpy.typing as npt
|
|
import requests
|
|
import zmq
|
|
from aiohttp import web
|
|
|
|
from sglang.srt.disaggregation.base.conn import BaseKVSender, KVArgs, KVPoll
|
|
from sglang.srt.disaggregation.common.conn import (
|
|
CommonKVBootstrapServer,
|
|
CommonKVManager,
|
|
CommonKVReceiver,
|
|
)
|
|
from sglang.srt.disaggregation.common.utils import group_concurrent_contiguous
|
|
from sglang.srt.disaggregation.utils import DisaggregationMode
|
|
from sglang.srt.server_args import ServerArgs
|
|
from sglang.srt.utils import (
|
|
format_tcp_address,
|
|
get_local_ip_auto,
|
|
is_valid_ipv6_address,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
GUARD = "NixlMsgGuard".encode("ascii")
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class TransferInfo:
|
|
"""Contains indices for a transfer, sent by KVReceiver. Received by prefill bootstrap thread."""
|
|
|
|
room: int
|
|
endpoint: str
|
|
dst_port: int
|
|
agent_name: str
|
|
dst_kv_indices: npt.NDArray[np.int32]
|
|
dst_aux_index: int
|
|
required_dst_info_num: int
|
|
|
|
def is_dummy(self):
|
|
return self.dst_kv_indices.size == 0
|
|
|
|
@classmethod
|
|
def from_zmq(cls, msg: List[bytes]):
|
|
return cls(
|
|
room=int(msg[0].decode("ascii")),
|
|
endpoint=msg[1].decode("ascii"),
|
|
dst_port=int(msg[2].decode("ascii")),
|
|
agent_name=msg[3].decode("ascii"),
|
|
dst_kv_indices=np.frombuffer(msg[4], dtype=np.int32),
|
|
dst_aux_index=int(msg[5].decode("ascii")),
|
|
required_dst_info_num=int(msg[6].decode("ascii")),
|
|
)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class KVArgsRegisterInfo:
|
|
"""Contains base pointers and other info which only needs to be sent once by KVReceiver. Received by prefill bootstrap thread."""
|
|
|
|
room: str
|
|
endpoint: str
|
|
dst_port: int
|
|
agent_name: str
|
|
agent_metadata: bytes
|
|
dst_kv_ptrs: list[int]
|
|
dst_aux_ptrs: list[int]
|
|
gpu_id: int
|
|
decode_tp_size: int
|
|
decode_tp_rank: int
|
|
dst_kv_item_len: 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")),
|
|
agent_name=msg[3].decode("ascii"),
|
|
agent_metadata=msg[4],
|
|
dst_kv_ptrs=list(struct.unpack(f"{len(msg[5])//8}Q", msg[5])),
|
|
dst_aux_ptrs=list(struct.unpack(f"{len(msg[6])//8}Q", msg[6])),
|
|
gpu_id=int(msg[7].decode("ascii")),
|
|
decode_tp_size=int(msg[8].decode("ascii")),
|
|
decode_tp_rank=int(msg[9].decode("ascii")),
|
|
dst_kv_item_len=int(msg[10].decode("ascii")),
|
|
)
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class TransferStatus:
|
|
"""Used by KV Receiver to know when a transfer is done."""
|
|
|
|
# KV chunk IDs that have been received.
|
|
received_kvs: Set[int] = dataclasses.field(default_factory=set)
|
|
# Number of kv chunks to expect, will know this after last chunk is received.
|
|
num_kvs_expected: Optional[int] = None
|
|
# Whether aux data has been received.
|
|
received_aux: bool = False
|
|
|
|
def is_done(self):
|
|
if self.num_kvs_expected is None:
|
|
return False
|
|
return self.num_kvs_expected == len(self.received_kvs) and self.received_aux
|
|
|
|
|
|
class NixlKVManager(CommonKVManager):
|
|
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)
|
|
try:
|
|
from nixl._api import nixl_agent
|
|
except ImportError as e:
|
|
raise ImportError(
|
|
"Please install NIXL by following the instructions at "
|
|
"https://github.com/ai-dynamo/nixl/blob/main/README.md "
|
|
"to run SGLang with NixlTransferEngine."
|
|
) from e
|
|
self.agent = nixl_agent(str(uuid.uuid4()))
|
|
self.local_ip = get_local_ip_auto()
|
|
self.server_socket = zmq.Context().socket(zmq.PULL)
|
|
if is_valid_ipv6_address(self.local_ip):
|
|
self.server_socket.setsockopt(zmq.IPV6, 1)
|
|
self.register_buffer_to_engine()
|
|
|
|
if self.disaggregation_mode == DisaggregationMode.PREFILL:
|
|
self.request_status: Dict[int, KVPoll] = {}
|
|
self.transfer_infos: Dict[int, Dict[str, TransferInfo]] = {}
|
|
self.decode_kv_args_table: Dict[str, KVArgsRegisterInfo] = {}
|
|
self._start_bootstrap_thread()
|
|
elif self.disaggregation_mode == DisaggregationMode.DECODE:
|
|
self.transfer_statuses: Dict[int, TransferStatus] = defaultdict(
|
|
TransferStatus
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"Unsupported DisaggregationMode: {self.disaggregation_mode}"
|
|
)
|
|
|
|
def check_status(self, bootstrap_room: int):
|
|
return self.request_status[bootstrap_room]
|
|
|
|
def update_status(self, bootstrap_room: int, status: KVPoll):
|
|
if bootstrap_room not in self.request_status:
|
|
self.request_status[bootstrap_room] = status
|
|
else:
|
|
# NOTE: The prefill engine could recv bootstrapping first
|
|
self.request_status[bootstrap_room] = max(
|
|
self.request_status[bootstrap_room], status
|
|
)
|
|
|
|
def register_buffer_to_engine(self):
|
|
kv_addrs = []
|
|
for kv_data_ptr, kv_data_len in zip(
|
|
self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens
|
|
):
|
|
kv_addrs.append((kv_data_ptr, kv_data_len, self.kv_args.gpu_id, ""))
|
|
self.kv_descs = self.agent.register_memory(kv_addrs, "VRAM")
|
|
logger.debug(f"Register kv tensors, len(kv_addr)= {len(kv_addrs)}")
|
|
if not self.kv_descs:
|
|
raise Exception("NIXL memory registration failed for kv tensors")
|
|
aux_addrs = []
|
|
for aux_data_ptr, aux_data_len in zip(
|
|
self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens
|
|
):
|
|
aux_addrs.append((aux_data_ptr, aux_data_len, 0, ""))
|
|
self.aux_descs = self.agent.register_memory(aux_addrs, "DRAM")
|
|
logger.debug(f"Register aux tensors, len(aux_addrs)= {len(aux_addrs)}")
|
|
if not self.aux_descs:
|
|
raise Exception("NIXL memory registration failed for aux tensors")
|
|
|
|
def _add_remote_peer(self, decode_kv_args: KVArgsRegisterInfo):
|
|
agent_name = decode_kv_args.agent_name
|
|
if agent_name in self.decode_kv_args_table:
|
|
logger.info(f"Peer {agent_name} was already registered, ignoring.")
|
|
return
|
|
self.decode_kv_args_table[agent_name] = decode_kv_args
|
|
self.agent.add_remote_agent(decode_kv_args.agent_metadata)
|
|
|
|
def send_kvcache(
|
|
self,
|
|
peer_name: str,
|
|
prefill_kv_indices: npt.NDArray[np.int32],
|
|
dst_kv_ptrs: list[int],
|
|
dst_kv_indices: npt.NDArray[np.int32],
|
|
dst_gpu_id: int,
|
|
notif: str,
|
|
):
|
|
# group by indices
|
|
prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous(
|
|
prefill_kv_indices, dst_kv_indices
|
|
)
|
|
|
|
logger.debug(f"sending kvcache to {peer_name} with notif {notif}")
|
|
# Make descs
|
|
num_layers = len(self.kv_args.kv_data_ptrs)
|
|
src_addrs = []
|
|
dst_addrs = []
|
|
for layer_id in range(num_layers):
|
|
src_ptr = self.kv_args.kv_data_ptrs[layer_id]
|
|
dst_ptr = dst_kv_ptrs[layer_id]
|
|
item_len = self.kv_args.kv_item_lens[layer_id]
|
|
|
|
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)
|
|
src_addrs.append((src_addr, length, self.kv_args.gpu_id))
|
|
dst_addrs.append((dst_addr, length, dst_gpu_id))
|
|
|
|
logger.debug(
|
|
f"len(src_addrs): before group: {len(prefill_kv_indices)}, after group: {len(src_addrs)}"
|
|
)
|
|
src_descs = self.agent.get_xfer_descs(src_addrs, "VRAM")
|
|
dst_descs = self.agent.get_xfer_descs(dst_addrs, "VRAM")
|
|
# Transfer data
|
|
xfer_handle = self.agent.initialize_xfer(
|
|
"WRITE",
|
|
src_descs,
|
|
dst_descs,
|
|
peer_name,
|
|
notif.encode("ascii"), # type: ignore
|
|
)
|
|
if not xfer_handle:
|
|
raise Exception("KVSender failed to create transfer")
|
|
state = self.agent.transfer(xfer_handle)
|
|
if state == "ERR":
|
|
raise Exception("KVSender failed to post transfer")
|
|
return xfer_handle
|
|
|
|
def send_kvcache_slice(
|
|
self,
|
|
peer_name: str,
|
|
prefill_kv_indices: npt.NDArray[np.int32],
|
|
dst_kv_ptrs: list[int],
|
|
dst_kv_indices: npt.NDArray[np.int32],
|
|
dst_gpu_id: int,
|
|
notif: str,
|
|
prefill_tp_size: int,
|
|
decode_tp_size: int,
|
|
decode_tp_rank: int,
|
|
dst_kv_item_len: int,
|
|
):
|
|
# Get configuration from kv_args
|
|
local_tp_rank_in_group = self.kv_args.engine_rank % prefill_tp_size
|
|
dst_tp_rank_in_group = decode_tp_rank % decode_tp_size
|
|
num_kv_heads = self.kv_args.kv_head_num
|
|
|
|
# Calculate head distribution
|
|
src_heads_per_rank = num_kv_heads
|
|
dst_heads_per_rank = num_kv_heads * prefill_tp_size // decode_tp_size
|
|
|
|
src_kv_item_len = self.kv_args.kv_item_lens[0]
|
|
page_size = self.kv_args.page_size
|
|
|
|
bytes_per_head_slice_to_send = (
|
|
dst_kv_item_len // page_size // dst_heads_per_rank
|
|
)
|
|
|
|
# Determine which heads to send
|
|
if prefill_tp_size > decode_tp_size:
|
|
# Multiple prefill ranks to one decode rank
|
|
src_head_start_offset = 0
|
|
num_heads_to_send = src_heads_per_rank
|
|
dst_head_start_offset = local_tp_rank_in_group * src_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
|
|
|
|
# Create transfer descriptors
|
|
src_addrs = []
|
|
dst_addrs = []
|
|
|
|
bytes_per_token_on_prefill = src_kv_item_len // page_size
|
|
bytes_per_token_on_decode = dst_kv_item_len // page_size
|
|
|
|
num_kv_layers = len(self.kv_args.kv_data_ptrs) // 2
|
|
src_k_ptrs = self.kv_args.kv_data_ptrs[:num_kv_layers]
|
|
src_v_ptrs = self.kv_args.kv_data_ptrs[num_kv_layers:]
|
|
dst_k_ptrs = dst_kv_ptrs[0 : len(src_k_ptrs)]
|
|
dst_v_ptrs = dst_kv_ptrs[num_kv_layers : num_kv_layers + len(src_v_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
|
|
|
|
src_dst_ptr_pairs = [
|
|
(
|
|
src_k_ptrs[layer_id],
|
|
dst_k_ptrs[layer_id],
|
|
)
|
|
for layer_id in range(len(src_k_ptrs))
|
|
] + [
|
|
(
|
|
src_v_ptrs[layer_id],
|
|
dst_v_ptrs[layer_id],
|
|
)
|
|
for layer_id in range(len(src_v_ptrs))
|
|
]
|
|
|
|
src_addrs = []
|
|
dst_addrs = []
|
|
|
|
# Calculate strides for a single token slot
|
|
bytes_per_token_on_prefill = src_kv_item_len // page_size
|
|
bytes_per_token_on_decode = dst_kv_item_len // page_size
|
|
|
|
for src_ptr, dst_ptr in src_dst_ptr_pairs:
|
|
for i in range(len(prefill_kv_indices)):
|
|
prefill_page_idx = int(prefill_kv_indices[i])
|
|
decode_page_idx = int(dst_kv_indices[i])
|
|
|
|
# Get the starting addresses for the current src and dst pages
|
|
src_page_start_addr = src_ptr + prefill_page_idx * src_kv_item_len
|
|
dst_page_start_addr = dst_ptr + decode_page_idx * dst_kv_item_len
|
|
|
|
# Iterate through each valid token slot within the current page
|
|
for token_slot_in_page in range(page_size):
|
|
# Calculate the start address of the current token slot
|
|
src_token_slot_start_addr = (
|
|
src_page_start_addr
|
|
+ token_slot_in_page * bytes_per_token_on_prefill
|
|
)
|
|
dst_token_slot_start_addr = (
|
|
dst_page_start_addr
|
|
+ token_slot_in_page * bytes_per_token_on_decode
|
|
)
|
|
|
|
# Calculate final src and dst addresses by applying head-slice offsets
|
|
src_slice_addr = src_token_slot_start_addr + src_head_slice_offset
|
|
dst_slice_addr = dst_token_slot_start_addr + dst_head_slice_offset
|
|
|
|
src_addrs.append(
|
|
(
|
|
src_slice_addr,
|
|
heads_bytes_per_token_to_send,
|
|
self.kv_args.gpu_id,
|
|
)
|
|
)
|
|
dst_addrs.append(
|
|
(dst_slice_addr, heads_bytes_per_token_to_send, dst_gpu_id)
|
|
)
|
|
|
|
# Use NIXL agent for transfer
|
|
src_descs = self.agent.get_xfer_descs(src_addrs, "VRAM")
|
|
dst_descs = self.agent.get_xfer_descs(dst_addrs, "VRAM")
|
|
|
|
xfer_handle = self.agent.initialize_xfer(
|
|
"WRITE", src_descs, dst_descs, peer_name, notif.encode("ascii")
|
|
)
|
|
if not xfer_handle:
|
|
raise Exception("Failed to create sliced KV transfer")
|
|
|
|
state = self.agent.transfer(xfer_handle)
|
|
if state == "ERR":
|
|
raise Exception("Failed to post sliced KV transfer")
|
|
|
|
return xfer_handle
|
|
|
|
def send_aux(
|
|
self,
|
|
peer_name: str,
|
|
prefill_aux_index: int,
|
|
dst_aux_ptrs: list[int],
|
|
dst_aux_index: int,
|
|
notif: str,
|
|
):
|
|
# Make descs
|
|
aux_item_len = self.kv_args.aux_item_lens[0]
|
|
prefill_aux_addr = (
|
|
self.kv_args.aux_data_ptrs[0] + prefill_aux_index * aux_item_len
|
|
)
|
|
decode_aux_addr = dst_aux_ptrs[0] + dst_aux_index * aux_item_len
|
|
src_addrs = [(prefill_aux_addr, aux_item_len, 0)]
|
|
dst_addrs = [(decode_aux_addr, aux_item_len, 0)]
|
|
src_descs = self.agent.get_xfer_descs(src_addrs, "DRAM")
|
|
dst_descs = self.agent.get_xfer_descs(dst_addrs, "DRAM")
|
|
# Transfer data
|
|
xfer_handle = self.agent.initialize_xfer(
|
|
"WRITE",
|
|
src_descs,
|
|
dst_descs,
|
|
peer_name,
|
|
notif.encode("ascii"), # type: ignore
|
|
)
|
|
if not xfer_handle:
|
|
raise Exception("KVSender failed to create transfer")
|
|
state = self.agent.transfer(xfer_handle)
|
|
if state == "ERR":
|
|
raise Exception("KVSender failed to post transfer")
|
|
return xfer_handle
|
|
|
|
def add_transfer_request(
|
|
self,
|
|
bootstrap_room: int,
|
|
kv_indices: npt.NDArray[np.int32],
|
|
index_slice: slice,
|
|
is_last: bool,
|
|
chunk_id: int,
|
|
aux_index: Optional[int] = None,
|
|
):
|
|
assert self.disaggregation_mode == DisaggregationMode.PREFILL
|
|
assert not is_last or (is_last and aux_index is not None)
|
|
|
|
reqs_to_be_processed = self.transfer_infos[bootstrap_room].values()
|
|
handles = []
|
|
for req in reqs_to_be_processed:
|
|
assert bootstrap_room == req.room
|
|
if req.is_dummy():
|
|
continue
|
|
|
|
chunked_dst_kv_indice = req.dst_kv_indices[index_slice]
|
|
assert len(chunked_dst_kv_indice) == len(kv_indices)
|
|
assert req.agent_name in self.decode_kv_args_table
|
|
|
|
notif = "_".join([str(req.room), "kv", str(chunk_id), str(int(is_last))])
|
|
decode_tp_size = self.decode_kv_args_table[req.agent_name].decode_tp_size
|
|
|
|
if decode_tp_size == self.tp_size:
|
|
kv_xfer_handle = self.send_kvcache(
|
|
req.agent_name,
|
|
kv_indices,
|
|
self.decode_kv_args_table[req.agent_name].dst_kv_ptrs,
|
|
chunked_dst_kv_indice,
|
|
self.decode_kv_args_table[req.agent_name].gpu_id,
|
|
notif,
|
|
)
|
|
else:
|
|
kv_xfer_handle = self.send_kvcache_slice(
|
|
req.agent_name,
|
|
kv_indices,
|
|
self.decode_kv_args_table[req.agent_name].dst_kv_ptrs,
|
|
chunked_dst_kv_indice,
|
|
self.decode_kv_args_table[req.agent_name].gpu_id,
|
|
notif,
|
|
prefill_tp_size=self.tp_size,
|
|
decode_tp_size=decode_tp_size,
|
|
decode_tp_rank=self.decode_kv_args_table[
|
|
req.agent_name
|
|
].decode_tp_rank,
|
|
dst_kv_item_len=self.decode_kv_args_table[
|
|
req.agent_name
|
|
].dst_kv_item_len,
|
|
)
|
|
|
|
handles.append(kv_xfer_handle)
|
|
# Only the last chunk we need to send the aux data.
|
|
if is_last:
|
|
assert aux_index is not None
|
|
aux_xfer_handle = self.send_aux(
|
|
req.agent_name,
|
|
aux_index,
|
|
self.decode_kv_args_table[req.agent_name].dst_aux_ptrs,
|
|
req.dst_aux_index,
|
|
str(req.room) + "_aux",
|
|
)
|
|
handles.append(aux_xfer_handle)
|
|
if is_last:
|
|
del self.transfer_infos[bootstrap_room]
|
|
return handles
|
|
|
|
def update_transfer_status(self):
|
|
# Process notifications from received transfers.
|
|
notif_map = self.agent.get_new_notifs()
|
|
for peer_name, messages in notif_map.items():
|
|
# We could also check that self.bootstrap_info['agent_name'] matches
|
|
# the message sender. But the bootstrap room alone should be
|
|
# sufficient to map the status.
|
|
for msg in messages:
|
|
components = msg.decode("ascii").split("_")
|
|
room = int(components[0])
|
|
if components[1] == "kv":
|
|
chunk_id = int(components[2])
|
|
is_last = bool(int(components[3]))
|
|
self.transfer_statuses[room].received_kvs.add(chunk_id)
|
|
if is_last:
|
|
self.transfer_statuses[room].num_kvs_expected = chunk_id + 1
|
|
elif components[1] == "aux":
|
|
self.transfer_statuses[room].received_aux = True
|
|
|
|
def check_transfer_done(self, room: int):
|
|
if room not in self.transfer_statuses:
|
|
return False
|
|
return self.transfer_statuses[room].is_done()
|
|
|
|
def _bind_server_socket(self):
|
|
self.server_socket.bind(format_tcp_address(self.local_ip, self.rank_port))
|
|
|
|
def _start_bootstrap_thread(self):
|
|
self._bind_server_socket()
|
|
|
|
def bootstrap_thread():
|
|
"""This thread recvs transfer info from the decode engine"""
|
|
while True:
|
|
waiting_req_bytes = self.server_socket.recv_multipart()
|
|
logger.debug(
|
|
f"Received multipart with total byte size {sum(len(x) for x in waiting_req_bytes)}"
|
|
)
|
|
assert (
|
|
waiting_req_bytes[0] == GUARD
|
|
), f"First message should be {GUARD}. Foreign traffic?"
|
|
waiting_req_bytes = waiting_req_bytes[1:]
|
|
room = waiting_req_bytes[0].decode("ascii")
|
|
agent_name = waiting_req_bytes[3].decode("ascii")
|
|
if room == "None":
|
|
# Register new peer and save KV base pointers.
|
|
self._add_remote_peer(
|
|
KVArgsRegisterInfo.from_zmq(waiting_req_bytes)
|
|
)
|
|
logger.debug(f"Register KVArgs from {agent_name} successfully")
|
|
continue
|
|
room = int(room)
|
|
if room not in self.transfer_infos:
|
|
self.transfer_infos[room] = {}
|
|
self.transfer_infos[room][agent_name] = TransferInfo.from_zmq(
|
|
waiting_req_bytes
|
|
)
|
|
required_dst_info_num = self.transfer_infos[room][
|
|
agent_name
|
|
].required_dst_info_num
|
|
logger.debug(f"got info {room=} {agent_name=} {required_dst_info_num=}")
|
|
if len(self.transfer_infos[room]) == required_dst_info_num:
|
|
logger.debug(f"{room=} is bootstrapped")
|
|
self.update_status(room, KVPoll.WaitingForInput)
|
|
|
|
threading.Thread(target=bootstrap_thread).start()
|
|
|
|
|
|
class NixlKVSender(BaseKVSender):
|
|
|
|
def __init__(
|
|
self,
|
|
mgr: NixlKVManager,
|
|
bootstrap_addr: str,
|
|
bootstrap_room: int,
|
|
dest_tp_ranks: List[int],
|
|
pp_rank: int,
|
|
):
|
|
self.kv_mgr = mgr
|
|
self.bootstrap_room = bootstrap_room
|
|
self.aux_index = None
|
|
self.bootstrap_server_url = bootstrap_addr
|
|
self.xfer_handles = []
|
|
self.has_sent = False
|
|
self.chunk_id = 0
|
|
self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping)
|
|
# inner state
|
|
self.curr_idx = 0
|
|
|
|
def init(self, num_kv_indices: int, aux_index: Optional[int] = None):
|
|
self.num_kv_indices = num_kv_indices
|
|
self.aux_index = aux_index
|
|
|
|
def send(
|
|
self,
|
|
kv_indices: npt.NDArray[np.int32],
|
|
):
|
|
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
|
|
|
|
new_xfer_handles = self.kv_mgr.add_transfer_request(
|
|
self.bootstrap_room,
|
|
kv_indices,
|
|
index_slice,
|
|
is_last,
|
|
self.chunk_id,
|
|
self.aux_index,
|
|
)
|
|
self.xfer_handles.extend(new_xfer_handles)
|
|
self.chunk_id += 1
|
|
if is_last:
|
|
self.has_sent = True
|
|
del self.kv_mgr.request_status[self.bootstrap_room]
|
|
|
|
def poll(self) -> KVPoll:
|
|
if not self.has_sent:
|
|
return self.kv_mgr.check_status(self.bootstrap_room)
|
|
states = [self.kv_mgr.agent.check_xfer_state(x) for x in self.xfer_handles]
|
|
if all([x == "DONE" for x in states]):
|
|
return KVPoll.Success # type: ignore
|
|
if any([x == "ERR" for x in states]):
|
|
raise Exception("KVSender transfer encountered an error.")
|
|
return KVPoll.WaitingForInput # type: ignore
|
|
|
|
def failure_exception(self):
|
|
raise Exception("Fake KVSender Exception")
|
|
|
|
|
|
class NixlKVReceiver(CommonKVReceiver):
|
|
def __init__(
|
|
self,
|
|
mgr: NixlKVManager,
|
|
bootstrap_addr: str,
|
|
bootstrap_room: Optional[int] = None,
|
|
prefill_dp_rank: Optional[int] = None,
|
|
):
|
|
self.started_transfer = False
|
|
self.conclude_state = None
|
|
super().__init__(mgr, bootstrap_addr, bootstrap_room, prefill_dp_rank)
|
|
|
|
def init(self, kv_indices: npt.NDArray[np.int32], aux_index: Optional[int] = None):
|
|
for bootstrap_info in self.bootstrap_infos:
|
|
logger.debug(
|
|
f"Fetched bootstrap info: {bootstrap_info} for engine rank: {self.kv_mgr.kv_args.engine_rank}"
|
|
)
|
|
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
|
|
is_dummy = bootstrap_info["is_dummy"]
|
|
logger.debug(
|
|
f"Sending to prefill server with bootstrap room {self.bootstrap_room} {is_dummy=}"
|
|
)
|
|
with lock:
|
|
sock.send_multipart(
|
|
[
|
|
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.agent.name.encode("ascii"),
|
|
kv_indices.tobytes() if not is_dummy else b"",
|
|
str(aux_index).encode("ascii"),
|
|
str(self.required_dst_info_num).encode("ascii"),
|
|
]
|
|
)
|
|
|
|
self.started_transfer = True
|
|
|
|
def poll(self) -> KVPoll:
|
|
if self.conclude_state is not None:
|
|
return self.conclude_state
|
|
if not self.started_transfer:
|
|
return KVPoll.WaitingForInput # type: ignore
|
|
|
|
self.kv_mgr.update_transfer_status()
|
|
if self.kv_mgr.check_transfer_done(self.bootstrap_room): # type: ignore
|
|
self.conclude_state = KVPoll.Success
|
|
del self.kv_mgr.transfer_statuses[self.bootstrap_room]
|
|
return KVPoll.Success # type: ignore
|
|
return KVPoll.WaitingForInput # type: ignore
|
|
|
|
def _register_kv_args(self):
|
|
for bootstrap_info in self.bootstrap_infos:
|
|
sock, lock = self._connect_to_bootstrap_server(bootstrap_info)
|
|
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
|
|
)
|
|
|
|
with lock:
|
|
sock.send_multipart(
|
|
[
|
|
GUARD,
|
|
"None".encode("ascii"),
|
|
self.kv_mgr.local_ip.encode("ascii"),
|
|
str(self.kv_mgr.rank_port).encode("ascii"),
|
|
self.kv_mgr.agent.name.encode("ascii"),
|
|
self.kv_mgr.agent.get_agent_metadata(),
|
|
packed_kv_data_ptrs,
|
|
packed_aux_data_ptrs,
|
|
str(self.kv_mgr.kv_args.gpu_id).encode("ascii"),
|
|
str(self.kv_mgr.kv_args.decode_tp_size).encode("ascii"),
|
|
str(self.kv_mgr.kv_args.engine_rank).encode("ascii"),
|
|
str(self.kv_mgr.kv_args.kv_item_lens[0]).encode("ascii"),
|
|
]
|
|
)
|
|
|
|
def failure_exception(self):
|
|
raise Exception("Fake KVReceiver Exception")
|
|
|
|
|
|
class NixlKVBootstrapServer(CommonKVBootstrapServer):
|
|
pass
|