from __future__ import annotations import dataclasses import logging import struct import threading import time import uuid from collections import defaultdict from typing import Dict, List, Optional, Set import numpy as np import numpy.typing as npt import requests 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 from sglang.srt.environ import envs from sglang.srt.server_args import ServerArgs 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 dst_state_indices: List[int] def is_dummy(self): return self.dst_kv_indices.size == 0 @classmethod def from_zmq(cls, msg: List[bytes]): # Parse state_indices from msg[7] if present if len(msg) > 7 and msg[7] != b"": dst_state_indices = list(np.frombuffer(msg[7], dtype=np.int32)) else: dst_state_indices = [] 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")), dst_state_indices=dst_state_indices, ) @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] dst_state_data_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]): # Parse state_data_ptrs from msg[7] if present if len(msg) > 7 and msg[7] != b"": dst_state_data_ptrs = list(struct.unpack(f"{len(msg[7]) // 8}Q", msg[7])) else: dst_state_data_ptrs = [] 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])), dst_state_data_ptrs=dst_state_data_ptrs, gpu_id=int(msg[8].decode("ascii")), decode_tp_size=int(msg[9].decode("ascii")), decode_tp_rank=int(msg[10].decode("ascii")), dst_kv_item_len=int(msg[11].decode("ascii")), ) @dataclasses.dataclass class TransferStatus: """Used by KV Receiver to know when a transfer is done.""" # KV chunks received per pp_rank: {pp_rank: set of chunk_ids} received_kvs_per_pp: Dict[int, Set[int]] = dataclasses.field( default_factory=lambda: defaultdict(set) ) # Expected chunk count per pp_rank (set when is_last=True): {pp_rank: expected_count} expected_kvs_per_pp: Dict[int, int] = dataclasses.field(default_factory=dict) # Number of PP ranks expected to send data. num_pp_ranks_expected: Optional[int] = None # Whether aux data has been received. received_aux: bool = False # PP ranks that have sent state data (state is layer-specific, each PP rank sends its portion). received_state_per_pp: Set[int] = dataclasses.field(default_factory=set) # Whether state data is expected (set based on state_type). expects_state: bool = False # Mark as failed is_failure: bool = False def is_done(self): if self.is_failure: return True if self.num_pp_ranks_expected is None or not self.received_aux: return False # If state data is expected, check all PP ranks have sent it if ( self.expects_state and len(self.received_state_per_pp) < self.num_pp_ranks_expected ): return False # All PP ranks must have reported their expected count if len(self.expected_kvs_per_pp) < self.num_pp_ranks_expected: return False # Each PP rank must have received all expected chunks for pp_rank, expected in self.expected_kvs_per_pp.items(): if len(self.received_kvs_per_pp[pp_rank]) != expected: return False return True def is_failed(self): return self.is_failure 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, nixl_agent_config 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 backend = envs.SGLANG_DISAGGREGATION_NIXL_BACKEND.get() agent_config = nixl_agent_config( backends=[backend], num_threads=(8 if disaggregation_mode == DisaggregationMode.PREFILL else 0), ) self.agent = nixl_agent(str(uuid.uuid4()), agent_config) available_plugins = self.agent.get_plugin_list() if backend not in available_plugins: raise ValueError( f"NIXL backend '{backend}' not found. Available: {available_plugins}. " f"Please install the required NIXL plugin or choose from: {available_plugins}" ) logger.info(f"NIXL KVManager initialized with backend: {backend}") self.register_buffer_to_engine() if self.disaggregation_mode == DisaggregationMode.PREFILL: self._start_bootstrap_thread() elif self.disaggregation_mode == DisaggregationMode.DECODE: self.transfer_statuses: Dict[int, TransferStatus] = defaultdict( TransferStatus ) self.heartbeat_failures = {} self.session_pool = defaultdict(requests.Session) self.session_pool_lock = threading.Lock() self.addr_to_rooms_tracker = defaultdict(set) self.connection_lock = threading.Lock() # Heartbeat interval should be at least 2 seconds self.heartbeat_interval = max( envs.SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL.get(), 2.0 ) # Heartbeat failure should be at least 1 self.max_failures = max( envs.SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE.get(), 1 ) self.waiting_timeout = envs.SGLANG_DISAGGREGATION_WAITING_TIMEOUT.get() self._start_heartbeat_checker_thread() else: raise ValueError( f"Unsupported DisaggregationMode: {self.disaggregation_mode}" ) def _start_heartbeat_checker_thread(self): """ Start the heartbeat checker thread for Decode worker. TODO (smor): unite nixl heartbeat checker with mooncake's. """ def heartbeat_checker(): while True: time.sleep(self.heartbeat_interval) with self.connection_lock: addresses = list(self.prefill_dp_size_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 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=heartbeat_checker, daemon=True).start() def _handle_node_failure(self, failed_bootstrap_addr): """Handle failure of a prefill node.""" 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] if failed_bootstrap_addr in self.prefill_attn_tp_size_table: del self.prefill_attn_tp_size_table[failed_bootstrap_addr] if failed_bootstrap_addr in self.prefill_dp_size_table: del self.prefill_dp_size_table[failed_bootstrap_addr] if failed_bootstrap_addr in self.prefill_pp_size_table: del self.prefill_pp_size_table[failed_bootstrap_addr] possible_affected_rooms = self.addr_to_rooms_tracker.get( failed_bootstrap_addr, [] ) if failed_bootstrap_addr in self.addr_to_rooms_tracker: del self.addr_to_rooms_tracker[failed_bootstrap_addr] # Mark all pending transfers associated with the failed node as failed affected_rooms = [] for room in possible_affected_rooms: if ( room in self.transfer_statuses and not self.transfer_statuses[room].is_done() ): # Mark the transfer as failed self.transfer_statuses[room].is_failure = True affected_rooms.append(room) logger.error( f"Lost connection with prefill instance (bootstrap_addr: {failed_bootstrap_addr}), " f"{len(affected_rooms)} transfers affected" ) for room in possible_affected_rooms: logger.error(f"Let room {room} be failed due to prefill down") self.update_status(room, KVPoll.Failed) 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: status is only allowed to be incremented unless it is KVPoll.Failed if status == KVPoll.Failed: self.request_status[bootstrap_room] = KVPoll.Failed else: self.request_status[bootstrap_room] = max( self.request_status[bootstrap_room], status ) def record_failure(self, bootstrap_room: int, failure_reason: str): pass 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") # Register state/extra pool data buffers if present if self.kv_args.state_data_ptrs and self.kv_args.state_data_lens: state_addrs = [] for state_data_ptr, state_data_len in zip( self.kv_args.state_data_ptrs, self.kv_args.state_data_lens ): state_addrs.append( (state_data_ptr, state_data_len, self.kv_args.gpu_id, "") ) self.state_descs = self.agent.register_memory(state_addrs, "VRAM") logger.debug( f"Register state tensors, len(state_addrs)= {len(state_addrs)}" ) if not self.state_descs: raise Exception("NIXL memory registration failed for state 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_generic( self, peer_name: 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], dst_gpu_id: int, notif: str, ): """Generic KV cache transfer supporting both MHA and MLA architectures. Used by both send_kvcache and maybe_send_extra.""" # group by indices prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous( prefill_data_indices, dst_data_indices ) logger.debug(f"sending kvcache to {peer_name} with notif {notif}") # Make descs 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) ) layers_params = [ ( src_k_ptrs[layer_id], dst_k_ptrs[layer_id], item_lens[layer_id], ) for layer_id in range(layers_current_pp_stage) ] + [ ( src_v_ptrs[layer_id], dst_v_ptrs[layer_id], item_lens[layer_id], ) for layer_id in range(layers_current_pp_stage) ] src_addrs = [] src_lens = [] dst_addrs = [] dst_lens = [] # Precompute block starts/lengths to reduce Python-level loops. prefill_starts = np.fromiter( (block[0] for block in prefill_kv_blocks), dtype=np.int64 ) dst_starts = np.fromiter((block[0] for block in dst_kv_blocks), dtype=np.int64) block_lens = np.fromiter( (len(block) for block in prefill_kv_blocks), dtype=np.int64 ) for src_ptr, dst_ptr, item_len in layers_params: lengths = item_len * block_lens src_addrs.append(src_ptr + prefill_starts * item_len) src_lens.append(lengths) dst_addrs.append(dst_ptr + dst_starts * item_len) dst_lens.append(lengths) def make_req_array(addr_chunks, len_chunks, gpu): if not addr_chunks: return np.empty((0, 3), dtype=np.int64) flat_addrs = np.concatenate(addr_chunks) flat_lens = np.concatenate(len_chunks) return np.column_stack( ( flat_addrs, flat_lens, np.full_like(flat_addrs, gpu), ) ) src_reqs = make_req_array(src_addrs, src_lens, self.kv_args.gpu_id) dst_reqs = make_req_array(dst_addrs, dst_lens, dst_gpu_id) logger.debug( f"len(src_addrs): before group: {len(prefill_data_indices)}, after group: {len(src_addrs)}" ) src_descs = self.agent.get_xfer_descs(src_reqs, "VRAM") dst_descs = self.agent.get_xfer_descs(dst_reqs, "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( 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, ): return self._send_kvcache_generic( peer_name=peer_name, 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, dst_gpu_id=dst_gpu_id, notif=notif, ) 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 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 src_dst_ptr_pairs = [ ( src_k_ptrs[layer_id], dst_k_ptrs[layer_id], ) for layer_id in range(layers_current_pp_stage) ] + [ ( src_v_ptrs[layer_id], dst_v_ptrs[layer_id], ) for layer_id in range(layers_current_pp_stage) ] prefill_indices = np.asarray(prefill_kv_indices, dtype=np.int64) dst_indices = np.asarray(dst_kv_indices, dtype=np.int64) bytes_per_token_prefill = src_kv_item_len // page_size bytes_per_token_decode = dst_kv_item_len // page_size token_offsets = np.arange(page_size, dtype=np.int64) src_addrs = [] dst_addrs = [] for src_ptr, dst_ptr in src_dst_ptr_pairs: src_page_bases = src_ptr + prefill_indices * src_kv_item_len dst_page_bases = dst_ptr + dst_indices * dst_kv_item_len src_all = ( src_page_bases[:, None] + token_offsets[None, :] * bytes_per_token_prefill + src_head_slice_offset ).ravel() dst_all = ( dst_page_bases[:, None] + token_offsets[None, :] * bytes_per_token_decode + dst_head_slice_offset ).ravel() src_addrs.append(src_all) dst_addrs.append(dst_all) def make_req_array(addr_chunks, size, gpu): if not addr_chunks: return np.empty((0, 3), dtype=np.int64) flat_addrs = np.concatenate(addr_chunks) return np.column_stack( ( flat_addrs, np.full_like(flat_addrs, size), np.full_like(flat_addrs, gpu), ) ) src_reqs = make_req_array( src_addrs, heads_bytes_per_token_to_send, self.kv_args.gpu_id ) dst_reqs = make_req_array(dst_addrs, heads_bytes_per_token_to_send, dst_gpu_id) # Use NIXL agent for transfer src_descs = self.agent.get_xfer_descs(src_reqs, "VRAM") dst_descs = self.agent.get_xfer_descs(dst_reqs, "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, ): src_addrs = [] dst_addrs = [] prefill_aux_ptrs = self.kv_args.aux_data_ptrs prefill_aux_item_lens = self.kv_args.aux_item_lens for i, _ 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 * dst_aux_index src_addrs.append((src_addr, length, 0)) dst_addrs.append((dst_addr, length, 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 _send_mamba_state( self, peer_name: str, prefill_state_indices: List[int], dst_state_data_ptrs: list[int], dst_state_indices: List[int], dst_gpu_id: int, notif: str, ): """Transfer Mamba states via RDMA.""" assert len(prefill_state_indices) == 1, "Mamba should have single state index" assert len(dst_state_indices) == len( prefill_state_indices ), "State indices count mismatch between Prefill and Decode" src_addrs = [] dst_addrs = [] 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_state_indices[0] ) dst_addr = dst_state_ptr + length * int(dst_state_indices[0]) src_addrs.append((src_addr, length, self.kv_args.gpu_id)) dst_addrs.append((dst_addr, length, dst_gpu_id)) 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 Mamba state transfer") state = self.agent.transfer(xfer_handle) if state == "ERR": raise Exception("Failed to post Mamba state transfer") return xfer_handle def maybe_send_extra( self, peer_name: str, prefill_state_indices: List[int], dst_state_data_ptrs: list[int], dst_state_indices: List[int], dst_gpu_id: int, notif: str, decode_tp_size: int, ): """Send state or extra pool data with type-specific handling.""" state_type = getattr(self.kv_args, "state_type", "none") if state_type == "mamba": if self.attn_tp_size != decode_tp_size: raise RuntimeError( "PD Disaggregation does NOT support PD different TP sizes for hybrid mamba models yet." ) return self._send_mamba_state( peer_name, prefill_state_indices, dst_state_data_ptrs, dst_state_indices, dst_gpu_id, notif, ) elif state_type in ["swa", "nsa"]: if not self.is_mla_backend and self.attn_tp_size != decode_tp_size: raise RuntimeError( f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {state_type.upper()} hybrid models yet." ) if len(prefill_state_indices) != len(dst_state_indices): raise RuntimeError( f"State index length mismatch: prefill={len(prefill_state_indices)}, " f"dst={len(dst_state_indices)}" ) return self._send_kvcache_generic( peer_name=peer_name, src_data_ptrs=self.kv_args.state_data_ptrs, dst_data_ptrs=dst_state_data_ptrs, item_lens=self.kv_args.state_item_lens, prefill_data_indices=np.array(prefill_state_indices, dtype=np.int32), dst_data_indices=np.array(dst_state_indices, dtype=np.int32), dst_gpu_id=dst_gpu_id, notif=notif, ) else: if state_type != "none": raise RuntimeError( f"PD Disaggregation via NIXL does NOT support {state_type} hybrid models yet." ) return None 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, state_indices: Optional[List[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 = f"{req.room}_kv_{chunk_id}_{int(is_last)}_{self.kv_args.pp_rank}" decode_tp_size = self.decode_kv_args_table[req.agent_name].decode_tp_size if self.is_mla_backend or (decode_tp_size == self.attn_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.attn_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: if state_indices is not None: dst_info = self.decode_kv_args_table[req.agent_name] state_xfer_handle = self.maybe_send_extra( req.agent_name, state_indices, dst_info.dst_state_data_ptrs, req.dst_state_indices, dst_info.gpu_id, f"{req.room}_state_{self.kv_args.pp_rank}", decode_tp_size, ) if state_xfer_handle is not None: handles.append(state_xfer_handle) 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, f"{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("_", 4) room = int(components[0]) if components[1] == "kv": chunk_id = int(components[2]) is_last = bool(int(components[3])) pp_rank = int(components[4]) if len(components) > 4 else 0 # Track received chunks per pp_rank self.transfer_statuses[room].received_kvs_per_pp[pp_rank].add( chunk_id ) if is_last: # Record expected chunk count for this pp_rank self.transfer_statuses[room].expected_kvs_per_pp[pp_rank] = ( chunk_id + 1 ) # Set num_pp_ranks_expected from table (or default to 1) if self.transfer_statuses[room].num_pp_ranks_expected is None: self.transfer_statuses[room].num_pp_ranks_expected = ( self.required_prefill_response_num_table.get(room, 1) ) elif components[1] == "aux": self.transfer_statuses[room].received_aux = True elif components[1] == "state": pp_rank = int(components[2]) if len(components) > 2 else 0 self.transfer_statuses[room].received_state_per_pp.add(pp_rank) def check_transfer_done(self, room: int): if room not in self.transfer_statuses: return False return self.transfer_statuses[room].is_done() def _start_bootstrap_thread(self): 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(CommonKVSender): def __init__( self, mgr: NixlKVManager, 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.xfer_handles = [] self.has_sent = False self.chunk_id = 0 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 new_xfer_handles = self.kv_mgr.add_transfer_request( self.bootstrap_room, kv_indices, index_slice, is_last, self.chunk_id, self.aux_index, state_indices, ) 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 RuntimeError("NIXL 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) # Track this room with its bootstrap address for heartbeat monitoring if hasattr(self.kv_mgr, "addr_to_rooms_tracker"): self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add( self.bootstrap_room ) self.init_time = None 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: logger.error( 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: 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"), ( np.array(state_indices, dtype=np.int32).tobytes() if not is_dummy and state_indices is not None else b"" ), ] ) # Mark that we expect state data if state_indices was provided if state_indices is not None: self.kv_mgr.transfer_statuses[self.bootstrap_room].expects_state = True self.started_transfer = True 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 not self.started_transfer: return KVPoll.WaitingForInput # type: ignore now = time.time() elapsed = now - self.init_time if elapsed >= self.kv_mgr.waiting_timeout: logger.error(f"Request {self.bootstrap_room} waiting_timeout") self.kv_mgr.record_failure( self.bootstrap_room, f"Request {self.bootstrap_room} timed out after {elapsed:.1f}s in KVPoll.WaitingForInput", ) self.conclude_state = KVPoll.Failed return KVPoll.Failed self.kv_mgr.update_transfer_status() if self.kv_mgr.check_transfer_done(self.bootstrap_room): # type: ignore self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].discard( self.bootstrap_room ) # Check if the transfer failed if self.kv_mgr.transfer_statuses[self.bootstrap_room].is_failed(): self.conclude_state = KVPoll.Failed logger.error( f"Transfer for room {self.bootstrap_room} failed due to node failure" ) else: self.conclude_state = KVPoll.Success del self.kv_mgr.transfer_statuses[self.bootstrap_room] return self.conclude_state # 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 ) packed_state_data_ptrs = b"".join( struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.state_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, packed_state_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 RuntimeError("NIXL KVReceiver Exception") class NixlKVBootstrapServer(CommonKVBootstrapServer): pass