From 62c2e091f6ba7181d69b660e57280ec98c0e55ad Mon Sep 17 00:00:00 2001 From: Niko Ma Date: Sat, 9 May 2026 07:07:22 +0800 Subject: [PATCH] [PD] MORI-IO: Add state transfer, inline transfer model, and high-concurrency fixes (#22665) --- python/sglang/srt/disaggregation/mori/conn.py | 747 ++++++++++++++---- .../test_mori_transfer_engine_e2e.py | 179 +++++ 2 files changed, 766 insertions(+), 160 deletions(-) create mode 100644 test/registered/amd/disaggregation/test_mori_transfer_engine_e2e.py diff --git a/python/sglang/srt/disaggregation/mori/conn.py b/python/sglang/srt/disaggregation/mori/conn.py index 4db7cb0cd..523b3d823 100644 --- a/python/sglang/srt/disaggregation/mori/conn.py +++ b/python/sglang/srt/disaggregation/mori/conn.py @@ -7,11 +7,13 @@ import os import struct import threading import time +import uuid from typing import Dict, List, Optional, Tuple import msgspec import numpy as np import numpy.typing as npt +import zmq from mori.cpp import TransferStatus from mori.io import ( BackendType, @@ -44,6 +46,14 @@ logger = logging.getLogger(__name__) MORI_GUARD = b"MoriMsgGuard" +def _normalize_state_indices( + state_indices, +) -> Optional[npt.NDArray[np.int32]]: + if state_indices is None: + return None + return np.asarray(state_indices, dtype=np.int32) + + def _pack_mem_desc_list(mems: List[MemoryDesc]) -> bytes: if not mems: return b"" @@ -66,6 +76,7 @@ class TransferInfo: engine_key: str dst_kv_indices: npt.NDArray[np.int32] dst_aux_index: int + dst_state_indices: npt.NDArray[np.int32] required_dst_info_num: int is_dummy: bool @@ -86,6 +97,11 @@ class TransferInfo: else: dst_aux_index = -1 + if len(payload) > 6 and payload[6]: + dst_state_indices = np.frombuffer(payload[6], dtype=np.int32) + else: + dst_state_indices = np.array([], dtype=np.int32) + required_dst_info_num = ( int(payload[7].decode("ascii")) if len(payload) > 7 else 1 ) @@ -97,6 +113,7 @@ class TransferInfo: engine_key=engine_key, dst_kv_indices=dst_kv_indices, dst_aux_index=dst_aux_index, + dst_state_indices=dst_state_indices, required_dst_info_num=required_dst_info_num, is_dummy=is_dummy, ) @@ -114,6 +131,8 @@ class KVArgsRegisterInfo: decode_tp_size: int decode_tp_rank: int dst_kv_item_len: int + dst_state_item_lens: List[int] + dst_state_dim_per_tensor: List[int] @property def engine_key(self) -> str: @@ -131,6 +150,16 @@ class KVArgsRegisterInfo: decode_tp_size = int(payload[8].decode("ascii")) decode_tp_rank = int(payload[9].decode("ascii")) dst_kv_item_len = int(payload[10].decode("ascii")) + dst_state_item_lens = ( + list(struct.unpack(f"{len(payload[11]) // 4}I", payload[11])) + if len(payload) > 11 and len(payload[11]) > 0 + else [] + ) + dst_state_dim_per_tensor = ( + list(struct.unpack(f"{len(payload[12]) // 4}I", payload[12])) + if len(payload) > 12 and len(payload[12]) > 0 + else [] + ) return cls( endpoint=endpoint, dst_port=dst_port, @@ -142,6 +171,8 @@ class KVArgsRegisterInfo: decode_tp_size=decode_tp_size, decode_tp_rank=decode_tp_rank, dst_kv_item_len=dst_kv_item_len, + dst_state_item_lens=dst_state_item_lens, + dst_state_dim_per_tensor=dst_state_dim_per_tensor, ) @@ -173,6 +204,48 @@ class TPSliceConfig: heads_bytes_per_token_to_send: int +@dataclasses.dataclass(frozen=True) +class GroupedIndexPlan: + src_starts: List[int] + dst_starts: List[int] + counts: List[int] + + @classmethod + def from_groups( + cls, src_groups: List[List[int]], dst_groups: List[List[int]] + ) -> GroupedIndexPlan: + if len(src_groups) != len(dst_groups): + raise ValueError("Source and destination groups must have the same length") + return cls( + src_starts=[int(group[0]) for group in src_groups], + dst_starts=[int(group[0]) for group in dst_groups], + counts=[len(group) for group in src_groups], + ) + + def materialize(self, item_len: int) -> BatchTransferPlan: + return BatchTransferPlan( + local_offsets=[start * item_len for start in self.src_starts], + remote_offsets=[start * item_len for start in self.dst_starts], + sizes=[count * item_len for count in self.counts], + ) + + +@dataclasses.dataclass(frozen=True) +class BatchTransferPlan: + local_offsets: List[int] + remote_offsets: List[int] + sizes: List[int] + + def empty(self) -> bool: + return not self.sizes + + +@dataclasses.dataclass(frozen=True) +class TransferTarget: + info: TransferInfo + peer_info: KVArgsRegisterInfo + + class MoriKVManager(CommonKVManager): AUX_DATA_HEADER = b"AUX_DATA" @@ -190,6 +263,13 @@ class MoriKVManager(CommonKVManager): self.aux_mem_descs: List[MemoryDesc] = [] self.state_mem_descs: List[MemoryDesc] = [] self.transfer_lock = threading.Lock() + self._zmq_ctx = zmq.Context() + self._socket_local = threading.local() + # Send CPU-resident AUX data via RDMA instead of ZMQ TCP. + # Default: TCP. Set SGLANG_MORI_SEND_AUX_RDMA=1 to use RDMA. + self._send_aux_rdma = os.environ.get( + "SGLANG_MORI_SEND_AUX_RDMA", "" + ).lower() in ("1", "true") self._register_local_buffers() if self.disaggregation_mode == DisaggregationMode.PREFILL: self._start_bootstrap_thread() @@ -207,7 +287,8 @@ class MoriKVManager(CommonKVManager): engine_key = ( f"io-{self.disaggregation_mode.value}-" f"dp{self.system_dp_rank}-tp{self.attn_tp_rank}-" - f"pid{os.getpid()}-{self.local_ip}" + f"pid{os.getpid()}-{self.local_ip}-" + f"{uuid.uuid4().hex[:8]}" ) engine = IOEngine(engine_key, config) @@ -215,8 +296,8 @@ class MoriKVManager(CommonKVManager): # Number of RDMA Queue Pairs (QPs) used per transfer operation. # Higher values can increase parallelism and bandwidth utilization. - # Default: 1 - qp_per_transfer = get_int_env_var("SGLANG_MORI_QP_PER_TRANSFER", 1) + # Default: 4 + qp_per_transfer = get_int_env_var("SGLANG_MORI_QP_PER_TRANSFER", 4) # Number of RDMA work requests posted in a single batch to each QP. # Larger batch sizes reduce per-operation overhead and improve throughput @@ -229,8 +310,8 @@ class MoriKVManager(CommonKVManager): # Each worker handles RDMA operations on a separate CPU core (with affinity). # More workers can improve parallelism for large batch transfers across # multiple QPs, but excessive threads may cause contention. - # Default: 1 - num_worker_threads = get_int_env_var("SGLANG_MORI_NUM_WORKERS", 1) + # Default: 4 + num_worker_threads = get_int_env_var("SGLANG_MORI_NUM_WORKERS", 4) rdma_cfg = RdmaBackendConfig( qp_per_transfer, @@ -281,6 +362,41 @@ class MoriKVManager(CommonKVManager): ) self.state_mem_descs.append(desc) + def update_status(self, bootstrap_room: int, status: KVPoll): + current = self.request_status.get(bootstrap_room) + if current is None: + # Room not yet created or already cleared. + # Only allow initial creation: Bootstrapping (normal) or + # WaitingForInput (dummy CP rank, see CommonKVSender.__init__). + if status not in (KVPoll.Bootstrapping, KVPoll.WaitingForInput): + return + elif current == KVPoll.Failed and status != KVPoll.Failed: + # Failed is terminal — never overwrite with non-Failed. + return + super().update_status(bootstrap_room, status) + + def _connect_threadsafe(self, endpoint: str, is_ipv6: bool = False): + """Thread-local ZMQ socket cache with shared Context. + + Each worker thread gets its own PUSH socket (ZMQ sockets are not + thread-safe), but all sockets share a single process-level Context + to avoid creating excessive I/O threads and TCP connections. + """ + cache = getattr(self._socket_local, "socket_cache", None) + if cache is None: + cache = {} + self._socket_local.socket_cache = cache + if endpoint not in cache: + sock = self._zmq_ctx.socket(zmq.PUSH) + sock.setsockopt(zmq.SNDHWM, 0) + sock.setsockopt(zmq.SNDTIMEO, 5000) + sock.setsockopt(zmq.LINGER, 0) + if is_ipv6: + sock.setsockopt(zmq.IPV6, 1) + sock.connect(endpoint) + cache[endpoint] = sock + return cache[endpoint] + def _handle_register_message(self, payload: List[bytes]) -> None: try: register_info = KVArgsRegisterInfo.from_zmq(payload) @@ -291,16 +407,30 @@ class MoriKVManager(CommonKVManager): def _handle_transfer_message(self, payload: List[bytes]) -> None: try: transfer_info = TransferInfo.from_zmq(payload) - infos = self.transfer_infos.setdefault(transfer_info.room, {}) - infos[transfer_info.engine_key] = transfer_info + with self.transfer_lock: + # Accept metadata when room is not yet created (None) or + # in Bootstrapping. Reject for active/terminal states where + # the worker may already be using transfer_infos. + # None is allowed because metadata can arrive from decode + # before the prefill scheduler creates the MoriKVSender. + current = self.request_status.get(transfer_info.room) + if current is not None and current != KVPoll.Bootstrapping: + logger.debug( + "Ignoring stale transfer info for room %s (status=%s)", + transfer_info.room, + current, + ) + return + infos = self.transfer_infos.setdefault(transfer_info.room, {}) + infos[transfer_info.engine_key] = transfer_info - if len(infos) >= transfer_info.required_dst_info_num: - logger.debug( - "Bootstrap room %s got enough transfer info (%s)", - transfer_info.room, - len(infos), - ) - self.update_status(transfer_info.room, KVPoll.WaitingForInput) + if len(infos) >= transfer_info.required_dst_info_num: + logger.debug( + "Bootstrap room %s got enough transfer info (%s)", + transfer_info.room, + len(infos), + ) + self.update_status(transfer_info.room, KVPoll.WaitingForInput) except Exception: logger.exception("Failed to parse transfer info message") @@ -395,6 +525,16 @@ class MoriKVManager(CommonKVManager): threading.Thread(target=decode_worker, daemon=True).start() + def _compute_prefill_unique_rank(self) -> int: + """Unique id per prefill sender, encoding TP/PP/CP ranks. + Must match Mooncake's formula so decode's response set size matches + expected_response_num when multiple CP ranks participate.""" + return ( + self.attn_tp_rank * (self.pp_size * self.attn_cp_size) + + self.pp_rank * self.attn_cp_size + + self.attn_cp_rank + ) + def notify_decode_status( self, infos: List[TransferInfo], @@ -408,13 +548,13 @@ class MoriKVManager(CommonKVManager): MORI_GUARD, str(bootstrap_room).encode("ascii"), str(int(status)).encode("ascii"), - str(self.attn_tp_rank * self.pp_size + self.pp_rank).encode("ascii"), + str(self._compute_prefill_unique_rank()).encode("ascii"), failure_reason.encode("utf-8") if failure_reason else b"", ] for info in infos: try: na = NetworkAddress(info.endpoint, info.dst_port) - socket = self._connect(na.to_tcp(), is_ipv6=na.is_ipv6) + socket = self._connect_threadsafe(na.to_tcp(), is_ipv6=na.is_ipv6) socket.send_multipart(payload) except Exception: logger.exception( @@ -479,32 +619,33 @@ class MoriKVManager(CommonKVManager): dst_slice = dst_mem_descs[start_layer:end_layer] return src_descs, dst_slice, num_local_layers - def _issue_layer_transfers( + def _submit_batch_transfer_plan( self, src_desc: MemoryDesc, dst_desc: MemoryDesc, - kv_item_len: int, - src_groups: List[List[int]], - dst_groups: List[List[int]], + plan: BatchTransferPlan, ) -> List[TransferStatus]: - if not src_groups: + if plan.empty(): return [] - local_offsets = [int(src_group[0]) * kv_item_len for src_group in src_groups] - remote_offsets = [int(dst_group[0]) * kv_item_len for dst_group in dst_groups] - sizes = [len(src_group) * kv_item_len for src_group in src_groups] transfer_uid = self.engine.allocate_transfer_uid() statuses = self.engine.batch_write( [src_desc], - [local_offsets], + [plan.local_offsets], [dst_desc], - [remote_offsets], - [sizes], + [plan.remote_offsets], + [plan.sizes], [transfer_uid], ) return statuses + def _build_contiguous_transfer_plan( + self, grouped_plan: GroupedIndexPlan, item_len: int + ) -> BatchTransferPlan: + # Reuse grouped indices across all layers/tensors that share the same item length. + return grouped_plan.materialize(item_len) + def _build_tp_slice_config(self, peer_info: KVArgsRegisterInfo) -> TPSliceConfig: page_size = self.kv_args.page_size @@ -517,23 +658,27 @@ class MoriKVManager(CommonKVManager): prefill_tp_size = self.attn_tp_size decode_tp_size = peer_info.decode_tp_size - num_kv_heads = self.kv_args.kv_head_num - src_heads_per_rank = num_kv_heads - dst_heads_per_rank = num_kv_heads * prefill_tp_size // decode_tp_size - if dst_heads_per_rank == 0: - raise ValueError("Destination heads per rank evaluates to zero") + total_kv_heads = getattr(self.kv_args, "total_kv_head_num", 0) + if total_kv_heads <= 0: + total_kv_heads = self.kv_args.kv_head_num * prefill_tp_size + + src_heads_per_rank = max(1, total_kv_heads // prefill_tp_size) + dst_heads_per_rank = max(1, total_kv_heads // decode_tp_size) bytes_per_head_slice = bytes_per_token_dst // dst_heads_per_rank if bytes_per_head_slice == 0: raise ValueError("Head slice size evaluates to zero") + src_replication = max(1, prefill_tp_size // total_kv_heads) + local_tp_rank = self.kv_args.engine_rank % prefill_tp_size dst_tp_rank = peer_info.decode_tp_rank % decode_tp_size if prefill_tp_size > decode_tp_size: src_head_start = 0 num_heads_to_send = src_heads_per_rank - dst_head_start = local_tp_rank * src_heads_per_rank + unique_head_idx = local_tp_rank // src_replication + dst_head_start = (unique_head_idx * src_heads_per_rank) % dst_heads_per_rank else: src_head_start = (dst_tp_rank * dst_heads_per_rank) % src_heads_per_rank num_heads_to_send = dst_heads_per_rank @@ -559,20 +704,18 @@ class MoriKVManager(CommonKVManager): heads_bytes_per_token_to_send=heads_bytes_per_token, ) - def _issue_tp_slice_transfers( + def _build_tp_slice_transfer_plan( self, - src_desc: MemoryDesc, - dst_desc: MemoryDesc, kv_indices: npt.NDArray[np.int32], dst_indices: npt.NDArray[np.int32], tp_cfg: TPSliceConfig, - ) -> List[TransferStatus]: + ) -> BatchTransferPlan: if kv_indices.size == 0 or dst_indices.size == 0: - return [] + return BatchTransferPlan([], [], []) limit = min(kv_indices.size, dst_indices.size) if not limit: - return [] + return BatchTransferPlan([], [], []) src_pages = kv_indices[:limit].astype(np.int64) dst_pages = dst_indices[:limit].astype(np.int64) @@ -607,18 +750,13 @@ class MoriKVManager(CommonKVManager): sizes = [tp_cfg.heads_bytes_per_token_to_send] * num_transfers if not local_offsets: - return [] + return BatchTransferPlan([], [], []) - transfer_uid = self.engine.allocate_transfer_uid() - statuses = self.engine.batch_write( - [src_desc], - [local_offsets], - [dst_desc], - [remote_offsets], - [sizes], - [transfer_uid], + return BatchTransferPlan( + local_offsets=local_offsets, + remote_offsets=remote_offsets, + sizes=sizes, ) - return statuses def send_kvcache( self, @@ -626,82 +764,76 @@ class MoriKVManager(CommonKVManager): prefill_kv_indices: npt.NDArray[np.int32], dst_kv_indices: npt.NDArray[np.int32], ) -> List[TransferStatus]: - src_groups, dst_groups = group_concurrent_contiguous( - prefill_kv_indices, dst_kv_indices + grouped_plan = GroupedIndexPlan.from_groups( + *group_concurrent_contiguous( + prefill_kv_indices, + dst_kv_indices, + ) ) - statuses = [] + statuses: List[TransferStatus] = [] kv_item_len = self.kv_args.kv_item_lens[0] + if self.is_mla_backend: - ( - src_descs, - dst_descs, - layers_current_pp_stage, - ) = self._get_mla_mem_desc_slices(peer_info.dst_kv_mem_descs) + layer_plan = self._build_contiguous_transfer_plan(grouped_plan, kv_item_len) + src_descs, dst_descs, layers_current_pp_stage = ( + self._get_mla_mem_desc_slices(peer_info.dst_kv_mem_descs) + ) for layer_id in range(layers_current_pp_stage): statuses.extend( - self._issue_layer_transfers( + self._submit_batch_transfer_plan( src_descs[layer_id], dst_descs[layer_id], - kv_item_len, - src_groups, - dst_groups, + layer_plan, ) ) - else: - tp_mismatch = peer_info.decode_tp_size != self.attn_tp_size - ( - src_k_descs, - src_v_descs, - dst_k_descs, - dst_v_descs, - layers_current_pp_stage, - ) = self._get_mha_mem_desc_slices(peer_info.dst_kv_mem_descs) + return statuses - if tp_mismatch: - tp_cfg = self._build_tp_slice_config(peer_info) - for layer_id in range(layers_current_pp_stage): - statuses.extend( - self._issue_tp_slice_transfers( - src_k_descs[layer_id], - dst_k_descs[layer_id], - prefill_kv_indices, - dst_kv_indices, - tp_cfg, - ) + ( + src_k_descs, + src_v_descs, + dst_k_descs, + dst_v_descs, + layers_current_pp_stage, + ) = self._get_mha_mem_desc_slices(peer_info.dst_kv_mem_descs) + + if peer_info.decode_tp_size != self.attn_tp_size: + tp_cfg = self._build_tp_slice_config(peer_info) + slice_plan = self._build_tp_slice_transfer_plan( + prefill_kv_indices, dst_kv_indices, tp_cfg + ) + for layer_id in range(layers_current_pp_stage): + statuses.extend( + self._submit_batch_transfer_plan( + src_k_descs[layer_id], + dst_k_descs[layer_id], + slice_plan, ) - statuses.extend( - self._issue_tp_slice_transfers( - src_v_descs[layer_id], - dst_v_descs[layer_id], - prefill_kv_indices, - dst_kv_indices, - tp_cfg, - ) - ) - else: - src_groups, dst_groups = group_concurrent_contiguous( - prefill_kv_indices, dst_kv_indices ) - for layer_id in range(layers_current_pp_stage): - statuses.extend( - self._issue_layer_transfers( - src_k_descs[layer_id], - dst_k_descs[layer_id], - kv_item_len, - src_groups, - dst_groups, - ) - ) - statuses.extend( - self._issue_layer_transfers( - src_v_descs[layer_id], - dst_v_descs[layer_id], - kv_item_len, - src_groups, - dst_groups, - ) + statuses.extend( + self._submit_batch_transfer_plan( + src_v_descs[layer_id], + dst_v_descs[layer_id], + slice_plan, ) + ) + return statuses + layer_plan = self._build_contiguous_transfer_plan(grouped_plan, kv_item_len) + for layer_id in range(layers_current_pp_stage): + statuses.extend( + self._submit_batch_transfer_plan( + src_k_descs[layer_id], + dst_k_descs[layer_id], + layer_plan, + ) + ) + statuses.extend( + self._submit_batch_transfer_plan( + src_v_descs[layer_id], + dst_v_descs[layer_id], + layer_plan, + ) + ) return statuses def send_aux( @@ -711,8 +843,42 @@ class MoriKVManager(CommonKVManager): dst_aux_index: int, room: int, ) -> List[TransferStatus]: + if self._send_aux_rdma: + return self.send_aux_rdma(peer_info, prefill_aux_index, dst_aux_index, room) return self.send_aux_tcp(peer_info, prefill_aux_index, dst_aux_index, room) + def send_aux_rdma( + self, + peer_info: KVArgsRegisterInfo, + prefill_aux_index: int, + dst_aux_index: int, + room: int, + ) -> List[TransferStatus]: + if not self.aux_mem_descs or len(self.aux_mem_descs) != len( + peer_info.dst_aux_mem_descs + ): + return self.send_aux_tcp(peer_info, prefill_aux_index, dst_aux_index, room) + + src_descs: List[MemoryDesc] = [] + dst_descs: List[MemoryDesc] = [] + local_offsets: List[List[int]] = [] + remote_offsets: List[List[int]] = [] + sizes: List[List[int]] = [] + uids = [] + for i in range(len(self.aux_mem_descs)): + item_len = self.kv_args.aux_item_lens[i] + src_descs.append(self.aux_mem_descs[i]) + dst_descs.append(peer_info.dst_aux_mem_descs[i]) + local_offsets.append([prefill_aux_index * item_len]) + remote_offsets.append([dst_aux_index * item_len]) + sizes.append([item_len]) + uids.append(self.engine.allocate_transfer_uid()) + return list( + self.engine.batch_write( + src_descs, local_offsets, dst_descs, remote_offsets, sizes, uids + ) + ) + def send_aux_tcp( self, peer_info: KVArgsRegisterInfo, @@ -720,15 +886,11 @@ class MoriKVManager(CommonKVManager): dst_aux_index: int, room: int, ) -> List[TransferStatus]: - prefill_aux_ptrs = self.kv_args.aux_data_ptrs - prefill_aux_item_lens = self.kv_args.aux_item_lens - - for i in range(len(prefill_aux_ptrs)): - length = prefill_aux_item_lens[i] - src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index + for i in range(len(self.kv_args.aux_data_ptrs)): + length = self.kv_args.aux_item_lens[i] + src_addr = self.kv_args.aux_data_ptrs[i] + length * prefill_aux_index data = AuxDataCodec.serialize_data_from_buffer(src_addr, length) - - self.send_aux_data_to_endpoint( + self._send_aux_data_to_endpoint( remote=peer_info.endpoint, dst_port=peer_info.dst_port, room=room, @@ -736,21 +898,13 @@ class MoriKVManager(CommonKVManager): aux_index=dst_aux_index, data=data, ) + return [] # TCP path has no TransferStatus to poll - return [] - - def send_aux_data_to_endpoint( - self, - remote: str, - dst_port: int, - room: int, - buffer_index: int, - aux_index: int, - data: bytes, + def _send_aux_data_to_endpoint( + self, remote, dst_port, room, buffer_index, aux_index, data ): na = NetworkAddress(remote, dst_port) - socket = self._connect(na.to_tcp(), is_ipv6=na.is_ipv6) - + socket = self._connect_threadsafe(na.to_tcp(), is_ipv6=na.is_ipv6) socket.send_multipart( [ MoriKVManager.AUX_DATA_HEADER, @@ -762,8 +916,201 @@ class MoriKVManager(CommonKVManager): ] ) + def send_state( + self, + peer_info: KVArgsRegisterInfo, + src_state_indices: npt.NDArray[np.int32], + dst_state_indices: npt.NDArray[np.int32], + ) -> List[TransferStatus]: + # Guard: no local state tensors -> no-op (e.g. SWA layers=0 on this PP rank) + if not self.state_mem_descs: + return [] + + state_type = getattr(self.kv_args, "state_type", "none") + + if state_type == "none": + raise RuntimeError( + "PD state transfer failed: state_type is 'none' but state_indices were provided" + ) + + if not peer_info.dst_state_mem_descs: + raise RuntimeError( + f"PD state transfer failed: remote peer has no state descriptors " + f"(state_type={state_type}, prefill_tp_size={self.attn_tp_size}, " + f"decode_tp_size={peer_info.decode_tp_size})" + ) + + if len(peer_info.dst_state_mem_descs) != len(self.state_mem_descs): + raise RuntimeError( + f"PD state transfer failed: state descriptor count mismatch " + f"(local={len(self.state_mem_descs)}, remote={len(peer_info.dst_state_mem_descs)}), " + f"likely PP configuration mismatch (state_type={state_type})" + ) + + if len(self.kv_args.state_item_lens) != len(self.state_mem_descs): + raise RuntimeError( + f"PD state transfer failed: local state_item_lens count " + f"({len(self.kv_args.state_item_lens)}) does not match state descriptor " + f"count ({len(self.state_mem_descs)}) (state_type={state_type})" + ) + + if state_type == "mamba": + return self._send_mamba_state( + peer_info, src_state_indices, dst_state_indices + ) + elif state_type in ("swa", "nsa"): + return self._send_swa_nsa_state( + peer_info, src_state_indices, dst_state_indices, state_type + ) + else: + raise RuntimeError( + f"PD state transfer failed: unknown state_type={state_type}" + ) + + def _send_mamba_state( + self, + peer_info: KVArgsRegisterInfo, + src_state_indices: npt.NDArray[np.int32], + dst_state_indices: npt.NDArray[np.int32], + ) -> List[TransferStatus]: + if len(src_state_indices) != 1 or len(dst_state_indices) != 1: + raise RuntimeError( + f"PD state transfer failed: mamba requires single state index, " + f"got src={len(src_state_indices)}, dst={len(dst_state_indices)}" + ) + + tp_mismatch = peer_info.decode_tp_size != self.attn_tp_size + src_state_dim_per_tensor = getattr(self.kv_args, "state_dim_per_tensor", []) + dst_state_dim_per_tensor = peer_info.dst_state_dim_per_tensor + + # If dim info missing, silently degrade to whole-item copy (Mooncake compat) + if tp_mismatch and ( + not src_state_dim_per_tensor or not dst_state_dim_per_tensor + ): + tp_mismatch = False + + if tp_mismatch: + logger.warning_once( + "Using Mamba state slice transfer for different TP sizes between prefill and decode. " + f"Prefill attn_tp_size={self.attn_tp_size}, Decode attn_tp_size={peer_info.decode_tp_size}. " + "Performance may be affected." + ) + + src_idx = int(src_state_indices[0]) + dst_idx = int(dst_state_indices[0]) + statuses = [] + + local_tp_rank = self.kv_args.engine_rank % self.attn_tp_size + dst_tp_rank = peer_info.decode_tp_rank % peer_info.decode_tp_size + + for i in range(len(self.state_mem_descs)): + src_desc = self.state_mem_descs[i] + dst_desc = peer_info.dst_state_mem_descs[i] + src_item_len = self.kv_args.state_item_lens[i] + + if not tp_mismatch: + # same-TP: whole item copy + src_offset = src_idx * src_item_len + dst_offset = dst_idx * src_item_len + size = src_item_len + else: + # TP mismatch slice copy + dst_item_len = peer_info.dst_state_item_lens[i] + src_dim = src_state_dim_per_tensor[i] + dst_dim = dst_state_dim_per_tensor[i] + + src_bytes_per_dim = src_item_len // src_dim + + if self.attn_tp_size > peer_info.decode_tp_size: + src_dim_start = 0 + num_dims_to_send = src_dim + writers_per_decode = self.attn_tp_size // peer_info.decode_tp_size + local_writer_idx = local_tp_rank % writers_per_decode + dst_dim_start = local_writer_idx * src_dim + else: + src_dim_start = (dst_tp_rank * dst_dim) % src_dim + num_dims_to_send = dst_dim + dst_dim_start = 0 + + dst_bytes_per_dim = dst_item_len // dst_dim + src_dim_offset = src_dim_start * src_bytes_per_dim + dst_dim_offset = dst_dim_start * dst_bytes_per_dim + bytes_to_send = num_dims_to_send * src_bytes_per_dim + + src_offset = src_idx * src_item_len + src_dim_offset + dst_offset = dst_idx * dst_item_len + dst_dim_offset + size = bytes_to_send + + transfer_uid = self.engine.allocate_transfer_uid() + batch_statuses = self.engine.batch_write( + [src_desc], + [[src_offset]], + [dst_desc], + [[dst_offset]], + [[size]], + [transfer_uid], + ) + statuses.extend(batch_statuses) + + return statuses + + def _send_swa_nsa_state( + self, + peer_info: KVArgsRegisterInfo, + src_state_indices: npt.NDArray[np.int32], + dst_state_indices: npt.NDArray[np.int32], + state_type: str, + ) -> List[TransferStatus]: + # TP mismatch check for non-MLA SWA + if ( + state_type == "swa" + and not self.is_mla_backend + and peer_info.decode_tp_size != self.attn_tp_size + ): + raise RuntimeError( + f"PD state transfer does not support TP-mismatched non-MLA SWA models " + f"(prefill_tp_size={self.attn_tp_size}, decode_tp_size={peer_info.decode_tp_size})" + ) + + common_len = min(len(src_state_indices), len(dst_state_indices)) + if common_len == 0 and max(len(src_state_indices), len(dst_state_indices)) > 0: + raise RuntimeError( + f"No overlapping state indices for state_type={state_type}" + ) + if len(src_state_indices) != len(dst_state_indices): + logger.warning( + "State index length mismatch for %s: src=%d dst=%d; truncating to common prefix=%d", + state_type, + len(src_state_indices), + len(dst_state_indices), + common_len, + ) + src_state_indices = src_state_indices[:common_len] + dst_state_indices = dst_state_indices[:common_len] + + # Group contiguous indices and issue per-tensor transfers + grouped_plan = GroupedIndexPlan.from_groups( + *group_concurrent_contiguous(src_state_indices, dst_state_indices) + ) + + statuses = [] + for i in range(len(self.state_mem_descs)): + src_desc = self.state_mem_descs[i] + dst_desc = peer_info.dst_state_mem_descs[i] + state_item_len = self.kv_args.state_item_lens[i] + + statuses.extend( + self._submit_batch_transfer_plan( + src_desc, + dst_desc, + self._build_contiguous_transfer_plan(grouped_plan, state_item_len), + ) + ) + + return statuses + def _handle_aux_data(self, msg: List[bytes]): - """Handle AUX_DATA messages received by the decode thread.""" + """Handle AUX_DATA messages received by the decode thread (legacy TCP path).""" room = int(msg[1].decode("ascii")) buffer_index = int(msg[2].decode("ascii")) aux_index = int(msg[3].decode("ascii")) @@ -778,10 +1125,6 @@ class MoriKVManager(CommonKVManager): self.kv_args, buffer_index, aux_index, data ) - logger.debug( - f"Received AUX_DATA for bootstrap_room {room} with length:{len(data)}" - ) - def add_transfer_request( self, bootstrap_room: int, @@ -792,31 +1135,59 @@ class MoriKVManager(CommonKVManager): state_indices: Optional[npt.NDArray[np.int32]] = None, ) -> Tuple[List[TransferStatus], Optional[List[TransferInfo]]]: assert self.disaggregation_mode == DisaggregationMode.PREFILL - transfer_infos = self.transfer_infos.get(bootstrap_room) - if not transfer_infos: - raise RuntimeError( - f"No transfer info found for bootstrap_room={bootstrap_room}" - ) - result_statuses = [] + + if ( + bootstrap_room not in self.request_status + or self.request_status.get(bootstrap_room) == KVPoll.Failed + ): + return [], None + + targets: List[TransferTarget] = [] target_infos_snapshot: Optional[List[TransferInfo]] = None with self.transfer_lock: + transfer_infos = self.transfer_infos.get(bootstrap_room) + if not transfer_infos: + reason = f"No transfer info found for bootstrap_room={bootstrap_room}" + self.record_failure(bootstrap_room, reason) + self.update_status(bootstrap_room, KVPoll.Failed) + return [], None + self.update_status(bootstrap_room, KVPoll.Transferring) for info in transfer_infos.values(): peer_info = self.decode_kv_args_table.get(info.engine_key) if not peer_info: - self.record_failure( - bootstrap_room, - f"Peer info missing for engine {info.engine_key}", - ) - raise RuntimeError( - f"Missing decode peer info for {info.engine_key}" - ) + reason = f"Peer info missing for engine {info.engine_key}" + self.record_failure(bootstrap_room, reason) + self.update_status(bootstrap_room, KVPoll.Failed) + return [], list(transfer_infos.values()) + targets.append(TransferTarget(info=info, peer_info=peer_info)) + if is_last: + target_infos_snapshot = list(transfer_infos.values()) + + result_statuses: List[TransferStatus] = [] + try: + for target in targets: + info = target.info + peer_info = target.peer_info + if not info.is_dummy: dst_indices_chunk = info.dst_kv_indices[index_slice] - statuses = self.send_kvcache( - peer_info, kv_indices, dst_indices_chunk + result_statuses.extend( + self.send_kvcache(peer_info, kv_indices, dst_indices_chunk) ) - result_statuses.extend(statuses) + + if ( + is_last + and state_indices is not None + and not info.is_dummy + and self.state_mem_descs + ): + result_statuses.extend( + self.send_state( + peer_info, state_indices, info.dst_state_indices + ) + ) + if ( is_last and aux_index is not None @@ -828,10 +1199,23 @@ class MoriKVManager(CommonKVManager): peer_info, aux_index, info.dst_aux_index, bootstrap_room ) ) - if is_last: + except Exception as e: + reason = f"Transfer submission failed: {e}" + with self.transfer_lock: + self.record_failure(bootstrap_room, reason) + self.update_status(bootstrap_room, KVPoll.Failed) + logger.exception( + "Mori KV transfer submission failed for bootstrap_room=%s", + bootstrap_room, + ) + return result_statuses, target_infos_snapshot + + if is_last: + with self.transfer_lock: + # Keep transfer_infos alive until sender.clear() so abort/failure + # paths can still recover notification targets after posting. self.update_status(bootstrap_room, KVPoll.Success) - target_infos_snapshot = list(transfer_infos.values()) - self.transfer_infos.pop(bootstrap_room, None) + return result_statuses, target_infos_snapshot @@ -874,24 +1258,40 @@ class MoriKVSender(CommonKVSender): else: self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Success) return + + normalized_state = _normalize_state_indices(state_indices) if is_last 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, + 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 - self.sent_last_chunk = True + if is_last: + self.sent_last_chunk = True + self._maybe_finalize_if_room_failed() + + def _maybe_finalize_if_room_failed(self) -> None: + if self.conclude_state is not None: + return + if self.kv_mgr.request_status.get(self.bootstrap_room) == KVPoll.Failed: + self._finalize_failure() def poll(self) -> KVPoll: if self.conclude_state is not None: return self.conclude_state + if self.bootstrap_room not in self.kv_mgr.request_status: + self._finalize_failure() + return KVPoll.Failed + 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: @@ -909,6 +1309,10 @@ class MoriKVSender(CommonKVSender): self._finalize_failure() return KVPoll.Failed + if status == KVPoll.Success and self.kv_mgr.is_dummy_cp_rank: + self.conclude_state = KVPoll.Success + return KVPoll.Success + transfers_done = self._all_transfers_finished() if transfers_done: if self._has_transfer_error(): @@ -943,9 +1347,16 @@ class MoriKVSender(CommonKVSender): ) -> None: if self.status_notified: return - if self.pending_infos: + + infos = self.pending_infos + if infos is None: + with self.kv_mgr.transfer_lock: + room_infos = self.kv_mgr.transfer_infos.get(self.bootstrap_room) + if room_infos is not None: + infos = list(room_infos.values()) + if infos: self.kv_mgr.notify_decode_status( - self.pending_infos, self.bootstrap_room, status, failure_reason + infos, self.bootstrap_room, status, failure_reason ) self.status_notified = True @@ -970,8 +1381,10 @@ class MoriKVSender(CommonKVSender): raise RuntimeError(failure_reason) def abort(self): - super().abort() + self.kv_mgr.record_failure(self.bootstrap_room, "Aborted by AbortReq.") + self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) self._notify_decode(KVPoll.Failed, "Aborted by AbortReq.") + self.conclude_state = KVPoll.Failed class MoriKVReceiver(CommonKVReceiver): @@ -1005,6 +1418,14 @@ class MoriKVReceiver(CommonKVReceiver): decode_tp_size = str(self.kv_mgr.attn_tp_size).encode("ascii") decode_tp_rank = str(self.kv_mgr.kv_args.engine_rank).encode("ascii") kv_item_len = str(self.kv_mgr.kv_args.kv_item_lens[0]).encode("ascii") + packed_state_item_lens = b"".join( + struct.pack("I", item_len) + for item_len in self.kv_mgr.kv_args.state_item_lens + ) + state_dim_per_tensor = getattr(self.kv_mgr.kv_args, "state_dim_per_tensor", []) + packed_state_dim_per_tensor = b"".join( + struct.pack("I", dim) for dim in state_dim_per_tensor + ) for bootstrap_info in self.bootstrap_infos: sock, lock = self._connect_to_bootstrap_server(bootstrap_info) @@ -1023,6 +1444,8 @@ class MoriKVReceiver(CommonKVReceiver): decode_tp_size, decode_tp_rank, kv_item_len, + packed_state_item_lens, + packed_state_dim_per_tensor, ] ) @@ -1040,11 +1463,15 @@ class MoriKVReceiver(CommonKVReceiver): np.asarray(kv_indices, dtype=np.int32).tobytes() if kv_indices.size else b"" ) aux_bytes = str(aux_index).encode("ascii") if aux_index is not None else b"" - state_bytes = b"" + normalized_state = _normalize_state_indices(state_indices) for bootstrap_info in self.bootstrap_infos: sock, lock = self._connect_to_bootstrap_server(bootstrap_info) is_dummy = bootstrap_info.get("is_dummy", False) + if not is_dummy and normalized_state is not None: + state_bytes = normalized_state.tobytes() + else: + state_bytes = b"" with lock: sock.send_multipart( [ diff --git a/test/registered/amd/disaggregation/test_mori_transfer_engine_e2e.py b/test/registered/amd/disaggregation/test_mori_transfer_engine_e2e.py new file mode 100644 index 000000000..3a9e1d459 --- /dev/null +++ b/test/registered/amd/disaggregation/test_mori_transfer_engine_e2e.py @@ -0,0 +1,179 @@ +import os +import unittest + +import requests + +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.server_fixtures.disaggregation_fixture import ( + PDDisaggregationServerBase, +) +from sglang.test.test_utils import ( + DEFAULT_SMALL_MODEL_NAME_FOR_TEST, + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + popen_launch_pd_server, + try_cached_model, +) + +register_amd_ci(est_time=300, suite="stage-b-test-large-8-gpu-35x-disaggregation-amd") + + +class MoriTransferEngineBase(PDDisaggregationServerBase): + port_delta = 0 + prefill_tp = 1 + decode_tp = 1 + decode_base_gpu_id = 1 + required_gpus = 2 + + @classmethod + def setUpClass(cls): + try: + import torch + + if not torch.cuda.is_available(): + raise unittest.SkipTest("torch.cuda is not available.") + if torch.cuda.device_count() < cls.required_gpus: + raise unittest.SkipTest( + f"MORI PD smoke test requires >= {cls.required_gpus} visible GPUs." + ) + except Exception as e: + raise unittest.SkipTest(f"torch is not available/usable: {e}") + + super().setUpClass() + + cls._old_use_aiter = os.environ.get("SGLANG_USE_AITER") + os.environ["SGLANG_USE_AITER"] = "1" + + # The shared fixture defaults to Mooncake in CI; pin Mori explicitly here. + cls.transfer_backend = ["--disaggregation-transfer-backend", "mori"] + + rdma_env = os.environ.get("SGLANG_TEST_RDMA_DEVICE") + if rdma_env: + cls.rdma_devices = ["--disaggregation-ib-device", rdma_env] + print(f"Found RDMA devices in env: {rdma_env}") + else: + print("SGLANG_TEST_RDMA_DEVICE is not set! Running without RDMA.") + cls.rdma_devices = [] + + cls._shift_ports() + cls.model = try_cached_model( + os.environ.get( + "SGLANG_MORI_E2E_TEST_MODEL", + DEFAULT_SMALL_MODEL_NAME_FOR_TEST, + ) + ) + + cls.start_prefill() + cls.start_decode() + + cls.wait_server_ready( + cls.prefill_url + "/health", + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + process=cls.process_prefill, + ) + cls.wait_server_ready( + cls.decode_url + "/health", + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + process=cls.process_decode, + ) + cls.launch_lb() + + @classmethod + def tearDownClass(cls): + if getattr(cls, "_old_use_aiter", None) is None: + os.environ.pop("SGLANG_USE_AITER", None) + else: + os.environ["SGLANG_USE_AITER"] = cls._old_use_aiter + super().tearDownClass() + + @classmethod + def _shift_ports(cls): + if cls.port_delta == 0: + return + + cls.lb_port = str(int(cls.lb_port) + cls.port_delta) + cls.prefill_port = str(int(cls.prefill_port) + cls.port_delta) + cls.decode_port = str(int(cls.decode_port) + cls.port_delta) + cls.bootstrap_port = str(int(cls.bootstrap_port) + cls.port_delta) + cls.prefill_url = f"http://{cls.base_host}:{cls.prefill_port}" + cls.decode_url = f"http://{cls.base_host}:{cls.decode_port}" + cls.lb_url = f"http://{cls.base_host}:{cls.lb_port}" + cls.base_url = cls.lb_url + + @classmethod + def start_prefill(cls): + prefill_args = [ + "--trust-remote-code", + "--disaggregation-mode", + "prefill", + "--disaggregation-bootstrap-port", + cls.bootstrap_port, + "--tp", + str(cls.prefill_tp), + "--attention-backend", + "aiter", + ] + prefill_args += cls.transfer_backend + cls.rdma_devices + cls.process_prefill = popen_launch_pd_server( + cls.model, + cls.prefill_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=prefill_args, + ) + + @classmethod + def start_decode(cls): + decode_args = [ + "--trust-remote-code", + "--disaggregation-mode", + "decode", + "--disaggregation-bootstrap-port", + cls.bootstrap_port, + "--tp", + str(cls.decode_tp), + "--base-gpu-id", + str(cls.decode_base_gpu_id), + "--attention-backend", + "aiter", + ] + decode_args += cls.transfer_backend + cls.rdma_devices + cls.process_decode = popen_launch_pd_server( + cls.model, + cls.decode_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=decode_args, + ) + + def _assert_generate_smoke(self): + resp = requests.post( + self.lb_url + "/generate", + json={ + "text": "Hello", + "sampling_params": {"temperature": 0, "max_new_tokens": 8}, + }, + timeout=120, + ) + self.assertEqual(resp.status_code, 200, resp.text) + out = resp.json() + self.assertIn("text", out) + self.assertIsInstance(out["text"], str) + self.assertGreater(len(out["text"]), 0) + + +class TestMoriTransferEngineE2E(MoriTransferEngineBase): + def test_generate_smoke(self): + self._assert_generate_smoke() + + +class TestMoriTransferEngineTPMismatchE2E(MoriTransferEngineBase): + port_delta = 10 + prefill_tp = 2 + decode_tp = 4 + decode_base_gpu_id = 2 + required_gpus = 6 + + def test_generate_smoke_tp_mismatch(self): + self._assert_generate_smoke() + + +if __name__ == "__main__": + unittest.main()