sglang 0.5.2rc2__py3-none-any.whl → 0.5.3__py3-none-any.whl

sglang/srt/layers/moe/token_dispatcher/deepep.py

@@ -1,17 +1,20 @@
 from __future__ import annotations
 
 import logging
+from contextlib import nullcontext
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, NamedTuple, Optional, Tuple, Union
 
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
-from sglang.srt.layers.moe import DeepEPMode, get_deepep_config, is_tbo_enabled
-from sglang.srt.layers.moe.token_dispatcher.base_dispatcher import (
+from sglang.srt.layers.moe.token_dispatcher.base import (
     BaseDispatcher,
     BaseDispatcherConfig,
+    CombineInput,
+    CombineInputFormat,
     DispatchOutput,
     DispatchOutputFormat,
 )
+from sglang.srt.layers.moe.utils import DeepEPMode, get_deepep_config, is_tbo_enabled
 from sglang.srt.layers.quantization import deep_gemm_wrapper
 from sglang.srt.utils import (
     get_bool_env_var,
@@ -23,6 +26,9 @@ from sglang.srt.utils import (
 
 _is_npu = is_npu()
 
+if TYPE_CHECKING:
+    from sglang.srt.single_batch_overlap import CombineOverlapArgs
+
 try:
     from deep_ep import Buffer, Config
 
@@ -40,11 +46,6 @@ from enum import Enum, IntEnum, auto
 import torch
 import torch.distributed as dist
 
-from sglang.srt.layers.moe.ep_moe.kernels import (
-    deepep_permute_triton_kernel,
-    deepep_post_reorder_triton_kernel,
-    deepep_run_moe_deep_preprocess,
-)
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 
 _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip()
@@ -56,6 +57,7 @@ class DeepEPNormalOutput(NamedTuple):
     """DeepEP normal dispatch output."""
 
     hidden_states: torch.Tensor | Tuple[torch.Tensor, torch.Tensor]
+    # hidden_states_scale
     topk_idx: torch.Tensor
     topk_weights: torch.Tensor
     num_recv_tokens_per_expert: List[int]
@@ -79,24 +81,32 @@ class DeepEPLLOutput(NamedTuple):
         return DispatchOutputFormat.DEEPEP_LL
 
 
-class AscendDeepEPLLOutput(NamedTuple):
-    """AscendDeepEP low latency dispatch output."""
+assert isinstance(DeepEPNormalOutput, DispatchOutput)
+assert isinstance(DeepEPLLOutput, DispatchOutput)
 
-    hidden_states_fp8: Tuple[torch.Tensor, torch.Tensor]
-    topk_idx: torch.Tensor
-    topk_weights: torch.Tensor
-    masked_m: torch.Tensor
-    seg_indptr: torch.Tensor
-    expected_m: int
+
+class DeepEPNormalCombineInput(NamedTuple):
+    """DeepEP normal combine input."""
+
+    pass
 
     @property
-    def format(self) -> DispatchOutputFormat:
-        return DispatchOutputFormat.ASCENT_LL
+    def format(self) -> CombineInputFormat:
+        return CombineInputFormat.DEEPEP_NORMAL
 
 
-assert isinstance(DeepEPNormalOutput, DispatchOutput)
-assert isinstance(DeepEPLLOutput, DispatchOutput)
-assert isinstance(AscendDeepEPLLOutput, DispatchOutput)
+class DeepEPLLCombineInput(NamedTuple):
+    """DeepEP low latency combine input."""
+
+    pass
+
+    @property
+    def format(self) -> CombineInputFormat:
+        return CombineInputFormat.DEEPEP_LL
+
+
+assert isinstance(DeepEPNormalCombineInput, CombineInput)
+assert isinstance(DeepEPLLCombineInput, CombineInput)
 
 
 class DeepEPDispatchMode(IntEnum):
@@ -158,10 +168,19 @@ class DeepEPBuffer:
                 num_rdma_bytes,
             )
 
+        # We should calculate num_qps_per_rank consistently with DeepEP's test script logic:
         if deepep_mode == DeepEPMode.NORMAL:
-            num_qps_per_rank = DeepEPConfig.get_instance().num_sms // 2
-        elif deepep_mode in [DeepEPMode.LOW_LATENCY, DeepEPMode.AUTO]:
+            # refer: https://github.com/deepseek-ai/DeepEP/blob/main/tests/test_internode.py#L235
+            num_qps_per_rank = DeepEPConfig.get_instance().num_sms
+        elif deepep_mode == DeepEPMode.LOW_LATENCY:
+            # refer: https://github.com/deepseek-ai/DeepEP/blob/main/tests/test_low_latency.py#L176
             num_qps_per_rank = num_experts // group.size()
+        elif deepep_mode == DeepEPMode.AUTO:
+            # low-latency and normal mode all need run
+            # refer: https://github.com/deepseek-ai/DeepEP/blob/main/tests/test_internode.py#L235
+            num_qps_per_rank = max(
+                DeepEPConfig.get_instance().num_sms, num_experts // group.size()
+            )
         else:
             raise NotImplementedError
 
@@ -272,12 +291,16 @@ class _DeepEPDispatcherImplBase:
         self.num_max_dispatch_tokens_per_rank = get_int_env_var(
             "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK", 128
         )
+        # DeepEP internode_ll dispatch uses FINISHED_SUM_TAG=1024
+        # and the logic requires num-tokens-sent-from-one-rank-to-another-rank less than it
+        assert self.num_max_dispatch_tokens_per_rank <= 1024
 
         self.handle = None
 
     def dispatch_a(
         self,
         hidden_states: torch.Tensor,
+        input_global_scale: Optional[torch.Tensor],
         topk_idx: torch.Tensor,
         topk_weights: torch.Tensor,
     ):
@@ -291,6 +314,7 @@ class _DeepEPDispatcherImplBase:
         hidden_states: torch.Tensor,
         topk_idx: torch.Tensor,
         topk_weights: torch.Tensor,
+        overlap_args: Optional["CombineOverlapArgs"],
     ):
         raise NotImplementedError
 
@@ -311,6 +335,7 @@ class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase):
     def dispatch_a(
         self,
         hidden_states: torch.Tensor,
+        input_global_scale: Optional[torch.Tensor],
         topk_idx: torch.Tensor,
         topk_weights: torch.Tensor,
     ):
