From aeca7d348c6bce7af5ae8eb198f07be495f49295 Mon Sep 17 00:00:00 2001 From: Tamir Baydasov <41994229+TamirBaydasov@users.noreply.github.com> Date: Tue, 17 Feb 2026 21:38:27 +0300 Subject: [PATCH] [3/N] Quantization Refactor: ModelSlim MoE schemes (#17993) Co-authored-by: ronnie_zheng --- .../quantization/modelslim/modelslim.py | 139 +++++++++++++-- .../modelslim/schemes/__init__.py | 9 +- .../modelslim/schemes/modelslim_scheme.py | 62 ++++++- .../modelslim/schemes/modelslim_w4a4_int4.py | 4 +- .../modelslim_w4a8_int8_moe.py} | 164 +----------------- .../modelslim/schemes/modelslim_w8a8_int8.py | 4 +- .../schemes/modelslim_w8a8_int8_moe.py | 139 +++++++++++++++ 7 files changed, 338 insertions(+), 183 deletions(-) rename python/sglang/srt/layers/quantization/modelslim/{modelslim_moe.py => schemes/modelslim_w4a8_int8_moe.py} (56%) create mode 100644 python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8_moe.py diff --git a/python/sglang/srt/layers/quantization/modelslim/modelslim.py b/python/sglang/srt/layers/quantization/modelslim/modelslim.py index 01a96672f..2106d6568 100644 --- a/python/sglang/srt/layers/quantization/modelslim/modelslim.py +++ b/python/sglang/srt/layers/quantization/modelslim/modelslim.py @@ -2,7 +2,7 @@ from __future__ import annotations import logging from types import MappingProxyType -from typing import Any, Dict, List, Mapping, Optional, Tuple, Union, cast +from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Tuple, Union, cast import torch @@ -10,19 +10,31 @@ from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import ( _NPULinearMethodBase, ) from sglang.srt.layers.quantization.base_config import ( + FusedMoEMethodBase, QuantizationConfig, - QuantizeMethodBase, ) from sglang.srt.layers.quantization.compressed_tensors.utils import should_ignore_layer -from sglang.srt.layers.quantization.modelslim.modelslim_moe import ModelSlimMoEMethod from sglang.srt.layers.quantization.modelslim.schemes import ( - ModelSlimScheme, ModelSlimW4A4Int4, + ModelSlimW4A8Int8MoE, ModelSlimW8A8Int8, + ModelSlimW8A8Int8MoE, ) from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod from sglang.srt.utils import apply_module_patch +if TYPE_CHECKING: + from sglang.srt.layers.moe import MoeRunnerConfig + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + StandardDispatchOutput, + ) + from sglang.srt.layers.quantization.base_config import QuantizeMethodBase + from sglang.srt.layers.quantization.modelslim.schemes import ( + ModelSlimLinearScheme, + ModelSlimMoEScheme, + ) + logger = logging.getLogger(__name__) @@ -150,17 +162,20 @@ class ModelSlimConfig(QuantizationConfig): if self.is_layer_skipped(prefix, packed_modules_mapping_subset): return UnquantizedLinearMethod() - scheme = self.get_scheme(layer=layer, layer_name=prefix_in_quant_config) + scheme = self.get_linear_scheme( + layer=layer, layer_name=prefix_in_quant_config + ) layer.scheme = scheme return ModelSlimLinearMethod(self) elif isinstance(layer, FusedMoE): - return ModelSlimMoEMethod.get_moe_method(self, layer, prefix) + layer.scheme = self.get_moe_scheme(layer, prefix) + return ModelSlimFusedMoEMethod(self) return None def _get_scheme_from_parts( self, layer_name: str, - ) -> ModelSlimScheme: + ) -> ModelSlimLinearScheme: quant_type = self.quant_description.get(layer_name + ".weight", "") if quant_type == "W8A8_DYNAMIC" or quant_type == "W8A8": @@ -173,9 +188,9 @@ class ModelSlimConfig(QuantizationConfig): ) raise NotImplementedError("No modelslim compatible scheme was found.") - def get_scheme( + def get_linear_scheme( self, layer: torch.nn.Module, layer_name: Optional[str] = None - ) -> Optional[ModelSlimScheme]: + ) -> Optional[ModelSlimLinearScheme]: """ get_scheme method adjusted for modelslim, taken from python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py @@ -188,6 +203,37 @@ class ModelSlimConfig(QuantizationConfig): logger.debug("Using scheme: %s for %s", scheme.__class__.__name__, layer_name) return scheme + def get_moe_scheme( + self, + layer: torch.nn.Module, + prefix: str, + ) -> Optional[ModelSlimMoEScheme]: + # TODO: @dsikka: refactor this to use schemes as other kernels + # are supported + check if the layer is being ignored. + + prefix_in_quant_config = prefix + ".0.down_proj.weight" + is_moe_w4a8_dynamic = ( + self.quant_description.get(prefix_in_quant_config, "STATIC") + == "W4A8_DYNAMIC" + ) + is_moe_w8a8_dynamic = ( + self.quant_description.get(prefix_in_quant_config, "STATIC") + == "W8A8_DYNAMIC" + ) + if is_moe_w4a8_dynamic: + logger.info_once("Using ModelSlimW4A8Int8MoE") + return ModelSlimW4A8Int8MoE(self) + elif is_moe_w8a8_dynamic: + logger.info_once("Using ModelSlimW8A8Int8MoE") + return ModelSlimW8A8Int8MoE(self) + else: + logger.warning( + f"Unsupported FusedMoe modelslim scheme: " + f"{self.quant_description.get(prefix_in_quant_config.strip())} " + f"in layer: {prefix}" + ) + return None + def is_layer_skipped( self, prefix: str, fused_mapping: Mapping[str, List[str]] = MappingProxyType({}) ): @@ -251,7 +297,7 @@ class ModelSlimLinearMethod(_NPULinearMethodBase): **extra_weight_attrs, ): """ - Use the ModelSlimScheme associated with each layer to create + Use the ModelSlimLinearScheme associated with the layer to create the necessary parameters for the layer. See LinearMethodBase for param details """ @@ -273,7 +319,7 @@ class ModelSlimLinearMethod(_NPULinearMethodBase): bias: Optional[torch.Tensor] = None, ): """ - Use the output of create_weights and the CompressedTensorsScheme + Use the output of create_weights and the ModelSlimLinearScheme associated with the layer to apply the forward pass with the layer input. See LinearMethodBase for param details @@ -283,3 +329,74 @@ class ModelSlimLinearMethod(_NPULinearMethodBase): if scheme is None: raise ValueError("A scheme must be defined for each layer") return scheme.apply_weights(layer, x, bias=bias) + + +class ModelSlimFusedMoEMethod(FusedMoEMethodBase): + + def __init__(self, quantization_config: ModelSlimConfig): + 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 ModelSlimMoEScheme 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 + ): + return layer.scheme.create_moe_runner(layer, moe_runner_config) + + def apply( + self, + layer: torch.nn.Module, + dispatch_output: StandardDispatchOutput, + ) -> CombineInput: + """ + Use the output of create_weights and the ModelSlimMoEScheme + associated with the layer to apply the forward pass with the + layer input. 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) + + def apply_without_routing_weights( + self, + layer, + hidden_states, + hidden_states_scale, + group_list_type, + group_list, + output_dtype, + ): + return layer.scheme.apply_without_routing_weights( + layer, + hidden_states, + hidden_states_scale, + group_list_type, + group_list, + output_dtype, + ) diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py b/python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py index 551b862a4..14005eb3d 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/__init__.py @@ -1,11 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 -from .modelslim_scheme import ModelSlimScheme +from .modelslim_scheme import ModelSlimLinearScheme, ModelSlimMoEScheme from .modelslim_w4a4_int4 import ModelSlimW4A4Int4 +from .modelslim_w4a8_int8_moe import ModelSlimW4A8Int8MoE from .modelslim_w8a8_int8 import ModelSlimW8A8Int8 +from .modelslim_w8a8_int8_moe import ModelSlimW8A8Int8MoE __all__ = [ - "ModelSlimScheme", + "ModelSlimLinearScheme", + "ModelSlimMoEScheme", "ModelSlimW8A8Int8", "ModelSlimW4A4Int4", + "ModelSlimW4A8Int8MoE", + "ModelSlimW8A8Int8MoE", ] diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_scheme.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_scheme.py index 1d09c384c..26f958e7b 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_scheme.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_scheme.py @@ -1,18 +1,24 @@ # Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/quantization/compressed_tensors # 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__ = ["ModelSlimScheme"] +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__ = ["ModelSlimLinearScheme", "ModelSlimMoEScheme"] -class ModelSlimScheme(ABC): +class ModelSlimLinearScheme(BaseLinearScheme): """ Abstract class used to describe the weight creation and forward pass - of different quantization schemes supported by CompressedTensors. + of different quantization schemes supported by ModelSlim. """ @abstractmethod @@ -23,6 +29,14 @@ class ModelSlimScheme(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] @@ -39,6 +53,21 @@ class ModelSlimScheme(ABC): """ raise NotImplementedError + +class ModelSlimMoEScheme(BaseMoEScheme): + """ + Abstract class used to describe the weight creation and forward pass + of different quantization schemes supported by ModelSlim. + """ + + @abstractmethod + def create_weights(self, *args, **kwargs): + """ + Weight creation for the particular scheme. Inputs to this function + + """ + raise NotImplementedError + @abstractmethod def process_weights_after_loading(self, layer: torch.nn.Module): """ @@ -46,3 +75,26 @@ class ModelSlimScheme(ABC): needs to occur. """ raise NotImplementedError + + def create_moe_runner( + self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig" + ): + raise NotImplementedError + + @abstractmethod + def apply_weights( + self, + layer, + 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/modelslim/schemes/modelslim_w4a4_int4.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4.py index 8e7f08277..1152aa804 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4.py @@ -9,11 +9,11 @@ from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import ( NPU_W4A4DynamicLinearMethod, ) from sglang.srt.layers.parameter import PerTensorScaleParameter -from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimScheme +from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimLinearScheme from sglang.srt.utils import set_weight_attrs -class ModelSlimW4A4Int4(ModelSlimScheme): +class ModelSlimW4A4Int4(ModelSlimLinearScheme): def __init__( self, diff --git a/python/sglang/srt/layers/quantization/modelslim/modelslim_moe.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a8_int8_moe.py similarity index 56% rename from python/sglang/srt/layers/quantization/modelslim/modelslim_moe.py rename to python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a8_int8_moe.py index c1674249f..4c3cd20f3 100644 --- a/python/sglang/srt/layers/quantization/modelslim/modelslim_moe.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a8_int8_moe.py @@ -1,5 +1,3 @@ -# Adapted from https://github.com/vllm-project/vllm/tree/v0.8.2/vllm/model_executor/layers/quantization/compressed_tensors -# SPDX-License-Identifier: Apache-2.0 from __future__ import annotations import logging @@ -9,9 +7,8 @@ import torch from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( NPUW4A8Int8DynamicMoEMethod, - NPUW8A8Int8DynamicMoEMethod, ) -from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase +from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimMoEScheme from sglang.srt.utils import set_weight_attrs if TYPE_CHECKING: @@ -20,56 +17,15 @@ if TYPE_CHECKING: CombineInput, StandardDispatchOutput, ) - from sglang.srt.layers.quantization.modelslim.modelslim import ModelSlimConfig logger = logging.getLogger(__name__) - __all__ = [ - "ModelSlimMoEMethod", "ModelSlimW4A8Int8MoE", - "ModelSlimW8A8Int8MoE", ] -class ModelSlimMoEMethod(FusedMoEMethodBase): - def __new__(cls, *args, **kwargs): - if cls is ModelSlimMoEMethod: - return super().__new__(cls) - return super().__new__(cls) - - @staticmethod - def get_moe_method( - quant_config: ModelSlimConfig, - layer: torch.nn.Module, - prefix: str, - ) -> "ModelSlimMoEMethod": - # TODO: @dsikka: refactor this to use schemes as other kernels - # are supported + check if the layer is being ignored. - - prefix_in_quant_config = prefix + ".0.down_proj.weight" - is_moe_w4a8_dynamic = ( - quant_config.quant_description.get(prefix_in_quant_config, "STATIC") - == "W4A8_DYNAMIC" - ) - is_moe_w8a8_dynamic = ( - quant_config.quant_description.get(prefix_in_quant_config, "STATIC") - == "W8A8_DYNAMIC" - ) - if is_moe_w4a8_dynamic: - logger.info_once("Using ModelSlimW4A8Int8MoE") - return ModelSlimW4A8Int8MoE(quant_config) - elif is_moe_w8a8_dynamic: - logger.info_once("Using ModelSlimW8A8Int8MoE") - return ModelSlimW8A8Int8MoE(quant_config) - else: - logger.warning(f"Unsupported FusedMoe modelslim scheme: \ - {quant_config.quant_description.get(prefix_in_quant_config.strip())} \ - in layer: {prefix}") - return None - - -class ModelSlimW4A8Int8MoE(ModelSlimMoEMethod): +class ModelSlimW4A8Int8MoE(ModelSlimMoEScheme): def __init__( self, @@ -234,7 +190,7 @@ class ModelSlimW4A8Int8MoE(ModelSlimMoEMethod): ): self.moe_runner_config = moe_runner_config - def apply( + def apply_weights( self, layer, dispatch_output: "StandardDispatchOutput", @@ -259,117 +215,3 @@ class ModelSlimW4A8Int8MoE(ModelSlimMoEMethod): group_list, output_dtype, ) - - -class ModelSlimW8A8Int8MoE(ModelSlimMoEMethod): - - def __init__( - self, - quant_config: Dict[str, Any], - prefix: str = None, - ): - self.quant_config = quant_config - self.kernel = NPUW8A8Int8DynamicMoEMethod() - - 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, - ) -> None: - from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported - - self.num_experts = num_experts - extra_weight_attrs.update( - {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} - ) - - # weight - w13_weight = torch.nn.Parameter( - torch.empty( - num_experts, - 2 * intermediate_size_per_partition, - hidden_size, - dtype=torch.int8, - ), - 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, - dtype=torch.int8, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight", w2_weight) - set_weight_attrs(w2_weight, extra_weight_attrs) - # scale - w13_weight_scale = torch.nn.Parameter( - torch.empty( - num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_scale", w13_weight_scale) - set_weight_attrs(w13_weight_scale, extra_weight_attrs) - w2_weight_scale = torch.nn.Parameter( - torch.empty(num_experts, hidden_size, 1, dtype=torch.float32), - requires_grad=False, - ) - layer.register_parameter("w2_weight_scale", w2_weight_scale) - set_weight_attrs(w2_weight_scale, extra_weight_attrs) - # offset - w13_weight_offset = torch.nn.Parameter( - torch.empty( - num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_offset", w13_weight_offset) - set_weight_attrs(w13_weight_offset, extra_weight_attrs) - w2_weight_offset = torch.nn.Parameter( - torch.empty(num_experts, hidden_size, 1, dtype=torch.float32), - requires_grad=False, - ) - layer.register_parameter("w2_weight_offset", w2_weight_offset) - set_weight_attrs(w2_weight_offset, extra_weight_attrs) - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - self.kernel.process_weights_after_loading(layer) - - def create_moe_runner( - self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig" - ): - self.moe_runner_config = moe_runner_config - - def apply( - self, - layer, - dispatch_output: "StandardDispatchOutput", - ) -> "CombineInput": - return self.kernel.apply(layer, dispatch_output) - - def apply_without_routing_weights( - self, - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, - ): - return self.kernel.apply_without_routing_weights( - layer, - hidden_states, - hidden_states_scale, - group_list_type, - group_list, - output_dtype, - ) diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8.py index 16c62d551..3770320ae 100644 --- a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8.py +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8.py @@ -14,10 +14,10 @@ from sglang.srt.layers.parameter import ( ModelWeightParameter, PerTensorScaleParameter, ) -from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimScheme +from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimLinearScheme -class ModelSlimW8A8Int8(ModelSlimScheme): +class ModelSlimW8A8Int8(ModelSlimLinearScheme): def __init__( self, diff --git a/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8_moe.py b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8_moe.py new file mode 100644 index 000000000..b226797f3 --- /dev/null +++ b/python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w8a8_int8_moe.py @@ -0,0 +1,139 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, Any, Dict + +import torch + +from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import ( + NPUW8A8Int8DynamicMoEMethod, +) +from sglang.srt.layers.quantization.modelslim.schemes import ModelSlimMoEScheme +from sglang.srt.utils import set_weight_attrs + +if TYPE_CHECKING: + from sglang.srt.layers.moe import MoeRunnerConfig + from sglang.srt.layers.moe.token_dispatcher import ( + CombineInput, + StandardDispatchOutput, + ) + +logger = logging.getLogger(__name__) + +__all__ = [ + "ModelSlimW8A8Int8MoE", +] + + +class ModelSlimW8A8Int8MoE(ModelSlimMoEScheme): + + def __init__( + self, + quant_config: Dict[str, Any], + prefix: str = None, + ): + self.quant_config = quant_config + self.kernel = NPUW8A8Int8DynamicMoEMethod() + + 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, + ) -> None: + from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + + self.num_experts = num_experts + extra_weight_attrs.update( + {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} + ) + + # weight + w13_weight = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size, + dtype=torch.int8, + ), + 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, + dtype=torch.int8, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight", w2_weight) + set_weight_attrs(w2_weight, extra_weight_attrs) + # scale + w13_weight_scale = torch.nn.Parameter( + torch.empty( + num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + set_weight_attrs(w13_weight_scale, extra_weight_attrs) + w2_weight_scale = torch.nn.Parameter( + torch.empty(num_experts, hidden_size, 1, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + set_weight_attrs(w2_weight_scale, extra_weight_attrs) + # offset + w13_weight_offset = torch.nn.Parameter( + torch.empty( + num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32 + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_offset", w13_weight_offset) + set_weight_attrs(w13_weight_offset, extra_weight_attrs) + w2_weight_offset = torch.nn.Parameter( + torch.empty(num_experts, hidden_size, 1, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w2_weight_offset", w2_weight_offset) + set_weight_attrs(w2_weight_offset, extra_weight_attrs) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + self.kernel.process_weights_after_loading(layer) + + 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, + dispatch_output: "StandardDispatchOutput", + ) -> "CombineInput": + return self.kernel.apply(layer, dispatch_output) + + def apply_without_routing_weights( + self, + layer, + hidden_states, + hidden_states_scale, + group_list_type, + group_list, + output_dtype, + ): + return self.kernel.apply_without_routing_weights( + layer, + hidden_states, + hidden_states_scale, + group_list_type, + group_list, + output_dtype, + )