sglang 0.5.1.post3__py3-none-any.whl → 0.5.2__py3-none-any.whl

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

@@ -23,8 +23,13 @@ from sglang.srt.layers.moe import (
     get_moe_runner_backend,
     should_use_flashinfer_trtllm_moe,
 )
+from sglang.srt.layers.moe.token_dispatcher.standard import (
+    CombineInput,
+    StandardDispatcher,
+)
 from sglang.srt.layers.moe.topk import TopKOutput, TopKOutputChecker
 from sglang.srt.layers.quantization.base_config import (
+    FusedMoEMethodBase,
     QuantizationConfig,
     QuantizeMethodBase,
 )
@@ -152,16 +157,6 @@ class FusedMoE(torch.nn.Module):
         self.expert_map_cpu = None
         self.expert_map_gpu = None
 
-        self.moe_runner_config = MoeRunnerConfig(
-            activation=activation,
-            apply_router_weight_on_input=apply_router_weight_on_input,
-            inplace=inplace,
-            no_combine=no_combine,
-            routed_scaling_factor=routed_scaling_factor,
-            gemm1_alpha=gemm1_alpha,
-            gemm1_clamp_limit=gemm1_clamp_limit,
-        )
-
         enable_flashinfer_cutlass_moe = get_moe_runner_backend().is_flashinfer_cutlass()
 
         if enable_flashinfer_cutlass_moe and quant_config is None:
@@ -175,6 +170,8 @@ class FusedMoE(torch.nn.Module):
         self.moe_tp_rank = get_moe_tensor_parallel_rank()
         assert num_experts % self.moe_ep_size == 0
         self.num_local_experts = num_experts // self.moe_ep_size
+        self.start_expert_id = self.moe_ep_rank * self.num_local_experts
+        self.end_expert_id = self.start_expert_id + self.num_local_experts - 1
         if self.moe_ep_size > 1:
             # TODO(ch-wan): support shared experts fusion
             # Create a tensor of size num_experts filled with -1
@@ -194,13 +191,6 @@ class FusedMoE(torch.nn.Module):
         self.use_presharded_weights = use_presharded_weights
 
         self.use_triton_kernels = get_moe_runner_backend().is_triton_kernel()
-        if quant_config is None:
-            self.quant_method: Optional[QuantizeMethodBase] = UnquantizedFusedMoEMethod(
-                self.use_triton_kernels
-            )
-        else:
-            self.quant_method = quant_config.get_quant_method(self, prefix)
-        assert self.quant_method is not None
 
         self.quant_config = quant_config
         self.use_flashinfer_mxfp4_moe = get_moe_runner_backend().is_flashinfer_mxfp4()
@@ -211,12 +201,40 @@ class FusedMoE(torch.nn.Module):
             and self.use_flashinfer_mxfp4_moe
         ):
             hidden_size = round_up(hidden_size, 256)
