diff --git a/python/sglang/srt/layers/moe/moe_runner/aiter.py b/python/sglang/srt/layers/moe/moe_runner/aiter.py new file mode 100644 index 000000000..36478e07e --- /dev/null +++ b/python/sglang/srt/layers/moe/moe_runner/aiter.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.srt.layers.moe.moe_runner.base import ( + MoeQuantInfo, + MoeRunnerConfig, + register_fused_func, +) + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher.standard import ( + StandardCombineInput, + StandardDispatchOutput, + ) + + +class AiterQuantType(str, Enum): + NONE = "No" + PER_TOKEN = "per_Token" + PER_128X128 = "per_128x128" + PER_1X32 = "per_1x32" + + +@dataclass +class AiterMoeQuantInfo(MoeQuantInfo): + w13_weight: torch.Tensor + w2_weight: torch.Tensor + quant_type: AiterQuantType = AiterQuantType.NONE + w13_scale: Optional[torch.Tensor] = None + w2_scale: Optional[torch.Tensor] = None + a13_scale: Optional[torch.Tensor] = None + a2_scale: Optional[torch.Tensor] = None + b13: Optional[torch.Tensor] = None + b2: Optional[torch.Tensor] = None + expert_mask: Optional[torch.Tensor] = None + doweight_stage1: bool = False + hidden_pad: int = 0 + intermediate_pad: int = 0 + + +_AITER_ACTIVATIONS = {"silu": "Silu", "swiglu": "Swiglu"} + + +@register_fused_func("none", "aiter") +def fused_experts_none_to_aiter( + dispatch_output: StandardDispatchOutput, + quant_info: AiterMoeQuantInfo, + runner_config: MoeRunnerConfig, +) -> StandardCombineInput: + from aiter import ActivationType, QuantType + from aiter.fused_moe import fused_moe + + from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput + + assert not runner_config.no_combine, "no_combine=True is not supported by AITER" + + hidden_states = dispatch_output.hidden_states + topk_weights, topk_ids, _ = dispatch_output.topk_output + topk_weights = topk_weights.to(torch.float32) + + if runner_config.apply_router_weight_on_input and not quant_info.doweight_stage1: + # Pre-scale at the Python level for kernels that don't honor doweight_stage1. + assert ( + topk_weights.dim() == 2 and topk_weights.shape[-1] == 1 + ), "apply_router_weight_on_input requires topk=1" + hidden_states = hidden_states * topk_weights.to(hidden_states.dtype) + topk_weights = torch.ones_like(topk_weights) + + activation = runner_config.activation + output = fused_moe( + hidden_states=hidden_states, + w1=quant_info.w13_weight, + w2=quant_info.w2_weight, + topk_weight=topk_weights, + topk_ids=topk_ids.to(torch.int32), + quant_type=getattr(QuantType, quant_info.quant_type.value), + activation=getattr(ActivationType, _AITER_ACTIVATIONS.get(activation, "Gelu")), + w1_scale=quant_info.w13_scale, + w2_scale=quant_info.w2_scale, + a1_scale=quant_info.a13_scale, + a2_scale=quant_info.a2_scale, + bias1=quant_info.b13, + bias2=quant_info.b2, + expert_mask=quant_info.expert_mask, + doweight_stage1=quant_info.doweight_stage1, + hidden_pad=quant_info.hidden_pad, + intermediate_pad=quant_info.intermediate_pad, + ) + return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/layers/moe/moe_runner/runner.py b/python/sglang/srt/layers/moe/moe_runner/runner.py index 09b04c250..d53f57cc1 100644 --- a/python/sglang/srt/layers/moe/moe_runner/runner.py +++ b/python/sglang/srt/layers/moe/moe_runner/runner.py @@ -43,6 +43,11 @@ class MoeRunner: self.runner_core = TritonKernelsRunnerCore(config) elif runner_backend.is_deep_gemm(): self.runner_core = DeepGemmRunnerCore(config) + elif runner_backend.is_aiter(): + # Side-effect import: registers the ("none", "aiter") fused func. + from sglang.srt.layers.moe.moe_runner import aiter # noqa: F401 + + self.runner_core = None # AITER only supports fused path elif runner_backend.is_marlin(): if lora_enabled: from sglang.srt.lora.lora_moe_runner_marlin import MarlinLoraRunnerCore diff --git a/python/sglang/srt/layers/moe/utils.py b/python/sglang/srt/layers/moe/utils.py index 59cdce5d1..e05167da9 100644 --- a/python/sglang/srt/layers/moe/utils.py +++ b/python/sglang/srt/layers/moe/utils.py @@ -78,6 +78,7 @@ class MoeRunnerBackend(Enum): FLASHINFER_CUTEDSL = "flashinfer_cutedsl" CUTLASS = "cutlass" MARLIN = "marlin" + AITER = "aiter" def is_auto(self): return self == MoeRunnerBackend.AUTO @@ -112,6 +113,9 @@ class MoeRunnerBackend(Enum): def is_marlin(self): return self == MoeRunnerBackend.MARLIN + def is_aiter(self): + return self == MoeRunnerBackend.AITER + class DeepEPMode(Enum): diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py index 6e00b51b7..22ae9283f 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py @@ -13,6 +13,7 @@ from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( ) from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo from sglang.srt.layers.moe.utils import ( + get_moe_a2a_backend, get_moe_runner_backend, get_moe_weight_sizes, ) @@ -41,8 +42,6 @@ _is_hip = is_hip() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip if _use_aiter: - from aiter import ActivationType, QuantType - from aiter.fused_moe import fused_moe from aiter.ops.shuffle import shuffle_weight @@ -347,8 +346,25 @@ class CompressedTensorsW8A8Fp8MoE(CompressedTensorsMoEScheme): self.moe_runner_config = moe_runner_config moe_runner_backend = get_moe_runner_backend() if moe_runner_backend.is_auto(): - moe_runner_backend = MoeRunnerBackend.TRITON - self.runner = MoeRunner(moe_runner_backend, moe_runner_config) + if ( + _use_aiter + and self.weight_quant.strategy == QuantizationStrategy.CHANNEL + and get_moe_a2a_backend().is_none() + ): + moe_runner_backend = MoeRunnerBackend.AITER + else: + moe_runner_backend = MoeRunnerBackend.TRITON + + if ( + moe_runner_backend.is_aiter() + or moe_runner_backend.is_triton() + or moe_runner_backend.is_flashinfer_trtllm() + or moe_runner_backend.is_flashinfer_trtllm_routed() + ): + self.runner = MoeRunner(moe_runner_backend, moe_runner_config) + else: + # TODO(cwan): refactor other backends + pass def apply_weights( self, @@ -356,46 +372,28 @@ class CompressedTensorsW8A8Fp8MoE(CompressedTensorsMoEScheme): dispatch_output: StandardDispatchOutput, ) -> CombineInput: - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - x = dispatch_output.hidden_states topk_output = dispatch_output.topk_output moe_runner_config = self.moe_runner_config - if _use_aiter and self.weight_quant.strategy == QuantizationStrategy.CHANNEL: + if self.runner.runner_backend.is_aiter(): + from sglang.srt.layers.moe.moe_runner.aiter import ( + AiterMoeQuantInfo, + AiterQuantType, + ) + assert not moe_runner_config.no_combine, "unsupported" - topk_weights, topk_ids, _ = topk_output - if moe_runner_config.apply_router_weight_on_input: - assert ( - topk_weights.dim() == 2 - ), "`topk_weights` should be in shape (num_tokens, topk)" - _, topk = topk_weights.shape - assert ( - topk == 1 - ), "Only support topk=1 when `apply_router_weight_on_input` is True" - x = x * topk_weights.to(x.dtype) - topk_weights = torch.ones_like( - topk_weights, dtype=torch.float32 - ) # topk_weights must be FP32 (float32) - output = fused_moe( - x, - layer.w13_weight, - layer.w2_weight, - topk_weights, - topk_ids, - activation=( - ActivationType.Silu - if moe_runner_config.activation == "silu" - else ActivationType.Gelu - ), - quant_type=QuantType.per_Token, - w1_scale=layer.w13_weight_scale, + quant_info = AiterMoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + quant_type=AiterQuantType.PER_TOKEN, + w13_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, - a1_scale=layer.w13_input_scale, + a13_scale=layer.w13_input_scale, a2_scale=layer.w2_input_scale, ) - return StandardCombineInput(hidden_states=output) + return self.runner.run(dispatch_output, quant_info) elif self.weight_quant.strategy == QuantizationStrategy.BLOCK: if self.use_flashinfer_trtllm: from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index ac2c31b6d..c85776d5b 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -28,6 +28,7 @@ from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import ( from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo from sglang.srt.layers.moe.utils import ( RoutingMethodType, + get_moe_a2a_backend, get_moe_padding_size, get_moe_runner_backend, get_moe_weight_sizes, @@ -97,8 +98,8 @@ from sglang.srt.utils import ( ) if TYPE_CHECKING: + from sglang.srt.layers.moe.moe_runner.aiter import AiterMoeQuantInfo from sglang.srt.layers.moe.token_dispatcher import CombineInput, DispatchOutput - from sglang.srt.layers.moe.topk import TopKOutput from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config from sglang.srt.models.utils import WeightsMapper @@ -113,8 +114,6 @@ _use_hip_int4 = get_bool_env_var("SGLANG_INT4_WEIGHT") and _is_hip _use_aiter = envs.SGLANG_USE_AITER.get() and _is_hip if _use_aiter or _use_hip_int4: - from aiter import ActivationType, QuantType - from aiter.fused_moe import fused_moe from aiter.ops.shuffle import shuffle_weight if _use_aiter: @@ -1521,11 +1520,21 @@ class Fp8MoEMethod(FusedMoEMethodBase): if moe_runner_backend.is_auto(): if self.is_deepgemm_moe_runner_backend_enabled(): moe_runner_backend = MoeRunnerBackend.DEEP_GEMM + elif ( + _is_hip + and (_use_aiter or _use_hip_int4) + and get_moe_a2a_backend().is_none() + ): + # *EPMoE backends bypass self.runner via run_moe_core, and the + # AITER fused func is only registered for ("none", "aiter"). + moe_runner_backend = MoeRunnerBackend.AITER else: moe_runner_backend = MoeRunnerBackend.TRITON + if ( moe_runner_backend.is_deep_gemm() or moe_runner_backend.is_triton() + or moe_runner_backend.is_aiter() or moe_runner_backend.is_flashinfer_trtllm() or moe_runner_backend.is_flashinfer_trtllm_routed() ): @@ -1590,16 +1599,17 @@ class Fp8MoEMethod(FusedMoEMethodBase): ) return StandardCombineInput(hidden_states=output) - if _is_hip: - ret = self.maybe_apply_hip_fused_experts( + if ( + _is_hip + and getattr(self, "runner", None) is not None + and self.runner.runner_backend.is_aiter() + ): + quant_info = self.maybe_get_hip_aiter_quant_info( layer, - x, - dispatch_output.topk_output, - moe_runner_config.activation, moe_runner_config.no_combine, ) - if ret is not None: - return StandardCombineInput(hidden_states=ret) + if quant_info is not None: + return self.runner.run(dispatch_output, quant_info) if get_moe_runner_backend().is_cutlass(): from sglang.srt.layers.moe.cutlass_moe import cutlass_fused_experts_fp8 @@ -1790,69 +1800,36 @@ class Fp8MoEMethod(FusedMoEMethodBase): self._cutlass_buffers_ready = True - def maybe_apply_hip_fused_experts( + def maybe_get_hip_aiter_quant_info( self, layer: torch.nn.Module, - x: torch.Tensor, - topk_output: TopKOutput, - activation: str = "silu", no_combine: bool = False, - ) -> Optional[torch.Tensor]: - topk_weights, topk_ids, _ = topk_output - if _use_hip_int4: - # TODO: add triton kernel and add check _use_aiter - assert not no_combine, f"{no_combine=} is not supported." - return fused_moe( - x, - layer.w13_weight, - layer.w2_weight, - topk_weights, - topk_ids, - quant_type=QuantType.per_Token, - w1_scale=layer.w13_weight_scale1, - w2_scale=layer.w2_weight_scale1, - activation=( - ActivationType.Silu if activation == "silu" else ActivationType.Gelu - ), - ) + ) -> Optional["AiterMoeQuantInfo"]: + if not (_use_aiter or _use_hip_int4): + return None + assert not no_combine, f"{no_combine=} is not supported." - if _use_aiter: - assert not no_combine, f"{no_combine=} is not supported." - if self.block_quant: - return fused_moe( - x, - layer.w13_weight, - layer.w2_weight, - topk_weights, - topk_ids, - w1_scale=layer.w13_weight_scale_inv, - w2_scale=layer.w2_weight_scale_inv, - quant_type=QuantType.per_128x128, - activation=( - ActivationType.Silu - if activation == "silu" - else ActivationType.Gelu - ), - expert_mask=layer.dispatcher.expert_mask_gpu, - ) - else: - return fused_moe( - x, - layer.w13_weight, - layer.w2_weight, - topk_weights, - topk_ids, - quant_type=QuantType.per_Token, - w1_scale=layer.w13_weight_scale1, - w2_scale=layer.w2_weight_scale1, - activation=( - ActivationType.Silu - if activation == "silu" - else ActivationType.Gelu - ), - expert_mask=layer.dispatcher.expert_mask_gpu, - ) - return None + from sglang.srt.layers.moe.moe_runner.aiter import ( + AiterMoeQuantInfo, + AiterQuantType, + ) + + if _use_aiter and self.block_quant: + quant_type = AiterQuantType.PER_128X128 + w13_scale = layer.w13_weight_scale_inv + w2_scale = layer.w2_weight_scale_inv + else: + quant_type = AiterQuantType.PER_TOKEN + w13_scale = layer.w13_weight_scale1 + w2_scale = layer.w2_weight_scale1 + return AiterMoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + quant_type=quant_type, + w13_scale=w13_scale, + w2_scale=w2_scale, + expert_mask=layer.dispatcher.expert_mask_gpu if _use_aiter else None, + ) class Fp8KVCacheMethod(BaseKVCacheMethod): diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py index d6435da3a..d5ad4d403 100644 --- a/python/sglang/srt/layers/quantization/mxfp4.py +++ b/python/sglang/srt/layers/quantization/mxfp4.py @@ -16,6 +16,7 @@ from __future__ import annotations +from dataclasses import replace from typing import TYPE_CHECKING, List, Optional import torch @@ -28,7 +29,7 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.layers.dp_attention import is_allocation_symmetric from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo -from sglang.srt.layers.moe.utils import get_moe_runner_backend +from sglang.srt.layers.moe.utils import get_moe_a2a_backend, get_moe_runner_backend from sglang.srt.layers.quantization.base_config import ( FusedMoEMethodBase, QuantizationConfig, @@ -116,8 +117,6 @@ _is_shuffle_moe_mxfp4 = is_gfx95_supported() if _is_hip: # import aiter try: - from aiter import ActivationType, QuantType - from aiter.fused_moe import fused_moe from aiter.ops.shuffle import ( shuffle_scale_a16w4, shuffle_weight, @@ -126,9 +125,7 @@ if _is_hip: from aiter.ops.triton.quant import dynamic_mxfp4_quant from aiter.utility.fp4_utils import e8m0_shuffle except ImportError as err: - ActivationType = QuantType = fused_moe = dynamic_mxfp4_quant = e8m0_shuffle = ( - err - ) + dynamic_mxfp4_quant = e8m0_shuffle = err def _swizzle_mxfp4(quant_tensor, scale, num_warps): @@ -743,12 +740,26 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig ): self.moe_runner_config = moe_runner_config - backend = ( - MoeRunnerBackend.TRITON_KERNELS - if self.use_triton_kernels - else MoeRunnerBackend.TRITON - ) - self.runner = MoeRunner(backend, moe_runner_config) + moe_runner_backend = get_moe_runner_backend() + if moe_runner_backend.is_auto(): + # Must match apply() priority: _use_aiter before use_triton_kernels. + if _use_aiter and get_moe_a2a_backend().is_none(): + moe_runner_backend = MoeRunnerBackend.AITER + elif self.use_triton_kernels: + moe_runner_backend = MoeRunnerBackend.TRITON_KERNELS + else: + moe_runner_backend = MoeRunnerBackend.TRITON + + if moe_runner_backend.is_aiter(): + # MXFP4 hard-codes Swiglu in the AITER kernel path. + self.runner = MoeRunner( + moe_runner_backend, replace(moe_runner_config, activation="swiglu") + ) + elif moe_runner_backend.is_triton_kernels() or moe_runner_backend.is_triton(): + self.runner = MoeRunner(moe_runner_backend, moe_runner_config) + else: + # TODO(cwan): refactor other backends + pass def apply( self, @@ -832,7 +843,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): )[0] return StandardCombineInput(hidden_states=trtllm_gen_output) if _use_aiter: - topk_weights, topk_ids, _ = topk_output + from sglang.srt.layers.moe.moe_runner.aiter import ( + AiterMoeQuantInfo, + AiterQuantType, + ) if hasattr(torch, "float4_e2m1fn_x2"): w13_weight = layer.w13_weight.view(torch.float4_e2m1fn_x2) @@ -841,33 +855,25 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): w13_weight = layer.w13_weight w2_weight = layer.w2_weight - origi_hidden_size = self.hidden_size - self.hidden_pad - - x = torch.nn.functional.pad( - x, - (0, self.hidden_pad), - mode="constant", - value=0.0, + x_padded = torch.nn.functional.pad( + x, (0, self.hidden_pad), mode="constant", value=0.0 ) - - output = fused_moe( - x, - w13_weight, - w2_weight, - topk_weights, - topk_ids, - expert_mask=layer.dispatcher.expert_mask_gpu, - activation=ActivationType.Swiglu, - quant_type=QuantType.per_1x32, - w1_scale=layer.w13_weight_scale, + quant_info = AiterMoeQuantInfo( + w13_weight=w13_weight, + w2_weight=w2_weight, + quant_type=AiterQuantType.PER_1X32, + w13_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, + b13=layer.w13_weight_bias, + b2=layer.w2_weight_bias, + expert_mask=layer.dispatcher.expert_mask_gpu, doweight_stage1=self.moe_runner_config.apply_router_weight_on_input, hidden_pad=self.hidden_pad, intermediate_pad=self.intermediate_pad, - bias1=layer.w13_weight_bias, - bias2=layer.w2_weight_bias, ) - return StandardCombineInput(hidden_states=output) + return self.runner.run( + dispatch_output._replace(hidden_states=x_padded), quant_info + ) backend = self.runner.runner_backend if backend.is_triton_kernels(): @@ -1002,22 +1008,25 @@ class Mxfp4DynamicQuantMoEMethod(FusedMoEMethodBase): self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig ): self.moe_runner_config = moe_runner_config + moe_runner_backend = get_moe_runner_backend() + if moe_runner_backend.is_auto() and get_moe_a2a_backend().is_none(): + moe_runner_backend = MoeRunnerBackend.AITER + + if moe_runner_backend.is_aiter(): + self.runner = MoeRunner(moe_runner_backend, moe_runner_config) + else: + # TODO(cwan): refactor other backends + pass def apply( self, layer: torch.nn.Module, dispatch_output: StandardDispatchOutput, ) -> CombineInput: - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - - x = dispatch_output.hidden_states - topk_output = dispatch_output.topk_output - - topk_weights, topk_ids, _ = topk_output - if _is_hip: - topk_weights = topk_weights.to( - torch.float32 - ) # aiter's moe_sorting requires topk_weights to be FP32 + from sglang.srt.layers.moe.moe_runner.aiter import ( + AiterMoeQuantInfo, + AiterQuantType, + ) if hasattr(torch, "float4_e2m1fn_x2"): w13_weight = layer.w13_weight.view(torch.float4_e2m1fn_x2) @@ -1030,21 +1039,12 @@ class Mxfp4DynamicQuantMoEMethod(FusedMoEMethodBase): w13_weight.is_shuffled = True w2_weight.is_shuffled = True - output = fused_moe( - x, - w13_weight, - w2_weight, - topk_weights, - topk_ids, - quant_type=QuantType.per_1x32, - w1_scale=layer.w13_weight_scale, + quant_info = AiterMoeQuantInfo( + w13_weight=w13_weight, + w2_weight=w2_weight, + quant_type=AiterQuantType.PER_1X32, + w13_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, - activation=( - ActivationType.Silu - if self.moe_runner_config.activation == "silu" - else ActivationType.Gelu - ), - doweight_stage1=False, expert_mask=layer.dispatcher.expert_mask_gpu, ) - return StandardCombineInput(hidden_states=output) + return self.runner.run(dispatch_output, quant_info) diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py index 715d782a5..2640d2319 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py +++ b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py @@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Any import torch -from sglang.srt.layers.moe import MoeRunnerConfig +from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig from sglang.srt.layers.moe.utils import get_moe_weight_sizes from sglang.srt.layers.quantization.quark.schemes import QuarkMoEScheme from sglang.srt.utils import ( @@ -32,8 +32,6 @@ __all__ = ["QuarkW4A4MXFp4MoE"] _is_hip = is_hip() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip if _use_aiter: - from aiter import ActivationType, QuantType - from aiter.fused_moe import fused_moe from aiter.ops.shuffle import shuffle_weight from aiter.utility.fp4_utils import e8m0_shuffle @@ -182,24 +180,31 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme): def create_moe_runner( self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig ): + from sglang.srt.layers.moe.utils import ( + get_moe_a2a_backend, + get_moe_runner_backend, + ) + self.moe_runner_config = moe_runner_config + moe_runner_backend = get_moe_runner_backend() + if moe_runner_backend.is_auto() and get_moe_a2a_backend().is_none(): + moe_runner_backend = MoeRunnerBackend.AITER + + if moe_runner_backend.is_aiter(): + self.runner = MoeRunner(moe_runner_backend, moe_runner_config) + else: + # TODO(cwan): refactor other backends + pass def apply_weights( self, layer: torch.nn.Module, dispatch_output: StandardDispatchOutput, ) -> CombineInput: - - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput - - x = dispatch_output.hidden_states - topk_output = dispatch_output.topk_output - moe_runner_config = self.moe_runner_config - topk_weights, topk_ids, _ = topk_output - if _is_hip: - topk_weights = topk_weights.to( - torch.float32 - ) # aiter's moe_sorting requires topk_weights to be FP32 + from sglang.srt.layers.moe.moe_runner.aiter import ( + AiterMoeQuantInfo, + AiterQuantType, + ) if hasattr(torch, "float4_e2m1fn_x2"): w13_weight = layer.w13_weight.view(torch.float4_e2m1fn_x2) @@ -212,21 +217,12 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme): w13_weight.is_shuffled = True w2_weight.is_shuffled = True - output = fused_moe( - x, - w13_weight, - w2_weight, - topk_weights, - topk_ids, - quant_type=QuantType.per_1x32, - w1_scale=layer.w13_weight_scale, + quant_info = AiterMoeQuantInfo( + w13_weight=w13_weight, + w2_weight=w2_weight, + quant_type=AiterQuantType.PER_1X32, + w13_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, - activation=( - ActivationType.Silu - if moe_runner_config.activation == "silu" - else ActivationType.Gelu - ), - doweight_stage1=False, expert_mask=layer.dispatcher.expert_mask_gpu, ) - return StandardCombineInput(hidden_states=output) + return self.runner.run(dispatch_output, quant_info) diff --git a/python/sglang/srt/layers/quantization/quark_int4fp8_moe.py b/python/sglang/srt/layers/quantization/quark_int4fp8_moe.py index 7fbbc6113..0e9204d19 100644 --- a/python/sglang/srt/layers/quantization/quark_int4fp8_moe.py +++ b/python/sglang/srt/layers/quantization/quark_int4fp8_moe.py @@ -11,7 +11,7 @@ from sglang.srt.layers.int4fp8_utils import ( quantize_fp8_scale_tensorwise, quantize_int4_scale_columnwise, ) -from sglang.srt.layers.moe import MoeRunnerConfig +from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig from sglang.srt.layers.quantization.base_config import ( FusedMoEMethodBase, QuantizationConfig, @@ -27,8 +27,6 @@ _is_hip = is_hip() if _is_hip: - from aiter import ActivationType, QuantType - from aiter.fused_moe import fused_moe from aiter.ops.shuffle import shuffle_weight ON_GFX950 = "gfx950" in torch.cuda.get_device_properties("cuda").gcnArchName @@ -405,18 +403,32 @@ class QuarkInt4Fp8MoEMethod(FusedMoEMethodBase): def create_moe_runner( self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig ): + from sglang.srt.layers.moe.utils import ( + get_moe_a2a_backend, + get_moe_runner_backend, + ) + self.moe_runner_config = moe_runner_config + moe_runner_backend = get_moe_runner_backend() + if moe_runner_backend.is_auto() and get_moe_a2a_backend().is_none(): + moe_runner_backend = MoeRunnerBackend.AITER + + if moe_runner_backend.is_aiter(): + self.runner = MoeRunner(moe_runner_backend, moe_runner_config) + else: + # TODO(cwan): refactor other backends + pass def apply( self, layer: torch.nn.Module, dispatch_output: "DispatchOutput", ) -> torch.Tensor: - # TODO: fix circular imports issues in sglang forcing us to import here instead of at - # the top of file. - from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput + from sglang.srt.layers.moe.moe_runner.aiter import ( + AiterMoeQuantInfo, + AiterQuantType, + ) - topk_output = dispatch_output.topk_output moe_runner_config = self.moe_runner_config # TODO: add triton kernel and add check get_bool_env_var("CK_MOE") @@ -424,20 +436,11 @@ class QuarkInt4Fp8MoEMethod(FusedMoEMethodBase): not moe_runner_config.no_combine ), f"no_combine={moe_runner_config.no_combine} is not supported." - output = fused_moe( - dispatch_output.hidden_states, - layer.w13_weight, - layer.w2_weight, - topk_output.topk_weights, - topk_output.topk_ids, - quant_type=QuantType.per_Token, - w1_scale=layer.w13_int4_scale, + quant_info = AiterMoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + quant_type=AiterQuantType.PER_TOKEN, + w13_scale=layer.w13_int4_scale, w2_scale=layer.w2_int4_scale, - activation=( - ActivationType.Silu - if moe_runner_config.activation == "silu" - else ActivationType.Gelu - ), ) - - return StandardCombineInput(hidden_states=output) + return self.runner.run(dispatch_output, quant_info) diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index bd9eb8b2a..a3a381171 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -17,6 +17,7 @@ from sglang.srt.layers.moe import ( MoeRunner, MoeRunnerBackend, MoeRunnerConfig, + get_moe_a2a_backend, get_moe_runner_backend, ) from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo @@ -52,8 +53,6 @@ _is_npu = is_npu() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip if _use_aiter: - from aiter import ActivationType - from aiter.fused_moe import fused_moe from aiter.ops.shuffle import shuffle_weight from aiter.tuned_gemm import tgemm @@ -232,9 +231,9 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): set_weight_attrs(w2_weight_bias, extra_weight_attrs) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - # Skip aiter weight shuffle when using non-auto MoE backend (e.g., triton, triton_kernels) - # because aiter CK kernels don't support all GEMM dimensions - _should_use_aiter_moe = _use_aiter and get_moe_runner_backend().is_auto() + _should_use_aiter_moe = _use_aiter and ( + get_moe_runner_backend().is_auto() or get_moe_runner_backend().is_aiter() + ) if _should_use_aiter_moe: copy_or_rebind_param( layer, "w13_weight", shuffle_weight(layer.w13_weight.data, (16, 16)) @@ -382,6 +381,18 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): backend = MoeRunnerBackend.TRITON self.runner = MoeRunner(backend, moe_runner_config) + # Separate runner so CK-shape errors fall back to self.runner on every call. + self._aiter_runner: Optional[MoeRunner] = None + if ( + _use_aiter + and ( + get_moe_runner_backend().is_auto() + or get_moe_runner_backend().is_aiter() + ) + and get_moe_a2a_backend().is_none() + ): + self._aiter_runner = MoeRunner(MoeRunnerBackend.AITER, moe_runner_config) + @property def load_up_proj_weight_first(self) -> bool: # FlashInfer CUTLASS kernel assumes [Up, Gate] Proj as W13 @@ -456,43 +467,19 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): ) return self.runner.run(dispatch_output, quant_info) else: - # Skip aiter fused_moe when using non-auto MoE backend (e.g., triton, triton_kernels) - # because aiter CK kernels don't support all GEMM dimensions - _should_use_aiter_moe = _use_aiter and get_moe_runner_backend().is_auto() - if _should_use_aiter_moe: - assert not moe_runner_config.no_combine, "unsupported" - topk_weights, topk_ids, _ = topk_output - if moe_runner_config.apply_router_weight_on_input: - assert ( - topk_weights.dim() == 2 - ), "`topk_weights` should be in shape (num_tokens, topk)" - _, topk = topk_weights.shape - assert ( - topk == 1 - ), "Only support topk=1 when `apply_router_weight_on_input` is True" - x = x * topk_weights.to(x.dtype) - topk_weights = torch.ones_like( - topk_weights, dtype=torch.float32 - ) # topk_weights must be FP32 (float32) + if self._aiter_runner is not None: + from sglang.srt.layers.moe.moe_runner.aiter import AiterMoeQuantInfo + try: - output = fused_moe( - x, - layer.w13_weight, - layer.w2_weight, - topk_weights, - topk_ids, - activation=( - ActivationType.Silu - if moe_runner_config.activation == "silu" - else ActivationType.Gelu - ), + quant_info = AiterMoeQuantInfo( + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, expert_mask=layer.dispatcher.expert_mask_gpu, ) - return StandardCombineInput(hidden_states=output) + return self._aiter_runner.run(dispatch_output, quant_info) except RuntimeError as e: # AITER CK fused_moe may not support all GEMM dimensions - # (e.g. Gemma4 MoE with 128 experts × 704 intermediate size). - # Fall through to Triton MoE runner below. + # (e.g. Gemma4 MoE with 128 experts x 704 intermediate size) logger.warning_once( f"AITER CK fused_moe failed ({e}), " "falling back to Triton MoE runner." diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 183607441..35be1c510 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -183,6 +183,7 @@ MOE_RUNNER_BACKEND_CHOICES = [ "flashinfer_mxfp4", "flashinfer_cutedsl", "cutlass", + "aiter", ] MOE_A2A_BACKEND_CHOICES = [