sglang 0.5.3rc2__py3-none-any.whl → 0.5.4.post1__py3-none-any.whl

sglang/srt/entrypoints/openai/serving_classify.py

@@ -0,0 +1,204 @@
+from __future__ import annotations
+
+import logging
+import time
+import uuid
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
+
+import torch
+import torch.nn.functional as F
+from fastapi import Request
+from fastapi.responses import ORJSONResponse
+
+from sglang.srt.entrypoints.openai.protocol import (
+    ClassifyRequest,
+    ClassifyResponse,
+    ErrorResponse,
+)
+from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase
+from sglang.srt.managers.io_struct import EmbeddingReqInput
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.template_manager import TemplateManager
+    from sglang.srt.managers.tokenizer_manager import TokenizerManager
+
+logger = logging.getLogger(__name__)
+
+
+class OpenAIServingClassify(OpenAIServingBase):
+    """Handler for v1/classify requests"""
+
+    def __init__(
+        self,
+        tokenizer_manager: TokenizerManager,
+        template_manager: TemplateManager,
+    ):
+        super().__init__(tokenizer_manager)
+        self.template_manager = template_manager
+        self.id2label = self._get_id2label_mapping()
+        self.model_name = (
+            self.tokenizer_manager.served_model_name
+            if self.tokenizer_manager.served_model_name
+            else self.tokenizer_manager.server_args.model_path
+        )
+        if not self.id2label:
+            raise ValueError("id2label mapping is missing")
+
+    def _request_id_prefix(self) -> str:
+        return "classify-"
+
+    def _convert_to_internal_request(
+        self,
+        request: ClassifyRequest,
+        raw_request: Request = None,
+    ) -> tuple[EmbeddingReqInput, ClassifyRequest]:
+        """Convert OpenAI embedding request to internal format"""
+        prompt = request.input
+
+        if isinstance(prompt, str):
+            # Single string input
+            prompt_kwargs = {"text": prompt}
+        elif isinstance(prompt, list):
+            if len(prompt) > 0 and isinstance(prompt[0], str):
+                prompt_kwargs = {"text": prompt}
+            else:
+                # List of integers (token IDs) or empty list
+                prompt_kwargs = {"input_ids": prompt}
+        else:
+            # Other types (should not happen but handle gracefully)
+            prompt_kwargs = {"input_ids": prompt}
+
+        adapted_request = EmbeddingReqInput(
+            **prompt_kwargs,
+            rid=request.rid,
+            priority=request.priority,
+        )
+
+        return adapted_request, request
+
+    def _validate_request(self, request: ClassifyRequest) -> Optional[str]:
+        """Validate that the input is not empty or whitespace only."""
+        if not (input := request.input):
+            return "Input cannot be empty"
+
+        # Handle single string
+        if isinstance(input, str):
+            if not input.strip():
+                return "Input cannot be empty or whitespace only"
+            return None
+
+        # Handle list inputs
+        if isinstance(input, list):
+            # Check first element to determine type
+            first_item = input[0]
+
+            if isinstance(first_item, str):
+                # List of strings
+                for i, item in enumerate(input):
+                    if not isinstance(item, str):
+                        return f"All items in input list must be strings"
+                    if not item.strip():
+                        return f"Input at index {i} cannot be empty or whitespace only"
+            elif isinstance(first_item, int):
+                # List of integers (token IDs)
+                for i, item in enumerate(input):
+                    if not isinstance(item, int):
+                        return f"All items in input list must be integers"
+                    if item < 0:
+                        return f"Token ID at index {i} must be non-negative"
+        return None
+
+    def _get_id2label_mapping(self) -> Optional[Dict[int, str]]:
+        """Get id2label mapping from model config."""
+        try:
+            hf_config = self.tokenizer_manager.model_config.hf_config
+            # Check for id2label in hf_config
+            if hf_config.id2label:
+                return hf_config.id2label
+            # Check for num_labels and create default mapping if needed
+            if hasattr(hf_config, "num_labels") and hf_config.num_labels:
+                num_labels = hf_config.num_labels
+                # Create default mapping: {0: "LABEL_0", 1: "LABEL_1", ...}
+                return {i: f"LABEL_{i}" for i in range(num_labels)}
+
+        except Exception as e:
+            logger.warning(f"Failed to get id2label mapping: {e}")
+
+        return None
+
+    async def _handle_non_streaming_request(
+        self,
+        adapted_request: EmbeddingReqInput,
+        request: ClassifyRequest,
+        raw_request: Request,
+    ) -> Union[ClassifyResponse, ErrorResponse, ORJSONResponse]:
+        """Handle non-streaming classification request."""
+        # Generate request ID
+
+        try:
+            ret = await self.tokenizer_manager.generate_request(
+                adapted_request, raw_request
+            ).__anext__()
+        except ValueError as e:
+            return self.create_error_response(str(e))
+
+        if not isinstance(ret, list):
+            ret = [ret]
+
+        response = self._build_classify_response(ret)
+        return response
+
+    def _build_classify_response(self, ret: List[Dict[str, Any]]) -> ClassifyResponse:
+        request_id = f"{self._request_id_prefix()}{uuid.uuid4().hex}"
+        created_time = int(time.time())
+        classify_objects = []
+        prompt_tokens = 0
+        total_latency = 0.0
+
+        for i, item in enumerate(ret):
+            embedding = item.get("embedding", [])
+            meta_info = item.get("meta_info", {})
+
+            prompt_tokens += meta_info.get("prompt_tokens", 0)
+            total_latency += meta_info.get("e2e_latency", 0.0)
+
+            if embedding:
+                try:
+                    embedding_tensor = torch.tensor(embedding, dtype=torch.float32)
+                    probs = F.softmax(embedding_tensor, dim=0).tolist()
+
+                    predicted_class = torch.argmax(embedding_tensor).item()
+
+                    label = self.id2label[predicted_class]
+
+                except Exception as e:
+                    logger.error(f"Error processing embedding for item {i}: {e}")
+                    probs = [1.0]
+                    label = "Default"
+            else:
+                probs = [1.0]
+                label = "Default"
+
+            classify_obj = {
+                "index": i,
+                "label": label,
+                "probs": probs,
+                "num_classes": len(probs),
+            }
+            classify_objects.append(classify_obj)
+
+        response = {
+            "id": request_id,
+            "object": "list",
+            "created": created_time,
+            "model": self.model_name,
+            "data": classify_objects,
+            "usage": {
+                "prompt_tokens": prompt_tokens,
+                "total_tokens": prompt_tokens,
+                "completion_tokens": 0,
+                "prompt_tokens_details": None,
+            },
+        }
+
+        return ClassifyResponse(**response)

