Support users to pass in expert distributions in init_expert_location

docs/backend/native_api.ipynb

@@ -408,19 +408,7 @@
     "print_highlight(response)\n",
     "\n",
     "response = requests.post(f\"http://localhost:{port}/dump_expert_distribution_record\")\n",
-    "print_highlight(response)\n",
-    "\n",
-    "import glob\n",
-    "\n",
-    "output_file = glob.glob(\"expert_distribution_*.csv\")[0]\n",
-    "with open(output_file, \"r\") as f:\n",
-    "    print_highlight(\"\\n| Layer ID | Expert ID | Count |\")\n",
-    "    print_highlight(\"|----------|-----------|--------|\")\n",
-    "    next(f)\n",
-    "    for i, line in enumerate(f):\n",
-    "        if i < 9:\n",
-    "            layer_id, expert_id, count = line.strip().split(\",\")\n",
-    "            print_highlight(f\"| {layer_id:8} | {expert_id:9} | {count:6} |\")"
+    "print_highlight(response)"
    ]
   },
   {

python/sglang/bench_one_batch.py

@@ -129,16 +129,7 @@ def load_model(server_args, port_args, tp_rank):
     suppress_other_loggers()
     rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None
 
-    model_config = ModelConfig(
-        server_args.model_path,
-        trust_remote_code=server_args.trust_remote_code,
-        revision=server_args.revision,
-        context_length=server_args.context_length,
-        model_override_args=server_args.json_model_override_args,
-        is_embedding=server_args.is_embedding,
-        dtype=server_args.dtype,
-        quantization=server_args.quantization,
-    )
+    model_config = ModelConfig.from_server_args(server_args)
     model_runner = ModelRunner(
         model_config=model_config,
         mem_fraction_static=server_args.mem_fraction_static,

python/sglang/srt/configs/model_config.py

@@ -24,6 +24,7 @@
 
 from sglang.srt.hf_transformers_utils import get_config, get_context_length
 from sglang.srt.layers.quantization import QUANTIZATION_METHODS
+from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import get_bool_env_var, is_hip
 
 logger = logging.getLogger(__name__)
@@ -171,6 +172,19 @@ def __init__(
         self.hf_eos_token_id = self.get_hf_eos_token_id()
         self.image_token_id = getattr(self.hf_config, "image_token_id", None)
 
+    @staticmethod
+    def from_server_args(server_args: ServerArgs, model_path: str = None):
+        return ModelConfig(
+            model_path=model_path or server_args.model_path,
+            trust_remote_code=server_args.trust_remote_code,
+            revision=server_args.revision,
+            context_length=server_args.context_length,
+            model_override_args=server_args.json_model_override_args,
+            is_embedding=server_args.is_embedding,
+            dtype=server_args.dtype,
+            quantization=server_args.quantization,
+        )
+
     # adapted from https://github.com/vllm-project/vllm/blob/main/vllm/config.py#L289
     def get_total_num_kv_heads(self) -> int:
         """Returns the total number of KV heads."""

python/sglang/srt/entrypoints/engine.py

@@ -20,6 +20,7 @@
 import asyncio
 import atexit
 import dataclasses
+import json
 import logging
 import multiprocessing as mp
 import os
@@ -31,6 +32,10 @@
 import zmq.asyncio
 from PIL.Image import Image
 
+from sglang.srt.configs.model_config import ModelConfig
+from sglang.srt.managers.eplb_manager import EPLBManager
+from sglang.srt.managers.expert_location import ExpertLocationMetadata
+
 # Fix a bug of Python threading
 setattr(threading, "_register_atexit", lambda *args, **kwargs: None)
 
@@ -51,6 +56,7 @@
     ResumeMemoryOccupationReqInput,
     RpcReqInput,
     RpcReqOutput,
+    UpdateExpertLocationReqInput,
     UpdateWeightFromDiskReqInput,
     UpdateWeightsFromDistributedReqInput,
     UpdateWeightsFromTensorReqInput,
@@ -278,6 +284,10 @@ def __exit__(self, exc_type, exc_value, traceback):
         self.shutdown()
         return False
 
+    def flush_cache(self):
+        loop = asyncio.get_event_loop()
+        return loop.run_until_complete(self.tokenizer_manager.flush_cache())
+
     def start_profile(self):
         loop = asyncio.get_event_loop()
         loop.run_until_complete(self.tokenizer_manager.start_profile())
@@ -354,10 +364,30 @@ def update_weights_from_tensor(
             self.tokenizer_manager.update_weights_from_tensor(obj, None)
         )
 
+    def eplb_rebalance(self):
+        loop = asyncio.get_event_loop()
+        return loop.run_until_complete(self.tokenizer_manager.eplb_rebalance())
+
+    def eplb_save_expert_distribution(self):
+        loop = asyncio.get_event_loop()
+        return loop.run_until_complete(
+            self.tokenizer_manager.eplb_save_expert_distribution()
+        )
+
+    def update_expert_location(self, expert_location_metadata: ExpertLocationMetadata):
+        obj = UpdateExpertLocationReqInput(
+            expert_location_metadata=expert_location_metadata,
+        )
+        loop = asyncio.get_event_loop()
+        return loop.run_until_complete(
+            self.tokenizer_manager.update_expert_location(obj)
+        )
+
     def update_weights_from_disk(
         self,
         model_path: str,
         load_format: Optional[str] = None,
+        param_categories: Optional[List[str]] = None,
     ):
         """Update the weights from disk inplace without re-launching the engine.
 
@@ -368,6 +398,7 @@ def update_weights_from_disk(
         obj = UpdateWeightFromDiskReqInput(
             model_path=model_path,
             load_format=load_format,
+            param_categories=param_categories,
         )
 
         loop = asyncio.get_event_loop()
@@ -495,6 +526,11 @@ def _launch_subprocesses(
         server_args.model_path, server_args.tokenizer_path
     )
 
+    eplb_manager = EPLBManager(server_args) if server_args.enable_eplb else None
+    expert_location_metadata = _compute_initial_expert_location_metadata(
+        server_args, eplb_manager
+    )
+
     scheduler_procs = []
     if server_args.dp_size == 1:
         # Launch tensor parallel scheduler processes
@@ -516,7 +552,15 @@ def _launch_subprocesses(
             )
             proc = mp.Process(
                 target=run_scheduler_process,
-                args=(server_args, port_args, gpu_id, tp_rank, None, writer),
+                args=(
+                    server_args,
+                    port_args,
+                    expert_location_metadata,
+                    gpu_id,
+                    tp_rank,
+                    None,
+                    writer,
+                ),
             )
             with memory_saver_adapter.configure_subprocess():
                 proc.start()
@@ -528,7 +572,7 @@ def _launch_subprocesses(
         scheduler_pipe_readers = [reader]
         proc = mp.Process(
             target=run_data_parallel_controller_process,
-            args=(server_args, port_args, writer),
+            args=(server_args, port_args, expert_location_metadata, writer),
         )
         proc.start()
         scheduler_procs.append(proc)
@@ -565,7 +609,9 @@ def _launch_subprocesses(
     detoken_proc.start()
 
     # Launch tokenizer process
-    tokenizer_manager = TokenizerManager(server_args, port_args)
+    tokenizer_manager = TokenizerManager(
+        server_args, port_args, expert_location_metadata, eplb_manager
+    )
     if server_args.chat_template:
         load_chat_template_for_openai_api(
             tokenizer_manager, server_args.chat_template, server_args.model_path
@@ -597,3 +643,22 @@ def _launch_subprocesses(
     scheduler_info = scheduler_infos[0]
     tokenizer_manager.max_req_input_len = scheduler_info["max_req_input_len"]
     return tokenizer_manager, scheduler_info
+
+
+def _compute_initial_expert_location_metadata(
+    server_args: ServerArgs, eplb_manager: EPLBManager
+) -> ExpertLocationMetadata:
+    if (data := server_args.init_expert_location) is not None:
+        data_dict = json.loads(data)
+        if "physical_to_logical_map" in data_dict:
+            # TODO We may want to allow users to not provide `logical_to_all_physical_map` if this API is frequently used
+            return ExpertLocationMetadata.init_by_mapping(server_args, **data_dict)
+        elif "logical_count" in data_dict:
+            return ExpertLocationMetadata.init_by_eplb(server_args, **data_dict)
+        else:
+            raise NotImplementedError(
+                f"Unknown init_expert_location format ({list(data_dict.keys())=})"
+            )
+    if server_args.enable_eplb:
+        return eplb_manager.compute_expert_location_metadata()
+    return ExpertLocationMetadata.init_trivial(server_args)

python/sglang/srt/entrypoints/http_server.py

@@ -310,11 +310,11 @@ async def classify_request(obj: EmbeddingReqInput, request: Request):
 @app.api_route("/flush_cache", methods=["GET", "POST"])
 async def flush_cache():
     """Flush the radix cache."""
-    _global_state.tokenizer_manager.flush_cache()
+    ret = await _global_state.tokenizer_manager.flush_cache()
     return Response(
         content="Cache flushed.\nPlease check backend logs for more details. "
         "(When there are running or waiting requests, the operation will not be performed.)\n",
-        status_code=200,
+        status_code=200 if ret.success else HTTPStatus.BAD_REQUEST,
     )
 
 
@@ -366,11 +366,20 @@ async def stop_expert_distribution_record_async():
 @app.api_route("/dump_expert_distribution_record", methods=["GET", "POST"])
 async def dump_expert_distribution_record_async():
     """Dump expert distribution record."""
-    await _global_state.tokenizer_manager.dump_expert_distribution_record()
-    return Response(
-        content="Dump expert distribution record.\n",
-        status_code=200,
-    )
+    content = await _global_state.tokenizer_manager.dump_expert_distribution_record()
+    return ORJSONResponse(content, status_code=200)
+
+
+@app.post("/eplb_rebalance")
+async def eplb_rebalance():
+    await _global_state.tokenizer_manager.eplb_rebalance()
+    return ORJSONResponse({}, status_code=200)
+
+
+@app.post("/eplb_save_expert_distribution")
+async def eplb_save_expert_distribution():
+    await _global_state.tokenizer_manager.eplb_save_expert_distribution()
+    return ORJSONResponse({}, status_code=200)
 
 
 @app.post("/update_weights_from_disk")

python/sglang/srt/layers/moe/ep_moe/layer.py

@@ -3,6 +3,8 @@
 
 import torch
 
+from sglang.srt.managers.schedule_batch import get_global_expert_location_metadata
+
 try:
     from deep_gemm import (
         get_col_major_tma_aligned_tensor,
@@ -131,6 +133,7 @@ def __init__(
         top_k: int,
         hidden_size: int,
         intermediate_size: int,
+        layer_id: int,
         params_dtype: Optional[torch.dtype] = None,
         renormalize: bool = True,
         use_grouped_topk: bool = False,
@@ -153,6 +156,7 @@ def __init__(
         )
         self.tp_rank = get_tensor_model_parallel_rank()
 
+        self.layer_id = layer_id
         self.num_experts = num_experts
         assert self.num_experts % self.tp_size == 0
         self.num_experts_per_partition = self.num_experts // self.tp_size
@@ -221,6 +225,9 @@ def forward(self, hidden_states: torch.Tensor, router_logits: torch.Tensor):
             num_expert_group=self.num_expert_group,
             correction_bias=self.correction_bias,
             custom_routing_function=self.custom_routing_function,
+            expert_logical_to_rank_dispatch_physical_map=get_global_expert_location_metadata().logical_to_rank_dispatch_physical_map[
+                self.tp_rank, self.layer_id, :
+            ],
         )
 
         reorder_topk_ids, src2dst, seg_indptr = run_moe_ep_preproess(
@@ -409,6 +416,28 @@ def weight_loader(
         weight_name: str,
         shard_id: str,
         expert_id: int,
+    ) -> None:
+        physical_expert_ids = (
+            get_global_expert_location_metadata().logical_to_all_physical(
+                self.layer_id, expert_id
+            )
+        )
+        for physical_expert_id in physical_expert_ids:
+            self._weight_loader_physical(
+                param=param,
+                loaded_weight=loaded_weight,
+                weight_name=weight_name,
+                shard_id=shard_id,
+                expert_id=physical_expert_id,
+            )
+
+    def _weight_loader_physical(
+        self,
+        param: torch.nn.Parameter,
+        loaded_weight: torch.Tensor,
+        weight_name: str,
+        shard_id: str,
+        expert_id: int,
     ) -> None:
         if expert_id < self.start_expert_id or expert_id > self.end_expert_id:
             return
@@ -802,6 +831,7 @@ def __init__(
         top_k: int,
         hidden_size: int,
         intermediate_size: int,
+        layer_id: int,
         params_dtype: Optional[torch.dtype] = None,
         renormalize: bool = True,
         use_grouped_topk: bool = False,
@@ -820,6 +850,7 @@ def __init__(
             top_k,
             hidden_size,
             intermediate_size,
+            layer_id,
             params_dtype,
             renormalize,
             use_grouped_topk,

python/sglang/srt/layers/moe/ep_moe/token_dispatcher.py

@@ -1,3 +1,4 @@
+from sglang.srt.managers.expert_distribution import expert_distribution_recorder
 from sglang.srt.utils import DeepEPMode
 
 try:
@@ -248,7 +249,7 @@ def _dispatch_core(
             recv_x,
             recv_topk_idx,
             recv_topk_weights,
-            _,  # num_recv_tokens_per_expert_list
+            num_recv_tokens_per_expert_list,
             self.handle,
             event,
         ) = buffer.dispatch(
@@ -264,6 +265,10 @@ def _dispatch_core(
             allocate_on_comm_stream=(previous_event is not None) and self.async_finish,
         )
 
+        expert_distribution_recorder.on_deepep_dispatch_normal(
+            num_recv_tokens_per_expert_list
+        )
+
         return (
             recv_x,
             recv_topk_idx,

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -268,6 +268,7 @@ def __init__(
         top_k: int,
         hidden_size: int,
         intermediate_size: int,
+        layer_id: Optional[int] = None,
         params_dtype: Optional[torch.dtype] = None,
         reduce_results: bool = False,
         renormalize: bool = True,