chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[4/N] Quantization Refactor: Quark MoE schemes (#18252)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Peng Zhang <aniz1905@gmail.com>
Co-authored-by: ronnie_zheng <zl19940307@163.com>
This commit is contained in:
Tamir Baydasov
2026-02-18 19:44:30 +03:00
committed by GitHub
parent 9c5aae4df5
commit 150ed881be
8 changed files with 396 additions and 243 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/layers/moe/ep_moe/layer.py
View File

@@ -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(

106
python/sglang/srt/layers/quantization/quark/quark.py
View File

@@ -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.

13
python/sglang/srt/layers/quantization/quark/schemes/__init__.py
View File

@@ -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",
]

69
python/sglang/srt/layers/quantization/quark/schemes/quark_scheme.py
View File

@@ -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

4
python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4.py
View File

@@ -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]

213
python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py Normal file
View File

@@ -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)

4
python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8.py
View File

@@ -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]]

226
python/sglang/srt/layers/quantization/quark/quark_moe.py → python/sglang/srt/layers/quantization/quark/schemes/quark_w8a8_fp8_moe.py
View File

@@ -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,