sglang/srt/entrypoints/openai/serving_completions.py

@@ -93,6 +93,17 @@ class OpenAIServingCompletion(OpenAIServingBase):
         # Extract custom labels from raw request headers
         custom_labels = self.extract_custom_labels(raw_request)
 
+        # Resolve LoRA adapter from model parameter or explicit lora_path
+        lora_path = self._resolve_lora_path(request.model, request.lora_path)
+        if lora_path:
+            first_adapter = (
+                lora_path
+                if isinstance(lora_path, str)
+                else next((a for a in lora_path if a), None)
+            )
+            if first_adapter:
+                self._validate_lora_enabled(first_adapter)
+
         adapted_request = GenerateReqInput(
             **prompt_kwargs,
             sampling_params=sampling_params,
@@ -101,7 +112,7 @@ class OpenAIServingCompletion(OpenAIServingBase):
             logprob_start_len=logprob_start_len,
             return_text_in_logprobs=True,
             stream=request.stream,
-            lora_path=request.lora_path,
+            lora_path=lora_path,
             bootstrap_host=request.bootstrap_host,
             bootstrap_port=request.bootstrap_port,
             bootstrap_room=request.bootstrap_room,
@@ -110,6 +121,7 @@ class OpenAIServingCompletion(OpenAIServingBase):
             extra_key=self._compute_extra_key(request),
             priority=request.priority,
             custom_labels=custom_labels,
+            custom_logit_processor=request.custom_logit_processor,
         )
 
         return adapted_request, request
