From a241659d188306bc5e67904502a0209ee741798b Mon Sep 17 00:00:00 2001 From: Shangming Cai Date: Sat, 23 May 2026 10:41:20 +0800 Subject: [PATCH] [PD] Consolidate shared logic into common backend (#25979) Signed-off-by: Shangming Cai --- .../sglang/srt/disaggregation/common/conn.py | 164 +++++++++++++- .../sglang/srt/disaggregation/common/utils.py | 37 +++- .../srt/disaggregation/mooncake/conn.py | 189 +++------------- python/sglang/srt/disaggregation/mori/conn.py | 91 +++----- python/sglang/srt/disaggregation/nixl/conn.py | 206 +++--------------- 5 files changed, 284 insertions(+), 403 deletions(-) diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 693e7fffe..94aa7d569 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -25,7 +25,10 @@ from sglang.srt.disaggregation.base.conn import ( KVPoll, KVTransferMetric, ) -from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.disaggregation.utils import ( + DisaggregationMode, + filter_kv_indices_for_cp_rank, +) from sglang.srt.distributed import get_pp_group, get_world_group from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import ( @@ -594,6 +597,92 @@ class CommonKVManager(BaseKVManager): return src_kv_ptrs, sliced_dst + def _start_heartbeat_checker_thread(self): + """Start the heartbeat checker thread for Decode worker.""" + + 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 + self._on_heartbeat_success(bootstrap_addr) + 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 _on_heartbeat_success(self, bootstrap_addr: str): + """Hook called on successful heartbeat. Override for backend-specific cleanup.""" + pass + + def _handle_node_failure(self, failed_bootstrap_addr: str): + """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] + self.prefill_info_table.pop(failed_bootstrap_addr, None) + + possible_affected_rooms = self.addr_to_rooms_tracker.get( + failed_bootstrap_addr, [] + ) + self.addr_to_rooms_tracker.pop(failed_bootstrap_addr, None) + + 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"Lost connection with prefill instance (bootstrap_addr: {failed_bootstrap_addr})", + ) + self.update_status(room, KVPoll.Failed) + affected_rooms.append(room) + + logger.error( + f"Lost connection with prefill instance (bootstrap_addr: {failed_bootstrap_addr}), " + f"{len(affected_rooms)} requests affected" + ) + class CommonKVSender(BaseKVSender): def __init__( @@ -614,6 +703,7 @@ class CommonKVSender(BaseKVSender): self._transfer_num_state_indices = 0 # inner state self.curr_idx = 0 + self.init_time: Optional[float] = None if self.kv_mgr.is_dummy_cp_rank: # Non-authoritative CP ranks are dummy participants. self.kv_mgr.update_status(self.bootstrap_room, KVPoll.WaitingForInput) @@ -667,10 +757,10 @@ class CommonKVSender(BaseKVSender): ) def pop_decode_prefix_len(self) -> int: - return 0 + return self.kv_mgr.req_to_decode_prefix_len.pop(self.bootstrap_room, 0) def should_send_kv_chunk(self, num_pages: int, last_chunk: bool) -> bool: - return num_pages > 0 + return num_pages > 0 or last_chunk def get_transfer_metric(self) -> KVTransferMetric: total_bytes = self._transfer_num_kv_indices * self.kv_mgr.kv_item_lens_sum @@ -691,6 +781,36 @@ class CommonKVSender(BaseKVSender): if component_indices is not None: self._transfer_num_state_indices += len(component_indices) + def _prepare_send_indices( + self, + kv_indices: npt.NDArray[np.int32], + state_indices: Optional[List] = None, + ) -> Tuple[npt.NDArray[np.int32], slice, bool, bool]: + """Common pre-processing for send(): index tracking and CP-rank handling. + + Returns: + (kv_indices, index_slice, is_last_chunk, should_skip) + If should_skip is True, the caller should return immediately. + """ + 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 + + 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_chunk: + return kv_indices, index_slice, is_last_chunk, True + else: + self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Success) + return kv_indices, index_slice, is_last_chunk, True + + return kv_indices, index_slice, is_last_chunk, False + def send( self, kv_indices: npt.NDArray[np.int32], @@ -698,6 +818,25 @@ class CommonKVSender(BaseKVSender): ): pass + def _check_bootstrap_timeout(self) -> Optional[KVPoll]: + if self.init_time is None: + return None + elapsed = time.time() - self.init_time + if elapsed < self.kv_mgr.bootstrap_timeout: + return None + 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 " + f"in KVPoll.Bootstrapping", + ) + self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) + return KVPoll.Failed + def poll(self) -> KVPoll: pass @@ -737,6 +876,7 @@ class CommonKVReceiver(BaseKVReceiver): self.kv_mgr = mgr self.conclude_state: Optional[KVPoll] = None self.require_staging: bool = False + self.init_time: Optional[float] = None self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room) self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Bootstrapping) @@ -906,6 +1046,24 @@ class CommonKVReceiver(BaseKVReceiver): ): raise NotImplementedError + def _check_waiting_timeout(self) -> Optional[KVPoll]: + if self.init_time is None: + return None + elapsed = time.time() - self.init_time + if elapsed < self.kv_mgr.waiting_timeout: + return None + 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 " + f"in KVPoll.WaitingForInput", + ) + self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) + return KVPoll.Failed + def failure_exception(self): raise Exception("Fake KVReceiver Exception") diff --git a/python/sglang/srt/disaggregation/common/utils.py b/python/sglang/srt/disaggregation/common/utils.py index 4e5e96c6f..1084b7536 100644 --- a/python/sglang/srt/disaggregation/common/utils.py +++ b/python/sglang/srt/disaggregation/common/utils.py @@ -1,12 +1,27 @@ +import ctypes +import dataclasses import struct import threading from collections import deque -from typing import List, Tuple +from typing import List, Optional, Tuple import numpy as np import numpy.typing as npt +@dataclasses.dataclass +class TransferKVChunk: + """Work unit for KV cache transfer from prefill to decode.""" + + room: int + prefill_kv_indices: npt.NDArray[np.int32] + index_slice: slice + is_last_chunk: bool + prefill_aux_index: Optional[int] + state_indices: Optional[List] + chunk_id: Optional[int] = None + + def pack_list_of_buffers(buffers: List[bytes]) -> bytes: if not buffers: return b"" @@ -59,6 +74,26 @@ class FastQueue: return self._buf.popleft() +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 + + def group_concurrent_contiguous( src_indices: npt.NDArray[np.int32], dst_indices: npt.NDArray[np.int32] ) -> Tuple[List[npt.NDArray[np.int32]], List[npt.NDArray[np.int32]]]: diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index ef7405624..e3951a777 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -1,7 +1,6 @@ from __future__ import annotations import concurrent.futures -import ctypes import dataclasses import logging import os @@ -29,7 +28,9 @@ from sglang.srt.disaggregation.common.staging_handler import ( StagingTransferInfo, ) from sglang.srt.disaggregation.common.utils import ( + AuxDataCodec, FastQueue, + TransferKVChunk, group_concurrent_contiguous, pack_int_lists, unpack_int_lists, @@ -37,10 +38,7 @@ from sglang.srt.disaggregation.common.utils import ( 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, -) +from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine from sglang.srt.environ import envs from sglang.srt.server_args import ServerArgs @@ -64,17 +62,6 @@ class KVTransferError(Exception): 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: slice - is_last_chunk: bool - prefill_aux_index: Optional[int] - state_indices: Optional[List] - - # decode @dataclasses.dataclass class TransferInfo: @@ -162,26 +149,6 @@ class KVArgsRegisterInfo: ) -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" @@ -1478,62 +1445,8 @@ class MooncakeKVManager(CommonKVManager): ) 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() + self._start_heartbeat_checker_thread() def add_transfer_request( self, @@ -1583,6 +1496,13 @@ class MooncakeKVManager(CommonKVManager): def get_session_id(self): return self.engine.get_session_id() + def _on_heartbeat_success(self, bootstrap_addr: str): + 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) + def _run_one_probe_pass(self) -> None: with self.session_lock: snapshot = list(self.failed_sessions) @@ -1666,34 +1586,16 @@ class MooncakeKVSender(CommonKVSender): self.conclude_state = None self.init_time = time.time() - def pop_decode_prefix_len(self) -> int: - return self.kv_mgr.req_to_decode_prefix_len.pop(self.bootstrap_room, 0) - - def should_send_kv_chunk(self, num_pages: int, last_chunk: bool) -> bool: - return num_pages > 0 or last_chunk - def send( self, kv_indices: npt.NDArray[np.int32], state_indices: Optional[List] = 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 - - # 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_chunk: - return - else: - self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Success) - return + kv_indices, index_slice, is_last_chunk, should_skip = ( + self._prepare_send_indices(kv_indices, state_indices) + ) + if should_skip: + return if not is_last_chunk: self.kv_mgr.add_transfer_request( @@ -1719,21 +1621,9 @@ class MooncakeKVSender(CommonKVSender): 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 + timeout_result = self._check_bootstrap_timeout() + if timeout_result is not None: + return timeout_result return status else: @@ -1819,12 +1709,6 @@ class MooncakeKVReceiver(CommonKVReceiver): ] ) - def init( - self, - prefill_dp_rank: int, - ): - super().init(prefill_dp_rank) - def send_metadata( self, kv_indices: npt.NDArray[np.int32], @@ -1874,33 +1758,20 @@ class MooncakeKVReceiver(CommonKVReceiver): 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", - ) - self.conclude_state = KVPoll.Failed - return KVPoll.Failed - - return status - - else: + if self.conclude_state is not None: return self.conclude_state + status = self.kv_mgr.check_status(self.bootstrap_room) + if status in (KVPoll.Success, KVPoll.Failed): + self.conclude_state = status + elif status == KVPoll.WaitingForInput: + timeout_result = self._check_waiting_timeout() + if timeout_result is not None: + return timeout_result + + return status + 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 diff --git a/python/sglang/srt/disaggregation/mori/conn.py b/python/sglang/srt/disaggregation/mori/conn.py index 45bdb6b50..408d0a980 100644 --- a/python/sglang/srt/disaggregation/mori/conn.py +++ b/python/sglang/srt/disaggregation/mori/conn.py @@ -1,6 +1,5 @@ from __future__ import annotations -import ctypes import dataclasses import logging import os @@ -33,11 +32,11 @@ from sglang.srt.disaggregation.common.conn import ( 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.disaggregation.common.utils import ( + AuxDataCodec, + group_concurrent_contiguous, ) +from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.server_args import ServerArgs from sglang.srt.utils.common import get_int_env_var from sglang.srt.utils.network import NetworkAddress, get_local_ip_auto @@ -176,22 +175,6 @@ class KVArgsRegisterInfo: ) -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 @@ -1132,7 +1115,7 @@ class MoriKVManager(CommonKVManager): bootstrap_room: int, kv_indices: npt.NDArray[np.int32], index_slice: slice, - is_last: bool, + is_last_chunk: bool, aux_index: Optional[int] = None, state_indices: Optional[npt.NDArray[np.int32]] = None, ) -> Tuple[List[TransferStatus], Optional[List[TransferInfo]]]: @@ -1163,7 +1146,7 @@ class MoriKVManager(CommonKVManager): self.update_status(bootstrap_room, KVPoll.Failed) return [], list(transfer_infos.values()) targets.append(TransferTarget(info=info, peer_info=peer_info)) - if is_last: + if is_last_chunk: target_infos_snapshot = list(transfer_infos.values()) result_statuses: List[TransferStatus] = [] @@ -1179,7 +1162,7 @@ class MoriKVManager(CommonKVManager): ) if ( - is_last + is_last_chunk and state_indices is not None and not info.is_dummy and self.state_mem_descs @@ -1191,7 +1174,7 @@ class MoriKVManager(CommonKVManager): ) if ( - is_last + is_last_chunk and aux_index is not None and info.dst_aux_index >= 0 and self.pp_group.is_last_rank @@ -1212,7 +1195,7 @@ class MoriKVManager(CommonKVManager): ) return result_statuses, target_infos_snapshot - if is_last: + if is_last_chunk: with self.transfer_lock: # Keep transfer_infos alive until sender.clear() so abort/failure # paths can still recover notification targets after posting. @@ -1243,38 +1226,28 @@ class MoriKVSender(CommonKVSender): kv_indices: npt.NDArray[np.int32], state_indices: Optional[List] = 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 + kv_indices, index_slice, is_last_chunk, should_skip = ( + self._prepare_send_indices(kv_indices, state_indices) + ) + if should_skip: + return - # 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 - - normalized_state = _normalize_state_indices(state_indices) if is_last else None + normalized_state = ( + _normalize_state_indices(state_indices) if is_last_chunk else None + ) 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, + is_last_chunk, + aux_index=self.aux_index if is_last_chunk else None, state_indices=normalized_state, ) self.transfer_statuses.extend(statuses) self._record_transfer_indices(kv_indices, None) if infos is not None: self.pending_infos = infos - if is_last: + if is_last_chunk: self.sent_last_chunk = True self._maybe_finalize_if_room_failed() @@ -1295,15 +1268,9 @@ class MoriKVSender(CommonKVSender): 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) + timeout_result = self._check_bootstrap_timeout() + if timeout_result is not None: + self._finalize_failure() return KVPoll.Failed return status @@ -1499,14 +1466,10 @@ class MoriKVReceiver(CommonKVReceiver): 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 + if status == KVPoll.WaitingForInput: + timeout_result = self._check_waiting_timeout() + if timeout_result is not None: + return timeout_result return status diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 6dc92555d..19f2dd627 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -26,14 +26,12 @@ from sglang.srt.disaggregation.common.conn import ( from sglang.srt.disaggregation.common.staging_handler import StagingRegisterInfo from sglang.srt.disaggregation.common.utils import ( FastQueue, + TransferKVChunk, group_concurrent_contiguous, pack_int_lists, unpack_int_lists, ) -from sglang.srt.disaggregation.utils import ( - DisaggregationMode, - filter_kv_indices_for_cp_rank, -) +from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.environ import envs from sglang.srt.server_args import ServerArgs @@ -104,17 +102,6 @@ class TransferInfo: ) -@dataclasses.dataclass -class TransferKVChunk: - room: int - prefill_kv_indices: npt.NDArray[np.int32] - index_slice: slice - is_last: bool - chunk_id: int - prefill_aux_index: Optional[int] - state_indices: Optional[List] - - @dataclasses.dataclass class KVArgsRegisterInfo: """Contains base pointers and other info which only needs to be sent once by KVReceiver. Received by prefill bootstrap thread.""" @@ -176,7 +163,7 @@ class TransferStatus: 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 chunk count per pp_rank (set when is_last_chunk=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 @@ -186,12 +173,8 @@ class TransferStatus: 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 @@ -209,9 +192,6 @@ class TransferStatus: return False return True - def is_failed(self): - return self.is_failure - class NixlKVManager(CommonKVManager): def __init__( @@ -471,92 +451,6 @@ class NixlKVManager(CommonKVManager): ) self._staging_ctx.prefetched_rooms.add(room) - 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_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 - - 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] - self.prefill_info_table.pop(failed_bootstrap_addr, None) - - possible_affected_rooms = self.addr_to_rooms_tracker.get( - failed_bootstrap_addr, [] - ) - self.addr_to_rooms_tracker.pop(failed_bootstrap_addr, None) - - # 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.get(bootstrap_room, KVPoll.WaitingForInput) @@ -606,7 +500,7 @@ class NixlKVManager(CommonKVManager): # Skip KV RDMA transfer when there are no pages to send # (e.g., decode-side radix cache matched the entire prefix). - # Aux data is still sent below when is_last=True. + # Aux data is still sent below when is_last_chunk=True. if len(kv_chunk.prefill_kv_indices) > 0: chunked_dst_kv_indice = req.dst_kv_indices[kv_chunk.index_slice] @@ -659,7 +553,7 @@ class NixlKVManager(CommonKVManager): if kv_xfer_handle is None: notif = ( f"{req.room}_kv_{kv_chunk.chunk_id}" - f"_{int(kv_chunk.is_last)}_{self.kv_args.engine_rank}" + f"_{int(kv_chunk.is_last_chunk)}_{self.kv_args.engine_rank}" ) if self.is_mla_backend or ( decode_tp_size == self.attn_tp_size @@ -688,7 +582,7 @@ class NixlKVManager(CommonKVManager): handles.append(kv_xfer_handle) - if kv_chunk.is_last: + if kv_chunk.is_last_chunk: dst_info = self.decode_kv_args_table[req.agent_name] if kv_chunk.state_indices: state_xfer_handles = self.maybe_send_extra( @@ -739,7 +633,7 @@ class NixlKVManager(CommonKVManager): break time.sleep(0) - if kv_chunk.is_last: + if kv_chunk.is_last_chunk: self.update_status(room, KVPoll.Success) # Drop per-room state on Success (parity with mooncake # transfer_worker; staging prefetch sets are NIXL-only). @@ -1265,7 +1159,7 @@ class NixlKVManager(CommonKVManager): return (None, True) notif_tag = ( - f"{req.room}_stg_{kv_chunk.chunk_id}_{int(kv_chunk.is_last)}" + f"{req.room}_stg_{kv_chunk.chunk_id}_{int(kv_chunk.is_last_chunk)}" f"_{self.kv_args.engine_rank}_{chunk_idx}" f"_{page_start}_{num_pages}_{req.agent_name}" ) @@ -1574,13 +1468,13 @@ class NixlKVManager(CommonKVManager): bootstrap_room: int, kv_indices: npt.NDArray[np.int32], index_slice: slice, - is_last: bool, + is_last_chunk: bool, chunk_id: int, aux_index: Optional[int] = None, state_indices: Optional[List] = None, ): assert self.disaggregation_mode == DisaggregationMode.PREFILL - assert not is_last or (is_last and aux_index is not None) + assert not is_last_chunk or (is_last_chunk and aux_index is not None) # Prefetch STAGING_REQ to decode before enqueueing so decode has # already allocated staging by the time the worker picks up the @@ -1601,7 +1495,7 @@ class NixlKVManager(CommonKVManager): room=bootstrap_room, prefill_kv_indices=kv_indices, index_slice=index_slice, - is_last=is_last, + is_last_chunk=is_last_chunk, chunk_id=chunk_id, prefill_aux_index=aux_index, state_indices=state_indices, @@ -1628,9 +1522,9 @@ class NixlKVManager(CommonKVManager): tag = components[1] if tag == "kv": chunk_id = int(components[2]) - is_last = bool(int(components[3])) + is_last_chunk = bool(int(components[3])) pp_rank = int(components[4]) if len(components) > 4 else 0 - self._track_kv_arrival(room, chunk_id, is_last, pp_rank) + self._track_kv_arrival(room, chunk_id, is_last_chunk, pp_rank) elif tag == "stg": self._handle_stg_notification(components, room) elif tag == "aux": @@ -1647,13 +1541,13 @@ class NixlKVManager(CommonKVManager): Format: {room}_stg_{chunk_id}_{is_last}_{pp_rank}_{chunk_idx}_{page_start}_{num_pages}_{agent_name} """ chunk_id = int(components[2]) - is_last = bool(int(components[3])) + is_last_chunk = bool(int(components[3])) pp_rank = int(components[4]) chunk_idx = int(components[5]) page_start = int(components[6]) num_pages = int(components[7]) agent_name = components[8] if len(components) > 8 else "" - self._track_kv_arrival(room, chunk_id, is_last, pp_rank) + self._track_kv_arrival(room, chunk_id, is_last_chunk, pp_rank) self._handle_staging_chunk_arrived( room, chunk_idx, page_start, num_pages, agent_name ) @@ -1683,10 +1577,12 @@ class NixlKVManager(CommonKVManager): ): self._maybe_submit_last_scatter(room) - def _track_kv_arrival(self, room: int, chunk_id: int, is_last: bool, pp_rank: int): + def _track_kv_arrival( + self, room: int, chunk_id: int, is_last_chunk: bool, pp_rank: int + ): """Update transfer status tracking for a kv chunk arrival.""" self.transfer_statuses[room].received_kvs_per_pp[pp_rank].add(chunk_id) - if is_last: + if is_last_chunk: self.transfer_statuses[room].expected_kvs_per_pp[pp_rank] = chunk_id + 1 if self.transfer_statuses[room].num_pp_ranks_expected is None: self.transfer_statuses[room].num_pp_ranks_expected = ( @@ -1827,12 +1723,6 @@ class NixlKVSender(CommonKVSender): self._send_error: Optional[Exception] = None self._transfer_start_time: Optional[float] = None - def pop_decode_prefix_len(self) -> int: - return self.kv_mgr.req_to_decode_prefix_len.pop(self.bootstrap_room, 0) - - def should_send_kv_chunk(self, num_pages: int, last_chunk: bool) -> bool: - return num_pages > 0 or last_chunk - def send( self, kv_indices: npt.NDArray[np.int32], @@ -1841,23 +1731,11 @@ class NixlKVSender(CommonKVSender): if self._send_failed: return - 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 + kv_indices, index_slice, is_last_chunk, should_skip = ( + self._prepare_send_indices(kv_indices, state_indices) + ) + if should_skip: + return if self._transfer_start_time is None and ( len(kv_indices) > 0 or state_indices is not None @@ -1868,14 +1746,14 @@ class NixlKVSender(CommonKVSender): self.bootstrap_room, kv_indices, index_slice, - is_last, + is_last_chunk, self.chunk_id, self.aux_index, state_indices, ) self._record_transfer_indices(kv_indices, state_indices) self.chunk_id += 1 - if is_last: + if is_last_chunk: self.has_sent = True def poll(self) -> KVPoll: @@ -1892,9 +1770,6 @@ class NixlKVSender(CommonKVSender): ) return status - def clear(self): - super().clear() - def failure_exception(self): if self._send_error is not None: raise self._send_error @@ -1915,12 +1790,6 @@ class NixlKVReceiver(CommonKVReceiver): super().__init__(mgr, bootstrap_addr, bootstrap_room) self.init_time = None - def init( - self, - prefill_dp_rank: int, - ): - super().init(prefill_dp_rank) - def send_metadata( self, kv_indices: npt.NDArray[np.int32], @@ -1997,31 +1866,16 @@ class NixlKVReceiver(CommonKVReceiver): if not self.started_transfer: return status - 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 + timeout_result = self._check_waiting_timeout() + if timeout_result is not None: + return timeout_result 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 + 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