diff --git a/python/sglang/srt/hardware_backend/xpu/quantization/__init__.py b/python/sglang/srt/hardware_backend/xpu/quantization/__init__.py new file mode 100644 index 000000000..1c37f034e --- /dev/null +++ b/python/sglang/srt/hardware_backend/xpu/quantization/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""XPU (Intel GPU) quantization kernels.""" diff --git a/python/sglang/srt/hardware_backend/xpu/quantization/awq_kernels.py b/python/sglang/srt/hardware_backend/xpu/quantization/awq_kernels.py new file mode 100644 index 000000000..4dfa74113 --- /dev/null +++ b/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, + ) diff --git a/python/sglang/srt/hardware_backend/xpu/quantization/gptq_kernels.py b/python/sglang/srt/hardware_backend/xpu/quantization/gptq_kernels.py new file mode 100644 index 000000000..a7928b260 --- /dev/null +++ b/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, + ) diff --git a/python/sglang/srt/hardware_backend/xpu/quantization/int4pack_utils.py b/python/sglang/srt/hardware_backend/xpu/quantization/int4pack_utils.py new file mode 100644 index 000000000..eac5b84d4 --- /dev/null +++ b/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) diff --git a/python/sglang/srt/layers/quantization/__init__.py b/python/sglang/srt/layers/quantization/__init__.py index 970afb35b..0ee2ac623 100644 --- a/python/sglang/srt/layers/quantization/__init__.py +++ b/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( { diff --git a/python/sglang/srt/layers/quantization/awq/__init__.py b/python/sglang/srt/layers/quantization/awq/__init__.py index 87f0387ce..e60db0bc6 100644 --- a/python/sglang/srt/layers/quantization/awq/__init__.py +++ b/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", diff --git a/python/sglang/srt/layers/quantization/awq/awq.py b/python/sglang/srt/layers/quantization/awq/awq.py index 3fdcc7399..2bac4a10a 100644 --- a/python/sglang/srt/layers/quantization/awq/awq.py +++ b/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""" diff --git a/python/sglang/srt/layers/quantization/awq/schemes/__init__.py b/python/sglang/srt/layers/quantization/awq/schemes/__init__.py index cbf38c3a8..72ff7be0c 100644 --- a/python/sglang/srt/layers/quantization/awq/schemes/__init__.py +++ b/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", diff --git a/python/sglang/srt/layers/quantization/awq/schemes/awq_linear.py b/python/sglang/srt/layers/quantization/awq/schemes/awq_linear.py index 30f8bef76..957d26635 100644 --- a/python/sglang/srt/layers/quantization/awq/schemes/awq_linear.py +++ b/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) diff --git a/python/sglang/srt/layers/quantization/gptq/__init__.py b/python/sglang/srt/layers/quantization/gptq/__init__.py index f706e77ec..6d6f75009 100644 --- a/python/sglang/srt/layers/quantization/gptq/__init__.py +++ b/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", diff --git a/python/sglang/srt/layers/quantization/gptq/gptq.py b/python/sglang/srt/layers/quantization/gptq/gptq.py index 03f2f6c7d..e6a6c65cc 100644 --- a/python/sglang/srt/layers/quantization/gptq/gptq.py +++ b/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""" diff --git a/python/sglang/srt/layers/quantization/gptq/schemes/__init__.py b/python/sglang/srt/layers/quantization/gptq/schemes/__init__.py index 72682db95..3d89eecf7 100644 --- a/python/sglang/srt/layers/quantization/gptq/schemes/__init__.py +++ b/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", diff --git a/python/sglang/srt/layers/quantization/gptq/schemes/gptq_linear.py b/python/sglang/srt/layers/quantization/gptq/schemes/gptq_linear.py index bfcab0e74..8f83426c0 100644 --- a/python/sglang/srt/layers/quantization/gptq/schemes/gptq_linear.py +++ b/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) diff --git a/test/registered/xpu/test_xpu_int4_dense.py b/test/registered/xpu/test_xpu_int4_dense.py new file mode 100644 index 000000000..8f0902e16 --- /dev/null +++ b/test/registered/xpu/test_xpu_int4_dense.py @@ -0,0 +1,243 @@ +"""Numeric unit tests for the XPU int4 *dense* linear kernels (GPTQ / AWQ).""" + +import unittest + +import torch + +from sglang.srt.runtime_context import get_parallel +from sglang.srt.utils import is_xpu +from sglang.test.ci.ci_register import register_xpu_ci +from sglang.test.test_utils import CustomTestCase + +register_xpu_ci(est_time=20, suite="stage-b-test-1-gpu-xpu") + +DEV = "xpu" + +REL_TOL = {torch.float16: 1.5e-3, torch.bfloat16: 2e-2} +M_VALUES = (1, 8, 256) + +# (K, N, group_size); K % 8 == 0, N % 8 == 0, K % group_size == 0. +SHAPES = [ + (128, 64, 32), + (256, 128, 64), + (256, 128, 128), + (512, 256, 256), +] + +# AutoAWQ forward pack order (inverse of reverse [0, 4, 1, 5, 2, 6, 3, 7]). +AWQ_PACK_ORDER = [0, 2, 4, 6, 1, 3, 5, 7] + + +def _awq_pack(codes: torch.Tensor) -> torch.Tensor: + """``[R, C]`` codes (0..15) -> ``[R, C // 8]`` int32 in AutoAWQ order.""" + r, c = codes.shape + codes = codes.reshape(r, c // 8, 8)[:, :, AWQ_PACK_ORDER] + packed = torch.zeros(r, c // 8, dtype=torch.int32, device=codes.device) + for i in range(8): + packed |= codes[:, :, i].to(torch.int32) << (4 * i) + return packed + + +def _gptq_pack_qweight(codes: torch.Tensor) -> torch.Tensor: + """``[K, N]`` codes -> ``[K // 8, N]`` int32 (packed sequentially along K).""" + k, n = codes.shape + codes = codes.reshape(k // 8, 8, n) + packed = torch.zeros(k // 8, n, dtype=torch.int32, device=codes.device) + for i in range(8): + packed |= codes[:, i, :].to(torch.int32) << (4 * i) + return packed + + +def _gptq_pack_qzeros(zc: torch.Tensor) -> torch.Tensor: + """``[ng, N]`` codes -> ``[ng, N // 8]`` int32 (packed sequentially along N).""" + ng, n = zc.shape + zc = zc.reshape(ng, n // 8, 8) + packed = torch.zeros(ng, n // 8, dtype=torch.int32, device=zc.device) + for j in range(8): + packed |= zc[:, :, j].to(torch.int32) << (4 * j) + return packed + + +def _make_layer(): + """A bare ``LinearBase`` with only ``nn.Module`` machinery initialised.""" + from sglang.srt.layers.linear import LinearBase + + layer = LinearBase.__new__(LinearBase) + torch.nn.Module.__init__(layer) + return layer + + +def _awq_config(group_size: int): + from sglang.srt.layers.quantization.awq import AWQXPUConfig + + cfg = AWQXPUConfig.__new__(AWQXPUConfig) + cfg.group_size = group_size + cfg.weight_bits = 4 + cfg.pack_factor = 8 + cfg.zero_point = True + cfg.lm_head_quantized = False + cfg.modules_to_not_convert = [] + return cfg + + +def _gptq_config(group_size: int, desc_act: bool, fmt: str): + from sglang.srt.layers.quantization.gptq import GPTQXPUConfig + + cfg = GPTQXPUConfig.__new__(GPTQXPUConfig) + cfg.group_size = group_size + cfg.desc_act = desc_act + cfg.checkpoint_format = fmt + cfg.weight_bits = 4 + cfg.lm_head_quantized = False + cfg.dynamic = {} + return cfg + + +@unittest.skipIf(not is_xpu(), "XPU int4 dense UT requires an Intel XPU") +class TestXPUInt4DenseKernel(CustomTestCase): + """AWQ / GPTQ int4pack kernel numerics vs a pure-torch dequant reference.""" + + def test_awq_numeric(self): + for dtype in (torch.float16, torch.bfloat16): + for m in M_VALUES: + for k, n, gs in SHAPES: + with self.subTest(dtype=dtype, M=m, K=k, N=n, gs=gs): + self._run_awq(m, k, n, gs, dtype) + + def _run_awq(self, m, k, n, gs, dtype): + from sglang.srt.hardware_backend.xpu.quantization.awq_kernels import ( + AWQXPULinearKernel, + ) + + torch.manual_seed(0) + ng = k // gs + wcodes = torch.randint(0, 16, (k, n), device=DEV) + zcodes = torch.randint(0, 16, (ng, n), device=DEV) + scales = torch.rand(ng, n, device=DEV, dtype=dtype) * 0.05 + 0.005 + + gidx = torch.arange(k, device=DEV) // gs + w_ref = (wcodes.to(dtype) - zcodes[gidx].to(dtype)) * scales[gidx] + x = torch.randn(m, k, device=DEV, dtype=dtype) + ref = x @ w_ref + + layer = _make_layer() + layer.qweight = torch.nn.Parameter(_awq_pack(wcodes), requires_grad=False) + layer.qzeros = torch.nn.Parameter(_awq_pack(zcodes), requires_grad=False) + layer.scales = torch.nn.Parameter(scales, requires_grad=False) + + kernel = AWQXPULinearKernel(_awq_config(gs)) + kernel.process_weights_after_loading(layer) + out = kernel.apply(layer, x) + + self.assertEqual(tuple(out.shape), (m, n)) + self.assertTrue(torch.isfinite(out).all()) + rel = (out - ref).abs().max().item() / ref.abs().max().item() + self.assertLess(rel, REL_TOL[dtype], f"rel={rel:.2e}") + + def test_gptq_numeric(self): + for dtype in (torch.float16, torch.bfloat16): + for fmt in ("", "gptq_v2"): + for desc_act in (False, True): + for m in M_VALUES: + for k, n, gs in SHAPES: + with self.subTest( + dtype=dtype, + fmt=fmt or "gptq_v1", + desc_act=desc_act, + M=m, + K=k, + N=n, + gs=gs, + ): + self._run_gptq(m, k, n, gs, dtype, desc_act, fmt) + + def _run_gptq(self, m, k, n, gs, dtype, desc_act, fmt, tp_size=1): + from sglang.srt.hardware_backend.xpu.quantization.gptq_kernels import ( + GPTQXPULinearKernel, + ) + + torch.manual_seed(0) + ng = k // gs + qnat = torch.randint(0, 16, (k, n), device=DEV) + zc = torch.randint(0, 14, (ng, n), device=DEV) # room for v1 +1 + scales = torch.rand(ng, n, device=DEV, dtype=dtype) * 0.05 + 0.005 + + if desc_act: + base = torch.arange(k, device=DEV) // gs + g_idx = base[torch.randperm(k, device=DEV)].to(torch.int32) + else: + g_idx = (torch.arange(k, device=DEV) // gs).to(torch.int32) + + zp_eff = zc + (0 if fmt == "gptq_v2" else 1) + w_true = (qnat.to(dtype) - zp_eff[g_idx].to(dtype)) * scales[g_idx] + x = torch.randn(m, k, device=DEV, dtype=dtype) + ref = x @ w_true + + layer = _make_layer() + layer.qweight = torch.nn.Parameter( + _gptq_pack_qweight(qnat), requires_grad=False + ) + layer.qzeros = torch.nn.Parameter(_gptq_pack_qzeros(zc), requires_grad=False) + layer.scales = torch.nn.Parameter(scales, requires_grad=False) + layer.g_idx = torch.nn.Parameter(g_idx, requires_grad=False) + + kernel = GPTQXPULinearKernel(_gptq_config(gs, desc_act, fmt)) + with get_parallel().override(tp_size=tp_size): + kernel.process_weights_after_loading(layer) + out = kernel.apply(layer, x) + + self.assertEqual(tuple(out.shape), (m, n)) + self.assertTrue(torch.isfinite(out).all()) + rel = (out - ref).abs().max().item() / ref.abs().max().item() + self.assertLess(rel, REL_TOL[dtype], f"rel={rel:.2e}") + + def test_gptq_act_order_rejects_split_group_shard(self): + from sglang.srt.hardware_backend.xpu.quantization.gptq_kernels import ( + GPTQXPULinearKernel, + ) + + k, n, gs = 128, 64, 32 + # A row-parallel shard of a permuted K owns only part of every group, so + # after sorting each gs-block still straddles two groups. + g_idx = (torch.arange(2 * k, device=DEV) // gs)[0::2].to(torch.int32) + + # Only g_idx matters here; the payload is never reached. + layer = _make_layer() + layer.qweight = torch.nn.Parameter( + torch.zeros(k // 8, n, dtype=torch.int32, device=DEV), requires_grad=False + ) + layer.qzeros = torch.nn.Parameter( + torch.zeros(k // gs, n // 8, dtype=torch.int32, device=DEV), + requires_grad=False, + ) + layer.scales = torch.nn.Parameter( + torch.ones(k // gs, n, device=DEV, dtype=torch.float16), + requires_grad=False, + ) + layer.g_idx = torch.nn.Parameter(g_idx, requires_grad=False) + kernel = GPTQXPULinearKernel(_gptq_config(gs, True, "")) + + # The limit is representability, not TP: tp_size only decides whether the + # actionable --tp-size hint is appended. The layer is safe to reuse + # because the check fires before any weight is replaced. + for tp_size, pattern in ( + (1, r"K boundary\.$"), + (2, r"tp_size=2.*--tp-size 1"), + ): + with self.subTest(tp_size=tp_size): + with get_parallel().override(tp_size=tp_size): + with self.assertRaisesRegex(NotImplementedError, pattern): + kernel.process_weights_after_loading(layer) + + def test_gptq_group_aligned_shard_allows_tensor_parallel(self): + # Whole-group shards stay representable, so TP is only rejected when a + # group is actually split (act_order) -- not for TP as such. + for dtype in (torch.float16, torch.bfloat16): + for desc_act in (False, True): + for k, n, gs in SHAPES: + with self.subTest(dtype=dtype, desc_act=desc_act, K=k, N=n, gs=gs): + self._run_gptq(8, k, n, gs, dtype, desc_act, "", tp_size=2) + + +if __name__ == "__main__": + unittest.main()