@@ -123,6 +135,7 @@ class OpenAIServingCompletion(OpenAIServingBase):
             "min_new_tokens": request.min_tokens,
             "stop": request.stop,
             "stop_token_ids": request.stop_token_ids,
+            "stop_regex": request.stop_regex,
             "top_p": request.top_p,
             "top_k": request.top_k,
             "min_p": request.min_p,
@@ -137,6 +150,7 @@ class OpenAIServingCompletion(OpenAIServingBase):
             "ignore_eos": request.ignore_eos,
             "skip_special_tokens": request.skip_special_tokens,
             "logit_bias": request.logit_bias,
+            "custom_params": request.custom_params,
         }
 
         # Handle response_format constraints

sglang/srt/entrypoints/openai/serving_responses.py

@@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Any, AsyncGenerator, AsyncIterator, Optional,
 
 import jinja2
 import openai.types.responses as openai_responses_types
+import orjson
 from fastapi import Request
 from fastapi.responses import ORJSONResponse
 from openai.types.responses import (
@@ -778,7 +779,9 @@ class OpenAIServingResponses(OpenAIServingChat):
             # Update the status to "cancelled"
             response.status = "cancelled"
 
-        # Abort the request
+        # The response_id is the same as the rid used when submitting the request
+        self.tokenizer_manager.abort_request(rid=response_id)
+
         if task := self.background_tasks.get(response_id):
             task.cancel()
             try:
@@ -1061,7 +1064,7 @@ class OpenAIServingResponses(OpenAIServingChat):
                 ):
                     function_name = previous_item.recipient[len("browser.") :]
                     action = None
