chore: import upstream snapshot with attribution

python/sglang/srt/elastic_ep/elastic_ep.py

@@ -0,0 +1,203 @@
+from __future__ import annotations
+
+import logging
+import time
+from dataclasses import dataclass
+from typing import Iterator, List, Optional
+
+import torch
+
+from sglang.srt.distributed import parallel_state
+from sglang.srt.managers.schedule_batch import ServerArgs
+from sglang.srt.utils import is_cpu, is_cuda
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ElasticEPState:
+    active_ranks: Optional[torch.Tensor]
+    last_active_ranks: Optional[torch.Tensor]
+    active_ranks_cpu: Optional[torch.Tensor]
+
+    def is_active_equal_last(self) -> bool:
+        return torch.equal(self.active_ranks, self.last_active_ranks)
+
+    def sync_active_to_cpu(self):
+        if self.active_ranks is not None:
+            self.active_ranks_cpu = self.active_ranks.detach().cpu().clone()
+
+    def snapshot_active_to_last(self):
+        if self.active_ranks is not None:
+            self.last_active_ranks = self.active_ranks.clone()
+
+    def reset(self):
+        if self.active_ranks is not None:
+            self.active_ranks.fill_(1)
+            self.snapshot_active_to_last()
+            self.sync_active_to_cpu()
+
+
+class ElasticEPStateManager:
+    _instance: Optional[ElasticEPState] = None
+
+    @classmethod
+    def instance(cls) -> ElasticEPState:
+        return cls._instance
+
+    @classmethod
+    def init(cls, server_args: ServerArgs):
+        if cls._instance is not None:
+            return cls._instance
+
+        if server_args.elastic_ep_backend is not None:
+            cls._instance = cls._build_state(ep_size=None, device=None)
+            if server_args.elastic_ep_rejoin:
+                # Mask out peer ranks to perform cuda graph capture on its own
+                cls._instance.active_ranks.zero_()
+                cls._instance.active_ranks[torch.distributed.get_rank()] = 1
+                cls._instance.snapshot_active_to_last()
+                cls._instance.sync_active_to_cpu()
+
+        return cls._instance
+
+    @staticmethod
+    def _select_device() -> torch.device:
+        if is_cuda():
+            return torch.device("cuda")
+        elif is_cpu():
+            return torch.device("cpu")
+        else:
+            raise NotImplementedError("Only CUDA and CPU support elastic ep now.")
+
+    @classmethod
+    def _build_state(
+        cls, *, ep_size: Optional[int] = None, device: Optional[torch.device] = None
+    ) -> ElasticEPState:
+        active = cls.healthy_rank_state(ep_size=ep_size, device=device)
+        return ElasticEPState(
+            active_ranks=active,
+            last_active_ranks=active.clone(),
+            active_ranks_cpu=active.detach().cpu().clone(),
+        )
+
+    @classmethod
+    def healthy_rank_state(
+        cls, *, ep_size: Optional[int] = None, device: Optional[torch.device] = None
+    ) -> torch.Tensor:
+        size = ep_size if ep_size is not None else torch.distributed.get_world_size()
+        dev = device if device is not None else cls._select_device()
+
+        return torch.ones(size, dtype=torch.int32, device=dev)
+
+
+# ---------------------------------------------------------------------------
+# Helpers for elastic EP recovery
+# ---------------------------------------------------------------------------
+
+
+_PEER_STATE_POLL_INTERVAL_SEC = 0.01
+
+
+def _get_process_group_backend(process_group, device: str):
+    return process_group
+
+
+def _iter_live_parallel_groups() -> Iterator[parallel_state.GroupCoordinator]:
+    groups = []
+    for group_ref in parallel_state._groups.values():
+        group = group_ref()
+        if group is not None:
+            groups.append(group)
+    for group in sorted(groups, key=lambda x: x.unique_name):
+        yield group
+
+
+def _map_global_to_group_local_ranks(
+    group_ranks: List[int], global_ranks: List[int]
+) -> List[int]:
+    rank_to_local = {rank: idx for idx, rank in enumerate(group_ranks)}
+    return [rank_to_local[rank] for rank in global_ranks if rank in rank_to_local]
+
+
+def _wait_for_peer_state(mooncake_ep, backend, ranks: List[int]) -> None:
+    # Relaunched ranks become recoverable asynchronously, so we poll until the
+    # target backend reports all requested peers as ready.
+    while not all(mooncake_ep.get_peer_state(backend, ranks)):
+        time.sleep(_PEER_STATE_POLL_INTERVAL_SEC)
+
+
+def _maybe_create_message_queue(group) -> None:
+    if not group.use_message_queue_broadcaster or group.world_size <= 1:
+        return
+
+    from sglang.srt.distributed.device_communicators.shm_broadcast import MessageQueue
+
+    group.mq_broadcaster = MessageQueue.create_from_process_group(
+        group.cpu_group, 1 << 22, 6
+    )
+
+
+def _refresh_ep_members() -> None:
+    from sglang.srt.layers.moe.token_dispatcher.mooncake import EPBuffer
+
+    EPBuffer.get_existing_buffer().update_ep_member()
+
+
+def try_recover_ranks(global_ranks: List[int]) -> bool:
+    from mooncake import ep as mooncake_ep
+
+    world_backend = _get_process_group_backend(torch.distributed.group.WORLD, "cuda")
+    if not all(mooncake_ep.get_peer_state(world_backend, global_ranks)):
+        # The relaunched ranks have not finished initializing yet.
+        return False
+
+    # Recover the world backend first, then recover each derived process group
+    # using ranks mapped into that group's local rank space.
+    mooncake_ep.recover_ranks(world_backend, global_ranks)
+
+    for group in _iter_live_parallel_groups():
+        group_local_ranks = _map_global_to_group_local_ranks(group.ranks, global_ranks)
+        if not group_local_ranks:
+            continue
+
+        device_backend = _get_process_group_backend(group.device_group, "cuda")
+        _wait_for_peer_state(mooncake_ep, device_backend, group_local_ranks)
+        mooncake_ep.recover_ranks(device_backend, group_local_ranks)
+
+        cpu_backend = _get_process_group_backend(group.cpu_group, "cpu")
+        _wait_for_peer_state(mooncake_ep, cpu_backend, group_local_ranks)
+        mooncake_ep.recover_ranks(cpu_backend, group_local_ranks)
+        _maybe_create_message_queue(group)
+
+    _refresh_ep_members()
+    return True
+
+
+def join_process_groups():
+    from mooncake import ep as mooncake_ep
+
+    def join_backend(label: str, backend) -> None:
+        logger.info("Recovered rank joining Mooncake backend %s", label)
+        mooncake_ep.join_group(backend)
+
+    join_backend(
+        "default_world",
+        _get_process_group_backend(torch.distributed.group.WORLD, "cuda"),
+    )
+
+    for group in _iter_live_parallel_groups():
+        if group.world_size <= 1:
+            continue
+
+        join_backend(
+            f"{group.unique_name}:device",
+            _get_process_group_backend(group.device_group, "cuda"),
+        )
+        join_backend(
+            f"{group.unique_name}:cpu",
+            _get_process_group_backend(group.cpu_group, "cpu"),
+        )
+        _maybe_create_message_queue(group)
+
+    _refresh_ep_members()

