chore: import upstream snapshot with attribution

python/sglang/srt/layers/quantization/awq/schemes/__init__.py

@@ -0,0 +1,19 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from .awq_cpu import AWQIntelAMXLinearScheme, AWQIntelAMXMoEScheme
+from .awq_linear import AWQAscendLinearScheme, AWQLinearScheme
+from .awq_marlin import AWQMarlinLinearScheme
+from .awq_moe import AWQAscendMoEScheme, AWQMoEScheme
+from .awq_scheme import AWQLinearSchemeBase, AWQMoESchemeBase
+
+__all__ = [
+    "AWQLinearSchemeBase",
+    "AWQMoESchemeBase",
+    "AWQLinearScheme",
+    "AWQAscendLinearScheme",
+    "AWQIntelAMXLinearScheme",
+    "AWQMarlinLinearScheme",
+    "AWQMoEScheme",
+    "AWQAscendMoEScheme",
+    "AWQIntelAMXMoEScheme",
+]

python/sglang/srt/layers/quantization/awq/schemes/awq_cpu.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.hardware_backend.cpu.quantization.awq_kernels import (
+    AWQIntelAMXLinearKernel,
+    AWQIntelAMXMoEKernel,
+)
+from sglang.srt.layers.moe import MoeRunnerConfig
+
+from .awq_linear import AWQLinearScheme
+from .awq_moe import AWQMoEScheme
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.quantization.awq.awq import AWQConfig
+
+__all__ = ["AWQIntelAMXLinearScheme", "AWQIntelAMXMoEScheme"]
+
+
+class AWQIntelAMXLinearScheme(AWQLinearScheme):
+    """Linear scheme for AWQ on Intel CPU with AMX."""
+
+    def _init_kernel(self, quant_config: AWQConfig):
+        return AWQIntelAMXLinearKernel(quant_config)
+
+
+class AWQIntelAMXMoEScheme(AWQMoEScheme):
+    """MoE scheme for AWQ on Intel CPU with AMX."""
+
+    def _init_kernel(self, quant_config: AWQConfig):
+        return AWQIntelAMXMoEKernel(quant_config)
+
+    def create_moe_runner(
+        self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
+    ):
+        self.moe_runner_config = moe_runner_config
+        self.kernel.create_moe_runner(layer, moe_runner_config)

python/sglang/srt/layers/quantization/awq/schemes/awq_linear.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Optional
+
+import torch
+
+from sglang.srt.layers.parameter import GroupQuantScaleParameter, PackedvLLMParameter
+
+from .awq_scheme import AWQLinearSchemeBase
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.quantization.awq.awq import AWQConfig
+
+__all__ = ["AWQLinearScheme", "AWQAscendLinearScheme"]
+
+
+class AWQLinearScheme(AWQLinearSchemeBase):
+    def __init__(self, quant_config: AWQConfig):
+        self.quant_config = quant_config
+        self.kernel = self._init_kernel(quant_config)
+
+    def _init_kernel(self, quant_config: AWQConfig):
+        from sglang.srt.hardware_backend.gpu.quantization.awq_kernels import (
+            AWQLinearKernel,
+        )
+
+        return AWQLinearKernel(quant_config)
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: List[int],
+        params_dtype: torch.dtype,
+        weight_loader,
+        **kwargs,
+    ):
+        if input_size_per_partition % self.quant_config.group_size != 0:
+            raise ValueError(
+                "The input size is not aligned with the quantized "
+                "weight shape. This can be caused by too large "
+                "tensor parallel size."
+            )
+
+        output_size_per_partition = sum(output_partition_sizes)
+        if output_size_per_partition % self.quant_config.pack_factor != 0:
+            raise ValueError(
+                "The output size is not aligned with the quantized "
+                "weight shape. This can be caused by too large "
+                "tensor parallel size."
+            )
+
+        qweight = PackedvLLMParameter(
+            data=torch.empty(
+                input_size_per_partition,
+                output_size_per_partition // self.quant_config.pack_factor,
+                dtype=torch.int32,
+            ),
+            input_dim=0,
+            output_dim=1,
+            packed_dim=1,
+            packed_factor=self.quant_config.pack_factor,
+            weight_loader=weight_loader,
+        )
+
+        qzeros = PackedvLLMParameter(
+            data=torch.empty(
+                input_size_per_partition // self.quant_config.group_size,
+                output_size_per_partition // self.quant_config.pack_factor,
+                dtype=torch.int32,
+            ),
+            input_dim=0,
+            output_dim=1,
+            packed_dim=1,
+            packed_factor=self.quant_config.pack_factor,
+            weight_loader=weight_loader,
+        )
+
+        scales = GroupQuantScaleParameter(
+            data=torch.empty(
+                input_size_per_partition // self.quant_config.group_size,
+                output_size_per_partition,
+                dtype=params_dtype,
+            ),
+            input_dim=0,
+            output_dim=1,
+            weight_loader=weight_loader,
+        )
+
+        layer.register_parameter("qweight", qweight)
+        layer.register_parameter("qzeros", qzeros)
+        layer.register_parameter("scales", scales)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        self.kernel.process_weights_after_loading(layer)
+
+    def apply_weights(
+        self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor]
+    ):
+        return self.kernel.apply(layer, x, bias)
+
+
+class AWQAscendLinearScheme(AWQLinearScheme):
+    def _init_kernel(self, quant_config: AWQConfig):
+        from sglang.srt.hardware_backend.npu.quantization.awq_kernels import (
+            AWQAscendLinearKernel,
+        )
+
+        return AWQAscendLinearKernel(quant_config)

