From 150ed881be2cf9a9e64b636226b3c6189d341a7d Mon Sep 17 00:00:00 2001 From: Tamir Baydasov <41994229+TamirBaydasov@users.noreply.github.com> Date: Wed, 18 Feb 2026 19:44:30 +0300 Subject: [PATCH] [4/N] Quantization Refactor: Quark MoE schemes (#18252) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Peng Zhang Co-authored-by: ronnie_zheng --- python/sglang/srt/layers/moe/ep_moe/layer.py | 4 +- .../srt/layers/quantization/quark/quark.py | 106 +++++++- .../quantization/quark/schemes/__init__.py | 13 +- .../quark/schemes/quark_scheme.py | 69 +++++- .../quark/schemes/quark_w4a4_mxfp4.py | 4 +- .../quark/schemes/quark_w4a4_mxfp4_moe.py | 213 +++++++++++++++++ .../quark/schemes/quark_w8a8_fp8.py | 4 +- .../quark_w8a8_fp8_moe.py} | 226 +----------------- 8 files changed, 396 insertions(+), 243 deletions(-) create mode 100644 python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py rename python/sglang/srt/layers/quantization/quark/{quark_moe.py => schemes/quark_w8a8_fp8_moe.py} (60%) diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index 2280f4535..ebcc696ec 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -32,7 +32,7 @@ from sglang.srt.layers.quantization.compressed_tensors.schemes import ( ) from sglang.srt.layers.quantization.fp8 import Fp8Config, Fp8MoEMethod from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz -from sglang.srt.layers.quantization.quark.quark_moe import QuarkW4A4MXFp4MoEMethod +from sglang.srt.layers.quantization.quark.schemes import QuarkW4A4MXFp4MoE from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config, W4AFp8MoEMethod from sglang.srt.utils import get_bool_env_var, is_hip, is_npu @@ -600,7 +600,7 @@ class MoriEPMoE(DeepEPMoE): output_dtype = hidden_states.dtype scale = None is_fp8_quant = isinstance(self.quant_method, Fp8MoEMethod) - is_quark_w4a4 = isinstance(self.quant_method, QuarkW4A4MXFp4MoEMethod) + is_quark_w4a4 = isinstance(self.scheme, QuarkW4A4MXFp4MoE) # dispatch dispatch_output = self.dispatcher.dispatch( diff --git a/python/sglang/srt/layers/quantization/quark/quark.py b/python/sglang/srt/layers/quantization/quark/quark.py index 6abde9b4f..542b0605c 100644 --- a/python/sglang/srt/layers/quantization/quark/quark.py +++ b/python/sglang/srt/layers/quantization/quark/quark.py @@ -2,29 +2,36 @@ import fnmatch import logging -from typing import Any, List, Optional, cast +from typing import TYPE_CHECKING, Any, List, Optional, cast import torch from sglang.srt.layers.linear import LinearBase +from sglang.srt.layers.moe import MoeRunnerConfig from sglang.srt.layers.quantization.base_config import ( # noqa: E501 + FusedMoEMethodBase, LinearMethodBase, QuantizationConfig, QuantizeMethodBase, ) from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod -from sglang.srt.layers.quantization.quark.quark_moe import QuarkMoEMethod from sglang.srt.layers.quantization.quark.schemes import ( - QuarkScheme, + QuarkLinearScheme, + QuarkMoEScheme, QuarkW4A4MXFP4, + QuarkW4A4MXFp4MoE, QuarkW8A8Fp8, + QuarkW8A8FP8MoE, ) from sglang.srt.layers.quantization.quark.utils import deep_compare, should_ignore_layer from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.utils import get_device_capability -__all__ = ["QuarkLinearMethod"] +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput + +__all__ = ["QuarkLinearMethod", "QuarkFusedMoEMethod"] logger = logging.getLogger(__name__) @@ -77,7 +84,7 @@ class QuarkConfig(QuantizationConfig): return None if isinstance(layer, LinearBase): - scheme = self.get_scheme(layer=layer, layer_name=prefix) + scheme = self.get_linear_scheme(layer=layer, layer_name=prefix) layer.scheme = scheme return QuarkLinearMethod(self) @@ -87,7 +94,8 @@ class QuarkConfig(QuantizationConfig): from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE if isinstance(layer, FusedMoE): - return QuarkMoEMethod.get_moe_method(self, module=layer, layer_name=prefix) + layer.scheme = self.get_moe_scheme(layer, prefix) + return QuarkFusedMoEMethod(self) return None @@ -308,7 +316,7 @@ class QuarkConfig(QuantizationConfig): ) return global_quant_config - def _get_scheme_from_config(self, config: dict[str, Any]) -> "QuarkScheme": + def _get_scheme_from_config(self, config: dict[str, Any]) -> "QuarkLinearScheme": if config.get("output_tensors") or config.get("bias"): raise NotImplementedError( "Currently, Quark models with output_tensors " @@ -332,7 +340,9 @@ class QuarkConfig(QuantizationConfig): f"Input config: {input_config}" ) - def get_scheme(self, layer: torch.nn.Module, layer_name: str) -> "QuarkScheme": + def get_linear_scheme( + self, layer: torch.nn.Module, layer_name: str + ) -> "QuarkLinearScheme": layer_quant_config = self._find_matched_config(layer_name, layer) @@ -345,6 +355,29 @@ class QuarkConfig(QuantizationConfig): return scheme + def get_moe_scheme( + self, + module: torch.nn.Module, + layer_name: str, + ) -> "QuarkMoEScheme": + layer_quant_config = self._find_matched_config(layer_name, module) + + if layer_quant_config.get("output_tensors") or layer_quant_config.get("bias"): + raise NotImplementedError( + "Currently, Quark models with " + "output_tensors and bias " + "quantized are not supported" + ) + weight_config = layer_quant_config.get("weight") + input_config = layer_quant_config.get("input_tensors") + + if self._is_mx_fp4(weight_config, input_config): + return QuarkW4A4MXFp4MoE(weight_config, input_config) + elif self._is_fp8_w8a8(weight_config, input_config): + return QuarkW8A8FP8MoE(weight_config, input_config) + else: + raise RuntimeError("Unsupported FusedMoe scheme") + def get_scaled_act_names(self) -> List[str]: return [] @@ -368,7 +401,7 @@ class QuarkLinearMethod(LinearMethodBase): **extra_weight_attrs, ): """ - Use the CompressedTensorsScheme associated with each layer to create + Use the QuarkLinearScheme associated with the layer to create the necessary parameters for the layer. See LinearMethodBase for param details """ @@ -390,7 +423,7 @@ class QuarkLinearMethod(LinearMethodBase): bias: Optional[torch.Tensor] = None, ): """ - Use the output of create_weights and the CompressedTensorsScheme + Use the output of create_weights and the QuarkLinearScheme associated with the layer to apply the forward pass with the layer input. See LinearMethodBase for param details @@ -401,6 +434,59 @@ class QuarkLinearMethod(LinearMethodBase): return scheme.apply_weights(layer, x, bias=bias) +class QuarkFusedMoEMethod(FusedMoEMethodBase): + + def __init__(self, quantization_config: QuarkConfig): + self.quantization_config = quantization_config + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.scheme.process_weights_after_loading(layer) + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + """ + Use the QuarkMoEScheme associated with the layer to create + the necessary parameters for the layer. See FusedMoEMethodBase for param + details + """ + layer.scheme.create_weights( + layer=layer, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size_per_partition=intermediate_size_per_partition, + params_dtype=params_dtype, + **extra_weight_attrs, + ) + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + layer.scheme.create_moe_runner(layer, moe_runner_config) + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: "StandardDispatchOutput", + ): + """ + Use the output of create_weights and the QuarkMoEScheme + associated with the layer to apply the forward pass with the + fused MoE layer. See FusedMoEMethodBase for param details + + """ + scheme = layer.scheme + if scheme is None: + raise ValueError("A scheme must be defined for each layer") + return scheme.apply_weights(layer, dispatch_output) + + class QuarkKVCacheMethod(BaseKVCacheMethod): """ Supports loading kv-cache scaling factors from quark checkpoints. diff --git a/python/sglang/srt/layers/quantization/quark/schemes/__init__.py b/python/sglang/srt/layers/quantization/quark/schemes/__init__.py index 91ceb6a1e..0ce9e0c29 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/__init__.py +++ b/python/sglang/srt/layers/quantization/quark/schemes/__init__.py @@ -1,7 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 -from .quark_scheme import QuarkScheme +from .quark_scheme import QuarkLinearScheme, QuarkMoEScheme from .quark_w4a4_mxfp4 import QuarkW4A4MXFP4 +from .quark_w4a4_mxfp4_moe import QuarkW4A4MXFp4MoE from .quark_w8a8_fp8 import QuarkW8A8Fp8 +from .quark_w8a8_fp8_moe import QuarkW8A8FP8MoE -__all__ = ["QuarkScheme", "QuarkW4A4MXFP4", "QuarkW8A8Fp8"] +__all__ = [ + "QuarkLinearScheme", + "QuarkMoEScheme", + "QuarkW4A4MXFP4", + "QuarkW8A8Fp8", + "QuarkW4A4MXFp4MoE", + "QuarkW8A8FP8MoE", +] diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_scheme.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_scheme.py index aab5c9c1e..556de3152 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/quark_scheme.py +++ b/python/sglang/srt/layers/quantization/quark/schemes/quark_scheme.py @@ -1,14 +1,20 @@ # SPDX-License-Identifier: Apache-2.0 -from abc import ABC, abstractmethod -from typing import Optional +from abc import abstractmethod +from typing import TYPE_CHECKING, Optional import torch -__all__ = ["QuarkScheme"] +from sglang.srt.layers.moe import MoeRunnerConfig +from sglang.srt.layers.quantization.base_scheme import BaseLinearScheme, BaseMoEScheme + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput + +__all__ = ["QuarkLinearScheme", "QuarkMoEScheme"] -class QuarkScheme(ABC): +class QuarkLinearScheme(BaseLinearScheme): """ Abstract class used to describe the weight creation and forward pass of different quantization schemes supported by Quark. @@ -30,6 +36,14 @@ class QuarkScheme(ABC): """ raise NotImplementedError + @abstractmethod + def process_weights_after_loading(self, layer: torch.nn.Module): + """ + Called after weight loading is complete for any cleanup that + needs to occur. + """ + raise NotImplementedError + @abstractmethod def apply_weights( self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor] @@ -46,6 +60,35 @@ class QuarkScheme(ABC): """ raise NotImplementedError + +class QuarkMoEScheme(BaseMoEScheme): + """ + Abstract class used to describe the weight creation and forward pass + of different quantization schemes supported by Quark. + """ + + @classmethod + @abstractmethod + def get_min_capability(cls) -> int: + """ + Get minimum device capability. + """ + raise NotImplementedError + + @abstractmethod + def create_weights(self, *args, **kwargs): + """ + Weight creation for the particular scheme. Inputs to this function + + """ + raise NotImplementedError + + @abstractmethod + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + raise NotImplementedError + @abstractmethod def process_weights_after_loading(self, layer: torch.nn.Module): """ @@ -53,3 +96,21 @@ class QuarkScheme(ABC): needs to occur. """ raise NotImplementedError + + @abstractmethod + def apply_weights( + self, + layer: torch.nn.Module, + dispatch_output: "StandardDispatchOutput", + ): + """ + Run the forward pass for the particular scheme. This is where + scheme-specific dequant/quant steps/kernels should be applied. + + :param layer: torch.nn.Module with the registered weights and + other parameters relevant to the particular scheme. + :param x: input to the layer + :param bias: bias parameter + + """ + raise NotImplementedError diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py index ccb34f749..984e15696 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py +++ b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py @@ -5,7 +5,7 @@ from typing import Any, Callable, Optional import torch from sglang.srt.layers.parameter import GroupQuantScaleParameter, PackedvLLMParameter -from sglang.srt.layers.quantization.quark.schemes import QuarkScheme +from sglang.srt.layers.quantization.quark.schemes import QuarkLinearScheme from sglang.srt.utils import is_hip _is_hip = is_hip() @@ -20,7 +20,7 @@ __all__ = ["QuarkW4A4MXFP4"] OCP_MX_BLOCK_SIZE = 32 -class QuarkW4A4MXFP4(QuarkScheme): +class QuarkW4A4MXFP4(QuarkLinearScheme): def __init__( self, weight_quant_spec: dict[str, Any], input_quant_spec: dict[str, Any] 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 new file mode 100644 index 000000000..df6650be8 --- /dev/null +++ b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py @@ -0,0 +1,213 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any + +import torch + +from sglang.srt.layers.moe import MoeRunnerConfig +from sglang.srt.layers.quantization.quark.schemes import QuarkMoEScheme +from sglang.srt.utils import ( + get_bool_env_var, + is_gfx95_supported, + is_hip, + set_weight_attrs, +) + +if TYPE_CHECKING: + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + StandardDispatchOutput, + ) + +logger = logging.getLogger(__name__) + +_is_shuffle_moe_mxfp4 = is_gfx95_supported() + +__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 + +OCP_MX_BLOCK_SIZE = 32 + + +class QuarkW4A4MXFp4MoE(QuarkMoEScheme): + + def __init__(self, weight_config: dict[str, Any], input_config: dict[str, Any]): + self.weight_quant = weight_config + self.input_quant = input_config + + weight_qscheme = self.weight_quant.get("qscheme") + input_qscheme = self.input_quant.get("qscheme") + if not (weight_qscheme == "per_group" and input_qscheme == "per_group"): + raise ValueError( + "For MX(FP4) Fused MoE layers, only per-group scales " + "for weights and activations are supported. Found " + f"{weight_qscheme}, {input_qscheme}" + ) # noqa E501 + + self.static_input_scales = not self.input_quant.get("is_dynamic") + self.with_bias = False + + @classmethod + def get_min_capability(cls) -> int: + return 70 + + def create_weights( + self, + layer: torch.nn.Module, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + # Add the quantization method used (per tensor/grouped/channel) + # to ensure the weight scales are loaded in properly + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value} + ) + + params_dtype = torch.uint8 + + # WEIGHTS + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // 2, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + + set_weight_attrs(w13_weight, extra_weight_attrs) + + w2_weight = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition // 2, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + + set_weight_attrs(w2_weight, extra_weight_attrs) + + # WEIGHT_SCALES + w13_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // OCP_MX_BLOCK_SIZE, + dtype=params_dtype, + ), + requires_grad=False, + ) + w2_weight_scale = torch.nn.Parameter( + torch.ones( + num_experts, + hidden_size, + intermediate_size_per_partition // OCP_MX_BLOCK_SIZE, + dtype=params_dtype, + ), + requires_grad=False, + ) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + + layer.register_parameter("w13_weight_scale", w13_weight_scale) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + float_dtype = torch.get_default_dtype() + + # Pre-shuffle weight scales + s0, s1, _ = layer.w13_weight_scale.shape + w13_weight_scale = layer.w13_weight_scale.view(s0 * s1, -1) + w13_weight_scale = e8m0_shuffle(w13_weight_scale) + # layer.w13_weight_scale = torch.nn.Parameter(w13_weight_scale, requires_grad=False) + layer.w13_weight_scale.data = w13_weight_scale.view(s0, s1, -1) + + s0, s1, _ = layer.w2_weight_scale.shape + w2_weight_scale = layer.w2_weight_scale.view(s0 * s1, -1) + w2_weight_scale = e8m0_shuffle(w2_weight_scale) + # layer.w2_weight_scale = torch.nn.Parameter(w2_weight_scale, requires_grad=False) + layer.w2_weight_scale.data = w2_weight_scale.view(s0, s1, -1) + + # Pre-shuffle weight + if _is_shuffle_moe_mxfp4: + layer.w13_weight.data = shuffle_weight( + layer.w13_weight.contiguous(), (16, 16) + ) + layer.w2_weight.data = shuffle_weight( + layer.w2_weight.contiguous(), (16, 16) + ) + layer.w13_weight.is_shuffled = True + layer.w2_weight.is_shuffled = True + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig + ): + self.moe_runner_config = moe_runner_config + + 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 + + if hasattr(torch, "float4_e2m1fn_x2"): + w13_weight = layer.w13_weight.view(torch.float4_e2m1fn_x2) + w2_weight = layer.w2_weight.view(torch.float4_e2m1fn_x2) + else: + w13_weight = layer.w13_weight + w2_weight = layer.w2_weight + + if hasattr(layer.w13_weight, "is_shuffled"): + 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, + w2_scale=layer.w2_weight_scale, + activation=( + ActivationType.Silu + if moe_runner_config.activation == "silu" + else ActivationType.Gelu + ), + doweight_stage1=False, + expert_mask=layer.expert_mask_gpu, + ) + return StandardCombineInput(hidden_states=output) diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8.py index 53001842a..c8b1dde59 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8.py +++ b/python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8.py @@ -16,7 +16,7 @@ from sglang.srt.layers.quantization.fp8_utils import ( cutlass_fp8_supported, normalize_e4m3fn_to_e4m3fnuz, ) -from sglang.srt.layers.quantization.quark.schemes import QuarkScheme +from sglang.srt.layers.quantization.quark.schemes import QuarkLinearScheme from sglang.srt.layers.quantization.utils import requantize_with_max_scale from sglang.srt.utils import get_bool_env_var, is_hip, set_weight_attrs @@ -29,7 +29,7 @@ if _use_aiter: from aiter.ops.shuffle import shuffle_weight -class QuarkW8A8Fp8(QuarkScheme): +class QuarkW8A8Fp8(QuarkLinearScheme): def __init__( self, weight_config: dict[str, Any], input_config: Optional[dict[str, Any]] diff --git a/python/sglang/srt/layers/quantization/quark/quark_moe.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8_moe.py similarity index 60% rename from python/sglang/srt/layers/quantization/quark/quark_moe.py rename to python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8_moe.py index 186f6c5d0..d55c9816b 100644 --- a/python/sglang/srt/layers/quantization/quark/quark_moe.py +++ b/python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8_moe.py @@ -9,248 +9,32 @@ import torch from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo -from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz, scaled_fp8_quant from sglang.srt.layers.quantization.fp8_utils import normalize_e4m3fn_to_e4m3fnuz +from sglang.srt.layers.quantization.quark.schemes import QuarkMoEScheme from sglang.srt.layers.quantization.utils import all_close_1d, per_tensor_dequantize -from sglang.srt.utils import ( - get_bool_env_var, - is_gfx95_supported, - is_hip, - set_weight_attrs, -) +from sglang.srt.utils import get_bool_env_var, is_hip, set_weight_attrs if TYPE_CHECKING: from sglang.srt.layers.moe.token_dispatcher import ( CombineInput, StandardDispatchOutput, ) - from sglang.srt.layers.quantization.quark.quark import QuarkConfig logger = logging.getLogger(__name__) -_is_shuffle_moe_mxfp4 = is_gfx95_supported() - -__all__ = ["QuarkMoEMethod", "QuarkW4A4MXFp4MoEMethod"] +__all__ = ["QuarkW8A8FP8MoE"] _is_fp8_fnuz = is_fp8_fnuz() _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 from sglang.srt.layers.moe.rocm_moe_utils import rocm_fused_experts_tkw1 -OCP_MX_BLOCK_SIZE = 32 -if TYPE_CHECKING: - from sglang.srt.layers.quantization import QuarkConfig - - -class QuarkMoEMethod(FusedMoEMethodBase): - - def __init__(self, quant_config: QuarkConfig): - self.quant_config = quant_config - - @staticmethod - def get_moe_method( - quant_config: QuarkConfig, # type: ignore # noqa E501 # noqa F821 - module: torch.nn.Module, - layer_name: str, - ) -> "QuarkMoEMethod": - layer_quant_config = quant_config._find_matched_config(layer_name, module) - - if layer_quant_config.get("output_tensors") or layer_quant_config.get("bias"): - raise NotImplementedError( - "Currently, Quark models with " - "output_tensors and bias " - "quantized are not supported" - ) - weight_config = layer_quant_config.get("weight") - input_config = layer_quant_config.get("input_tensors") - - if quant_config._is_mx_fp4(weight_config, input_config): - return QuarkW4A4MXFp4MoEMethod(weight_config, input_config) - elif quant_config._is_fp8_w8a8(weight_config, input_config): - return QuarkW8A8FP8MoEMethod(weight_config, input_config) - else: - raise RuntimeError("Unsupported FusedMoe scheme") - - -class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod): - - def __init__(self, weight_config: dict[str, Any], input_config: dict[str, Any]): - self.weight_quant = weight_config - self.input_quant = input_config - - weight_qscheme = self.weight_quant.get("qscheme") - input_qscheme = self.input_quant.get("qscheme") - if not (weight_qscheme == "per_group" and input_qscheme == "per_group"): - raise ValueError( - "For MX(FP4) Fused MoE layers, only per-group scales " - "for weights and activations are supported. Found " - f"{weight_qscheme}, {input_qscheme}" - ) # noqa E501 - - self.static_input_scales = not self.input_quant.get("is_dynamic") - self.with_bias = False - - def create_weights( - self, - layer: torch.nn.Module, - num_experts: int, - hidden_size: int, - intermediate_size_per_partition: int, - params_dtype: torch.dtype, - **extra_weight_attrs, - ): - - from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported - - # Add the quantization method used (per tensor/grouped/channel) - # to ensure the weight scales are loaded in properly - extra_weight_attrs.update( - {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value} - ) - - params_dtype = torch.uint8 - - # WEIGHTS - w13_weight = torch.nn.Parameter( - torch.empty( - num_experts, - 2 * intermediate_size_per_partition, - hidden_size // 2, - dtype=params_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight", w13_weight) - - set_weight_attrs(w13_weight, extra_weight_attrs) - - w2_weight = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - intermediate_size_per_partition // 2, - dtype=params_dtype, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight", w2_weight) - - set_weight_attrs(w2_weight, extra_weight_attrs) - - # WEIGHT_SCALES - w13_weight_scale = torch.nn.Parameter( - torch.ones( - num_experts, - 2 * intermediate_size_per_partition, - hidden_size // OCP_MX_BLOCK_SIZE, - dtype=params_dtype, - ), - requires_grad=False, - ) - w2_weight_scale = torch.nn.Parameter( - torch.ones( - num_experts, - hidden_size, - intermediate_size_per_partition // OCP_MX_BLOCK_SIZE, - dtype=params_dtype, - ), - requires_grad=False, - ) - set_weight_attrs(w2_weight_scale, extra_weight_attrs) - set_weight_attrs(w13_weight_scale, extra_weight_attrs) - - layer.register_parameter("w13_weight_scale", w13_weight_scale) - layer.register_parameter("w2_weight_scale", w2_weight_scale) - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - float_dtype = torch.get_default_dtype() - - # Pre-shuffle weight scales - s0, s1, _ = layer.w13_weight_scale.shape - w13_weight_scale = layer.w13_weight_scale.view(s0 * s1, -1) - w13_weight_scale = e8m0_shuffle(w13_weight_scale) - # layer.w13_weight_scale = torch.nn.Parameter(w13_weight_scale, requires_grad=False) - layer.w13_weight_scale.data = w13_weight_scale.view(s0, s1, -1) - - s0, s1, _ = layer.w2_weight_scale.shape - w2_weight_scale = layer.w2_weight_scale.view(s0 * s1, -1) - w2_weight_scale = e8m0_shuffle(w2_weight_scale) - # layer.w2_weight_scale = torch.nn.Parameter(w2_weight_scale, requires_grad=False) - layer.w2_weight_scale.data = w2_weight_scale.view(s0, s1, -1) - - # Pre-shuffle weight - if _is_shuffle_moe_mxfp4: - layer.w13_weight.data = shuffle_weight( - layer.w13_weight.contiguous(), (16, 16) - ) - layer.w2_weight.data = shuffle_weight( - layer.w2_weight.contiguous(), (16, 16) - ) - layer.w13_weight.is_shuffled = True - layer.w2_weight.is_shuffled = True - - def create_moe_runner( - self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig - ): - self.moe_runner_config = moe_runner_config - - 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 - 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 - - if hasattr(torch, "float4_e2m1fn_x2"): - w13_weight = layer.w13_weight.view(torch.float4_e2m1fn_x2) - w2_weight = layer.w2_weight.view(torch.float4_e2m1fn_x2) - else: - w13_weight = layer.w13_weight - w2_weight = layer.w2_weight - - if hasattr(layer.w13_weight, "is_shuffled"): - 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, - w2_scale=layer.w2_weight_scale, - activation=( - ActivationType.Silu - if moe_runner_config.activation == "silu" - else ActivationType.Gelu - ), - doweight_stage1=False, - expert_mask=layer.expert_mask_gpu, - ) - return StandardCombineInput(hidden_states=output) - - -class QuarkW8A8FP8MoEMethod(QuarkMoEMethod): +class QuarkW8A8FP8MoE(QuarkMoEScheme): def __init__(self, weight_config: dict[str, Any], input_config: dict[str, Any]): self.is_static_input_scheme: bool = False @@ -479,7 +263,7 @@ class QuarkW8A8FP8MoEMethod(QuarkMoEMethod): self.moe_runner_config = moe_runner_config self.runner = MoeRunner(MoeRunnerBackend.TRITON, moe_runner_config) - def apply( + def apply_weights( self, layer: torch.nn.Module, dispatch_output: StandardDispatchOutput,