-                    parsed_args = json.loads(previous_item.content[0].text)
+                    parsed_args = orjson.loads(previous_item.content[0].text)
                     if function_name == "search":
                         action = openai_responses_types.response_function_web_search.ActionSearch(
                             type="search",

sglang/srt/entrypoints/openai/serving_tokenize.py

@@ -0,0 +1,144 @@
+import logging
+from http import HTTPStatus
+from typing import List, Union
+
+from fastapi import Request
+
+from sglang.srt.entrypoints.openai.protocol import (
+    DetokenizeRequest,
+    DetokenizeResponse,
+    ErrorResponse,
+    TokenizeRequest,
+    TokenizeResponse,
+)
+from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase
+
+logger = logging.getLogger(__name__)
+
+
+class OpenAIServingTokenize(OpenAIServingBase):
+    """Handler for /v1/tokenize requests"""
+
+    def _request_id_prefix(self) -> str:
+        return "tok-"
+
+    def _convert_to_internal_request(
+        self, request: TokenizeRequest, raw_request: Request
+    ) -> tuple[TokenizeRequest, TokenizeRequest]:
+        return request, request
+
+    async def _handle_non_streaming_request(
+        self,
+        adapted_request: TokenizeRequest,
+        request: TokenizeRequest,
+        raw_request: Request,
+    ) -> Union[TokenizeResponse, ErrorResponse]:
+        try:
+            tokenizer = self.tokenizer_manager.tokenizer
+            max_model_len = getattr(tokenizer, "model_max_length", -1)
+
+            if isinstance(request.prompt, str):
+                token_ids = tokenizer.encode(
+                    request.prompt,
+                    add_special_tokens=request.add_special_tokens,
+                )
+                tokens = token_ids
+                count = len(token_ids)
+            elif isinstance(request.prompt, list):
+                token_ids_list = [
+                    tokenizer.encode(
+                        text, add_special_tokens=request.add_special_tokens
+                    )
+                    for text in request.prompt
+                ]
+                tokens = token_ids_list
+                count = [len(ids) for ids in token_ids_list]
+            else:
+                return self.create_error_response(
+                    f"Invalid prompt type: {type(request.prompt)}. Expected str or List[str]."
+                )
+
+            return TokenizeResponse(
+                tokens=tokens, count=count, max_model_len=max_model_len
+            )
+        except Exception as e:
+            logger.error("Error during tokenization", exc_info=True)
+            return self.create_error_response(
+                f"Internal server error during tokenization: {e}",
+                err_type="InternalServerError",
+                status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
+            )
+
+
+class OpenAIServingDetokenize(OpenAIServingBase):
+    """Handler for /v1/detokenize requests"""
+
+    def _request_id_prefix(self) -> str:
+        return "detok-"
+
+    def _convert_to_internal_request(
+        self, request: DetokenizeRequest, raw_request: Request
+    ) -> tuple[DetokenizeRequest, DetokenizeRequest]:
+        return request, request
+
+    async def _handle_non_streaming_request(
+        self,
+        adapted_request: DetokenizeRequest,
+        request: DetokenizeRequest,
+        raw_request: Request,
+    ) -> Union[DetokenizeResponse, ErrorResponse]:
+        try:
+            tokenizer = self.tokenizer_manager.tokenizer
+
+            if (
+                isinstance(request.tokens, list)
+                and request.tokens
+                and isinstance(request.tokens[0], int)
+            ):
+                if not all(isinstance(t, int) for t in request.tokens):
+                    return self.create_error_response(
+                        "Invalid input: 'tokens' must be a list of integers."
+                    )
+                tokens_to_decode = [int(t) for t in request.tokens]
+                text = tokenizer.decode(
+                    tokens_to_decode, skip_special_tokens=request.skip_special_tokens
+                )
+                text_out: Union[str, List[str]] = text
+            elif (
+                isinstance(request.tokens, list)
+                and request.tokens
+                and isinstance(request.tokens[0], list)
+            ):
+                texts: List[str] = []
+                for token_list in request.tokens:
+                    if not all(isinstance(t, int) for t in token_list):
+                        return self.create_error_response(
+                            f"Invalid input: Sublist in 'tokens' must contain only integers. Found: {token_list}"
+                        )
+                    decoded_text = tokenizer.decode(
+                        [int(t) for t in token_list],
+                        skip_special_tokens=request.skip_special_tokens,
+                    )
+                    texts.append(decoded_text)
+                text_out = texts
+            elif isinstance(request.tokens, list) and not request.tokens:
+                text_out = ""
+            else:
+                return self.create_error_response(
+                    f"Invalid tokens type: {type(request.tokens)}. Expected List[int] or List[List[int]]."
+                )
+
+            return DetokenizeResponse(text=text_out)
+        except Exception as e:
+            logger.error("Error during detokenization", exc_info=True)
+            if "decode" in str(e).lower():
+                return self.create_error_response(
+                    f"Error decoding tokens: {e}. Input tokens might be invalid for the model.",
+                    err_type="DecodeError",
+                    status_code=HTTPStatus.BAD_REQUEST,
+                )
+            return self.create_error_response(
+                f"Internal server error during detokenization: {e}",
+                err_type="InternalServerError",
+                status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
+            )

sglang/srt/environ.py

@@ -111,25 +111,45 @@ class Envs:
     # Model & File Download
     SGLANG_USE_MODELSCOPE = EnvBool(False)
 
+    # Logging Options
+    SGLANG_LOG_GC = EnvBool(False)
+    SGLANG_LOG_FORWARD_ITERS = EnvBool(False)
+    SGLANG_DISABLE_REQUEST_LOGGING = EnvBool(False)
+
     # Test & Debug
     SGLANG_IS_IN_CI = EnvBool(False)
-    SGLANG_AMD_CI = EnvBool(False)
-    SGLANG_TEST_RETRACT = EnvBool(False)
+    SGLANG_IS_IN_CI_AMD = EnvBool(False)
     SGLANG_SET_CPU_AFFINITY = EnvBool(False)
     SGLANG_PROFILE_WITH_STACK = EnvBool(True)
     SGLANG_RECORD_STEP_TIME = EnvBool(False)
-    SGLANG_GC_LOG = EnvBool(False)
     SGLANG_FORCE_SHUTDOWN = EnvBool(False)
     SGLANG_DEBUG_MEMORY_POOL = EnvBool(False)
     SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False)
     SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False)