python/sglang/srt/layers/quantization/awq/schemes/awq_marlin.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.layers.parameter import GroupQuantScaleParameter, PackedvLLMParameter
+from sglang.srt.layers.quantization.marlin_utils import verify_marlin_supports_shape
+
+from .awq_scheme import AWQLinearSchemeBase
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.quantization.awq.awq import AWQMarlinConfig
+
+__all__ = ["AWQMarlinLinearScheme"]
+
+
+class AWQMarlinLinearScheme(AWQLinearSchemeBase):
+    def __init__(self, quant_config: AWQMarlinConfig):
+        self.quant_config = quant_config
+        self.kernel = self._init_kernel(quant_config)
+
+    def _init_kernel(self, quant_config: AWQMarlinConfig):
+        from sglang.srt.hardware_backend.gpu.quantization.awq_kernels import (
+            AWQMarlinLinearKernel,
+        )
+
+        return AWQMarlinLinearKernel(quant_config)
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        input_size_per_partition: int,
+        output_partition_sizes: list[int],
+        input_size: int,
+        params_dtype: torch.dtype,
+        weight_loader,
+        **kwargs,
+    ) -> None:
+        output_size_per_partition = sum(output_partition_sizes)
+
+        group_size = (
+            self.quant_config.group_size
+            if self.quant_config.group_size != -1
+            else input_size
+        )
+
+        verify_marlin_supports_shape(
+            output_size_per_partition=output_size_per_partition,
+            input_size_per_partition=input_size_per_partition,
+            input_size=input_size,
+            group_size=group_size,
+        )
+
+        qweight = PackedvLLMParameter(
+            data=torch.empty(
+                input_size_per_partition,
+                output_size_per_partition // self.quant_config.pack_factor,
+                dtype=torch.int32,
+            ),
+            input_dim=0,
+            output_dim=1,
+            packed_dim=1,
+            packed_factor=self.quant_config.pack_factor,
+            weight_loader=weight_loader,
+        )
+
+        num_groups = input_size_per_partition // group_size
+
+        qzeros = PackedvLLMParameter(
+            data=torch.empty(
+                num_groups,
+                output_size_per_partition // self.quant_config.pack_factor,
+                dtype=torch.int32,
+            ),
+            input_dim=0,
+            output_dim=1,
+            packed_dim=1,
+            packed_factor=self.quant_config.pack_factor,
+            weight_loader=weight_loader,
+        )
+
+        scales = GroupQuantScaleParameter(
+            data=torch.empty(
+                num_groups,
+                output_size_per_partition,
+                dtype=params_dtype,
+            ),
+            input_dim=0,
+            output_dim=1,
+            weight_loader=weight_loader,
+        )
+
+        layer.register_parameter("qweight", qweight)
+        layer.register_parameter("qzeros", qzeros)
+        layer.register_parameter("scales", scales)
+
+        layer.input_size_per_partition = input_size_per_partition
+        layer.output_size_per_partition = output_size_per_partition
+        layer.num_groups = num_groups
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        self.kernel.process_weights_after_loading(layer)
+
+    def apply_weights(
+        self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor]
+    ):
+        return self.kernel.apply(layer, x, bias)

python/sglang/srt/layers/quantization/awq/schemes/awq_moe.py

