[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>
This commit is contained in:
YangKai0616
co-authored by
gemini-code-assist[bot]
gemini-code-assist[bot]
parent
f6fff25756
commit
fbdec2855a
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""XPU (Intel GPU) quantization kernels."""
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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(
|
||||
{
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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"""
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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"""
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user