-    SGLANG_DISABLE_REQUEST_LOGGING = EnvBool(False)
     SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1)
     SGLANG_SIMULATE_ACC_METHOD = EnvStr("multinomial")
     SGLANG_TORCH_PROFILER_DIR = EnvStr("/tmp")
 
+    # Scheduler: memory leak test
+    SGLANG_TEST_RETRACT = EnvBool(False)
+    SGLANG_TEST_RETRACT_INTERVAL = EnvInt(3)
+    SGLANG_ENABLE_RUNTIME_MEM_LEAK_CHECK = EnvBool(False)
+
+    # Scheduler: new token ratio hyperparameters
+    SGLANG_INIT_NEW_TOKEN_RATIO = EnvFloat(0.7)
+    SGLANG_MIN_NEW_TOKEN_RATIO_FACTOR = EnvFloat(0.14)
+    SGLANG_NEW_TOKEN_RATIO_DECAY_STEPS = EnvInt(600)
+    SGLANG_RETRACT_DECODE_STEPS = EnvInt(20)
+
+    # Scheduler: others:
+    SGLANG_EMPTY_CACHE_INTERVAL = EnvFloat(-1)  # in seconds. Set if you observe high memory accumulation over a long serving period.
+    # Test: pd-disaggregation
+    SGLANG_TEST_PD_DISAGG_BACKEND = EnvStr("mooncake")
+    SGLANG_TEST_PD_DISAGG_DEVICES = EnvStr(None)
+
     # Model Parallel
     SGLANG_USE_MESSAGE_QUEUE_BROADCASTER = EnvBool(True)
+    SGLANG_ONE_VISIBLE_DEVICE_PER_PROCESS = EnvBool(False)
 
     # Constrained Decoding
     SGLANG_DISABLE_OUTLINES_DISK_CACHE = EnvBool(True)
@@ -145,6 +165,7 @@ class Envs:
     # AMD & ROCm
     SGLANG_USE_AITER = EnvBool(False)
     SGLANG_ROCM_FUSED_DECODE_MLA = EnvBool(False)
+    SGLANG_ROCM_DISABLE_LINEARQUANT = EnvBool(False)
 
     # Quantization
     SGLANG_INT4_WEIGHT = EnvBool(False)
@@ -155,6 +176,7 @@ class Envs:
     # Flashinfer
     SGLANG_IS_FLASHINFER_AVAILABLE = EnvBool(True)
     SGLANG_ENABLE_FLASHINFER_GEMM = EnvBool(False)
+    SGLANG_FLASHINFER_WORKSPACE_SIZE = EnvInt(384 * 1024 * 1024)
 
     # Triton
     SGLANG_TRITON_DECODE_ATTN_STATIC_KV_SPLITS = EnvBool(False)
@@ -167,6 +189,7 @@ class Envs:
     SGLANG_EXPERT_LOCATION_UPDATER_CANARY = EnvBool(False)
     SGLANG_EXPERT_LOCATION_UPDATER_LOG_METRICS = EnvBool(False)
     SGLANG_LOG_EXPERT_LOCATION_METADATA = EnvBool(False)
+    SGLANG_EXPERT_DISTRIBUTION_RECORDER_DIR = EnvStr("/tmp")
 
     # TBO
     SGLANG_TBO_DEBUG = EnvBool(False)
@@ -183,12 +206,12 @@ class Envs:
     # sgl-kernel
     SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK = EnvBool(False)
 
