[XPU] Support INT4 dense linear (AWQ/GPTQ) for XPU (#30236)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

python/sglang/srt/hardware_backend/xpu/quantization/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+"""XPU (Intel GPU) quantization kernels."""

python/sglang/srt/hardware_backend/xpu/quantization/awq_kernels.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+"""AWQ int4 dense linear for Intel XPU."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.hardware_backend.xpu.quantization.int4pack_utils import (
+    SUPPORTED_GROUP_SIZES,
+    pack_int4_to_uint8,
+    unpack_awq_to_codes,
+    xpu_int4pack_mm,
+)
+from sglang.srt.layers.quantization.utils import replace_parameter
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.quantization.base_config import QuantizationConfig
+
+
+class AWQXPULinearKernel:
+    def __init__(self, quant_config: Optional[QuantizationConfig] = None):
+        self.quant_config = quant_config
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        group_size = self.quant_config.group_size
+        if group_size not in SUPPORTED_GROUP_SIZES:
+            raise ValueError(
+                f"AWQ on XPU requires group_size in {SUPPORTED_GROUP_SIZES}, "
+                f"got {group_size}. The native XPU INT4 operator does not "
+                "support this group size (per-channel/-1 is out of scope)."
+            )
+
+        qweight = layer.qweight.data  # [K, N // 8] int32
+        qzeros = layer.qzeros.data  # [K // gs, N // 8] int32
+        scales = layer.scales.data  # [K // gs, N]
+
+        k = qweight.shape[0]
+        n = scales.shape[1]
+
+        # qweight -> [N, K // 2] uint8 (torch int4pack B layout)
+        codes = unpack_awq_to_codes(qweight, k)  # [K, N]
+        codes = codes.t().contiguous()  # [N, K]
+        qweight_uint8 = pack_int4_to_uint8(codes)  # [N, K // 2]
+        qweight_packed = torch.ops.aten._convert_weight_to_int4pack(
+            qweight_uint8, 8
+        )  # [N, K // 8] int32
+
+        # qzeros -> [K // gs, N] int8 zero-points expected by the native op.
+        zero_points = unpack_awq_to_codes(qzeros, scales.shape[0])
+
+        replace_parameter(layer, "qweight", qweight_packed)
+        layer.register_parameter(
+            "xpu_scales",
+            torch.nn.Parameter(scales.contiguous(), requires_grad=False),
+        )
+        layer.register_parameter(
+            "xpu_zero_points",
+            torch.nn.Parameter(zero_points.to(torch.int8), requires_grad=False),
+        )
+        del layer.qzeros
+        del layer.scales
+
+        layer.xpu_out_features = n
+        layer.xpu_group_size = group_size
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        return xpu_int4pack_mm(
+            x,
+            layer.qweight,
+            layer.xpu_group_size,
+            layer.xpu_scales,
+            layer.xpu_zero_points,
+            layer.xpu_out_features,
+            bias,
+        )

python/sglang/srt/hardware_backend/xpu/quantization/gptq_kernels.py

@@ -0,0 +1,127 @@
+# SPDX-License-Identifier: Apache-2.0
+"""GPTQ int4 dense linear for Intel XPU."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.hardware_backend.xpu.quantization.int4pack_utils import (
+    SUPPORTED_GROUP_SIZES,
+    pack_int4_to_uint8,
+    unpack_gptq_qweight,
+    unpack_gptq_qzeros,
+    xpu_int4pack_mm,
+)
+from sglang.srt.layers.quantization.utils import replace_parameter
+from sglang.srt.runtime_context import get_parallel
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.quantization.base_config import QuantizationConfig
+
+
+class GPTQXPULinearKernel:
+    def __init__(self, quant_config: Optional[QuantizationConfig] = None):
+        self.quant_config = quant_config
+        self.use_v2_format = getattr(quant_config, "checkpoint_format", "") == "gptq_v2"
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        group_size = self.quant_config.group_size
+        if group_size not in SUPPORTED_GROUP_SIZES:
+            raise ValueError(
+                f"GPTQ on XPU requires group_size in {SUPPORTED_GROUP_SIZES}, "
+                f"got {group_size}. The native XPU INT4 operator does not "
+                "support this group size (per-channel/-1 is out of scope)."
+            )
+
+        qweight = layer.qweight.data  # [K // 8, N] int32
+        qzeros = layer.qzeros.data  # [K // gs, N // 8] int32
+        scales = layer.scales.data  # [K // gs, N]
+        desc_act = bool(self.quant_config.desc_act)
+
+        codes = unpack_gptq_qweight(qweight)  # [K, N]
+        k, n = codes.shape
+
+        # qzeros -> [K // gs, N] effective zero-points (v1 is off-by-one).
+        zp = unpack_gptq_qzeros(qzeros).to(torch.int32)  # [num_groups, N]
+        if not self.use_v2_format:
+            zp = zp + 1
+
+        act_perm = None
+        if desc_act:
+            g_idx = layer.g_idx.data
+            if g_idx.numel() != k:
+                raise ValueError(
+                    "GPTQ act_order on XPU expects a per-channel g_idx of length "
+                    f"K={k}, got {g_idx.numel()}."
+                )
+            # Sort K by group id so groups become contiguous gs-blocks.
+            act_perm = torch.argsort(g_idx, stable=True).to(torch.int64)
+            codes = codes[act_perm, :]
+            sorted_g = g_idx[act_perm].to(torch.int64)
+            blocks = sorted_g.view(-1, group_size)
+            if not torch.equal(blocks, blocks[:, :1].expand_as(blocks)):
+                tp_size = get_parallel().tp_size
+                tp_hint = (
+                    f" Got tp_size={tp_size}; please use --tp-size 1."
+                    if tp_size > 1
+                    else ""
+                )
+                raise NotImplementedError(
+                    "GPTQ act_order on XPU requires each group_size block of "
+                    "input channels to map to a single group, but this shard "
+                    "splits a group across the K boundary." + tp_hint
+                )
+            # Reorder scales/zeros to follow the block group order.
+            block_gid = blocks[:, 0]  # [num_blocks]
+            scales = scales[block_gid]
+            zp = zp[block_gid]
+
+        codes = codes.t().contiguous()  # [N, K]
+        qweight_uint8 = pack_int4_to_uint8(codes)  # [N, K // 2]
+        qweight_packed = torch.ops.aten._convert_weight_to_int4pack(
+            qweight_uint8, 8
+        )  # [N, K // 8] int32
+
+        replace_parameter(layer, "qweight", qweight_packed)
+        layer.register_parameter(
+            "xpu_scales",
+            torch.nn.Parameter(scales.contiguous(), requires_grad=False),
+        )
+        layer.register_parameter(
+            "xpu_zero_points",
+            torch.nn.Parameter(zp.to(torch.int8).contiguous(), requires_grad=False),
+        )
+        if act_perm is not None:
+            layer.register_buffer(
+                "xpu_act_perm", act_perm.to(qweight_packed.device), persistent=False
+            )
+        else:
+            layer.xpu_act_perm = None
+        del layer.qzeros
+        del layer.scales
+        if hasattr(layer, "g_idx"):
+            del layer.g_idx
+
+        layer.xpu_out_features = n
+        layer.xpu_group_size = group_size
+
+    def apply(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        act_perm = getattr(layer, "xpu_act_perm", None)
+        if act_perm is not None:
+            x = x.index_select(-1, act_perm)
+        return xpu_int4pack_mm(
+            x,
+            layer.qweight,
+            layer.xpu_group_size,
+            layer.xpu_scales,
+            layer.xpu_zero_points,
+            layer.xpu_out_features,
+            bias,
+        )

python/sglang/srt/hardware_backend/xpu/quantization/int4pack_utils.py

@@ -0,0 +1,80 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Helpers for lowering GPTQ/AWQ int4 weights to the torch XPU int4pack layout."""
+
+from __future__ import annotations
+
+import torch
+
+# AutoAWQ packs 8 nibbles per int32 in this interleaved order; reversing it
+# recovers natural column order. Matches reverse_awq_pack_order used in
+# moe_wna16.convert_awq_tensor and awq_triton.
+AWQ_REVERSE_PACK_ORDER = [0, 4, 1, 5, 2, 6, 3, 7]
+
+# Group sizes accepted by _weight_int4pack_mm_with_scales_and_zeros on XPU.
+SUPPORTED_GROUP_SIZES = (32, 64, 128, 256)
+
+
+def pack_int4_to_uint8(q: torch.Tensor) -> torch.Tensor:
+    """Pack an ``[N, K]`` tensor of codes ``q in [0, 15]`` into ``[N, K // 2]``.
+
+    Low nibble holds even ``k``, high nibble holds odd ``k`` (torch int4pack B).
+    """
+    assert q.shape[-1] % 2 == 0, "K must be even to pack into int4 bytes"
+    q = q.to(torch.uint8)
+    low = q[..., 0::2]
+    high = q[..., 1::2]
+    return (low | (high << 4)).contiguous()
+
+
+def unpack_awq_to_codes(packed: torch.Tensor, rows: int) -> torch.Tensor:
+    """Deinterleave AWQ-packed int32 ``[rows, cols]`` into codes ``[rows, cols*8]``.
+
+    Codes are in ``[0, 15]`` and restored to natural (non-interleaved) order.
+    """
+    t = packed.contiguous().view(torch.uint8)  # [rows, cols * 4]
+    shifter = torch.tensor([0, 4], dtype=torch.uint8, device=t.device)
+    t = (t[:, :, None] >> shifter) & 0xF  # [rows, cols * 4, 2]
+    t = t.view(-1, 8)[:, AWQ_REVERSE_PACK_ORDER]  # undo interleave
+    return t.reshape(rows, -1)  # [rows, cols * 8]
+
+
+def _nibble_shifts(device: torch.device) -> torch.Tensor:
+    return torch.arange(0, 32, 4, device=device, dtype=torch.int32)
+
+
+def unpack_gptq_qweight(qweight: torch.Tensor) -> torch.Tensor:
+    """``[K // 8, N]`` int32 packed along K -> ``[K, N]`` codes in ``[0, 15]``."""
+    n = qweight.shape[1]
+    shifts = _nibble_shifts(qweight.device)  # [8]
+    # [K // 8, 8, N]; sub-index i selects k = row * 8 + i
+    codes = (qweight.unsqueeze(1) >> shifts.view(1, 8, 1)) & 0xF
+    return codes.reshape(-1, n)  # [K, N]
+
+
+def unpack_gptq_qzeros(qzeros: torch.Tensor) -> torch.Tensor:
+    """``[K // gs, N // 8]`` int32 packed along N -> ``[K // gs, N]`` codes."""
+    rows = qzeros.shape[0]
+    shifts = _nibble_shifts(qzeros.device)  # [8]
+    # [K // gs, N // 8, 8]; sub-index j selects n = col * 8 + j
+    codes = (qzeros.unsqueeze(-1) >> shifts.view(1, 1, 8)) & 0xF
+    return codes.reshape(rows, -1)  # [K // gs, N]
+
+
+def xpu_int4pack_mm(
+    x: torch.Tensor,
+    qweight_packed: torch.Tensor,
+    group_size: int,
+    scales: torch.Tensor,
+    zero_points: torch.Tensor,
+    out_features: int,
+    bias: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Run XPU int4pack MM with leading-dim flatten / restore and bias."""
+    out_shape = x.shape[:-1] + (out_features,)
+    reshaped_x = x.reshape(-1, x.shape[-1]).contiguous()
+    out = torch.ops.aten._weight_int4pack_mm_with_scales_and_zeros(
+        reshaped_x, qweight_packed, group_size, scales, zero_points
+    )
+    if bias is not None:
+        out = out + bias
+    return out.reshape(out_shape)

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

@@ -19,7 +19,12 @@ class DummyConfig:
 CompressedTensorsConfig = DummyConfig
 
 from sglang.srt.layers.quantization.auto_round import AutoRoundConfig
-from sglang.srt.layers.quantization.awq import AWQConfig, AWQCPUConfig, AWQMarlinConfig
+from sglang.srt.layers.quantization.awq import (
+    AWQConfig,
+    AWQCPUConfig,
+    AWQMarlinConfig,
+    AWQXPUConfig,
+)
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.layers.quantization.bitsandbytes import BitsAndBytesConfig
 from sglang.srt.layers.quantization.blockwise_int8 import BlockInt8Config
@@ -33,6 +38,7 @@ from sglang.srt.layers.quantization.gptq import (
     GPTQAscendConfig,
     GPTQConfig,
     GPTQMarlinConfig,
+    GPTQXPUConfig,
 )
 from sglang.srt.layers.quantization.humming import HummingConfig
 from sglang.srt.layers.quantization.mlx import MlxQuantizationConfig
@@ -61,6 +67,7 @@ from sglang.srt.utils import (
     is_gfx95_supported,
     is_mps,
     is_npu,
+    is_xpu,
 )
 
 _is_gfx95_supported = is_gfx95_supported()
@@ -121,6 +128,15 @@ if is_npu():
     )
 
 
+if is_xpu():
+    BASE_QUANTIZATION_METHODS.update(
+        {
+            "gptq": GPTQXPUConfig,
+            "awq": AWQXPUConfig,
+        }
+    )
+
+
 if is_mps():
     BASE_QUANTIZATION_METHODS.update(
         {

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

@@ -11,6 +11,7 @@ from .awq import (
     AWQLinearMethod,
     AWQMarlinConfig,
     AWQMoEMethod,
+    AWQXPUConfig,
 )
 from .schemes import (
     AWQAscendLinearScheme,
@@ -24,6 +25,7 @@ __all__ = [
     "AWQConfig",
     "AWQCPUConfig",
     "AWQMarlinConfig",
+    "AWQXPUConfig",
     "AWQLinearMethod",
     "AWQMoEMethod",
     "AWQLinearScheme",

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

@@ -33,6 +33,7 @@ from .schemes import (
     AWQLinearScheme,
     AWQMarlinLinearScheme,
     AWQMoEScheme,
+    AWQXPULinearScheme,
 )
 
 if TYPE_CHECKING:
@@ -213,6 +214,34 @@ class AWQCPUConfig(AWQConfig):
         return AWQIntelAMXMoEScheme(self)
 
 
+class AWQXPUConfig(AWQConfig):
+    """AWQ int4 dense linear on Intel XPU.
+
+    Lowers to torch's native ``_weight_int4pack_mm_with_scales_and_zeros`` op.
+    MoE is out of scope for the dense phase (mirrors ``GPTQXPUConfig``).
+    """
+
+    def get_supported_act_dtypes(self) -> List[torch.dtype]:
+        return [torch.float16, torch.bfloat16]
+
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> Optional[LinearMethodBase]:
+        from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
+
+        if isinstance(layer, FusedMoE):
+            raise NotImplementedError(
+                "AWQ MoE is not yet supported on XPU (dense-only phase)."
+            )
+        return super().get_quant_method(layer, prefix)
+
+    def get_linear_scheme(self, layer: torch.nn.Module):
+        from sglang.srt.layers.linear import LinearBase
+
+        assert isinstance(layer, LinearBase)
+        return AWQXPULinearScheme(self)
+
+
 class AWQMarlinConfig(QuantizationConfig):
     """Config class for AWQ Marlin"""
 

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

@@ -1,7 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from .awq_cpu import AWQIntelAMXLinearScheme, AWQIntelAMXMoEScheme
-from .awq_linear import AWQAscendLinearScheme, AWQLinearScheme
+from .awq_linear import (
+    AWQAscendLinearScheme,
+    AWQLinearScheme,
+    AWQXPULinearScheme,
+)
 from .awq_marlin import AWQMarlinLinearScheme
 from .awq_moe import AWQAscendMoEScheme, AWQMoEScheme
 from .awq_scheme import AWQLinearSchemeBase, AWQMoESchemeBase
@@ -11,6 +15,7 @@ __all__ = [
     "AWQMoESchemeBase",
     "AWQLinearScheme",
     "AWQAscendLinearScheme",
+    "AWQXPULinearScheme",
     "AWQIntelAMXLinearScheme",
     "AWQMarlinLinearScheme",
     "AWQMoEScheme",

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

@@ -12,7 +12,7 @@ from .awq_scheme import AWQLinearSchemeBase
 if TYPE_CHECKING:
     from sglang.srt.layers.quantization.awq.awq import AWQConfig
 
-__all__ = ["AWQLinearScheme", "AWQAscendLinearScheme"]
+__all__ = ["AWQLinearScheme", "AWQAscendLinearScheme", "AWQXPULinearScheme"]
 
 
 class AWQLinearScheme(AWQLinearSchemeBase):
@@ -108,3 +108,12 @@ class AWQAscendLinearScheme(AWQLinearScheme):
         )
 
         return AWQAscendLinearKernel(quant_config)
+
+
+class AWQXPULinearScheme(AWQLinearScheme):
+    def _init_kernel(self, quant_config: AWQConfig):
+        from sglang.srt.hardware_backend.xpu.quantization.awq_kernels import (
+            AWQXPULinearKernel,
+        )
+
+        return AWQXPULinearKernel(quant_config)

python/sglang/srt/layers/quantization/gptq/__init__.py

@@ -9,6 +9,7 @@ from .gptq import (
     GPTQMarlinLinearMethod,
     GPTQMarlinMoEMethod,
     GPTQMoEMethod,
+    GPTQXPUConfig,
     check_marlin_format,
 )
 from .schemes import (
@@ -19,6 +20,7 @@ from .schemes import (
     GPTQMarlinLinearScheme,
     GPTQMarlinMoEScheme,
     GPTQMoEAscendScheme,
+    GPTQXPULinearScheme,
 )
 
 __all__ = [
@@ -32,6 +34,8 @@ __all__ = [
     "GPTQMarlinMoEMethod",
     "GPTQLinearScheme",
     "GPTQAscendLinearScheme",
+    "GPTQXPULinearScheme",
+    "GPTQXPUConfig",
     "GPTQIntelAMXLinearScheme",
     "GPTQIntelAMXMoEScheme",
     "GPTQMarlinLinearScheme",

python/sglang/srt/layers/quantization/gptq/gptq.py

@@ -28,6 +28,7 @@ from .schemes import (
     GPTQMarlinLinearScheme,
     GPTQMarlinMoEScheme,
     GPTQMoEAscendScheme,
+    GPTQXPULinearScheme,
 )
 
 if TYPE_CHECKING:
@@ -261,6 +262,35 @@ class CPUGPTQConfig(GPTQConfig):
         return GPTQIntelAMXMoEScheme(self)
 
 
+class GPTQXPUConfig(GPTQConfig):
+    """Config class for GPTQ on Intel XPU.
+
+    Dense int4 GPTQ lowers to torch's native
+    ``_weight_int4pack_mm_with_scales_and_zeros`` op (no Marlin on XPU). MoE is
+    out of scope for the dense phase.
+    """
+
+    @classmethod
+    def get_supported_act_dtypes(cls) -> List[torch.dtype]:
+        return [torch.half, torch.bfloat16]
+
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> Optional[LinearMethodBase]:
+        from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
+
+        if isinstance(layer, FusedMoE):
+            raise NotImplementedError(
+                "GPTQ MoE is not yet supported on XPU (dense-only phase)."
+            )
+        return get_linear_quant_method(
+            self, layer, prefix=prefix, linear_method_cls=GPTQLinearMethod
+        )
+
+    def get_linear_scheme(self, layer: torch.nn.Module):
+        return GPTQXPULinearScheme(self)
+
+
 class GPTQMarlinConfig(QuantizationConfig):
     """Config class for GPTQ Marlin"""
 

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

@@ -1,7 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from .gptq_cpu import GPTQIntelAMXLinearScheme, GPTQIntelAMXMoEScheme
-from .gptq_linear import GPTQAscendLinearScheme, GPTQLinearScheme
+from .gptq_linear import (
+    GPTQAscendLinearScheme,
+    GPTQLinearScheme,
+    GPTQXPULinearScheme,
+)
 from .gptq_marlin import GPTQMarlinLinearScheme
 from .gptq_moe import GPTQMarlinMoEScheme, GPTQMoEAscendScheme
 from .gptq_scheme import GPTQLinearSchemeBase, GPTQMoESchemeBase
@@ -11,6 +15,7 @@ __all__ = [
     "GPTQMoESchemeBase",
     "GPTQLinearScheme",
     "GPTQAscendLinearScheme",
+    "GPTQXPULinearScheme",
     "GPTQIntelAMXLinearScheme",
     "GPTQMarlinLinearScheme",
     "GPTQMoEAscendScheme",

python/sglang/srt/layers/quantization/gptq/schemes/gptq_linear.py

@@ -19,7 +19,7 @@ from .gptq_scheme import GPTQLinearSchemeBase
 if TYPE_CHECKING:
     from sglang.srt.layers.quantization.gptq.gptq import GPTQConfig
 
-__all__ = ["GPTQLinearScheme", "GPTQAscendLinearScheme"]
+__all__ = ["GPTQLinearScheme", "GPTQAscendLinearScheme", "GPTQXPULinearScheme"]
 
 
 class GPTQLinearScheme(GPTQLinearSchemeBase):
@@ -169,3 +169,12 @@ class GPTQAscendLinearScheme(GPTQLinearScheme):
         super().create_weights(layer=layer, **kwargs)
         set_weight_attrs(layer.qzeros, {"pack_factor": self.quant_config.pack_factor})
         set_weight_attrs(layer.qweight, {"pack_factor": self.quant_config.pack_factor})
+
+
+class GPTQXPULinearScheme(GPTQLinearScheme):
+    def _init_kernel(self, quant_config: GPTQConfig):
+        from sglang.srt.hardware_backend.xpu.quantization.gptq_kernels import (
+            GPTQXPULinearKernel,
+        )
+
+        return GPTQXPULinearKernel(quant_config)