LPLB: linear-programming load balancer for MoE expert parallelism (#24515)

Co-authored-by: xutizhou <xutingz@nvidia.com>

python/sglang/srt/eplb/expert_location_dispatch.py

@@ -77,7 +77,9 @@ def transform_select_experts_inputs(
 
 
 def topk_ids_logical_to_physical(
-    topk_ids: torch.Tensor, info: Optional[ExpertLocationDispatchInfo]
+    topk_ids: torch.Tensor,
+    info: Optional[ExpertLocationDispatchInfo],
+    log2phy_prob: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     if info is None:
         return topk_ids
@@ -86,6 +88,13 @@ def topk_ids_logical_to_physical(
         return _topk_ids_logical_to_physical_static(topk_ids, info)
     if info.ep_dispatch_algorithm in ["dynamic", "fake"]:
         return _topk_ids_logical_to_physical_dynamic(topk_ids, info)
+    if info.ep_dispatch_algorithm == "lp":
+        if log2phy_prob is None:
+            raise RuntimeError(
+                "ep_dispatch_algorithm='lp' but log2phy_prob is None at dispatch "
+                f"time (topk_ids.shape={tuple(topk_ids.shape)})."
+            )
+        return _topk_ids_logical_to_physical_probability(topk_ids, info, log2phy_prob)
     raise NotImplementedError(f"Unknown algorithm {info.ep_dispatch_algorithm}")
 
 
@@ -115,3 +124,24 @@ def _topk_ids_logical_to_physical_dynamic(
 
     topk_ids = topk_ids.view(topk_ids_original_shape)
     return topk_ids
+
+
+def _topk_ids_logical_to_physical_probability(
+    topk_ids: torch.Tensor,
+    info: ExpertLocationDispatchInfo,
+    log2phy_prob: torch.Tensor,
+) -> torch.Tensor:
+    """Select physical experts via the JIT-compiled CUDA dispatch kernel.
+
+    Raises if ``topk_ids`` isn't on CUDA — the LP path requires the fused
+    kernel and there is no torch reference fallback at runtime.
+    """
+    if not topk_ids.is_cuda:
+        raise RuntimeError(
+            "LP dispatch requires CUDA tensors; got topk_ids on " f"{topk_ids.device}."
+        )
+    from sglang.jit_kernel.lplb import cuda_solver
+
+    return cuda_solver.dispatch_probability(
+        topk_ids, log2phy_prob, info.partial_logical_to_all_physical_map
+    )

python/sglang/srt/eplb/lplb_solver.py

@@ -0,0 +1,280 @@
+"""
+LPLBSolver — Linear-Programming Load Balancer for Expert Parallelism.
+
+Encapsulates LP matrix construction (offline, at init/rebalance) and
+per-batch solving (online, per MoE layer forward pass).
+
+Design for DP-attention:
+    Each EP rank counts its local tokens, then all ranks participate in an
+    all-reduce to obtain identical global counts.  Every rank then solves
+    the same LP independently, producing the same log2phy_prob — no
+    broadcast is needed.  Empty-token ranks contribute zeros in the
+    all-reduce so the collective never deadlocks.
+
+Usage:
+    solver = LPLBSolver(phy2log, log2phy, num_gpus, ep_group)
+    log2phy_prob = solver.solve(topk_ids)  # per batch
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+# Global per-layer LPLB solvers
+_global_lplb_solvers: dict[int, LPLBSolver] = {}
+
+
+# LP dispatch requires every EP rank to call solver.solve() on every forward
+# pass (including empty-topk ranks under DP-attention) — the all-reduce inside
+# would otherwise hang. Only the DeepSeek-v2 family and its subclasses route
+# empty-rank paths through solver.solve(); other MoE families would deadlock.
+_LPLB_SUPPORTED_MODEL_ARCHS: frozenset[str] = frozenset(
+    {
+        "DeepseekV2ForCausalLM",
+        "DeepseekV3ForCausalLM",
+        "DeepseekV32ForCausalLM",
+        "MistralLarge3ForCausalLM",
+        "MistralLarge3ForCausalLMEagle",
+        "Glm4MoeLiteForCausalLM",
+        "GlmMoeDsaForCausalLM",
+    }
+)
+
+
+def assert_lplb_supported_model(architecture: str) -> None:
+    if architecture not in _LPLB_SUPPORTED_MODEL_ARCHS:
+        supported = ", ".join(sorted(_LPLB_SUPPORTED_MODEL_ARCHS))
+        raise NotImplementedError(
+            f"{architecture} does not support --ep-dispatch-algorithm lp. "
+            f"Validated targets: {supported}. Other MoE families have "
+            "empty-token early returns that don't participate in the EP "
+            "all-reduce inside LPLBSolver.solve(), which would deadlock "
+            "under DP-attention."
+        )
+
+
+def get_global_lplb_solver(layer_id: int) -> Optional[LPLBSolver]:
+    return _global_lplb_solvers.get(layer_id)
+
+
+def set_global_lplb_solver(layer_id: int, solver: LPLBSolver):
+    _global_lplb_solvers[layer_id] = solver
+
+
+def clear_global_lplb_solvers():
+    _global_lplb_solvers.clear()
+
+
+class LPLBSolver:
+    """
+    Per-layer LPLB solver.
+
+    At init: pre-computes LP constraint matrices from expert-to-GPU mapping.
+    At solve: takes topk_ids, counts tokens, all-reduces, runs LP,
+              returns log2phy_prob for probability-based token dispatch.
+    """
+
+    def __init__(
+        self,
+        phy2log: torch.Tensor,
+        log2phy: torch.Tensor,
+        num_gpus: int,
+        ep_group=None,
+        logical_to_all_physical_map_num_valid=None,
+    ):
+        """
+        Args:
+            phy2log: (num_physical_experts,) physical-to-logical expert mapping.
+            log2phy: (num_logical_experts, max_copies) logical-to-physical mapping (-1 padded).
+            num_gpus: Number of GPUs in the EP group.
+            ep_group: GroupCoordinator for EP communication (all-reduce).
+            logical_to_all_physical_map_num_valid: (num_logical_experts,) number of valid physical copies.
+        """
+        device = phy2log.device
+        self.num_gpus = num_gpus
+        self.ep_group = ep_group
+        self._has_redundancy = False
+        if logical_to_all_physical_map_num_valid is not None:
+            self._has_redundancy = bool(
+                (logical_to_all_physical_map_num_valid > 1).any()
+            )
+
+        self.num_logical = log2phy.shape[0]
+        self.max_copies = log2phy.shape[1]
+        self.num_phy = phy2log.shape[0]
+        # B1/B2 GPU-assignment matrices below assume each rank owns a
+        # contiguous block of num_phy // num_gpus physical experts.
+        if self.num_phy % num_gpus != 0:
+            raise ValueError(
+                f"LPLBSolver requires num_phy ({self.num_phy}) to be divisible "
+                f"by num_gpus ({num_gpus}); per-rank-contiguous ownership is "
+                "currently the only supported allocation."
+            )
+        num_phy_per_gpu = self.num_phy // num_gpus
+
+        # Count copies per logical expert
+        logcnt = torch.bincount(phy2log, minlength=self.num_logical)
+
+        # Separate single-copy vs replicated experts.
+        # Stored as int64 so they can be used directly as index tensors in
+        # _solve without per-call .long() casts (Tier 1 optimization).
+        self.log_single = torch.nonzero(logcnt == 1).flatten().to(torch.int64)
+        self.phy_single = log2phy[self.log_single, 0].to(torch.int64)
+        self.log_replicated = torch.nonzero(logcnt > 1).flatten().to(torch.int64)
+        self.phy_replicated = (
+            torch.nonzero(logcnt[phy2log] > 1).flatten().to(torch.int64)
+        )
+
+        self.num_single = len(self.log_single)
+        self.num_red_log = len(self.log_replicated)
+        self.num_red_phy = len(self.phy_replicated)
+
+        # Build GPU assignment matrices
+        B_full = torch.zeros(
+            (num_gpus, self.num_phy), dtype=torch.float32, device=device
+        )
+        for i in range(num_gpus):
+            B_full[i, i * num_phy_per_gpu : (i + 1) * num_phy_per_gpu] = 1
+        self.B1 = B_full[:, self.phy_single].contiguous()
+        B2 = B_full[:, self.phy_replicated]
+
+        # Build C matrix (copy-to-logical mapping)
+        C = torch.zeros(
+            (self.num_red_log, self.num_red_phy), dtype=torch.float32, device=device
+        )
+        phy2log_rep = phy2log[self.phy_replicated]
+        for i in range(self.num_red_log):
+            C[i, phy2log_rep == self.log_replicated[i]] = 1.0
+
+        # Build A_base = [[C, 0, 0], [B2, I, -1]]  (without Big-M column)
+        zeros_top_g = torch.zeros(
+            (self.num_red_log, num_gpus), dtype=torch.float32, device=device
+        )
+        zeros_top_1 = torch.zeros(
+            (self.num_red_log, 1), dtype=torch.float32, device=device
+        )
+        I_g = torch.eye(num_gpus, dtype=torch.float32, device=device)
+        neg_ones = torch.full((num_gpus, 1), -1.0, dtype=torch.float32, device=device)
+
+        A_top = torch.hstack([C, zeros_top_g, zeros_top_1])
+        A_bottom = torch.hstack([B2, I_g, neg_ones])
+        self.A_base = torch.vstack([A_top, A_bottom]).contiguous()
+
+        # Objective: minimize M (second-to-last var), penalize Big-M auxiliary
+        nv = self.A_base.shape[1] + 1  # +1 for Big-M column
+        self.c_vec = torch.zeros(nv, dtype=torch.float32, device=device)
+        self.c_vec[-2] = 1.0
+        self.c_vec[-1] = 1000.0
+
+        # Store log2phy as int64 so it can be used directly as index tensor
+        # without per-call .long() casts (Tier 1 optimization).
+        self.log2phy = log2phy.to(torch.int64).contiguous()
+
+        # Pre-JIT-compile the fused IPM kernel for this (NC, NV) shape so the
+        # 20-40s compile cost happens once at startup rather than on the first
+        # real request. No-op when the fused backend is unavailable.
+        nc = self.A_base.shape[0]
+        nv = self.A_base.shape[1] + 1  # +1 for Big-M column added in solve()
+        from sglang.jit_kernel.lplb.torch_solver import warmup as _ipm_warmup
+
+        _ipm_warmup(nc, nv, num_iters=5, device=device)
+
+        # Pre-compute A_base row sum (used in every prep call).
+        self._A_base_row_sum = self.A_base.sum(dim=1).contiguous()  # (NC,)
+
+        # Pre-allocate the buffers the JIT CUDA prep / IPM / post kernels write
+        # into. All writes are contiguous full-tensor stores (no strided
+        # ``out=`` semantics), so the reuse is safe under high concurrency.
+        # Constructed lazily on the first solve() call (we don't know the
+        # device-side log2phy_prob shape until then) — see _solve.
+        self._A_full = torch.empty(nc, nv, dtype=torch.float32, device=device)
+        self._A_full[:, : nv - 1].copy_(self.A_base)
+        self._b = torch.empty(nc, dtype=torch.float32, device=device)
+        self._t1 = torch.empty(self.num_single, dtype=torch.float32, device=device)
+        self._x = torch.empty(nv, dtype=torch.float32, device=device)
+        self._log2phy_prob = torch.empty(
+            log2phy.shape, dtype=torch.float32, device=device
+        )
+
+    def solve(self, topk_ids: torch.Tensor) -> torch.Tensor:
+        """
+        Full LPLB pipeline: count -> all-reduce -> LP solve -> return log2phy_prob.
+
+        All EP ranks must call this method every MoE layer forward pass,
+        including empty-token ranks (which pass an empty topk_ids tensor).
+        This ensures the all-reduce collective does not deadlock under
+        DP-attention where different ranks may have different token counts.
+
+        Args:
+            topk_ids: (num_tokens, topk) int32 tensor of logical expert IDs.
+                      Can be empty (shape (0, topk)) for idle ranks.
+
+        Returns:
+            log2phy_prob: (num_logical, max_copies) float32 probability tensor.
+        """
+        device = topk_ids.device
+
+        # Step 1: Count local tokens per logical expert.
+        # topk_ids comes from the router and is by construction in
+        # [0, num_logical), so we can scatter_add directly without filtering.
+        # Boolean masking + numel() (the previous defensive form) forced a
+        # GPU->host sync on every forward pass via aten::nonzero and a
+        # tensor-shape read; scatter_add on the flattened tensor is async
+        # and a no-op when topk_ids is empty (DP-attention idle rank case).
+        local_counts = torch.zeros(self.num_logical, dtype=torch.int32, device=device)
+        flat_ids = topk_ids.flatten()
+        local_counts.scatter_add_(
+            0,
+            flat_ids.long(),
+            torch.ones_like(flat_ids, dtype=torch.int32),
+        )
+
+        # Step 2: All-reduce to get global counts across all EP ranks.
+        # All EP ranks must participate — empty-token ranks contribute zeros.
+        # After all-reduce, every rank has identical global_counts and solves
+        # the same LP independently, so no broadcast is needed.
+        # GroupCoordinator.all_reduce may be in-place (pynccl) or out-of-place
+        # (ca_comm / pymscclpp / ...) depending on tensor size; small tensors
+        # like ours (~num_logical * 4 B) typically take the out-of-place path,
+        # so we must capture the return value.
+        global_counts = local_counts.float()
+        if self.ep_group is not None:
+            global_counts = self.ep_group.all_reduce(global_counts)
+
+        # Step 3: Run LP solver
+        return self._solve(global_counts)
+
+    def _solve(self, global_counts: torch.Tensor) -> torch.Tensor:
+        """Three CUDA kernel launches replace ~14 torch ops.
+
+        Pipeline (all writes go into pre-allocated buffers from __init__):
+            prep_lp_inputs → solve_ipm → extract_log2phy_prob
+        Raises if the JIT CUDA backend is unavailable.
+        """
+        from sglang.jit_kernel.lplb import cuda_solver
+
+        cuda_solver.prep_lp_inputs(
+            self._A_full,
+            self._b,
+            self._t1,
+            global_counts,
+            self.log_single,
+            self.log_replicated,
+            self.B1,
+            self._A_base_row_sum,
+        )
+        cuda_solver.solve_ipm(self._A_full, self._b, self.c_vec, result=self._x)
+        cuda_solver.extract_log2phy_prob(
+            self._log2phy_prob,
+            self._x,
+            self._t1,
+            self.phy_single,
+            self.phy_replicated,
+            self.log2phy,
+        )
+        return self._log2phy_prob

python/sglang/srt/layers/moe/hash_topk.py

@@ -34,9 +34,10 @@ class HashTopK(nn.Module):
         scoring_func="sqrtsoftplus",
         routed_scaling_factor=1.5,
         apply_routed_scaling_factor_on_output=False,
+        layer_id: Optional[int] = None,
     ):
         super().__init__()
-        self.layer_id = None
+        self.layer_id = layer_id
         from sglang.srt.server_args import get_global_server_args
 
         self.enable_deepep_waterfill = (
@@ -80,8 +81,18 @@ class HashTopK(nn.Module):
         with torch.no_grad():
             self.tid2eid.copy_(tid2eid.to(self.tid2eid.dtype))
 
-    def empty_topk_output(self, device: torch.device):
+    def empty_topk_output(
+        self, device: torch.device, *, layer_id: Optional[int] = None
+    ):
         topk = self.topk - self.num_fused_shared_experts
+        if layer_id is not None:
+            from sglang.srt.eplb.lplb_solver import get_global_lplb_solver
+
+            lplb_solver = get_global_lplb_solver(layer_id)
+            if lplb_solver is not None:
+                lplb_solver.solve(
+                    torch.empty((0, topk), dtype=torch.int32, device=device)
+                )
         topk_weights = torch.empty((0, topk), dtype=torch.float32, device=device)
         topk_ids = torch.full((0, topk), -1, dtype=torch.int32, device=device)
         router_logits = torch.empty((0, topk), dtype=torch.float32, device=device)
@@ -175,7 +186,23 @@ class HashTopK(nn.Module):
         if is_hip():
             topk_weights = topk_weights.to(torch.float32)
 
-        topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info)
+        log2phy_prob = None
+        if (
+            expert_location_dispatch_info is not None
+            and getattr(expert_location_dispatch_info, "ep_dispatch_algorithm", None)
+            == "lp"
+        ):
+            if self.layer_id is None:
+                raise RuntimeError("HashTopK LP dispatch requires layer_id.")
+            from sglang.srt.eplb.lplb_solver import get_global_lplb_solver
+
+            lplb_solver = get_global_lplb_solver(self.layer_id)
+            if lplb_solver is not None:
+                log2phy_prob = lplb_solver.solve(topk_ids)
+
+        topk_ids = topk_ids_logical_to_physical(
+            topk_ids, expert_location_dispatch_info, log2phy_prob
+        )
         if is_hip():
             _zero_topk_weights_padded_region(topk_weights, num_token_non_padded)
         else:

python/sglang/srt/layers/moe/topk.py

@@ -552,7 +552,30 @@ class TopK(MultiPlatformOp):
             layer_id=self.layer_id,
         )
 
-    def empty_topk_output(self, device: torch.device) -> TopKOutput:
+    def empty_topk_output(
+        self, device: torch.device, *, layer_id: Optional[int] = None
+    ) -> TopKOutput:
+        """Return an empty topk output for a rank with zero tokens this forward.
+
+        When ``layer_id`` is provided and the active dispatch algorithm is LP,
+        also calls ``LPLBSolver.solve(empty)`` so that this rank participates
+        in the EP all-reduce. Without this, an empty rank would skip the
+        collective and deadlock under DP-attention.
+        """
+        if layer_id is not None:
+            # Skip the full ExpertLocationDispatchInfo allocation — we only
+            # need the per-layer solver to participate in the EP all-reduce.
+            from sglang.srt.eplb.lplb_solver import get_global_lplb_solver
+
+            lplb_solver = get_global_lplb_solver(layer_id)
+            if lplb_solver is not None:
+                lplb_solver.solve(
+                    torch.empty(
+                        (0, self.topk_config.top_k),
+                        dtype=torch.int32,
+                        device=device,
+                    )
+                )
         topk = self.topk_config.top_k - self.topk_config.num_fused_shared_experts
         with use_symmetric_memory(
             get_tp_group(), disabled=not is_allocation_symmetric()
@@ -1508,11 +1531,29 @@ def _post_process_topk_ids(
         )
     recorder_topk_ids = None
     if _is_cuda:
-        # When shared experts are fused (appended as extra columns in topk_ids),
-        # EPLB dispatch must only remap the routed expert columns.
-        # The shared expert column (value = n_routed_experts) would be out-of-bounds
-        # for the logical-to-physical dispatch table.
-        if num_fused_shared_experts > 0 and is_deepep_class_backend():
+        # LP path: solve LP outside torch.compile (the solver contains an
+        # EP all-reduce that can't run inside compiled regions).
+        log2phy_prob = None
+        if (
+            expert_location_dispatch_info is not None
+            and getattr(expert_location_dispatch_info, "ep_dispatch_algorithm", None)
+            == "lp"
+        ):
+            from sglang.srt.eplb.lplb_solver import get_global_lplb_solver
+
+            lplb_solver = get_global_lplb_solver(layer_id)
+            if lplb_solver is not None:
+                log2phy_prob = lplb_solver.solve(topk_ids)
+
+        if log2phy_prob is not None:
+            topk_ids = topk_ids_logical_to_physical(
+                topk_ids, expert_location_dispatch_info, log2phy_prob
+            )
+            _mask_topk_ids_padded_region(topk_ids, num_token_non_padded)
+        elif num_fused_shared_experts > 0 and is_deepep_class_backend():
+            # Shared experts appended as extra columns in topk_ids: their value
+            # would be out-of-bounds for the logical-to-physical dispatch table,
+            # so split, dispatch the routed cols, recombine.
             shared_cols = topk_ids[:, -num_fused_shared_experts:]
             routed_cols = topk_ids[:, :-num_fused_shared_experts]
             routed_cols = _biased_grouped_topk_postprocess(

python/sglang/srt/model_executor/model_runner.py

@@ -106,6 +106,12 @@ from sglang.srt.eplb.expert_location import (
     set_global_expert_location_metadata,
 )
 from sglang.srt.eplb.expert_location_updater import ExpertLocationUpdater
+from sglang.srt.eplb.lplb_solver import (
+    LPLBSolver,
+    assert_lplb_supported_model,
+    clear_global_lplb_solvers,
+    set_global_lplb_solver,
+)
 from sglang.srt.hardware_backend.npu.graph_runner.npu_graph_runner import NPUGraphRunner
 from sglang.srt.kv_canary.api import install_canary
 from sglang.srt.kv_canary.runner.canary_manager import context_tuple
@@ -691,6 +697,9 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 )
             )
 
+        if self.server_args.ep_dispatch_algorithm == "lp" and not self.is_draft_worker:
+            self._init_lplb_solvers()
+
         # Expert parallelism
         self.eplb_manager = (
             EPLBManager(self)
@@ -1612,6 +1621,35 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 logger, f"Prepared {num_prepared} DeepEP waterfill TopK modules."
             )
 
+    def _init_lplb_solvers(self):
+        """Initialize per-layer LPLB solvers from current expert location metadata."""
+        from sglang.srt.distributed import get_moe_ep_group
+
+        # Gate: refuse LP for non-DeepSeek MoE families whose empty-token paths
+        # don't participate in the EP all-reduce (would deadlock under DP-
+        # attention). Failure here happens before any forward pass.
+        architectures = getattr(self.model_config.hf_config, "architectures", None)
+        if architectures:
+            assert_lplb_supported_model(architectures[0])
+
+        metadata = get_global_expert_location_metadata()
+        if metadata is None:
+            return
+        clear_global_lplb_solvers()
+        ep_group = get_moe_ep_group()
+        for lid in range(metadata.num_layers):
+            solver = LPLBSolver(
+                phy2log=metadata.physical_to_logical_map[lid],
+                log2phy=metadata.logical_to_all_physical_map[lid],
+                num_gpus=metadata.ep_size,
+                ep_group=ep_group,
+                logical_to_all_physical_map_num_valid=(
+                    metadata.logical_to_all_physical_map_num_valid[lid]
+                ),
+            )
+            set_global_lplb_solver(lid, solver)
+        logger.info(f"Initialized LPLB solvers for {metadata.num_layers} layers")
+
     def update_expert_location(
         self,
         new_expert_location_metadata: ExpertLocationMetadata,
@@ -1654,6 +1692,10 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                     weight_name_filter=weight_name_filter,
                 )
 
+        # Re-init LPLB solvers after expert location update
+        if self.server_args.ep_dispatch_algorithm == "lp":
+            self._init_lplb_solvers()
+
     def maybe_recover_ep_ranks(self):
         # TODO(perf): `active_ranks.all()` on a CUDA tensor triggers host-device
         # synchronization, and this function is on the forward-path.

python/sglang/srt/models/deepseek_v2.py

@@ -631,6 +631,7 @@ class DeepseekV2MoE(nn.Module):
                 scoring_func=config.scoring_func,
                 routed_scaling_factor=self.routed_scaling_factor,
                 apply_routed_scaling_factor_on_output=self.experts.should_fuse_routed_scaling_factor_in_topk,
+                layer_id=self.layer_id,
             )
         else:
             # Default: grouped noaux_tc top-k. Covers V3/V3.2/GLM-5/Glm4MoeLite.
@@ -996,7 +997,9 @@ class DeepseekV2MoE(nn.Module):
             )
         else:
             shared_output = None
-            topk_output = self.topk.empty_topk_output(hidden_states.device)
+            topk_output = self.topk.empty_topk_output(
+                hidden_states.device, layer_id=self.layer_id
+            )
 
         if self._fuse_shared_experts_inside_sbo:
             shared_output = None
@@ -1168,7 +1171,19 @@ class DeepseekV2MoE(nn.Module):
                 **topk_kwargs,
             )
         else:
-            topk_output = self.topk.empty_topk_output(hidden_states.device)
+            topk_output = self.topk.empty_topk_output(
+                hidden_states.device, layer_id=self.layer_id
+            )
+            if is_deepep_class_backend() and self.num_fused_shared_experts > 0:
+                n = self.num_fused_shared_experts
+                topk_output = topk_output._replace(
+                    topk_ids=topk_output.topk_ids.new_empty(
+                        (0, topk_output.topk_ids.shape[-1] + n)
+                    ),
+                    topk_weights=topk_output.topk_weights.new_empty(
+                        (0, topk_output.topk_weights.shape[-1] + n)
+                    ),
+                )
 
         if sbo_overlap_dispatch_flag:
             shared_output = None
@@ -1391,7 +1406,9 @@ class DeepseekV2MoE(nn.Module):
                     ),
                 )
         else:
-            state.topk_output = self.topk.empty_topk_output(hidden_states.device)
+            state.topk_output = self.topk.empty_topk_output(
+                hidden_states.device, layer_id=self.layer_id
+            )
 
     def op_dispatch_a(self, state):
         if self.ep_size > 1:

python/sglang/srt/server_args.py

@@ -668,7 +668,7 @@ class ServerArgs:
         "auto"
     )
     ep_num_redundant_experts: int = 0
-    ep_dispatch_algorithm: Optional[Literal["static", "dynamic", "fake"]] = None
+    ep_dispatch_algorithm: Optional[Literal["static", "dynamic", "fake", "lp"]] = None
     init_expert_location: str = "trivial"
     enable_eplb: bool = False
     eplb_algorithm: str = "auto"