-    # vLLM dependencies
+    # vLLM dependencies (TODO: they have been deprecated, we can remove them safely)
     USE_VLLM_CUSTOM_ALLREDUCE = EnvBool(False)
     USE_VLLM_CUTLASS_W8A8_FP8_KERNEL = EnvBool(False)
 
     USE_TRITON_W8A8_FP8_KERNEL = EnvBool(False)
-    RETURN_ORIGINAL_LOGPROB = EnvBool(False)
+    SGLANG_RETURN_ORIGINAL_LOGPROB = EnvBool(False)
     SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN = EnvBool(False)
     SGLANG_MOE_PADDING = EnvBool(False)
     SGLANG_CUTLASS_MOE = EnvBool(False)
@@ -207,13 +230,42 @@ class Envs:
     SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE = EnvInt(4096)
     SGLANG_TRITON_DECODE_SPLIT_TILE_SIZE = EnvInt(256)
 
+    # Overlap Spec V2
+    SGLANG_ENABLE_SPEC_V2 = EnvBool(False)
+    SGLANG_ENABLE_OVERLAP_PLAN_STREAM = EnvBool(False)
+
+    # VLM
+    SGLANG_IMAGE_MAX_PIXELS = EnvInt(16384 * 28 * 28)
+    SGLANG_RESIZE_RESAMPLE = EnvStr("")
+
+    # Ktransformers
+    SGLANG_KT_MOE_NUM_GPU_EXPERTS = EnvInt(None)
+    SGLANG_KT_MOE_CPUINFER = EnvInt(None)
+    SGLANG_KT_THREADPOOL_COUNT = EnvInt(None)
+    SGLANG_KT_MOE_AMX_WEIGHT_PATH = EnvStr(None)
+    SGLANG_KT_AMX_METHOD = EnvStr(None)
+    SGLANG_KT_MOE_CHUNKED_PREFILL_SIZE = EnvInt(None)
+
+    # Sparse Embeddings
+    SGLANG_EMBEDDINGS_SPARSE_HEAD = EnvStr(None)
+
     # fmt: on
 
 
 envs = Envs()
 
 
+def _print_deprecated_env(new_name: str, old_name: str):
+    if old_name in os.environ:
+        warnings.warn(
+            f"Environment variable {old_name} will be deprecated, please use {new_name} instead"
+        )
+        os.environ[new_name] = os.environ[old_name]
+
+
 def _convert_SGL_to_SGLANG():
+    _print_deprecated_env("SGLANG_LOG_GC", "SGLANG_GC_LOG")
+
     for key, value in os.environ.items():
         if key.startswith("SGL_"):
             new_key = key.replace("SGL_", "SGLANG_", 1)

sglang/srt/eplb/eplb_algorithms/__init__.py

@@ -3,7 +3,8 @@ from typing import Optional
 
 import torch
 
-from sglang.srt.eplb.eplb_algorithms import deepseek, deepseek_vec
+from sglang.srt.elastic_ep.elastic_ep import ElasticEPStateManager
+from sglang.srt.eplb.eplb_algorithms import deepseek, deepseek_vec, elasticity_aware
 
 
 class EplbAlgorithm(Enum):
@@ -11,6 +12,7 @@ class EplbAlgorithm(Enum):
     deepseek_hierarchical = auto()
     deepseek_vec = auto()
     deepseek_vec_hierarchical = auto()
+    elasticity_aware = auto()
     # TODO may have more algorithm later
 
 
@@ -45,6 +47,21 @@ def rebalance_experts(
             enable_hierarchical=algorithm == EplbAlgorithm.deepseek_vec_hierarchical,
         )
 
+    if algorithm == EplbAlgorithm.elasticity_aware:
+        return elasticity_aware.rebalance_experts(
+            weight=tokens_per_expert.sum(dim=0),
+            num_replicas=num_physical_experts,
+            num_groups=num_groups,
+            num_nodes=num_nodes,
+            num_gpus=num_physical_experts // num_local_physical_experts,
+            enable_hierarchical=True,
+            active_ranks=(
+                ElasticEPStateManager.instance().active_ranks
+                if ElasticEPStateManager.instance() is not None
+                else ElasticEPStateManager.healthy_rank_state()
+            ),
+        )
+
     raise NotImplementedError
 
 

