minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[MoE] Add Aiter MoE runner backend and purge aiter.fused_moe from quant methods (#23597)

Browse Source 
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-04-30 19:50:52 -07:00
committed by GitHub
co-authored by Claude Opus 4.7
parent 0acc569edd
commit 108bfd8b6a
10 changed files with 316 additions and 251 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/moe_runner/aiter.py
+94
View File
@@ -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)
python/sglang/srt/layers/moe/moe_runner/runner.py
+5
View File
@@ -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
python/sglang/srt/layers/moe/utils.py
+4
View File
@@ -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):
python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8_moe.py
+33 -35
View File
@@ -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 (
python/sglang/srt/layers/quantization/fp8.py
+46 -69
View File
@@ -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):
python/sglang/srt/layers/quantization/mxfp4.py
+59 -59
View File
@@ -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)
python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py
+25 -29
View File
@@ -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)
python/sglang/srt/layers/quantization/quark_int4fp8_moe.py
+25 -22
View File
@@ -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)
python/sglang/srt/layers/quantization/unquant.py
+24 -37
View File
@@ -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."
python/sglang/srt/server_args.py
+1
View File
@@ -183,6 +183,7 @@ MOE_RUNNER_BACKEND_CHOICES = [
"flashinfer_mxfp4",
"flashinfer_cutedsl",
"cutlass",
"aiter",
]
MOE_A2A_BACKEND_CHOICES = [