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

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/deep_gemm.py

@@ -0,0 +1,304 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, List, Optional
+
+import torch
+
+from sglang.srt.layers.moe.moe_runner.base import (
+    MoeQuantInfo,
+    MoeRunnerConfig,
+    MoeRunnerCore,
+    RunnerInput,
+    RunnerOutput,
+    register_post_permute,
+    register_pre_permute,
+)
+from sglang.srt.layers.moe.utils import MoeRunnerBackend
+from sglang.srt.utils import dispose_tensor
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.token_dispatcher.standard import (
+        StandardCombineInput,
+        StandardDispatchOutput,
+    )
+
+
+# TODO(kaixih@nvidia): ideally we should merge this logic into
+# `fill_gateup_input_triton_kernel` to directly generate e8m0 scale.
+@torch.compile
+def _cast_to_e8m0_with_rounding_up(x: torch.Tensor) -> torch.Tensor:
+    temp = x.to(torch.float32).view(torch.int32)
+    exp = torch.bitwise_right_shift(temp, 23)
+    mant = torch.bitwise_and(temp, 0x7FFFFF)
+    is_ru = torch.logical_and(
+        torch.logical_and((mant > 0), (exp != 0xFE)),
+        ~torch.logical_and((exp == 0), (mant <= 0x400000)),
+    )
+    exp = torch.where(is_ru, exp + 1, exp)
+    new_x = exp.to(torch.uint8).view(torch.int)
+    return new_x.transpose(1, 2).contiguous().transpose(1, 2)
+
+
+@dataclass
+class DeepGemmRunnerInput(RunnerInput):
+    hidden_states: torch.Tensor
+    hidden_states_scale: torch.Tensor
+    masked_m: torch.Tensor
+    expected_m: int
+    use_masked_gemm: bool
+
+    @property
+    def runner_backend(self) -> MoeRunnerBackend:
+        return MoeRunnerBackend.DEEP_GEMM
+
+
+@dataclass
+class DeepGemmRunnerOutput(RunnerOutput):
+    hidden_states: torch.Tensor
+
+    @property
+    def runner_backend(self) -> MoeRunnerBackend:
+        return MoeRunnerBackend.DEEP_GEMM
+
+
+@dataclass
+class DeepGemmMoeQuantInfo(MoeQuantInfo):
+    w13_weight: torch.Tensor
+    w2_weight: torch.Tensor
+    use_fp8: bool
+    w13_scale: Optional[torch.Tensor] = None
+    w2_scale: Optional[torch.Tensor] = None
+    block_shape: Optional[List[int]] = None
+
+
+class DeepGemmRunnerCore(MoeRunnerCore):
+    def __init__(self, config: MoeRunnerConfig):
+        super().__init__(config)
+        assert self.config.activation == "silu"
+
+    def run(
+        self,
+        runner_input: DeepGemmRunnerInput,
+        quant_info: DeepGemmMoeQuantInfo,
+        running_state: dict,
+    ) -> DeepGemmRunnerOutput:
+
+        if runner_input.use_masked_gemm:
+            hidden_states = self._run_masked_gemm(
+                runner_input,
+                quant_info,
+                running_state,
+            )
+        else:
+            hidden_states = self._run_contiguous_gemm(
+                runner_input,
+                quant_info,
+                running_state,
+            )
+        return DeepGemmRunnerOutput(hidden_states=hidden_states)
+
+    def _run_masked_gemm(
+        self,
+        runner_input: DeepGemmRunnerInput,
+        quant_info: DeepGemmMoeQuantInfo,
+        running_state: dict,
+    ) -> torch.Tensor:
+
+        from sglang.srt.layers.moe.ep_moe.kernels import (
+            silu_and_mul_masked_post_quant_fwd,
+        )
+        from sglang.srt.layers.quantization import deep_gemm_wrapper
+
+        hidden_states = runner_input.hidden_states
+        hidden_states_scale = runner_input.hidden_states_scale
+        masked_m = runner_input.masked_m
+        expected_m = runner_input.expected_m
+
+        w13_weight = quant_info.w13_weight
+        w2_weight = quant_info.w2_weight
+        w13_scale = quant_info.w13_scale
+        w2_scale = quant_info.w2_scale
+
+        hidden_states_device = running_state["hidden_states_device"]
+
+        if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
+            b, s_mn, s_k = hidden_states_scale.shape
+            assert (
+                s_mn % 4 == 0 and s_k % 4 == 0
+            ), f"scales must be aligned to 4, but got ({b}, {s_mn}, {s_k})"
+
+        # GroupGemm-0
+        if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
+            hidden_states_scale = _cast_to_e8m0_with_rounding_up(hidden_states_scale)
+        else:
+            hidden_states_scale = deep_gemm_wrapper.get_mn_major_tma_aligned_tensor(
+                hidden_states_scale
+            )
+
+        num_groups, m, k = hidden_states.shape
+        n = w13_weight.size(1)
+        gateup_output = torch.empty(
+            (num_groups, m, n), device=hidden_states_device, dtype=torch.bfloat16
+        )
+        deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked(
+            (hidden_states, hidden_states_scale),
+            (w13_weight, w13_scale),
+            gateup_output,
+            masked_m,
+            expected_m,
+        )
+        dispose_tensor(hidden_states)
+
+        # Act
+        down_input = torch.empty(
+            (
+                gateup_output.shape[0],
+                gateup_output.shape[1],
+                gateup_output.shape[2] // 2,
+            ),
+            device=hidden_states_device,
+            dtype=torch.float8_e4m3fn,
+        )
+        scale_block_size = 128
+        down_input_scale = torch.empty(
+            (
+                gateup_output.shape[0],
+                gateup_output.shape[1],
+                gateup_output.shape[2] // 2 // scale_block_size,
+            ),
+            device=hidden_states_device,
+            dtype=torch.float32,
+        )
+        silu_and_mul_masked_post_quant_fwd(
+            gateup_output,
+            down_input,
+            down_input_scale,
+            scale_block_size,
+            masked_m,
+            scale_ue8m0=deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0,
+        )
+        del gateup_output
+
+        # GroupGemm-1
+        n = w2_weight.shape[1]
+
+        if not deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
+            down_input_scale = deep_gemm_wrapper.get_mn_major_tma_aligned_tensor(
+                down_input_scale
+            )
+
+        down_output = torch.empty(
+            (num_groups, m, n), device=hidden_states_device, dtype=torch.bfloat16
+        )
+        deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked(
+            (down_input, down_input_scale),
+            (w2_weight, w2_scale),
+            down_output,
+            masked_m,
+            expected_m,
+        )
+        del down_input
+
+        return down_output
+
+    def _run_contiguous_gemm(
+        self,
+        runner_input: DeepGemmRunnerInput,
+        quant_info: DeepGemmMoeQuantInfo,
+        running_state: dict,
+    ) -> torch.Tensor:
+        pass
+
+    @property
+    def runner_backend(self) -> MoeRunnerBackend:
+        return MoeRunnerBackend.DEEP_GEMM
+
+
+@register_pre_permute("standard", "deep_gemm")
+def pre_permute_standard_to_deep_gemm(
+    dispatch_output: StandardDispatchOutput,
+    quant_info: DeepGemmMoeQuantInfo,
+    runner_config: MoeRunnerConfig,
+    running_state: dict,
+) -> DeepGemmRunnerInput:
+    from sglang.srt.layers.moe.ep_moe.kernels import moe_ep_deepgemm_preprocess
+
+    hidden_states, topk_output = dispatch_output
+    topk_weights, topk_ids, _ = topk_output
+
+    hidden_states_shape = hidden_states.shape
+    hidden_states_dtype = hidden_states.dtype
+    hidden_states_device = hidden_states.device
+    hidden_states_ref = hidden_states
+
+    topk_weights, topk_ids = topk_weights, topk_ids
+
+    # PreReorder
+    masked_m, expected_m, src2dst, hidden_states, hidden_states_scale = (
+        moe_ep_deepgemm_preprocess(
+            topk_ids,
+            runner_config.num_local_experts,
+            hidden_states,
+            runner_config.top_k,
+            quant_info.block_shape,
+        )
+    )
+
+    dispose_tensor(hidden_states_ref)
+
+    running_state["topk_ids"] = topk_ids
+    running_state["topk_weights"] = topk_weights
+    running_state["hidden_states_shape"] = hidden_states_shape
+    running_state["hidden_states_dtype"] = hidden_states_dtype
+    running_state["hidden_states_device"] = hidden_states_device
+    running_state["src2dst"] = src2dst
+
+    return DeepGemmRunnerInput(
+        hidden_states=hidden_states,
+        hidden_states_scale=hidden_states_scale,
+        masked_m=masked_m,
+        expected_m=expected_m,
+        use_masked_gemm=True,
+    )
+
+
+@register_post_permute("deep_gemm", "standard")
+def post_permute_deep_gemm_to_standard(
+    runner_output: DeepGemmRunnerOutput,
+    quant_info: DeepGemmMoeQuantInfo,
+    runner_config: MoeRunnerConfig,
+    running_state: dict,
+) -> StandardCombineInput:
+    from sglang.srt.layers.moe.ep_moe.kernels import post_reorder_triton_kernel
+    from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
+
+    hidden_states_shape = running_state["hidden_states_shape"]
+    hidden_states_dtype = running_state["hidden_states_dtype"]
+    hidden_states_device = running_state["hidden_states_device"]
+    src2dst = running_state["src2dst"]
+    topk_ids = running_state["topk_ids"]
+    topk_weights = running_state["topk_weights"]
+
+    output = torch.empty(
+        hidden_states_shape, dtype=hidden_states_dtype, device=hidden_states_device
+    )
+    post_reorder_triton_kernel[(hidden_states_shape[0],)](
+        runner_output.hidden_states,
+        output,
+        src2dst,
+        topk_ids,
+        topk_weights,
+        runner_config.top_k,
+        hidden_states_shape[1],
+        BLOCK_SIZE=512,
+    )
+
+    dispose_tensor(runner_output.hidden_states)
+
+    if runner_config.routed_scaling_factor is not None:
+        output *= runner_config.routed_scaling_factor
+
+    return StandardCombineInput(
+        hidden_states=output,
+    )

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

@@ -0,0 +1,83 @@
+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.deep_gemm import DeepGemmRunnerCore
+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)
+        elif runner_backend.is_deep_gemm():
+            self.runner_core = DeepGemmRunnerCore(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