sglang/srt/eplb/eplb_algorithms/deepseek.py

@@ -3,8 +3,6 @@ from typing import Tuple
 
 import torch
 
-from sglang.srt.utils import get_bool_env_var
-
 
 def balanced_packing(
     weight: torch.Tensor, num_packs: int

sglang/srt/eplb/eplb_algorithms/elasticity_aware.py

@@ -0,0 +1,87 @@
+from typing import Tuple
+
+import torch
+
+from sglang.srt.eplb.eplb_algorithms.deepseek import rebalance_experts_hierarchical
+
+
+def rebalance_experts(
+    weight: torch.Tensor,
+    num_replicas: int,
+    num_groups: int,
+    num_nodes: int,
+    num_gpus: int,
+    enable_hierarchical: bool,
+    active_ranks: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Entry point for expert-parallelism load balancer.
+
+    Parameters:
+        weight: [layers, num_logical_experts], the load statistics for all logical experts
+        num_replicas: number of physical experts, must be a multiple of `num_gpus`
+        num_groups: number of expert groups
+        num_nodes: number of server nodes, where the intra-node network (e.g, NVLink) is faster
+        num_gpus: number of GPUs, must be a multiple of `num_nodes`
+
+    Returns:
+        physical_to_logical_map: [layers, num_replicas], the expert index of each replica
+        logical_to_physical_map: [layers, num_logical_experts, X], the replica indices for each expert
+        expert_count: [layers, num_logical_experts], number of physical replicas for each logical expert
+    """
+
+    num_layers, num_logical_experts = weight.shape
+    weight = weight.float().cpu()
+    num_active_ranks = active_ranks.sum().item()
+    num_local_experts = num_replicas // num_gpus
+    if num_active_ranks < num_gpus:
+        # Must fall back to global load-balance policy
+        # and fix some params
+        phy2log, phyrank, logcnt = rebalance_experts_hierarchical(
+            weight,
+            num_local_experts * num_active_ranks,
+            1,
+            1,
+            num_active_ranks,
+        )
+    elif enable_hierarchical:
+        # use hierarchical load-balance policy
+        phy2log, phyrank, logcnt = rebalance_experts_hierarchical(
+            weight, num_replicas, num_groups, num_nodes, num_gpus
+        )
+    else:
+        # use global load-balance policy
+        phy2log, phyrank, logcnt = rebalance_experts_hierarchical(
+            weight, num_replicas, 1, 1, num_gpus
+        )
+    maxlogcnt = logcnt.max().item()
+    log2phy: torch.Tensor = torch.full(
+        (num_layers, num_logical_experts, maxlogcnt),
+        -1,
+        dtype=torch.int64,
+        device=logcnt.device,
+    )
+    log2phy.view(num_layers, -1).scatter_(
+        -1,
+        phy2log * maxlogcnt + phyrank,
+        torch.arange(
+            num_local_experts * num_active_ranks,
+            dtype=torch.int64,
+            device=log2phy.device,
+        ).expand(num_layers, -1),
+    )
+    if num_active_ranks < num_gpus:
+        phy2log_slices = list(
+            phy2log.view(num_layers, num_active_ranks, -1).unbind(dim=1)
+        )
+        active_ranks_list = active_ranks.tolist()
+        for idx, active_rank in enumerate(active_ranks_list):
+            if not active_rank:
+                phy2log_slices.insert(idx, torch.zeros_like(phy2log_slices[0]))
+                log2phy = torch.where(
+                    log2phy >= idx * num_local_experts,
+                    log2phy + num_local_experts,
+                    log2phy,
+                )
+        phy2log = torch.stack(phy2log_slices, dim=1).contiguous().view(num_layers, -1)
+    return phy2log, log2phy, logcnt