+        self.hidden_size = hidden_size
+
+        self.moe_runner_config = MoeRunnerConfig(
+            num_experts=num_experts,
+            num_local_experts=self.num_local_experts,
+            hidden_size=hidden_size,
+            intermediate_size_per_partition=self.intermediate_size_per_partition,
+            layer_id=layer_id,
+            top_k=top_k,
+            num_fused_shared_experts=num_fused_shared_experts,
+            params_dtype=params_dtype,
+            activation=activation,
+            apply_router_weight_on_input=apply_router_weight_on_input,
+            inplace=inplace,
+            no_combine=no_combine,
+            routed_scaling_factor=routed_scaling_factor,
+            gemm1_alpha=gemm1_alpha,
+            gemm1_clamp_limit=gemm1_clamp_limit,
+        )
+
+        if quant_config is None:
+            self.quant_method: FusedMoEMethodBase = UnquantizedFusedMoEMethod(
+                self.use_triton_kernels
+            )
+        else:
+            self.quant_method: FusedMoEMethodBase = quant_config.get_quant_method(
+                self, prefix
+            )
+        assert self.quant_method is not None
+
         self.quant_method.create_weights(
             layer=self,
             num_experts=self.num_local_experts,
             hidden_size=hidden_size,
-            # FIXME: figure out which intermediate_size to use
-            intermediate_size=self.intermediate_size_per_partition,
             intermediate_size_per_partition=self.intermediate_size_per_partition,
             params_dtype=params_dtype,
             weight_loader=(
@@ -227,6 +245,9 @@ class FusedMoE(torch.nn.Module):
             with_bias=with_bias,
         )
 
+        self.quant_method.create_moe_runner(self, self.moe_runner_config)
+        self.dispatcher = StandardDispatcher()
+
     def _load_per_tensor_weight_scale(
         self,
         shard_id: str,
@@ -592,9 +613,12 @@ class FusedMoE(torch.nn.Module):
             loaded_weight = loaded_weight.to(param.data.device)
 
             if (
-                "compressed" in self.quant_method.__class__.__name__.lower()
-                and param.data[expert_id] != 1
-                and (param.data[expert_id] - loaded_weight).abs() > 1e-5
+                (
+                    "compressed" in self.quant_method.__class__.__name__.lower()
+                    or "w4afp8" in self.quant_config.get_name()
+                )
+                and (param.data[expert_id] != 1).any()
+                and ((param.data[expert_id] - loaded_weight).abs() > 1e-5).any()
             ):
                 raise ValueError(
                     "input_scales of w1 and w3 of a layer "
@@ -808,16 +832,17 @@ class FusedMoE(torch.nn.Module):
             elif TopKOutputChecker.format_is_triton_kernel(topk_output):
                 raise NotImplementedError()
 
-        # Matrix multiply.
-        with use_symmetric_memory(get_tp_group()) as sm:
+        dispatch_output = self.dispatcher.dispatch(
+            hidden_states=hidden_states, topk_output=topk_output
+        )
 
-            final_hidden_states = self.quant_method.apply(
-                layer=self,
-                x=hidden_states,
-                topk_output=topk_output,
-                moe_runner_config=self.moe_runner_config,
-            )
-            sm.tag(final_hidden_states)
+        # TODO: consider using symmetric memory
+        combine_input = self.quant_method.apply(
+            layer=self,
+            dispatch_output=dispatch_output,
+        )
+
+        final_hidden_states = self.dispatcher.combine(combine_input)
 
         final_hidden_states = final_hidden_states[
             ..., :origin_hidden_states_dim
@@ -952,7 +977,6 @@ class FlashInferFusedMoE(FusedMoE):
             layer=self,
             x=hidden_states,
             topk_output=topk_output,
-            moe_runner_config=self.moe_runner_config,
         )
 
         if self.reduce_results and (self.moe_tp_size > 1 or self.moe_ep_size > 1):
@@ -1052,16 +1076,3 @@ class FlashInferFP4MoE(FusedMoE):
         )[0]
 
         return result
-
-
-def get_fused_moe_impl_class():
-    """Factory function to get the appropriate FusedMoE implementation class."""
-    if should_use_flashinfer_trtllm_moe() and _is_fp4_quantization_enabled():
-        # Use FP4 variant when FP4 quantization is enabled
-        return FlashInferFP4MoE
-    elif should_use_flashinfer_trtllm_moe():
-        # Use regular FlashInfer variant for non-FP4 FlashInfer cases
-        return FlashInferFusedMoE
-    else:
-        # Default case
-        return FusedMoE

sglang/srt/layers/moe/fused_moe_triton/moe_align_block_size.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+from typing import Tuple
+
+import torch
+import triton
+
+from sglang.srt.utils import is_cuda, is_hip
+
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+
+if _is_cuda or _is_hip:
+    from sgl_kernel import moe_align_block_size as sgl_moe_align_block_size
+
+
+def moe_align_block_size(
+    topk_ids: torch.Tensor, block_size: int, num_experts: int
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Aligns the token distribution across experts to be compatible with block
+    size for matrix multiplication.
+
+    Parameters:
+    - topk_ids: A tensor of shape [total_tokens, top_k] representing the
+        top-k expert indices for each token.
+    - block_size: The block size used in block matrix multiplication.
+    - num_experts: The total number of experts.
+
+    Returns:
+    - sorted_token_ids: A tensor containing the sorted token indices according
+        to their allocated expert.
+    - expert_ids: A tensor indicating the assigned expert index for each block.
+    - num_tokens_post_padded: The total number of tokens after padding,
+        ensuring divisibility by block_size.
+
+    This function pads the number of tokens that each expert needs to process
+    so that it is divisible by block_size.
+    Padding ensures that during block matrix multiplication, the dimensions
+    align correctly.
+
+    Example:
+    Given topk_ids = [[2, 3, 4], [1, 2, 4], [1, 3, 4], [1, 2, 3]],
+    block_size = 4, and num_experts = 4:
+    - We initially have 12 tokens (after repeating 'top_k' times) and 4 experts,
+        with each expert needing to process 3 tokens.
+    - As block_size is 4, we pad 1 token for each expert.
+    - First, flatten topk_ids to [2, 3, 4, 1, 2, 4, 1, 3, 4, 1, 2, 3].
+    - Then append padding tokens [12, 12, 12, 12] for each block.
+    - After sorting by expert index, we obtain token_ids
+        [3, 6, 9, 12, 0, 4, 10, 12, 1, 7, 11, 12, 2, 5, 8, 12].
+        Tokens 12 are non-existent (padding) and are ignored in
+        the subsequent matrix multiplication.
+    - The padding ensures that the total number of tokens is now divisible
+        by block_size for proper block matrix operations.
+    """
+    max_num_tokens_padded = topk_ids.numel() + (num_experts + 1) * (block_size - 1)
+    sorted_ids = torch.empty(
+        (max_num_tokens_padded,), dtype=torch.int32, device=topk_ids.device
+    )
+    max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
+    expert_ids = torch.empty(
+        (max_num_m_blocks,), dtype=torch.int32, device=topk_ids.device
+    )
+    num_tokens_post_pad = torch.empty((1), dtype=torch.int32, device=topk_ids.device)
+
+    # In EP, expert_ids for filtered experts are -1. We have num_experts + 1 ids in total.
+    cumsum_buffer = torch.empty(
+        (num_experts + 2,), dtype=torch.int32, device=topk_ids.device
+    )
+
+    # Threshold based on benchmark results
+    fuse_sorted_ids_padding = sorted_ids.shape[0] <= 4096
+    if not fuse_sorted_ids_padding:
+        sorted_ids.fill_(topk_ids.numel())
+
+    sgl_moe_align_block_size(
+        topk_ids,
+        num_experts + 1,
+        block_size,
+        sorted_ids,
+        expert_ids,
+        num_tokens_post_pad,
+        cumsum_buffer,
+        fuse_sorted_ids_padding,
+    )
+    return sorted_ids, expert_ids, num_tokens_post_pad

sglang/srt/layers/moe/moe_runner/__init__.py

@@ -1,3 +1,4 @@
 from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
+from sglang.srt.layers.moe.moe_runner.runner import MoeRunner
 
-__all__ = ["MoeRunnerConfig"]
+__all__ = ["MoeRunnerConfig", "MoeRunner"]

sglang/srt/layers/moe/moe_runner/base.py

@@ -1,9 +1,41 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
 from dataclasses import dataclass
-from typing import Optional
+from typing import TYPE_CHECKING, Callable, Optional, Tuple, TypeGuard
+
+import torch
+
+from sglang.srt.layers.moe.utils import MoeA2ABackend, MoeRunnerBackend
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.moe_runner.triton import (
+        TritonRunnerCore,
+        TritonRunnerInput,
+        TritonRunnerOutput,
+    )
+    from sglang.srt.layers.moe.token_dispatcher import (
+        CombineInput,
+        CombineInputFormat,
+        DispatchOutput,
+        DispatchOutputFormat,
+    )
 
 
 @dataclass
 class MoeRunnerConfig:
+
+    # MoE parameters
+    num_experts: Optional[int] = None
+    num_local_experts: Optional[int] = None
+    hidden_size: Optional[int] = None
+    intermediate_size_per_partition: Optional[int] = None
+    layer_id: Optional[int] = None
+    top_k: Optional[int] = None
+    num_fused_shared_experts: Optional[int] = None
+    params_dtype: Optional[torch.dtype] = None
+
+    # Runner configuration
     activation: str = "silu"
     apply_router_weight_on_input: bool = False
     inplace: bool = True
@@ -11,3 +43,244 @@ class MoeRunnerConfig:
     routed_scaling_factor: Optional[float] = None
     gemm1_alpha: Optional[float] = None
     gemm1_clamp_limit: Optional[float] = None
+
+
+@dataclass
+class RunnerInput(ABC):
+
+    @property
+    @abstractmethod
+    def runner_backend(self) -> MoeRunnerBackend: ...
+
+    def runner_backend_is_triton(self) -> TypeGuard[TritonRunnerInput]:
+        return self.runner_backend == MoeRunnerBackend.TRITON
+
+
+class RunnerOutput(ABC):
+
+    @property
+    @abstractmethod
+    def runner_backend(self) -> MoeRunnerBackend: ...
+
+    def runner_backend_is_triton(self) -> TypeGuard[TritonRunnerOutput]:
+        return self.runner_backend == MoeRunnerBackend.TRITON
+
+
+@dataclass
+class MoeQuantInfo(ABC):
+    """Moe quantization data."""
+
+    pass
+
+
+class MoeRunnerCore(ABC):
+
+    def __init__(self, config: MoeRunnerConfig):
+        self.config = config
+
+    @abstractmethod
+    def run(
+        self, runner_input: RunnerInput, quant_info: MoeQuantInfo, running_state: dict
+    ) -> RunnerOutput:
+        pass
+
+    @property
+    @abstractmethod
+    def runner_backend(self) -> MoeRunnerBackend: ...
+
+    def runner_backend_is_triton(self) -> TypeGuard[TritonRunnerCore]:
+        return self.runner_backend == MoeRunnerBackend.TRITON
+
+
+class FusedOpPool:
+
+    _fused_funcs: dict[str, Callable] = {}
+
+    @classmethod
+    def register_fused_func(
+        cls, a2a_backend_name: str, runner_backend_name: str, fused_func: Callable
+    ):
+        key = (a2a_backend_name, runner_backend_name)
+        if key in cls._fused_funcs:
+            raise ValueError(
+                f"Fused function for {a2a_backend_name} to {runner_backend_name} is already registered."
+            )
+        assert MoeA2ABackend(
+            a2a_backend_name
+        ), f"Invalid dispatch name: {a2a_backend_name}"
+        assert MoeRunnerBackend(
+            runner_backend_name
+        ), f"Invalid runner name: {runner_backend_name}"
+        cls._fused_funcs[key] = fused_func
+
+    @classmethod
+    def get_fused_func(cls, dispatch_name: str, runner_name: str) -> Optional[Callable]:
+        key = (dispatch_name, runner_name)
+        fused_func = cls._fused_funcs.get(key)
+        return fused_func
+
+
+class PermuteMethodPool:
+
+    _pre_permute_methods: dict[
+        Tuple[DispatchOutputFormat, MoeRunnerBackend], Callable
+    ] = {}
+    _post_permute_methods: dict[
+        Tuple[MoeRunnerBackend, CombineInputFormat], Callable
+    ] = {}
+
+    @classmethod
+    def register_pre_permute(
+        cls,
+        dispatch_output_name: str,
+        runner_backend_name: str,
+        permute_func: Callable,
+    ):
+        """
+        Register a customized pre-permute function for the given DispatchOutputFormat and MoeRunnerBackend.
+
+        :param dispatch_output_name: The DispatchOutputFormat name.
+        :param runner_backend_name: The MoeRunnerBackend name.
+        :param permute_func: The permute function to register.
+        """
+        # TODO: check if registration is valid
+        key = (dispatch_output_name, runner_backend_name)
+        if key in cls._pre_permute_methods:
+            raise ValueError(
+                f"Pre-permute method for {dispatch_output_name} to {runner_backend_name} is already registered."
+            )
+        cls._pre_permute_methods[key] = permute_func
+
+    @classmethod
+    def register_post_permute(
+        cls,
+        runner_backend_name: str,
+        combine_input_name: str,
+        permute_func: Callable,
+    ):
+        """
+        Register a customized post-permute function for the given MoeRunnerBackend and CombineInputFormat.
+
+        :param runner_backend_name: The MoeRunnerBackend name.
+        :param combine_input_name: The CombineInputFormat name.
+        :param permute_func: The permute function to register.
+        """
+        # TODO: check if registration is valid
+        key = (runner_backend_name, combine_input_name)
+        if key in cls._post_permute_methods:
+            raise ValueError(
+                f"Post-permute method for {runner_backend_name} to {combine_input_name} is already registered."
+            )
+        cls._post_permute_methods[key] = permute_func
+
+    @classmethod
+    def get_pre_permute(
+        cls,
+        dispatch_output_format: DispatchOutputFormat,
+        runner_input_format: MoeRunnerBackend,
+    ) -> Callable:
+        """
+        Retrieve the pre-permute function for the given DispatchOutputFormat and MoeRunnerBackend.
+
+        :param dispatch_output_format: The DispatchOutputFormat type.
+        :param runner_input_format: The MoeRunnerBackend type.
+        :return: The registered permute function or None if not found.
+        """
+        key = (dispatch_output_format, runner_input_format)
+        pre_permute_func = cls._pre_permute_methods.get(key)
+        assert (
+            pre_permute_func is not None
+        ), f"Pre-permute function for {dispatch_output_format} to {runner_input_format} is not registered"
+        return pre_permute_func
+
+    @classmethod
+    def get_post_permute(
+        cls,
+        runner_output_format: MoeRunnerBackend,
+        combine_input_format: CombineInputFormat,
+    ) -> Callable:
+        """
+        Retrieve the post-permute function for the given MoeRunnerBackend and CombineInputFormat.
+
+        :param runner_output_format: The MoeRunnerBackend type.
+        :param combine_input_format: The CombineInputFormat type.
+        :return: The registered permute function or None if not found.
+        """
+        key = (runner_output_format, combine_input_format)
+        post_permute_func = cls._post_permute_methods.get(key)
+        assert (
+            post_permute_func is not None
+        ), f"Post-permute function for {runner_output_format} to {combine_input_format} is not registered"
+        return post_permute_func
+
+
+def register_fused_func(
+    a2a_backend_name: str,
+    runner_backend_name: str,
+) -> Callable:
+    """
+    Decorator to register a fused function for the given DispatchOutputFormat and MoeRunnerBackend.
+
+    :param a2a_backend_name: The A2A backend name.
+    :param runner_backend_name: The MoeRunnerBackend name.
+    :return: The decorator function.
+    """
+
+    def decorator(fused_func: Callable):
+        FusedOpPool.register_fused_func(
+            a2a_backend_name, runner_backend_name, fused_func
+        )
+        return fused_func
+
+    return decorator
+
+
+def register_pre_permute(
+    dispatch_output_name: str,
+    runner_backend_name: str,
+) -> Callable:
+    """
+    Decorator to register a pre-permute function for the given DispatchOutputFormat and MoeRunnerBackend.
+
+    :param dispatch_output_name: The DispatchOutputFormat name.
+    :param runner_backend_name: The MoeRunnerBackend name.
+    :return: The decorator function.
+    """
+
+    def decorator(
+        permute_func: Callable[
+            [DispatchOutput, MoeQuantInfo, MoeRunnerConfig, dict], RunnerInput
+        ]
+    ) -> Callable:
+
+        PermuteMethodPool.register_pre_permute(
+            dispatch_output_name, runner_backend_name, permute_func
+        )
+        return permute_func
+
+    return decorator
+
+
+def register_post_permute(
+    runner_backend_name: str,
+    combine_input_name: str,
+) -> Callable:
+    """
+    Decorator to register a post-permute function for the given MoeRunnerBackend and CombineInputFormat.
+
+    :param runner_backend_name: The MoeRunnerBackend name.
+    :param combine_input_name: The CombineInputFormat name.
+    :return: The decorator function.
+    """
+
+    def decorator(
+        permute_func: Callable[
+            [RunnerOutput, MoeQuantInfo, MoeRunnerConfig, dict], CombineInput
+        ]
+    ) -> Callable:
+        PermuteMethodPool.register_post_permute(
+            runner_backend_name, combine_input_name, permute_func
+        )
+        return permute_func
+
+    return decorator

sglang/srt/layers/moe/moe_runner/runner.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import logging
+import os
+from typing import TYPE_CHECKING
+
+from sglang.srt.layers.moe.moe_runner.base import (
+    FusedOpPool,
+    MoeRunnerConfig,
+    PermuteMethodPool,
+)
+from sglang.srt.layers.moe.moe_runner.triton import TritonRunnerCore
+from sglang.srt.layers.moe.utils import get_moe_a2a_backend
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.moe_runner.base import MoeQuantInfo
+    from sglang.srt.layers.moe.token_dispatcher.base import CombineInput, DispatchOutput
+    from sglang.srt.layers.moe.utils import MoeRunnerBackend
+
+logger = logging.getLogger(__name__)
+
+
+class MoeRunner:
+
+    def __init__(self, runner_backend: MoeRunnerBackend, config: MoeRunnerConfig):
+        self.runner_backend = runner_backend
+        self.config = config
+
+        self.fused_func = None
+
+        if runner_backend.is_triton():
+            self.runner_core = TritonRunnerCore(config)
+        else:
+            raise NotImplementedError(f"Unsupported runner backend: {runner_backend}")
+
+        a2a_backend_name = get_moe_a2a_backend().value
+        runner_backend_name = runner_backend.value
+
+        self.fused_func = FusedOpPool.get_fused_func(
+            a2a_backend_name, runner_backend_name
+        )
+
+        SGLANG_CI_DISABLE_MOE_FUSED_FUNC = os.environ.get(
+            "SGLANG_CI_DISABLE_MOE_FUSED_FUNC", "0"
+        )
+        if SGLANG_CI_DISABLE_MOE_FUSED_FUNC == "1":
+            logger.info(
+                "SGLANG_CI_DISABLE_MOE_FUSED_FUNC is set to 1, disabling fused func"
+            )
+            self.fused_func = None
+
+    def run(
+        self, dispatch_output: DispatchOutput, quant_info: MoeQuantInfo
+    ) -> CombineInput:
+
+        if self.fused_func is not None:
+            return self.fused_func(dispatch_output, quant_info, self.config)
+
+        dispatch_format = dispatch_output.format.value
+        runner_format = self.runner_core.runner_backend.value
+        self.pre_permute_func = PermuteMethodPool.get_pre_permute(
+            dispatch_format, runner_format
+        )
+
+        running_state = {}
+        runner_input = self.pre_permute_func(
+            dispatch_output, quant_info, self.config, running_state
+        )
+        runner_output = self.runner_core.run(runner_input, quant_info, running_state)
+
+        runner_format = self.runner_core.runner_backend.value
+        combine_format = dispatch_output.format.value
+        self.post_permute_func = PermuteMethodPool.get_post_permute(
+            runner_format, combine_format
+        )
+        combine_input = self.post_permute_func(
+            runner_output, quant_info, self.config, running_state
+        )
+
+        return combine_input