@@ -408,8 +433,13 @@ class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase):
         hidden_states: torch.Tensor,
         topk_idx: torch.Tensor,
         topk_weights: torch.Tensor,
+        overlap_args: Optional["CombineOverlapArgs"],
     ):
-        if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM or _use_aiter:
+        from sglang.srt.layers.moe.ep_moe.kernels import (
+            deepep_post_reorder_triton_kernel,
+        )
+
+        if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM or _use_aiter or _is_npu:
             output = hidden_states
         else:
             if hidden_states.shape[0] > 0:
@@ -479,10 +509,12 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
         https://github.com/deepseek-ai/DeepEP?tab=readme-ov-file#example-use-in-inference-decoding
         """
         self.return_recv_hook = return_recv_hook
+        self.device_module = torch.get_device_module()
 
     def dispatch_a(
         self,
         hidden_states: torch.Tensor,
+        input_global_scale: Optional[torch.Tensor],
         topk_idx: torch.Tensor,
         topk_weights: torch.Tensor,
     ):
@@ -494,8 +526,8 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
         ) // self.num_experts
         hidden_states, masked_m, event, hook = self._dispatch_core(
             hidden_states,
+            input_global_scale,
             topk_idx,
-            use_fp8=True,
         )
         return (
             hidden_states,
@@ -523,39 +555,41 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
             masked_m
         )
 
-        if _is_npu:
-            deepep_output = AscendDeepEPLLOutput(
-                hidden_states,
-                topk_idx,
-                topk_weights,
-                masked_m,
-                self.handle[1],
-                expected_m,
-            )
-        else:
-            deepep_output = DeepEPLLOutput(
-                hidden_states,
-                topk_idx,
-                topk_weights,
-                masked_m,
-                expected_m,
-            )
+        deepep_output = DeepEPLLOutput(
+            hidden_states,
+            topk_idx,
+            topk_weights,
+            masked_m,
+            expected_m,
+        )
         return deepep_output
 
     def _dispatch_core(
         self,
         hidden_states: torch.Tensor,
+        input_global_scale: Optional[torch.Tensor],
         topk_idx: torch.Tensor,
-        use_fp8: bool = False,
     ):
+        use_nvfp4 = use_fp8 = False
+        if input_global_scale is not None:
+            use_nvfp4 = True
+        elif not get_bool_env_var("SGLANG_DEEPEP_BF16_DISPATCH"):
+            use_fp8 = True
+
         buffer = self._get_buffer()
-        packed_recv_hidden, packed_recv_count, self.handle, event, hook = (
+        packed_recv_hidden, self.packed_recv_count, self.handle, event, hook = (
             buffer.low_latency_dispatch(
                 hidden_states,
                 topk_idx,
                 self.num_max_dispatch_tokens_per_rank,
                 self.num_experts,
                 use_fp8=use_fp8,
+                **(dict(use_nvfp4=True) if use_nvfp4 else dict()),
+                **(
+                    dict(x_global_scale=input_global_scale)
+                    if input_global_scale is not None
+                    else dict()
+                ),
                 async_finish=not self.return_recv_hook,
                 return_recv_hook=self.return_recv_hook,
                 round_scale=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
@@ -564,23 +598,29 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
                 and deep_gemm_wrapper.DEEPGEMM_BLACKWELL,
             )
         )
-        return packed_recv_hidden, packed_recv_count, event, hook
+        return packed_recv_hidden, self.packed_recv_count, event, hook
 
     def combine_a(
         self,
         hidden_states: torch.Tensor,
         topk_idx: torch.Tensor,
         topk_weights: torch.Tensor,
+        overlap_args: Optional["CombineOverlapArgs"],
     ):
         hidden_states, event, hook = self._combine_core(
             hidden_states,
             topk_idx,
             topk_weights,
+            overlap_args=overlap_args,
         )
-        return hidden_states, event, hook
+        return hidden_states, event, hook, overlap_args
 
-    def combine_b(self, hidden_states, event, hook):
+    def combine_b(self, hidden_states, event, hook, overlap_args):
         hook() if self.return_recv_hook else event.current_stream_wait()
+
+        if overlap_args is not None:
+            self.device_module.current_stream().wait_stream(overlap_args.stream)
+
         return hidden_states
 
     def _combine_core(
@@ -588,17 +628,35 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
         hidden_states: torch.Tensor,
         topk_idx: torch.Tensor,
         topk_weights: torch.Tensor,
+        overlap_args: Optional["CombineOverlapArgs"],
     ):
         buffer = self._get_buffer()
-        combined_hidden_states, event, hook = buffer.low_latency_combine(
-            hidden_states,
-            topk_idx,
-            topk_weights,
-            self.handle,
-            async_finish=not self.return_recv_hook,
-            return_recv_hook=self.return_recv_hook,
-        )
-        self.handle = None
+
+        ctx = nullcontext()
+        if overlap_args is not None:
+            overlap_args.stream.wait_event(overlap_args.wait_event)
+            ctx = torch.cuda.stream(overlap_args.stream)
+
+        with ctx:
+            combined_hidden_states, event, hook = buffer.low_latency_combine(
+                x=hidden_states,
+                topk_idx=topk_idx,
+                topk_weights=topk_weights,
+                handle=self.handle,
+                async_finish=not self.return_recv_hook,
+                return_recv_hook=self.return_recv_hook,
+                **(
+                    dict(
+                        overlap=overlap_args.overlap,
+                        src_signals=overlap_args.signal,
+                        src_signal_expect_value=overlap_args.threshold,
+                    )
+                    if overlap_args is not None
+                    else {}
+                ),
+            )
+
+        self.packed_recv_count = self.handle = None
         return combined_hidden_states, event, hook
 
     def _get_buffer(self):
@@ -669,6 +727,7 @@ class DeepEPDispatcher(BaseDispatcher):
     def dispatch_a(
         self,
         hidden_states: torch.Tensor,
+        input_global_scale: Optional[torch.Tensor],
         topk_idx: torch.Tensor,
         topk_weights: torch.Tensor,
         forward_batch: ForwardBatch,
@@ -676,6 +735,7 @@ class DeepEPDispatcher(BaseDispatcher):
         self._update_stage(_Stage.INITIAL, _Stage.AFTER_DISPATCH_A)
         inner_state = self._get_impl(forward_batch).dispatch_a(
             hidden_states=hidden_states,
+            input_global_scale=input_global_scale,
             topk_idx=topk_idx,
             topk_weights=topk_weights,
         )
@@ -698,12 +758,14 @@ class DeepEPDispatcher(BaseDispatcher):
         topk_idx: torch.Tensor,
         topk_weights: torch.Tensor,
         forward_batch: ForwardBatch,
+        overlap_args: Optional["CombineOverlapArgs"] = None,
     ):
         self._update_stage(_Stage.AFTER_DISPATCH_B, _Stage.AFTER_COMBINE_A)
         inner_state = self._get_impl(forward_batch).combine_a(
             hidden_states=hidden_states,
             topk_idx=topk_idx,
             topk_weights=topk_weights,
+            overlap_args=overlap_args,
         )
         self._combine_intermediate_state = forward_batch, inner_state
 

sglang/srt/layers/moe/token_dispatcher/standard.py

@@ -1,19 +1,61 @@
 from __future__ import annotations
 
-from typing import NamedTuple
+from typing import TYPE_CHECKING, NamedTuple
 
-from sglang.srt.layers.moe.token_dispatcher.base_dispatcher import (
+import torch
+
+from sglang.srt.layers.moe.token_dispatcher.base import (
+    BaseDispatcher,
+    CombineInput,
+    CombineInputFormat,
     DispatchOutput,
     DispatchOutputFormat,
 )
 
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.topk import TopKOutput
+
 
 class StandardDispatchOutput(NamedTuple):
     """Standard dispatch output."""
 
+    hidden_states: torch.Tensor
+    topk_output: TopKOutput
+
     @property
     def format(self) -> DispatchOutputFormat:
         return DispatchOutputFormat.STANDARD
 
 
 assert isinstance(StandardDispatchOutput, DispatchOutput)
+
+
+class StandardCombineInput(NamedTuple):
+    """Standard combine input."""
+
+    hidden_states: torch.Tensor
+
+    @property
+    def format(self) -> CombineInputFormat:
+        return CombineInputFormat.STANDARD
+
+
+assert isinstance(StandardCombineInput, CombineInput)
+
+
+class StandardDispatcher(BaseDispatcher):
+
+    def dispatch(
+        self, hidden_states: torch.Tensor, topk_output: TopKOutput
+    ) -> DispatchOutput:
+        return StandardDispatchOutput(
+            hidden_states=hidden_states, topk_output=topk_output
+        )
+
+    def combine(self, combine_input: CombineInput) -> torch.Tensor:
+        if isinstance(combine_input, StandardCombineInput):
+            return combine_input.hidden_states
+        else:
+            # TODO: this branch should be removed in the future
+            assert isinstance(combine_input, torch.Tensor)
+            return combine_input

sglang/srt/layers/moe/topk.py

@@ -19,6 +19,7 @@ import math
 from dataclasses import dataclass
 from enum import Enum, auto
 from typing import (
+    TYPE_CHECKING,
     Callable,
     NamedTuple,
     Optional,
@@ -51,6 +52,9 @@ from sglang.srt.utils import (
     is_npu,
 )
 
+if TYPE_CHECKING:
+    from sglang.srt.layers.quantization import QuantizationConfig
+
 try:
     from triton_kernels.routing import GatherIndx, RoutingData, ScatterIndx, routing
 except ImportError:
@@ -94,6 +98,7 @@ class TopKConfig:
     torch_native: bool = False
     routed_scaling_factor: Optional[float] = None
     apply_routed_scaling_factor_on_output: bool = False
+    output_format: Optional[TopKOutputFormat] = None
 
 
 # -------------------------------- TopKOutput ---------------------------------------
@@ -196,9 +201,10 @@ class TopK(CustomOp):
         custom_routing_function: Optional[Callable] = None,
         scoring_func: str = "softmax",
         correction_bias: Optional[torch.Tensor] = None,
+        quant_config: Optional[QuantizationConfig] = None,
         routed_scaling_factor: Optional[float] = None,
         apply_routed_scaling_factor_on_output: Optional[bool] = False,
-        force_topk: bool = False,
+        output_format: Optional[TopKOutputFormat] = None,
     ):
         # NOTE: scoring_func is not used for now, but we keep it for future use
         # see https://github.com/sgl-project/sglang/pull/4505 for more details
@@ -218,11 +224,9 @@ class TopK(CustomOp):
             correction_bias=correction_bias,
             routed_scaling_factor=routed_scaling_factor,
             apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output,
+            output_format=output_format,
         )
 
-        self.use_triton_kernels = get_moe_runner_backend().is_triton_kernel()
-        self.force_topk = force_topk
-
     def forward_native(
         self,
         hidden_states: torch.Tensor,
@@ -248,7 +252,19 @@ class TopK(CustomOp):
         num_token_non_padded: Optional[torch.Tensor] = None,
         expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
     ) -> TopKOutput:
-        if self.use_triton_kernels:
+        if self.topk_config.output_format is not None:
+            output_format = self.topk_config.output_format
+        elif get_moe_runner_backend().is_triton_kernel():
+            output_format = TopKOutputFormat.TRITON_KERNEL
+        elif (
+            should_use_flashinfer_trtllm_moe()
+            or get_moe_runner_backend().is_flashinfer_mxfp4()
+        ):
+            output_format = TopKOutputFormat.BYPASSED
+        else:
+            output_format = TopKOutputFormat.STANDARD
+
+        if output_format == TopKOutputFormat.TRITON_KERNEL:
             # renormalize=True is equivalent to sm_first=False
             routing_data, gather_idx, scatter_idx = routing(
                 router_logits,
@@ -256,10 +272,7 @@ class TopK(CustomOp):
                 sm_first=not self.topk_config.renormalize,
             )
             return TritonKernelTopKOutput(routing_data, gather_idx, scatter_idx)
-        elif not self.force_topk and (
-            should_use_flashinfer_trtllm_moe()
-            or get_moe_runner_backend().is_flashinfer_mxfp4()
-        ):
+        elif output_format == TopKOutputFormat.BYPASSED:
             return BypassedTopKOutput(
                 hidden_states=hidden_states,
                 router_logits=router_logits,
@@ -330,6 +343,14 @@ class TopK(CustomOp):
                 )
                 topk_weights = topk_weights / topk_weights_sum
 
+            if expert_location_dispatch_info is not None:
+                topk_ids = topk_ids_logical_to_physical(
+                    topk_ids, expert_location_dispatch_info
+                )
+            get_global_expert_distribution_recorder().on_select_experts(
+                topk_ids=topk_ids
+            )
+
             return StandardTopKOutput(topk_weights, topk_ids, _)
         else:
             self.topk_config.torch_native = True

sglang/srt/layers/moe/utils.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import importlib.util
+import logging
 from enum import Enum
 from functools import lru_cache
 from typing import TYPE_CHECKING, Optional
@@ -12,11 +13,12 @@ from sglang.srt.layers.dp_attention import (
     get_attention_dp_size,
     is_dp_attention_enabled,
 )
-from sglang.srt.utils import logger
 
 if TYPE_CHECKING:
     from sglang.srt.server_args import ServerArgs
 
+logger = logging.getLogger(__name__)
+
 
 class MoeA2ABackend(Enum):
 
@@ -44,9 +46,10 @@ class MoeRunnerBackend(Enum):
     AUTO = "auto"
     TRITON = "triton"
     TRITON_KERNEL = "triton_kernel"
-    FLASHINFER = "flashinfer_trtllm"
+    FLASHINFER_TRTLLM = "flashinfer_trtllm"
     FLASHINFER_CUTLASS = "flashinfer_cutlass"
     FLASHINFER_MXFP4 = "flashinfer_mxfp4"
+    FLASHINFER_CUTEDSL = "flashinfer_cutedsl"
 
     def is_auto(self):
         return self == MoeRunnerBackend.AUTO
@@ -58,11 +61,14 @@ class MoeRunnerBackend(Enum):
         return self == MoeRunnerBackend.TRITON_KERNEL
 
     def is_flashinfer_trtllm(self):
-        return self == MoeRunnerBackend.FLASHINFER
+        return self == MoeRunnerBackend.FLASHINFER_TRTLLM
 
     def is_flashinfer_cutlass(self):
         return self == MoeRunnerBackend.FLASHINFER_CUTLASS
 
+    def is_flashinfer_cutedsl(self):
+        return self == MoeRunnerBackend.FLASHINFER_CUTEDSL
+
     def is_flashinfer_mxfp4(self):
         return self == MoeRunnerBackend.FLASHINFER_MXFP4
 
@@ -102,6 +108,7 @@ MOE_A2A_BACKEND: Optional[MoeA2ABackend] = None
 MOE_RUNNER_BACKEND: Optional[MoeRunnerBackend] = None
 DEEPEP_MODE: Optional[DeepEPMode] = None
 IS_TBO_ENABLED: Optional[bool] = None
+IS_SBO_ENABLED: Optional[bool] = None
 TBO_TOKEN_DISTRIBUTION_THRESHOLD: Optional[float] = None
 DEEPEP_CONFIG: Optional[str] = None
 DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER: Optional[bool] = None
@@ -113,6 +120,7 @@ def initialize_moe_config(server_args: ServerArgs):
     global DEEPEP_MODE
     global DEEPEP_CONFIG
     global IS_TBO_ENABLED
+    global IS_SBO_ENABLED
     global TBO_TOKEN_DISTRIBUTION_THRESHOLD
     global DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER
 
@@ -121,6 +129,7 @@ def initialize_moe_config(server_args: ServerArgs):
     DEEPEP_MODE = DeepEPMode(server_args.deepep_mode)
     DEEPEP_CONFIG = server_args.deepep_config or ""
     IS_TBO_ENABLED = server_args.enable_two_batch_overlap
+    IS_SBO_ENABLED = server_args.enable_single_batch_overlap
     TBO_TOKEN_DISTRIBUTION_THRESHOLD = server_args.tbo_token_distribution_threshold
     DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER = (
         server_args.disable_flashinfer_cutlass_moe_fp4_allgather
@@ -131,7 +140,7 @@ def get_moe_a2a_backend() -> MoeA2ABackend:
     global MOE_A2A_BACKEND
     if MOE_A2A_BACKEND is None:
         logger.warning("MOE_A2A_BACKEND is not initialized, using default backend")
-        MOE_A2A_BACKEND = MoeA2ABackend(None)
+        MOE_A2A_BACKEND = MoeA2ABackend.NONE
     return MOE_A2A_BACKEND
 
 
@@ -139,7 +148,7 @@ def get_moe_runner_backend() -> MoeRunnerBackend:
     global MOE_RUNNER_BACKEND
     if MOE_RUNNER_BACKEND is None:
         logger.warning("MOE_RUNNER_BACKEND is not initialized, using triton backend")
-        MOE_RUNNER_BACKEND = MoeRunnerBackend("triton")
+        MOE_RUNNER_BACKEND = MoeRunnerBackend.AUTO
     return MOE_RUNNER_BACKEND
 
 
@@ -147,7 +156,7 @@ def get_deepep_mode() -> DeepEPMode:
     global DEEPEP_MODE
     if DEEPEP_MODE is None:
         logger.warning("DEEPEP_MODE is not initialized, using auto mode")
-        DEEPEP_MODE = DeepEPMode("auto")
+        DEEPEP_MODE = DeepEPMode.AUTO
     return DEEPEP_MODE
 
 
@@ -166,6 +175,13 @@ def is_tbo_enabled() -> bool:
     return IS_TBO_ENABLED
 
 
+def is_sbo_enabled() -> bool:
+    global IS_SBO_ENABLED
+    if IS_SBO_ENABLED is None:
+        IS_SBO_ENABLED = False
+    return IS_SBO_ENABLED
+
+
 def get_tbo_token_distribution_threshold() -> float:
     global TBO_TOKEN_DISTRIBUTION_THRESHOLD
     if TBO_TOKEN_DISTRIBUTION_THRESHOLD is None:

sglang/srt/layers/parameter.py

@@ -7,6 +7,7 @@ from typing import Callable, Optional, Union
 import torch
 from torch.nn import Parameter
 
+from sglang.srt.layers.utils import pad_or_narrow_weight
 from sglang.srt.utils import is_cpu
 
 __all__ = [
@@ -156,9 +157,17 @@ class _ColumnvLLMParameter(BasevLLMParameter):
             )
         else:
             if not use_presharded_weights:
-                loaded_weight = loaded_weight.narrow(
-                    self.output_dim, tp_rank * shard_size, shard_size
-                )
+                # Padding for special case like qwen2_5_VL's mlp which is not 8-aligned
+                start_idx = tp_rank * shard_size
+                end_idx = start_idx + shard_size
+                if end_idx > loaded_weight.shape[self.output_dim]:
+                    loaded_weight = pad_or_narrow_weight(
+                        loaded_weight, self.output_dim, start_idx, shard_size
+                    )
+                else:
+                    loaded_weight = loaded_weight.narrow(
+                        self.output_dim, start_idx, shard_size
+                    )
 
         assert param_data.shape == loaded_weight.shape
         param_data.copy_(loaded_weight)
@@ -258,9 +267,17 @@ class RowvLLMParameter(BasevLLMParameter):
 
                 return
             else:
-                loaded_weight = loaded_weight.narrow(
-                    self.input_dim, tp_rank * shard_size, shard_size
-                )
+                # Padding for special case like qwen2_5_VL's mlp which is not 8-aligned
+                start_idx = tp_rank * shard_size
+                end_idx = start_idx + shard_size
+                if end_idx > loaded_weight.shape[self.input_dim]:
+                    loaded_weight = pad_or_narrow_weight(
+                        loaded_weight, self.input_dim, start_idx, shard_size
+                    )
+                else:
+                    loaded_weight = loaded_weight.narrow(
+                        self.input_dim, start_idx, shard_size
+                    )
 
         if len(loaded_weight.shape) == 0:
             loaded_weight = loaded_weight.reshape(1)