@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.srt.layers.linear import set_weight_attrs
+from sglang.srt.layers.moe import (
+    MoeRunner,
+    MoeRunnerBackend,
+    MoeRunnerConfig,
+    get_moe_runner_backend,
+)
+
+from .awq_scheme import AWQMoESchemeBase
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.token_dispatcher import StandardDispatchOutput
+    from sglang.srt.layers.quantization.awq.awq import AWQConfig, AWQMarlinConfig
+
+__all__ = ["AWQMoEScheme", "AWQAscendMoEScheme"]
+
+
+class AWQMoEScheme(AWQMoESchemeBase):
+    def __init__(self, quant_config: AWQMarlinConfig):
+        self.quant_config = quant_config
+        if self.quant_config.weight_bits != 4:
+            raise ValueError("AWQMoEScheme only supports 4bit now.")
+        self.kernel = self._init_kernel(quant_config)
+
+    def _init_kernel(self, quant_config: AWQMarlinConfig):
+        from sglang.srt.hardware_backend.gpu.quantization.awq_kernels import (
+            AWQMoEKernel,
+        )
+
+        return AWQMoEKernel(quant_config)
+
+    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
+
+        extra_weight_attrs.update(
+            {
+                "is_transposed": True,
+                "quant_method": FusedMoeWeightScaleSupported.GROUP.value,
+            }
+        )
+
+        w13_qweight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                hidden_size,
+                2 * intermediate_size_per_partition // self.quant_config.pack_factor,
+                dtype=torch.int32,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_qweight", w13_qweight)
+        set_weight_attrs(w13_qweight, extra_weight_attrs)
+
+        w2_qweight = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                intermediate_size_per_partition,
+                hidden_size // self.quant_config.pack_factor,
+                dtype=torch.int32,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_qweight", w2_qweight)
+        set_weight_attrs(w2_qweight, extra_weight_attrs)
+
+        num_groups_w13 = hidden_size // self.quant_config.group_size
+        num_groups_w2 = intermediate_size_per_partition // self.quant_config.group_size
+
+        w13_scales = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                num_groups_w13,
+                intermediate_size_per_partition * 2,
+                dtype=params_dtype,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_scales", w13_scales)
+        set_weight_attrs(w13_scales, extra_weight_attrs)
+
+        w2_scales = torch.nn.Parameter(
+            torch.empty(num_experts, num_groups_w2, hidden_size, dtype=params_dtype),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_scales", w2_scales)
+        set_weight_attrs(w2_scales, extra_weight_attrs)
+
+        w13_qzeros = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                num_groups_w13,
+                2 * intermediate_size_per_partition // self.quant_config.pack_factor,
+                dtype=torch.int32,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_qzeros", w13_qzeros)
+        set_weight_attrs(w13_qzeros, extra_weight_attrs)
+
+        w2_qzeros = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                num_groups_w2,
+                hidden_size // self.quant_config.pack_factor,
+                dtype=torch.int32,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_qzeros", w2_qzeros)
+        set_weight_attrs(w2_qzeros, 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
+    ):
+        assert get_moe_runner_backend().is_auto()
+        self.moe_runner_config = moe_runner_config
+        self.kernel.runner = MoeRunner(MoeRunnerBackend.MARLIN, moe_runner_config)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: StandardDispatchOutput,
+    ):
+        return self.kernel.apply(layer, dispatch_output)
+
+
+class AWQAscendMoEScheme(AWQMoEScheme):
+    def _init_kernel(self, quant_config: AWQConfig):
+        from sglang.srt.hardware_backend.npu.quantization.awq_kernels import (
+            AWQAscendMoEKernel,
+        )
+
+        return AWQAscendMoEKernel(quant_config)
+
+    def create_moe_runner(
+        self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
+    ):
+        self.moe_runner_config = moe_runner_config

python/sglang/srt/layers/quantization/awq/schemes/awq_scheme.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import abstractmethod
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+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__ = ["AWQLinearSchemeBase", "AWQMoESchemeBase"]
+
+
+class AWQLinearSchemeBase(BaseLinearScheme):
+    @abstractmethod
+    def create_weights(self, *args, **kwargs):
+        raise NotImplementedError
+
+    @abstractmethod
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        raise NotImplementedError
+
+    @abstractmethod
+    def apply_weights(
+        self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor]
+    ):
+        raise NotImplementedError
+
+
+class AWQMoESchemeBase(BaseMoEScheme):
+    @abstractmethod
+    def create_weights(self, *args, **kwargs):
+        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):
+        raise NotImplementedError
+
+    @abstractmethod
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        dispatch_output: "StandardDispatchOutput",
+    ):
+        raise NotImplementedError