diff --git a/docker/rocm.Dockerfile b/docker/rocm.Dockerfile index 51f225b03..44687afd4 100644 --- a/docker/rocm.Dockerfile +++ b/docker/rocm.Dockerfile @@ -59,6 +59,16 @@ ARG TILELANG_COMMIT="ebf4a7cb8881432165ae8760e99d209d905c704a" ARG FHT_REPO="https://github.com/jeffdaily/fast-hadamard-transform.git" ARG FHT_BRANCH="rocm" ARG FHT_COMMIT="46efb7d776d38638fc39f3c803eaee3dd7016bd1" + +ARG ENABLE_MORI=0 +ARG NIC_BACKEND=none + +ARG MORI_REPO="https://github.com/ROCm/mori.git" +ARG MORI_COMMIT="b0dce4beebeb1f26c784eee17d5fd9785ee9447f" + +# AMD AINIC apt repo settings +ARG AINIC_VERSION=1.117.5 +ARG UBUNTU_CODENAME=jammy USER root # Install some basic utilities @@ -283,6 +293,73 @@ RUN python3 -m pip install --no-cache-dir \ py-spy \ pre-commit +# ----------------------- +# MORI (optional) +RUN /bin/bash -lc 'set -euo pipefail; \ + if [ "${ENABLE_MORI}" != "1" ]; then \ + echo "[MORI] Skipping (ENABLE_MORI=${ENABLE_MORI})"; \ + exit 0; \ + fi; \ + echo "[MORI] Enabling MORI (NIC_BACKEND=${NIC_BACKEND})"; \ + \ + # Base deps for MORI build + apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ + g++ \ + jq \ + libopenmpi-dev \ + libpci-dev \ + initramfs-tools \ + && rm -rf /var/lib/apt/lists/*; \ + \ + # NIC backend deps + case "${NIC_BACKEND}" in \ + # default: mlx5 + none) \ + export USE_IONIC="OFF"; \ + export USE_BNXT="OFF"; \ + ;; \ + # AMD NIC + ainic) \ + export USE_IONIC="ON"; \ + export USE_BNXT="OFF"; \ + apt-get update && apt-get install -y --no-install-recommends ca-certificates curl gnupg apt-transport-https && \ + rm -rf /var/lib/apt/lists/* && mkdir -p /etc/apt/keyrings; \ + curl -fsSL https://repo.radeon.com/rocm/rocm.gpg.key | gpg --dearmor > /etc/apt/keyrings/amdainic.gpg; \ + echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/amdainic.gpg] https://repo.radeon.com/amdainic/pensando/ubuntu/${AINIC_VERSION} ${UBUNTU_CODENAME} main" \ + > /etc/apt/sources.list.d/amdainic.list; \ + apt-get update && apt-get install -y --no-install-recommends \ + libionic-dev \ + ionic-common \ + ; \ + rm -rf /var/lib/apt/lists/*; \ + ;; \ + # TODO: Add Broadcom bnxt packages/repos here later. + # bnxt) \ + # export USE_IONIC="OFF"; \ + # export USE_BNXT="ON"; \ + # echo "[MORI] NIC_BACKEND=bnxt: USE_BNXT=ON. Add Broadcom bnxt packages/repos here later."; \ + # ;; \ + *) \ + echo "ERROR: unknown NIC_BACKEND=${NIC_BACKEND}. Use one of: none, ainic"; \ + exit 2; \ + ;; \ + esac; \ + \ + # Build/install MORI + export MORI_GPU_ARCHS="${GPU_ARCH_LIST}"; \ + echo "[MORI] MORI_GPU_ARCHS=${MORI_GPU_ARCHS} USE_IONIC=${USE_IONIC} USE_BNXT=${USE_BNXT}"; \ + rm -rf /sgl-workspace/mori; \ + git clone "${MORI_REPO}" /sgl-workspace/mori; \ + cd /sgl-workspace/mori; \ + git checkout "${MORI_COMMIT}"; \ + git submodule update --init --recursive; \ + python3 setup.py develop; \ + python3 -c "import os, torch; print(os.path.join(os.path.dirname(torch.__file__), \"lib\"))" > /etc/ld.so.conf.d/torch.conf; \ + ldconfig; \ + echo "export PYTHONPATH=/sgl-workspace/mori:\${PYTHONPATH}" >> /etc/bash.bashrc; \ + echo "[MORI] Done."' + # ----------------------- # Performance environment variable. diff --git a/python/sglang/srt/disaggregation/mori/__init__.py b/python/sglang/srt/disaggregation/mori/__init__.py new file mode 100644 index 000000000..f537f00b4 --- /dev/null +++ b/python/sglang/srt/disaggregation/mori/__init__.py @@ -0,0 +1,6 @@ +from sglang.srt.disaggregation.mori.conn import ( + MoriKVBootstrapServer, + MoriKVManager, + MoriKVReceiver, + MoriKVSender, +) diff --git a/python/sglang/srt/disaggregation/mori/conn.py b/python/sglang/srt/disaggregation/mori/conn.py new file mode 100644 index 000000000..2aca028c2 --- /dev/null +++ b/python/sglang/srt/disaggregation/mori/conn.py @@ -0,0 +1,1136 @@ +from __future__ import annotations + +import ctypes +import dataclasses +import logging +import os +import struct +import threading +import time +from collections import defaultdict +from typing import Dict, List, Optional, Set, Tuple + +import msgpack +import numpy as np +import numpy.typing as npt +from mori.cpp import TransferStatus +from mori.io import ( + BackendType, + EngineDesc, + IOEngine, + IOEngineConfig, + MemoryDesc, + MemoryLocationType, + PollCqMode, + RdmaBackendConfig, +) + +from sglang.srt.disaggregation.base.conn import KVArgs, KVPoll +from sglang.srt.disaggregation.common.conn import ( + CommonKVBootstrapServer, + CommonKVManager, + CommonKVReceiver, + CommonKVSender, +) +from sglang.srt.disaggregation.common.utils import group_concurrent_contiguous +from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils.common import ( + format_tcp_address, + get_int_env_var, + get_local_ip_auto, + is_valid_ipv6_address, +) + +logger = logging.getLogger(__name__) +MORI_GUARD = b"MoriMsgGuard" + + +def _pack_mem_desc_list(mems: List[MemoryDesc]) -> bytes: + if not mems: + return b"" + packed_descs = [mem.pack() for mem in mems] + return msgpack.packb(packed_descs, use_bin_type=True) + + +def _unpack_mem_desc_list(blob: bytes) -> List[MemoryDesc]: + if not blob: + return [] + desc_blobs = msgpack.unpackb(blob, raw=False) + return [MemoryDesc.unpack(b) for b in desc_blobs] + + +@dataclasses.dataclass +class TransferInfo: + room: int + endpoint: str + dst_port: int + engine_key: str + dst_kv_indices: npt.NDArray[np.int32] + dst_aux_index: int + required_dst_info_num: int + is_dummy: bool + + @classmethod + def from_zmq(cls, payload: List[bytes]) -> TransferInfo: + room = int(payload[0].decode("ascii")) + endpoint = payload[1].decode("ascii") + dst_port = int(payload[2].decode("ascii")) + engine_key = payload[3].decode("ascii") + + if payload[4]: + dst_kv_indices = np.frombuffer(payload[4], dtype=np.int32) + else: + dst_kv_indices = np.array([], dtype=np.int32) + + if payload[5]: + dst_aux_index = int(payload[5].decode("ascii")) + else: + dst_aux_index = -1 + + required_dst_info_num = ( + int(payload[7].decode("ascii")) if len(payload) > 7 else 1 + ) + is_dummy = dst_kv_indices.size == 0 and dst_aux_index < 0 + return cls( + room=room, + endpoint=endpoint, + dst_port=dst_port, + engine_key=engine_key, + dst_kv_indices=dst_kv_indices, + dst_aux_index=dst_aux_index, + required_dst_info_num=required_dst_info_num, + is_dummy=is_dummy, + ) + + +@dataclasses.dataclass +class KVArgsRegisterInfo: + endpoint: str + dst_port: int + engine_desc: EngineDesc + dst_kv_mem_descs: List[MemoryDesc] + dst_aux_mem_descs: List[MemoryDesc] + dst_state_mem_descs: List[MemoryDesc] + gpu_id: int + decode_tp_size: int + decode_tp_rank: int + dst_kv_item_len: int + + @property + def engine_key(self) -> str: + return self.engine_desc.key + + @classmethod + def from_zmq(cls, payload: List[bytes]) -> KVArgsRegisterInfo: + endpoint = payload[1].decode("ascii") + dst_port = int(payload[2].decode("ascii")) + engine_desc = EngineDesc.unpack(payload[3]) + dst_kv_mem_descs = _unpack_mem_desc_list(payload[4]) + dst_aux_mem_descs = _unpack_mem_desc_list(payload[5]) + dst_state_mem_descs = _unpack_mem_desc_list(payload[6]) + gpu_id = int(payload[7].decode("ascii")) + decode_tp_size = int(payload[8].decode("ascii")) + decode_tp_rank = int(payload[9].decode("ascii")) + dst_kv_item_len = int(payload[10].decode("ascii")) + return cls( + endpoint=endpoint, + dst_port=dst_port, + engine_desc=engine_desc, + dst_kv_mem_descs=dst_kv_mem_descs, + dst_aux_mem_descs=dst_aux_mem_descs, + dst_state_mem_descs=dst_state_mem_descs, + gpu_id=gpu_id, + decode_tp_size=decode_tp_size, + decode_tp_rank=decode_tp_rank, + dst_kv_item_len=dst_kv_item_len, + ) + + +class AuxDataCodec: + @staticmethod + def serialize_data_from_buffer(src_addr, data_length): + buffer = (ctypes.c_byte * data_length).from_address(src_addr) + return bytes(buffer) + + @staticmethod + def deserialize_data_to_buffer(kv_args, buffer_index, aux_index, data): + dst_aux_ptr = kv_args.aux_data_ptrs[buffer_index] + item_len = kv_args.aux_item_lens[buffer_index] + dst_addr = dst_aux_ptr + item_len * aux_index + buffer = (ctypes.c_byte * len(data)).from_address(dst_addr) + buffer[:] = data + return + + +@dataclasses.dataclass +class TPSliceConfig: + page_size: int + src_item_len: int + dst_item_len: int + bytes_per_token_src: int + bytes_per_token_dst: int + src_head_slice_offset: int + dst_head_slice_offset: int + heads_bytes_per_token_to_send: int + + +class MoriKVManager(CommonKVManager): + AUX_DATA_HEADER = b"AUX_DATA" + + def __init__( + self, + args: KVArgs, + disaggregation_mode: DisaggregationMode, + server_args: ServerArgs, + is_mla_backend: Optional[bool] = False, + ): + super().__init__(args, disaggregation_mode, server_args, is_mla_backend) + self.engine = self._init_engine() + self.engine_desc = self.engine.get_engine_desc() + self.kv_mem_descs: List[MemoryDesc] = [] + self.aux_mem_descs: List[MemoryDesc] = [] + self.state_mem_descs: List[MemoryDesc] = [] + self.failure_records: Dict[int, str] = {} + self.failure_lock = threading.Lock() + self.transfer_lock = threading.Lock() + self._register_local_buffers() + if self.disaggregation_mode == DisaggregationMode.PREFILL: + self.bootstrap_timeout = get_int_env_var( + "SGLANG_DISAGGREGATION_BOOTSTRAP_TIMEOUT", 300 + ) + self._start_bootstrap_thread() + elif self.disaggregation_mode == DisaggregationMode.DECODE: + self.waiting_timeout = get_int_env_var( + "SGLANG_DISAGGREGATION_WAITING_TIMEOUT", 300 + ) + self.prefill_response_tracker: Dict[int, Set[int]] = defaultdict(set) + self.addr_to_rooms_tracker = defaultdict(set) + self.room_to_bootstrap_addr: Dict[int, str] = {} + self._start_decode_thread() + + def _init_engine(self) -> IOEngine: + if self.kv_args.ib_device: + os.environ["MORI_RDMA_DEVICES"] = self.kv_args.ib_device + + self.local_ip = get_local_ip_auto() + config = IOEngineConfig(host=self.local_ip, port=0) + + engine_key = ( + f"io-{self.disaggregation_mode.value}-" + f"dp{self.system_dp_rank}-tp{self.attn_tp_rank}-" + f"pid{os.getpid()}-{self.local_ip}" + ) + + engine = IOEngine(engine_key, config) + poll_mode = PollCqMode.POLLING + + # Number of RDMA Queue Pairs (QPs) used per transfer operation. + # Higher values can increase parallelism and bandwidth utilization. + # Default: 1 + qp_per_transfer = get_int_env_var("SGLANG_MORI_QP_PER_TRANSFER", 1) + + # Number of RDMA work requests posted in a single batch to each QP. + # Larger batch sizes reduce per-operation overhead and improve throughput + # at the cost of higher latency. Use -1 for automatic sizing based on + # the number of merged work requests and available endpoints. + # Default: -1 (automatic) + post_batch_size = get_int_env_var("SGLANG_MORI_POST_BATCH_SIZE", -1) + + # Number of worker threads in the RDMA executor thread pool. + # Each worker handles RDMA operations on a separate CPU core (with affinity). + # More workers can improve parallelism for large batch transfers across + # multiple QPs, but excessive threads may cause contention. + # Default: 1 + num_worker_threads = get_int_env_var("SGLANG_MORI_NUM_WORKERS", 1) + + rdma_cfg = RdmaBackendConfig( + qp_per_transfer, + post_batch_size, + num_worker_threads, + poll_mode, + False, + ) + engine.create_backend(BackendType.RDMA, rdma_cfg) + actual_port = engine.get_engine_desc().port + assert actual_port > 0, f"Failed to bind port for engine {engine_key}" + logger.debug( + "Initialized Mori IOEngine %s at %s:%s (qp_per_transfer=%s, workers=%s, poll_mode=%s)", + engine_key, + self.local_ip, + actual_port, + qp_per_transfer, + num_worker_threads, + poll_mode.name, + ) + return engine + + def _register_local_buffers(self) -> None: + for ptr, length in zip(self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens): + mem_desc = self.engine.register_memory( + ptr, + length, + self.kv_args.gpu_id, + MemoryLocationType.GPU, + ) + self.kv_mem_descs.append(mem_desc) + for ptr, length in zip(self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens): + desc = self.engine.register_memory( + ptr, + length, + -1, + MemoryLocationType.CPU, + ) + self.aux_mem_descs.append(desc) + for ptr, length in zip( + self.kv_args.state_data_ptrs, getattr(self.kv_args, "state_data_lens", []) + ): + desc = self.engine.register_memory( + ptr, + length, + self.kv_args.gpu_id, + MemoryLocationType.GPU, + ) + self.state_mem_descs.append(desc) + + def 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: + 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) -> None: + with self.failure_lock: + self.failure_records[bootstrap_room] = failure_reason + + def _handle_register_message(self, payload: List[bytes]) -> None: + try: + register_info = KVArgsRegisterInfo.from_zmq(payload) + self._add_remote_peer(register_info) + except Exception: + logger.exception("Failed to register remote peer") + + def _handle_transfer_message(self, payload: List[bytes]) -> None: + try: + transfer_info = TransferInfo.from_zmq(payload) + infos = self.transfer_infos.setdefault(transfer_info.room, {}) + infos[transfer_info.engine_key] = transfer_info + + if len(infos) >= transfer_info.required_dst_info_num: + logger.debug( + "Bootstrap room %s got enough transfer info (%s)", + transfer_info.room, + len(infos), + ) + self.update_status(transfer_info.room, KVPoll.WaitingForInput) + except Exception: + logger.exception("Failed to parse transfer info message") + + def _validate_message(self, msg: List[bytes]) -> Optional[List[bytes]]: + if not msg or msg[0] != MORI_GUARD: + logger.warning("Received malformed bootstrap message") + return None + payload = msg[1:] + if not payload: + return None + return payload + + def _start_bootstrap_thread(self) -> None: + def bootstrap_worker(): + while True: + try: + msg = self.server_socket.recv_multipart() + payload = self._validate_message(msg) + if payload is None: + continue + room = payload[0].decode("ascii") + + if room == "None": + self._handle_register_message(payload) + else: + self._handle_transfer_message(payload) + except Exception: + logger.exception("Bootstrap worker failed") + + threading.Thread(target=bootstrap_worker, daemon=True).start() + + def _cleanup_room_tracking(self, bootstrap_room: int) -> None: + bootstrap_addr = self.room_to_bootstrap_addr.pop(bootstrap_room, None) + if bootstrap_addr is not None: + rooms = self.addr_to_rooms_tracker.get(bootstrap_addr) + if rooms is not None: + rooms.discard(bootstrap_room) + if not rooms: + self.addr_to_rooms_tracker.pop(bootstrap_addr, None) + + def _start_decode_thread(self) -> None: + def decode_worker(): + while True: + try: + msg = self.server_socket.recv_multipart() + if msg and msg[0] == MoriKVManager.AUX_DATA_HEADER: + self._handle_aux_data(msg) + continue + + if not msg or msg[0] != MORI_GUARD: + logger.warning( + "Received malformed status message on decode worker" + ) + continue + payload = msg[1:] + if len(payload) < 3: + logger.warning("Incomplete status payload received") + continue + bootstrap_room = int(payload[0].decode("ascii")) + status_code = int(payload[1].decode("ascii")) + prefill_rank = int(payload[2].decode("ascii")) + failure_reason = ( + payload[3].decode("utf-8") + if len(payload) > 3 and payload[3] + else None + ) + + if status_code == KVPoll.Success: + tracker = self.prefill_response_tracker[bootstrap_room] + tracker.add(prefill_rank) + expected = self.required_prefill_response_num_table.get( + bootstrap_room, 1 + ) + if len(tracker) >= expected: + self.prefill_response_tracker.pop(bootstrap_room, None) + self.update_status(bootstrap_room, KVPoll.Success) + self._cleanup_room_tracking(bootstrap_room) + elif status_code == KVPoll.Failed: + if failure_reason: + self.record_failure(bootstrap_room, failure_reason) + self.prefill_response_tracker.pop(bootstrap_room, None) + self.update_status(bootstrap_room, KVPoll.Failed) + self._cleanup_room_tracking(bootstrap_room) + else: + logger.warning( + "Unknown status code %s received for room %s", + status_code, + bootstrap_room, + ) + except Exception: + logger.exception("Decode status worker failed") + + threading.Thread(target=decode_worker, daemon=True).start() + + def notify_decode_status( + self, + infos: List[TransferInfo], + bootstrap_room: int, + status: KVPoll, + failure_reason: Optional[str] = None, + ) -> None: + if not infos: + return + payload = [ + MORI_GUARD, + str(bootstrap_room).encode("ascii"), + str(int(status)).encode("ascii"), + str(self.attn_tp_rank * self.pp_size + self.pp_rank).encode("ascii"), + failure_reason.encode("utf-8") if failure_reason else b"", + ] + for info in infos: + try: + endpoint = format_tcp_address(info.endpoint, info.dst_port) + socket = self._connect( + endpoint, is_ipv6=is_valid_ipv6_address(info.endpoint) + ) + socket.send_multipart(payload) + except Exception: + logger.exception( + "Failed to sync status %s to decode endpoint %s:%s for room %s", + status, + info.endpoint, + info.dst_port, + bootstrap_room, + ) + + def _add_remote_peer(self, register_info: KVArgsRegisterInfo) -> None: + engine_key = register_info.engine_key + if engine_key in self.decode_kv_args_table: + logger.debug("Remote peer %s already registered. Skipping.", engine_key) + return + self.engine.register_remote_engine(register_info.engine_desc) + self.decode_kv_args_table[engine_key] = register_info + logger.debug( + "Registered decode peer %s (%s:%s)", + engine_key, + register_info.endpoint, + register_info.dst_port, + ) + + def _get_mha_mem_desc_slices( + self, dst_mem_descs: List[MemoryDesc] + ) -> tuple[ + List[MemoryDesc], List[MemoryDesc], List[MemoryDesc], List[MemoryDesc], int + ]: + src_descs = self.kv_mem_descs + if not src_descs: + raise RuntimeError("KV memory descriptors are empty on prefill side") + + num_local_layers = len(src_descs) // 2 + src_k_descs = src_descs[:num_local_layers] + src_v_descs = src_descs[num_local_layers:] + + start_layer = self.kv_args.prefill_start_layer + end_layer = start_layer + num_local_layers + dst_total_layers = len(dst_mem_descs) // 2 + if len(dst_mem_descs) < 2 or end_layer > dst_total_layers: + raise ValueError( + "Destination KV descriptors do not match prefill pp configuration" + ) + dst_k_descs = dst_mem_descs[start_layer:end_layer] + dst_v_descs = dst_mem_descs[ + dst_total_layers + start_layer : dst_total_layers + end_layer + ] + return src_k_descs, src_v_descs, dst_k_descs, dst_v_descs, num_local_layers + + def _get_mla_mem_desc_slices( + self, dst_mem_descs: List[MemoryDesc] + ) -> tuple[List[MemoryDesc], List[MemoryDesc], int]: + src_descs = self.kv_mem_descs + num_local_layers = len(src_descs) + start_layer = self.kv_args.prefill_start_layer + end_layer = start_layer + num_local_layers + if end_layer > len(dst_mem_descs): + raise ValueError( + "Destination MLA KV descriptors do not match prefill pp configuration" + ) + dst_slice = dst_mem_descs[start_layer:end_layer] + return src_descs, dst_slice, num_local_layers + + def _issue_layer_transfers( + self, + src_desc: MemoryDesc, + dst_desc: MemoryDesc, + kv_item_len: int, + src_groups: List[List[int]], + dst_groups: List[List[int]], + ) -> List[TransferStatus]: + if not src_groups: + return [] + local_offsets = [int(src_group[0]) * kv_item_len for src_group in src_groups] + remote_offsets = [int(dst_group[0]) * kv_item_len for dst_group in dst_groups] + sizes = [len(src_group) * kv_item_len for src_group in src_groups] + + transfer_uid = self.engine.allocate_transfer_uid() + + statuses = self.engine.batch_write( + [src_desc], + [local_offsets], + [dst_desc], + [remote_offsets], + [sizes], + [transfer_uid], + ) + return statuses + + def _build_tp_slice_config(self, peer_info: KVArgsRegisterInfo) -> TPSliceConfig: + page_size = self.kv_args.page_size + + src_item_len = self.kv_args.kv_item_lens[0] + dst_item_len = peer_info.dst_kv_item_len + + bytes_per_token_src = src_item_len // page_size + bytes_per_token_dst = dst_item_len // page_size + + prefill_tp_size = self.attn_tp_size + decode_tp_size = peer_info.decode_tp_size + + num_kv_heads = self.kv_args.kv_head_num + src_heads_per_rank = num_kv_heads + dst_heads_per_rank = num_kv_heads * prefill_tp_size // decode_tp_size + if dst_heads_per_rank == 0: + raise ValueError("Destination heads per rank evaluates to zero") + + bytes_per_head_slice = bytes_per_token_dst // dst_heads_per_rank + if bytes_per_head_slice == 0: + raise ValueError("Head slice size evaluates to zero") + + local_tp_rank = self.kv_args.engine_rank % prefill_tp_size + dst_tp_rank = peer_info.decode_tp_rank % decode_tp_size + + if prefill_tp_size > decode_tp_size: + src_head_start = 0 + num_heads_to_send = src_heads_per_rank + dst_head_start = local_tp_rank * src_heads_per_rank + else: + src_head_start = (dst_tp_rank * dst_heads_per_rank) % src_heads_per_rank + num_heads_to_send = dst_heads_per_rank + dst_head_start = 0 + + src_head_slice_offset = src_head_start * bytes_per_head_slice + dst_head_slice_offset = dst_head_start * bytes_per_head_slice + heads_bytes_per_token = num_heads_to_send * bytes_per_head_slice + + if heads_bytes_per_token > bytes_per_token_dst: + raise ValueError( + "Slice size exceeds destination token capacity for TP slice transfer" + ) + + return TPSliceConfig( + page_size=page_size, + src_item_len=src_item_len, + dst_item_len=dst_item_len, + bytes_per_token_src=bytes_per_token_src, + bytes_per_token_dst=bytes_per_token_dst, + src_head_slice_offset=src_head_slice_offset, + dst_head_slice_offset=dst_head_slice_offset, + heads_bytes_per_token_to_send=heads_bytes_per_token, + ) + + def _issue_tp_slice_transfers( + self, + src_desc: MemoryDesc, + dst_desc: MemoryDesc, + kv_indices: npt.NDArray[np.int32], + dst_indices: npt.NDArray[np.int32], + tp_cfg: TPSliceConfig, + ) -> List[TransferStatus]: + if kv_indices.size == 0 or dst_indices.size == 0: + return [] + + limit = min(kv_indices.size, dst_indices.size) + if not limit: + return [] + + src_pages = kv_indices[:limit].astype(np.int64) + dst_pages = dst_indices[:limit].astype(np.int64) + token_slots = np.arange(tp_cfg.page_size, dtype=np.int64) + + src_page_bases = src_pages * tp_cfg.src_item_len + dst_page_bases = dst_pages * tp_cfg.dst_item_len + + src_token_offsets = token_slots * tp_cfg.bytes_per_token_src + dst_token_offsets = token_slots * tp_cfg.bytes_per_token_dst + + local_offsets = ( + ( + src_page_bases[:, np.newaxis] + + src_token_offsets + + tp_cfg.src_head_slice_offset + ) + .flatten() + .tolist() + ) + remote_offsets = ( + ( + dst_page_bases[:, np.newaxis] + + dst_token_offsets + + tp_cfg.dst_head_slice_offset + ) + .flatten() + .tolist() + ) + + num_transfers = limit * tp_cfg.page_size + sizes = [tp_cfg.heads_bytes_per_token_to_send] * num_transfers + + if not local_offsets: + return [] + + transfer_uid = self.engine.allocate_transfer_uid() + statuses = self.engine.batch_write( + [src_desc], + [local_offsets], + [dst_desc], + [remote_offsets], + [sizes], + [transfer_uid], + ) + return statuses + + def send_kvcache( + self, + peer_info: KVArgsRegisterInfo, + prefill_kv_indices: npt.NDArray[np.int32], + dst_kv_indices: npt.NDArray[np.int32], + ) -> List[TransferStatus]: + src_groups, dst_groups = group_concurrent_contiguous( + prefill_kv_indices, dst_kv_indices + ) + statuses = [] + kv_item_len = self.kv_args.kv_item_lens[0] + if self.is_mla_backend: + ( + src_descs, + dst_descs, + layers_current_pp_stage, + ) = self._get_mla_mem_desc_slices(peer_info.dst_kv_mem_descs) + for layer_id in range(layers_current_pp_stage): + statuses.extend( + self._issue_layer_transfers( + src_descs[layer_id], + dst_descs[layer_id], + kv_item_len, + src_groups, + dst_groups, + ) + ) + else: + tp_mismatch = peer_info.decode_tp_size != self.attn_tp_size + ( + src_k_descs, + src_v_descs, + dst_k_descs, + dst_v_descs, + layers_current_pp_stage, + ) = self._get_mha_mem_desc_slices(peer_info.dst_kv_mem_descs) + + if tp_mismatch: + tp_cfg = self._build_tp_slice_config(peer_info) + for layer_id in range(layers_current_pp_stage): + statuses.extend( + self._issue_tp_slice_transfers( + src_k_descs[layer_id], + dst_k_descs[layer_id], + prefill_kv_indices, + dst_kv_indices, + tp_cfg, + ) + ) + statuses.extend( + self._issue_tp_slice_transfers( + src_v_descs[layer_id], + dst_v_descs[layer_id], + prefill_kv_indices, + dst_kv_indices, + tp_cfg, + ) + ) + else: + src_groups, dst_groups = group_concurrent_contiguous( + prefill_kv_indices, dst_kv_indices + ) + for layer_id in range(layers_current_pp_stage): + statuses.extend( + self._issue_layer_transfers( + src_k_descs[layer_id], + dst_k_descs[layer_id], + kv_item_len, + src_groups, + dst_groups, + ) + ) + statuses.extend( + self._issue_layer_transfers( + src_v_descs[layer_id], + dst_v_descs[layer_id], + kv_item_len, + src_groups, + dst_groups, + ) + ) + + return statuses + + def send_aux( + self, + peer_info: KVArgsRegisterInfo, + prefill_aux_index: int, + dst_aux_index: int, + room: int, + ) -> List[TransferStatus]: + return self.send_aux_tcp(peer_info, prefill_aux_index, dst_aux_index, room) + + def send_aux_tcp( + self, + peer_info: KVArgsRegisterInfo, + prefill_aux_index: int, + dst_aux_index: int, + room: int, + ) -> List[TransferStatus]: + prefill_aux_ptrs = self.kv_args.aux_data_ptrs + prefill_aux_item_lens = self.kv_args.aux_item_lens + + for i in range(len(prefill_aux_ptrs)): + length = prefill_aux_item_lens[i] + src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index + data = AuxDataCodec.serialize_data_from_buffer(src_addr, length) + + self.send_aux_data_to_endpoint( + remote=peer_info.endpoint, + dst_port=peer_info.dst_port, + room=room, + buffer_index=i, + aux_index=dst_aux_index, + data=data, + ) + + return [] + + def send_aux_data_to_endpoint( + self, + remote: str, + dst_port: int, + room: int, + buffer_index: int, + aux_index: int, + data: bytes, + ): + socket = self._connect( + format_tcp_address(remote, dst_port), is_ipv6=is_valid_ipv6_address(remote) + ) + + socket.send_multipart( + [ + MoriKVManager.AUX_DATA_HEADER, + str(room).encode("ascii"), + str(buffer_index).encode("ascii"), + str(aux_index).encode("ascii"), + struct.pack(">I", len(data)), + data, + ] + ) + + def _handle_aux_data(self, msg: List[bytes]): + """Handle AUX_DATA messages received by the decode thread.""" + room = int(msg[1].decode("ascii")) + buffer_index = int(msg[2].decode("ascii")) + aux_index = int(msg[3].decode("ascii")) + data_length = struct.unpack(">I", msg[4])[0] + data = msg[5] + + if len(data) != data_length: + logger.error(f"AUX_DATA length mismatch for bootstrap_room {room}") + return + + AuxDataCodec.deserialize_data_to_buffer( + self.kv_args, buffer_index, aux_index, data + ) + + logger.debug( + f"Received AUX_DATA for bootstrap_room {room} with length:{len(data)}" + ) + + def add_transfer_request( + self, + bootstrap_room: int, + kv_indices: npt.NDArray[np.int32], + index_slice: slice, + is_last: bool, + aux_index: Optional[int] = None, + state_indices: Optional[npt.NDArray[np.int32]] = None, + ) -> Tuple[List[TransferStatus], Optional[List[TransferInfo]]]: + assert self.disaggregation_mode == DisaggregationMode.PREFILL + transfer_infos = self.transfer_infos.get(bootstrap_room) + if not transfer_infos: + raise RuntimeError( + f"No transfer info found for bootstrap_room={bootstrap_room}" + ) + result_statuses = [] + target_infos_snapshot: Optional[List[TransferInfo]] = None + with self.transfer_lock: + self.update_status(bootstrap_room, KVPoll.Transferring) + for info in transfer_infos.values(): + peer_info = self.decode_kv_args_table.get(info.engine_key) + if not peer_info: + self.record_failure( + bootstrap_room, + f"Peer info missing for engine {info.engine_key}", + ) + raise RuntimeError( + f"Missing decode peer info for {info.engine_key}" + ) + if not info.is_dummy: + dst_indices_chunk = info.dst_kv_indices[index_slice] + statuses = self.send_kvcache( + peer_info, kv_indices, dst_indices_chunk + ) + result_statuses.extend(statuses) + if ( + is_last + and aux_index is not None + and info.dst_aux_index >= 0 + and self.pp_group.is_last_rank + ): + result_statuses.extend( + self.send_aux( + peer_info, aux_index, info.dst_aux_index, bootstrap_room + ) + ) + if is_last: + self.update_status(bootstrap_room, KVPoll.Success) + target_infos_snapshot = list(transfer_infos.values()) + self.transfer_infos.pop(bootstrap_room, None) + return result_statuses, target_infos_snapshot + + +class MoriKVSender(CommonKVSender): + def __init__( + self, + mgr: MoriKVManager, + bootstrap_addr: str, + bootstrap_room: int, + dest_tp_ranks: List[int], + pp_rank: int, + ): + super().__init__(mgr, bootstrap_addr, bootstrap_room, dest_tp_ranks, pp_rank) + self.transfer_statuses: List[TransferStatus] = [] + self.pending_infos: Optional[List[TransferInfo]] = None + self.sent_last_chunk = False + self.conclude_state: Optional[KVPoll] = None + self.status_notified = False + self.init_time = time.time() + + def send( + self, + kv_indices: npt.NDArray[np.int32], + state_indices: Optional[List[int]] = None, + ): + index_slice = slice(self.curr_idx, self.curr_idx + len(kv_indices)) + self.curr_idx += len(kv_indices) + is_last = self.curr_idx == self.num_kv_indices + + statuses, infos = self.kv_mgr.add_transfer_request( + self.bootstrap_room, + kv_indices, + index_slice, + is_last, + aux_index=self.aux_index if is_last else None, + ) + self.transfer_statuses.extend(statuses) + if infos is not None: + self.pending_infos = infos + self.sent_last_chunk = True + + def poll(self) -> KVPoll: + if self.conclude_state is not None: + return self.conclude_state + + status = self.kv_mgr.check_status(self.bootstrap_room) + if status == KVPoll.Bootstrapping: + elapsed = time.time() - self.init_time + if elapsed >= self.kv_mgr.bootstrap_timeout: + reason = ( + f"Request {self.bootstrap_room} timed out after {elapsed:.1f}s " + "waiting for decode handshake" + ) + self.kv_mgr.record_failure(self.bootstrap_room, reason) + self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) + self._finalize_failure(reason) + return KVPoll.Failed + return status + + if status == KVPoll.Failed: + self._finalize_failure() + return KVPoll.Failed + + transfers_done = self._all_transfers_finished() + if transfers_done: + if self._has_transfer_error(): + reason = self._collect_failure_reason() + self.kv_mgr.record_failure(self.bootstrap_room, reason) + self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) + self._finalize_failure(reason) + return KVPoll.Failed + self._notify_decode(KVPoll.Success) + self.conclude_state = KVPoll.Success + return KVPoll.Success + return KVPoll.Transferring if status == KVPoll.Success else status + + def _all_transfers_finished(self) -> bool: + if not self.sent_last_chunk: + return False + if not self.transfer_statuses: + return True + return all(not status.InProgress() for status in self.transfer_statuses) + + def _has_transfer_error(self) -> bool: + return any(status.Failed() for status in self.transfer_statuses) + + def _collect_failure_reason(self) -> str: + for status in self.transfer_statuses: + if status.Failed(): + return f"KV transfer failed: {status.Message()}" + return "KV transfer failed due to unknown reason" + + def _notify_decode( + self, status: KVPoll, failure_reason: Optional[str] = None + ) -> None: + if self.status_notified: + return + if self.pending_infos: + self.kv_mgr.notify_decode_status( + self.pending_infos, self.bootstrap_room, status, failure_reason + ) + self.status_notified = True + + def _finalize_failure(self, failure_reason: Optional[str] = None) -> None: + if self.conclude_state == KVPoll.Failed: + return + if failure_reason is None: + failure_reason = self.kv_mgr.failure_records.get( + self.bootstrap_room, "KV transfer failed" + ) + self._notify_decode(KVPoll.Failed, failure_reason) + self.conclude_state = KVPoll.Failed + + def clear(self) -> None: + self.kv_mgr.request_status.pop(self.bootstrap_room, None) + + def failure_exception(self): + if self.conclude_state is None: + self._finalize_failure() + self.clear() + with self.kv_mgr.failure_lock: + failure_reason = self.kv_mgr.failure_records.pop( + self.bootstrap_room, "KV transfer failed" + ) + raise RuntimeError(failure_reason) + + def abort(self): + reason = "Aborted by AbortReq." + self.kv_mgr.record_failure(self.bootstrap_room, reason) + self._notify_decode(KVPoll.Failed, reason) + self.conclude_state = KVPoll.Failed + + +class MoriKVReceiver(CommonKVReceiver): + + def __init__( + self, + mgr: MoriKVManager, + bootstrap_addr: str, + bootstrap_room: Optional[int] = None, + prefill_dp_rank: Optional[int] = None, + ): + super().__init__(mgr, bootstrap_addr, bootstrap_room, prefill_dp_rank) + self.conclude_state: Optional[KVPoll] = None + self.init_time: Optional[float] = None + if self.bootstrap_room is None or self.bootstrap_infos is None: + return + self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room) + self.kv_mgr.room_to_bootstrap_addr[self.bootstrap_room] = self.bootstrap_addr + self.kv_mgr.update_status(self.bootstrap_room, KVPoll.WaitingForInput) + self._register_kv_args() + + def _register_kv_args(self): + if self.bootstrap_infos is None: + return + engine_desc_blob = self.kv_mgr.engine_desc.pack() + packed_kv_descs = _pack_mem_desc_list(self.kv_mgr.kv_mem_descs) + packed_aux_descs = _pack_mem_desc_list(self.kv_mgr.aux_mem_descs) + packed_state_descs = _pack_mem_desc_list(self.kv_mgr.state_mem_descs) + gpu_id = str(self.kv_mgr.kv_args.gpu_id).encode("ascii") + decode_tp_size = str(self.kv_mgr.kv_args.decode_tp_size).encode("ascii") + decode_tp_rank = str(self.kv_mgr.kv_args.engine_rank).encode("ascii") + kv_item_len = str(self.kv_mgr.kv_args.kv_item_lens[0]).encode("ascii") + + for bootstrap_info in self.bootstrap_infos: + sock, lock = self._connect_to_bootstrap_server(bootstrap_info) + with lock: + sock.send_multipart( + [ + MORI_GUARD, + "None".encode("ascii"), + self.kv_mgr.local_ip.encode("ascii"), + str(self.kv_mgr.rank_port).encode("ascii"), + engine_desc_blob, + packed_kv_descs, + packed_aux_descs, + packed_state_descs, + gpu_id, + decode_tp_size, + decode_tp_rank, + kv_item_len, + ] + ) + + def init( + self, + kv_indices: npt.NDArray[np.int32], + aux_index: Optional[int] = None, + state_indices: Optional[List[int]] = None, + ): + if self.bootstrap_infos is None or self.bootstrap_room is None: + return + + kv_indices_bytes = ( + np.asarray(kv_indices, dtype=np.int32).tobytes() if kv_indices.size else b"" + ) + aux_bytes = str(aux_index).encode("ascii") if aux_index is not None else b"" + state_bytes = b"" + + for bootstrap_info in self.bootstrap_infos: + sock, lock = self._connect_to_bootstrap_server(bootstrap_info) + is_dummy = bootstrap_info.get("is_dummy", False) + with lock: + sock.send_multipart( + [ + MORI_GUARD, + str(self.bootstrap_room).encode("ascii"), + self.kv_mgr.local_ip.encode("ascii"), + str(self.kv_mgr.rank_port).encode("ascii"), + self.kv_mgr.engine_desc.key.encode("ascii"), + kv_indices_bytes if not is_dummy else b"", + aux_bytes if not is_dummy else b"", + state_bytes, + str(self.required_dst_info_num).encode("ascii"), + ] + ) + self.init_time = time.time() + + def poll(self) -> KVPoll: + if self.conclude_state is not None: + return self.conclude_state + + status = self.kv_mgr.check_status(self.bootstrap_room) + if status in (KVPoll.Success, KVPoll.Failed): + self.conclude_state = status + return status + + if status == KVPoll.WaitingForInput and self.init_time is not None: + elapsed = time.time() - self.init_time + if elapsed >= self.kv_mgr.waiting_timeout: + reason = f"Request {self.bootstrap_room} timed out after {elapsed:.1f}s waiting for KV transfer" + self.kv_mgr.record_failure(self.bootstrap_room, reason) + self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) + self.conclude_state = KVPoll.Failed + return KVPoll.Failed + + return status + + def clear(self) -> None: + if self.bootstrap_room is None: + return + self.kv_mgr.request_status.pop(self.bootstrap_room, None) + self.kv_mgr.required_prefill_response_num_table.pop(self.bootstrap_room, None) + self.kv_mgr.prefill_response_tracker.pop(self.bootstrap_room, None) + self.kv_mgr._cleanup_room_tracking(self.bootstrap_room) + + def failure_exception(self): + if self.conclude_state is None: + self.conclude_state = KVPoll.Failed + + self.clear() + with self.kv_mgr.failure_lock: + failure_reason = self.kv_mgr.failure_records.pop( + self.bootstrap_room, "KV transfer failed" + ) + raise RuntimeError(failure_reason) + + def abort(self): + if self.bootstrap_room is None: + return + reason = "Aborted by AbortReq." + self.kv_mgr.record_failure(self.bootstrap_room, reason) + self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) + self.conclude_state = KVPoll.Failed + self.clear() + + +class MoriKVBootstrapServer(CommonKVBootstrapServer): + pass diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index d660172de..690ca7b35 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -231,6 +231,7 @@ class MetadataBuffers: class TransferBackend(Enum): MOONCAKE = "mooncake" + MORI = "mori" NIXL = "nixl" ASCEND = "ascend" FAKE = "fake" @@ -266,6 +267,23 @@ def get_kv_class( KVClassType.BOOTSTRAP_SERVER: MooncakeKVBootstrapServer, } return class_mapping.get(class_type) + elif transfer_backend == TransferBackend.MORI: + from sglang.srt.disaggregation.base import KVArgs + from sglang.srt.disaggregation.mori import ( + MoriKVBootstrapServer, + MoriKVManager, + MoriKVReceiver, + MoriKVSender, + ) + + class_mapping = { + KVClassType.KVARGS: KVArgs, + KVClassType.MANAGER: MoriKVManager, + KVClassType.SENDER: MoriKVSender, + KVClassType.RECEIVER: (MoriKVReceiver), + KVClassType.BOOTSTRAP_SERVER: MoriKVBootstrapServer, + } + return class_mapping.get(class_type) elif transfer_backend == TransferBackend.ASCEND: from sglang.srt.disaggregation.ascend import ( AscendKVBootstrapServer, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 49975de64..6333bd26c 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -143,7 +143,7 @@ ATTENTION_BACKEND_CHOICES = [ LORA_BACKEND_CHOICES = ["triton", "csgmv", "ascend", "torch_native"] -DISAGG_TRANSFER_BACKEND_CHOICES = ["mooncake", "nixl", "ascend", "fake"] +DISAGG_TRANSFER_BACKEND_CHOICES = ["mooncake", "nixl", "ascend", "fake", "mori"] ENCODER_TRANSFER_BACKEND_CHOICES = ["zmq_to_scheduler", "zmq_to_tokenizer", "mooncake"] diff --git a/test/manual/test_mori_transfer_engine_e2e.py b/test/manual/test_mori_transfer_engine_e2e.py new file mode 100644 index 000000000..249d14420 --- /dev/null +++ b/test/manual/test_mori_transfer_engine_e2e.py @@ -0,0 +1,281 @@ +import os +import subprocess +import unittest + +import requests + +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, +) + + +class TestMoriTransferEngineE2E(PDDisaggregationServerBase): + """ + Run: + SGLANG_MORI_MANUAL_E2E=1 python3 test/manual/test_mori_transfer_engine_e2e.py + + Optional: + - SGLANG_MORI_E2E_TEST_MODEL: override model (defaults to a small test model) + - SGLANG_TEST_PD_DISAGG_DEVICES: RDMA devices string, e.g. "mlx5_roce0,mlx5_roce4" + """ + + @classmethod + def setUpClass(cls): + if os.environ.get("SGLANG_MORI_MANUAL_E2E", "") not in ("1", "true", "True"): + raise unittest.SkipTest( + "Set SGLANG_MORI_MANUAL_E2E=1 to run this manual MORI E2E test." + ) + + try: + import torch + + if not torch.cuda.is_available(): + raise unittest.SkipTest("torch.cuda is not available.") + except Exception as e: + raise unittest.SkipTest(f"torch is not available/usable: {e}") + + # Force the disaggregation fixture to use MORI backend in local/manual runs. + os.environ["SGLANG_TEST_PD_DISAGG_BACKEND"] = "mori" + + super().setUpClass() + + cls.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 + ) + cls.wait_server_ready( + cls.decode_url + "/health", timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH + ) + + cls.launch_lb() + + @classmethod + def tearDownClass(cls): + os.environ.pop("SGLANG_TEST_PD_DISAGG_BACKEND", None) + super().tearDownClass() + + @classmethod + def launch_lb(cls): + lb_command = [ + "python3", + "-m", + "sglang_router.launch_router", + "--pd-disaggregation", + "--mini-lb", + "--prefill", + cls.prefill_url, + "--decode", + cls.decode_url, + "--host", + cls.base_host, + "--port", + cls.lb_port, + ] + print("Starting load balancer:", " ".join(lb_command)) + cls.process_lb = subprocess.Popen(lb_command, stdout=None, stderr=None) + cls.wait_server_ready( + cls.lb_url + "/health", timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH + ) + + @classmethod + def start_prefill(cls): + prefill_args = [ + "--trust-remote-code", + "--disaggregation-mode", + "prefill", + "--tp", + "1", + ] + 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", + "--tp", + "1", + "--base-gpu-id", + "1", + ] + 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 test_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 TestMoriTransferEngineTPMismatchE2E(PDDisaggregationServerBase): + """Manual MORI PD-disaggregation E2E with TP mismatch. + + Scenario: + - prefill: tp=2 (GPU 0-1) + - decode: tp=4 (GPU 2-5) + + Manual-only and requires >= 6 visible GPUs. + """ + + _PORT_DELTA = 10 + + @classmethod + def setUpClass(cls): + if os.environ.get("SGLANG_MORI_MANUAL_E2E", "") not in ("1", "true", "True"): + raise unittest.SkipTest( + "Set SGLANG_MORI_MANUAL_E2E=1 to run this manual MORI E2E test." + ) + + try: + import torch + + if not torch.cuda.is_available(): + raise unittest.SkipTest("torch.cuda is not available.") + if torch.cuda.device_count() < 6: + raise unittest.SkipTest( + "TP-mismatch test requires >= 6 visible GPUs (prefill tp=2 + decode tp=4)." + ) + except Exception as e: + raise unittest.SkipTest(f"torch is not available/usable: {e}") + + os.environ["SGLANG_TEST_PD_DISAGG_BACKEND"] = "mori" + super().setUpClass() + + # Shift ports to avoid clashing with TestMoriTransferEngineE2E. + 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.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.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 + ) + cls.wait_server_ready( + cls.decode_url + "/health", timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH + ) + cls.launch_lb() + + @classmethod + def tearDownClass(cls): + os.environ.pop("SGLANG_TEST_PD_DISAGG_BACKEND", None) + super().tearDownClass() + + @classmethod + def launch_lb(cls): + lb_command = [ + "python3", + "-m", + "sglang_router.launch_router", + "--pd-disaggregation", + "--mini-lb", + "--prefill", + cls.prefill_url, + "--decode", + cls.decode_url, + "--host", + cls.base_host, + "--port", + cls.lb_port, + ] + print("Starting load balancer:", " ".join(lb_command)) + cls.process_lb = subprocess.Popen(lb_command, stdout=None, stderr=None) + cls.wait_server_ready( + cls.lb_url + "/health", timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH + ) + + @classmethod + def start_prefill(cls): + prefill_args = [ + "--trust-remote-code", + "--disaggregation-mode", + "prefill", + "--tp", + "2", + ] + 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", + "--tp", + "4", + "--base-gpu-id", + "2", + ] + 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 test_generate_smoke_tp_mismatch(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) + + +if __name__ == "__main__": + unittest.main()