python/sglang/srt/elastic_ep/expert_backup_client.py

@@ -0,0 +1,173 @@
+import logging
+import re
+import threading
+import time
+
+import torch
+import zmq
+
+from sglang.srt.distributed.parallel_state import (
+    get_world_group,
+    get_world_size,
+)
+from sglang.srt.environ import envs
+from sglang.srt.eplb.expert_location import get_global_expert_location_metadata
+from sglang.srt.managers.io_struct import UpdateExpertBackupReq, sock_recv, sock_send
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils.network import get_local_ip_auto
+
+PORT_BASE = envs.SGLANG_BACKUP_PORT_BASE.get()
+logger = logging.getLogger(__name__)
+
+
+def extract_layer_and_expert_id(param_name):
+    pattern = r"layers\.(\d+)\.mlp\.experts\.(\d+)\.(.+?)\."
+    match = re.search(pattern, param_name)
+    if match:
+        return int(match.group(1)), int(match.group(2)), match.group(3)
+    return -1, -1, ""
+
+
+class ExpertBackupClient:
+    def __init__(self, server_args: ServerArgs, model_runner):
+        context = zmq.Context(2)
+        self.server_args = server_args
+        self.engine_num = server_args.nnodes
+        self.engine_rank = server_args.node_rank
+        self.recv_list = [None] * self.engine_num
+        self.ready_sockets = [None] * self.engine_num
+        self.model_runner = model_runner
+        self.moe_ep_size = model_runner.moe_ep_size
+        self.model_config = model_runner.model_config
+        self.moe_ep_rank = model_runner.moe_ep_rank
+        self.dram_map_list = [None] * self.engine_num
+        self.session_id_list = [None] * self.engine_num
+        self.transfer_engine = None
+        self.gpu_buffer = None
+        self.buffer_size = 0
+        self.use_backup = False
+        local_ip = get_local_ip_auto()
+        all_ips = [None] * get_world_size()
+        torch.distributed.all_gather_object(
+            all_ips, local_ip, group=get_world_group().cpu_group
+        )
+        logger.info(f"all_ips: {all_ips}")
+
+        for i in range(self.engine_num):
+            self.recv_list[i] = context.socket(zmq.SUB)
+            self.recv_list[i].connect(
+                f"tcp://{all_ips[i * get_world_size() // server_args.nnodes]}:{PORT_BASE + i * 2 + 1}"
+            )
+            self.recv_list[i].setsockopt(zmq.SUBSCRIBE, b"")
+
+            # Synchronization channel to notify the manager when this client is ready.
+            self.ready_sockets[i] = context.socket(zmq.PUSH)
+            self.ready_sockets[i].connect(
+                f"tcp://{all_ips[i * get_world_size() // server_args.nnodes]}:{PORT_BASE + i * 2}"
+            )
+            sock_send(self.ready_sockets[i], UpdateExpertBackupReq())
+
+        self._receive_thread = threading.Thread(target=self._receive_loop, daemon=True)
+        self._receive_thread.start()
+
+    def _receive_loop(self):
+        cnt = 0
+        while cnt < self.engine_num:
+            response = sock_recv(self.recv_list[cnt])
+            self.dram_map_list[response.rank] = response.weight_pointer_map
+            self.session_id_list[response.rank] = response.session_id
+            self.buffer_size = max(self.buffer_size, response.buffer_size)
+            cnt += 1
+
+        self.use_backup = True
+        self.start_transfer_client()
+
+    def start_transfer_client(self):
+        from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine
+
+        self.transfer_engine = get_mooncake_transfer_engine()
+
+        self.params_dict = dict(self.model_runner.model.named_parameters())
+        for name, param in self.params_dict.items():
+            param_data = param.data
+            ret_value = self.transfer_engine.engine.register_memory(
+                param_data.data_ptr(), param_data.numel() * param_data.element_size()
+            )
+            if ret_value != 0:
+                self.use_backup = False
+                logger.warning("Register fails. Stop using expert weight backup!")
+                break
+
+    def update_weights(self, weight_name_filter=None):
+        global_expert_location_metadata = get_global_expert_location_metadata()
+        num_experts = (
+            self.model_config.hf_config.n_routed_experts
+            + self.server_args.ep_num_redundant_experts
+        )
+        num_local_experts = num_experts // self.moe_ep_size
+        for i in range(self.engine_num):
+            server_ptr_list = []
+            local_ptr_list = []
+            weight_size_list = []
+
+            for name, weight_info in self.dram_map_list[i].items():
+                if weight_name_filter is not None and not weight_name_filter(name):
+                    continue
+                layer_id, expert_id, weight_name = extract_layer_and_expert_id(name)
+                if layer_id >= self.model_config.hf_config.num_hidden_layers:
+                    continue
+
+                if weight_name == "gate_proj":
+                    shard_id = "w1"
+                    param_name = "experts.w13_"
+                elif weight_name == "down_proj":
+                    shard_id = "w2"
+                    param_name = "experts.w2_"
+                elif weight_name == "up_proj":
+                    shard_id = "w3"
+                    param_name = "experts.w13_"
+                else:
+                    raise RuntimeError(f"Unknown weight name {weight_name}")
+
+                name = name.replace(f"experts.{expert_id}.{weight_name}.", param_name)
+                weight_param = self.params_dict[name]
+
+                physical_expert_ids = (
+                    global_expert_location_metadata.logical_to_all_physical(
+                        layer_id, expert_id
+                    )
+                )
+                for physical_expert_id in physical_expert_ids:
+                    if physical_expert_id not in range(
+                        num_local_experts * self.moe_ep_rank,
+                        num_local_experts * (self.moe_ep_rank + 1),
+                    ):
+                        continue
+                    param = weight_param[physical_expert_id % num_local_experts]
+                    if shard_id == "w1":
+                        param = param.narrow(0, 0, param.shape[0] // 2)
+                    elif shard_id == "w3":
+                        param = param.narrow(
+                            0, param.shape[0] // 2, param.shape[0] // 2
+                        )
+                    server_ptr_list.append(weight_info["weight_ptr"])
+                    local_ptr_list.append(param.data_ptr())
+                    assert (
+                        param.numel() * param.element_size() == weight_info["byte_size"]
+                    )
+                    weight_size_list.append(weight_info["byte_size"])
+            before_transfer = time.time()
+            ret = self.transfer_engine.engine.batch_transfer_sync_read(
+                self.session_id_list[i],
+                local_ptr_list,
+                server_ptr_list,
+                weight_size_list,
+            )
+            after_transfer = time.time()
+            logger.info(f"transfer time = {after_transfer - before_transfer} s")
+
+            if ret != 0:
+                raise RuntimeError(
+                    f"Failed to read weights from backup, error code: {ret}"
+                )
+        return

python/sglang/srt/elastic_ep/expert_backup_manager.py

@@ -0,0 +1,186 @@
+import logging
+import multiprocessing as mp
+import re
+import signal
+
+import torch
+import zmq
+
+from sglang.srt.configs.load_config import LoadConfig
+from sglang.srt.configs.model_config import ModelConfig
+from sglang.srt.environ import envs
+from sglang.srt.managers.io_struct import BackupDramReq, sock_recv, sock_send
+from sglang.srt.model_loader.loader import DefaultModelLoader, get_model_loader
+from sglang.srt.model_loader.utils import set_default_torch_dtype
+from sglang.srt.server_args import (
+    PortArgs,
+    ServerArgs,
+    set_global_server_args_for_scheduler,
+)
+from sglang.srt.utils.network import get_local_ip_auto
+
+PORT_BASE = envs.SGLANG_BACKUP_PORT_BASE.get()
+logger = logging.getLogger(__name__)
+
+
+def extract_expert_id(param_name):
+    pattern = r"\.experts\.(\d+)\."
+    match = re.search(pattern, param_name)
+    if match:
+        return int(match.group(1))
+    return -1
+
+
+class ExpertBackupManager:
+    def __init__(self, server_args: ServerArgs, port_args: PortArgs):
+        self.load_format = server_args.load_format
+        self.model_config = ModelConfig.from_server_args(server_args)
+        self.continuous_buffer = None
+        self.weight_pointer_map = {}
+        self.transfer_engine = None
+        self.session_id = None
+        self.engine_num = server_args.nnodes
+        self.engine_rank = server_args.node_rank
+        self.expert_num = self.model_config.hf_config.n_routed_experts
+        self.idmn = (self.expert_num // self.engine_num) * self.engine_rank
+        self.idmx = (self.expert_num // self.engine_num) * (self.engine_rank + 1)
+        context = zmq.Context(2)
+        # Synchronization socket to avoid PUB/SUB slow joiner issues.
+        self.recv_from_expert_backup_client = context.socket(zmq.PULL)
+        self.recv_from_expert_backup_client.bind(
+            f"tcp://{get_local_ip_auto()}:{PORT_BASE + server_args.node_rank * 2}"
+        )
+        self.send_to_expert_backup_client = context.socket(zmq.PUB)
+        self.send_to_expert_backup_client.bind(
+            f"tcp://{get_local_ip_auto()}:{PORT_BASE + server_args.node_rank * 2 + 1}"
+        )
+        self.backup_weights_from_disk()
+        self.start_transfer_server()
+
+        # Block until all expert backup clients have reported readiness, to avoid
+        # losing the initial PUB message due to slow joiners.
+        num_ready_clients = 0
+
+        while num_ready_clients < server_args.tp_size:
+            sock_recv(self.recv_from_expert_backup_client)
+            num_ready_clients += 1
+
+        back_req = BackupDramReq(
+            rank=self.engine_rank,
+            weight_pointer_map=self.weight_pointer_map,
+            session_id=self.session_id,
+            buffer_size=self.continuous_buffer.numel()
+            * self.continuous_buffer.element_size(),
+        )
+        sock_send(self.send_to_expert_backup_client, back_req)
+
+        # Keep the manager subprocess alive until signals
+        signal.pause()
+
+    def backup_weights_from_disk(self):
+        load_config = LoadConfig(load_format=self.load_format)
+        loader = get_model_loader(load_config, self.model_config)
+
+        with set_default_torch_dtype(self.model_config.dtype):
+            iter = loader._get_weights_iterator(
+                DefaultModelLoader.Source.init_new(self.model_config, None)
+            )
+
+            total_bytes = 0
+            weight_info_dict = {}
+
+            for name, weight in iter:
+                expert_id = extract_expert_id(name)
+                if expert_id < self.idmx and expert_id >= self.idmn:
+                    numel = weight.numel()
+                    element_size = weight.element_size()
+                    byte_size = numel * element_size
+                    weight_info_dict[name] = {
+                        "name": name,
+                        "weight": weight,
+                        "numel": numel,
+                        "shape": weight.shape,
+                        "dtype": weight.dtype,
+                        "element_size": element_size,
+                        "byte_size": byte_size,
+                    }
+                    total_bytes += byte_size
+
+            if total_bytes == 0:
+                self.continuous_buffer = None
+                self.weight_pointer_map = {}
+                return
+
+            self.continuous_buffer = torch.empty(
+                total_bytes, dtype=torch.uint8, device="cpu"
+            )
+            buffer_base_ptr = self.continuous_buffer.data_ptr()
+            self.weight_pointer_map = {}
+            current_byte_offset = 0
+
+            for name in sorted(weight_info_dict.keys()):
+                weight_info = weight_info_dict[name]
+                weight = weight_info["weight"]
+                byte_size = weight_info["byte_size"]
+                weight_flat = weight.flatten().contiguous()
+                weight_bytes = weight_flat.view(torch.uint8)
+                start_byte = current_byte_offset
+                end_byte = current_byte_offset + byte_size
+                weight_ptr = buffer_base_ptr + current_byte_offset
+                self.continuous_buffer[start_byte:end_byte].copy_(weight_bytes)
+                self.weight_pointer_map[name] = {
+                    "name": name,
+                    "weight_ptr": weight_ptr,
+                    "shape": weight_info["shape"],
+                    "numel": weight_info["numel"],
+                    "dtype": weight_info["dtype"],
+                    "element_size": weight_info["element_size"],
+                    "byte_size": byte_size,
+                }
+
+                current_byte_offset = end_byte
+
+    def start_transfer_server(self):
+        from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine
+
+        self.transfer_engine = get_mooncake_transfer_engine()
+        self.session_id = self.transfer_engine.session_id
+        server_ptr = self.continuous_buffer.data_ptr()
+        server_len = (
+            self.continuous_buffer.numel() * self.continuous_buffer.element_size()
+        )
+
+        ret_value = self.transfer_engine.engine.register_memory(server_ptr, server_len)
+        if ret_value != 0:
+            raise RuntimeError("Mooncake memory registration failed.")
+
+
+def run_expert_backup_manager_process(
+    server_args: ServerArgs,
+    port_args: PortArgs,
+):
+    set_global_server_args_for_scheduler(server_args)
+    from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
+        init_mooncake_transfer_engine,
+    )
+
+    init_mooncake_transfer_engine(
+        hostname=get_local_ip_auto(),
+        gpu_id=0,
+        ib_device=(
+            server_args.disaggregation_ib_device or server_args.mooncake_ib_device
+        ),
+    )
+    manager = ExpertBackupManager(server_args, port_args)
+
+
+def run_expert_backup_manager(
+    server_args: ServerArgs,
+    port_args: PortArgs,
+):
+    proc = mp.Process(
+        target=run_expert_backup_manager_process,
+        args=(server_args, port_args),
+    )
+    proc.start()
+    return proc