wehub-resource-sync
232 lines
8.3 KiB
Python
232 lines
8.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
|
|
import torch
|
|
from torch.nn.parameter import Parameter
|
|
|
|
from vllm.logger import init_logger
|
|
from vllm.model_executor.layers.quantization.utils import replace_parameter
|
|
from vllm.platforms import current_platform
|
|
from vllm.scalar_type import scalar_types
|
|
|
|
from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
|
|
|
|
_XPUWNA16_SUPPORTED_QUANT_TYPES = (scalar_types.uint4, scalar_types.uint4b8)
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class XPUwNa16LinearKernel(MPLinearKernel):
|
|
@classmethod
|
|
def get_min_capability(cls) -> int:
|
|
return -1
|
|
|
|
@classmethod
|
|
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
|
|
if not current_platform.is_xpu():
|
|
return False, "XPUwNa16 only supported on XPU"
|
|
|
|
if c.act_type != torch.bfloat16 and c.act_type != torch.float16:
|
|
return False, "XPUwNa16 only supports BF16/FP16 activations"
|
|
|
|
if c.weight_type not in _XPUWNA16_SUPPORTED_QUANT_TYPES:
|
|
return (
|
|
False,
|
|
f"Quant type ({c.weight_type}) not supported by "
|
|
"XPUwNa16, supported types are: "
|
|
f"{_XPUWNA16_SUPPORTED_QUANT_TYPES}",
|
|
)
|
|
if c.group_size != -1 and c.group_size % 32 != 0:
|
|
return (
|
|
False,
|
|
f"Group size ({c.group_size}) not supported by "
|
|
"XPUwNa16, supported group sizes are multiples of 32",
|
|
)
|
|
|
|
if c.partition_weight_shape[0] % 32 != 0:
|
|
return (
|
|
False,
|
|
f"Input size ({c.partition_weight_shape[0]}) not supported by "
|
|
"XPUwNa16, supported sizes are multiples of 32",
|
|
)
|
|
|
|
return True, None
|
|
|
|
def process_weights_after_loading(self, layer: torch.nn.Module):
|
|
# Default names since marlin requires empty parameters for these,
|
|
# TODO: remove this requirement from marlin (allow optional tensors)
|
|
if self.w_gidx_name is None:
|
|
self.w_gidx_name = "g_idx"
|
|
if self.w_zp_name is None:
|
|
self.w_zp_name = "w_zp"
|
|
|
|
need_transpose = False
|
|
qweight_shape = getattr(layer, self.w_q_name).shape
|
|
scale_shape = getattr(layer, self.w_s_name).shape
|
|
# gptq marlin and compressed tensors wna16 expect different default
|
|
# layouts for weight and scale, so we check the shapes to determine
|
|
# if we need to transpose
|
|
if qweight_shape[0] != scale_shape[0]:
|
|
need_transpose = True
|
|
|
|
if need_transpose:
|
|
getattr(layer, self.w_q_name).data = (
|
|
getattr(layer, self.w_q_name).data.t().contiguous()
|
|
)
|
|
getattr(layer, self.w_s_name).data = getattr(layer, self.w_s_name).data
|
|
else:
|
|
getattr(layer, self.w_s_name).data = (
|
|
getattr(layer, self.w_s_name).data.t().contiguous()
|
|
)
|
|
|
|
if self.config.zero_points:
|
|
# (FIXME): maybe zero points should also be transposed.
|
|
getattr(layer, self.w_zp_name).data = (
|
|
getattr(layer, self.w_zp_name).data.t().contiguous()
|
|
)
|
|
else:
|
|
weight_zero_point = torch.Tensor([8]).to(torch.int8).to("xpu")
|
|
setattr(
|
|
layer, self.w_zp_name, Parameter(weight_zero_point, requires_grad=False)
|
|
)
|
|
if self.config.has_g_idx:
|
|
setattr(
|
|
layer,
|
|
self.w_gidx_name,
|
|
Parameter(
|
|
getattr(layer, self.w_gidx_name).data.t().contiguous(),
|
|
requires_grad=False,
|
|
),
|
|
)
|
|
else:
|
|
setattr(layer, self.w_gidx_name, None)
|
|
|
|
def apply_weights(
|
|
self,
|
|
layer: torch.nn.Module,
|
|
x: torch.Tensor,
|
|
bias: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
reshaped_x = x.reshape(-1, x.shape[-1])
|
|
w_q, w_s, w_zp, w_gidx = self._get_weight_params(layer)
|
|
out = torch.ops._xpu_C.int4_gemm_w4a16(
|
|
reshaped_x,
|
|
w_q.t(),
|
|
bias if bias is not None else None,
|
|
w_s,
|
|
w_zp,
|
|
self.config.group_size,
|
|
w_gidx,
|
|
)
|
|
return out
|
|
|
|
|
|
class XPUW4A8IntLinearKernel(MPLinearKernel):
|
|
"""XPU kernel for W4A8 integer quantization using oneDNN int4_gemm_w4a8.
|
|
|
|
Weights are symmetric group-quantized int4 packed as uint4.
|
|
Activations are dynamically quantized per-token to symmetric int8.
|
|
"""
|
|
|
|
@classmethod
|
|
def get_min_capability(cls) -> int:
|
|
return -1
|
|
|
|
@classmethod
|
|
def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
|
|
if not current_platform.is_xpu():
|
|
return False, "XPUW4A8Int only supported on XPU"
|
|
if c.act_type not in (torch.bfloat16, torch.float16):
|
|
return False, "XPUW4A8Int requires BF16/FP16 activations"
|
|
if c.weight_type != scalar_types.int4:
|
|
return (
|
|
False,
|
|
f"XPUW4A8Int requires int4 weights, got {c.weight_type}",
|
|
)
|
|
if c.zero_points:
|
|
return False, "XPUW4A8Int only supports symmetric weight quantization"
|
|
if c.group_size != -1 and c.group_size % 32 != 0:
|
|
return (
|
|
False,
|
|
f"Group size ({c.group_size}) not supported by XPUW4A8Int, "
|
|
"must be a multiple of 32",
|
|
)
|
|
in_size, out_size = c.partition_weight_shape
|
|
if in_size % 8 != 0 or out_size % 8 != 0:
|
|
return (
|
|
False,
|
|
f"in/out sizes ({in_size}, {out_size}) must be multiples of 8",
|
|
)
|
|
|
|
if c.act_type != torch.float16:
|
|
logger.warning_once(
|
|
"XPUW4A8IntLinearKernel is running with model dtype %s, "
|
|
"but int4_gemm_w4a8 produces float16 output. Recommend "
|
|
"setting --dtype float16 for best performance.",
|
|
c.act_type,
|
|
)
|
|
|
|
return True, None
|
|
|
|
def _pack_int4_weight(self, w: torch.Tensor) -> torch.Tensor:
|
|
# w is [N, K] int8 with values in [-8, 7]
|
|
w_u4 = w.to(torch.int32) + 8 # shift to [0, 15]
|
|
w_u4 = w_u4.reshape(w.shape[0], w.shape[1] // 8, 8) # [N, K/8, 8]
|
|
shifts = torch.arange(0, 32, 4, dtype=torch.int32, device=w.device)
|
|
packed = ((w_u4 & 0xF) << shifts[None, None, :]).sum(dim=2).to(torch.int32)
|
|
return packed
|
|
|
|
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
|
layer.weight_scale.data = layer.weight_scale.data.t().contiguous()
|
|
|
|
device = layer.weight_packed.device
|
|
# TODO: support asymmetric quantization
|
|
weight_zero_point = torch.tensor([8], dtype=torch.int8, device=device)
|
|
layer.weight_zero_point = Parameter(weight_zero_point, requires_grad=False)
|
|
|
|
# weight_packed is [out, in] int8, signed int4 values in [-8, 7]
|
|
w = layer.weight_packed.data # [out, in]
|
|
|
|
# TODO: implement asym case
|
|
packed = self._pack_int4_weight(w) # [out, in/8] packed uint4
|
|
|
|
replace_parameter(
|
|
layer,
|
|
self.w_q_name,
|
|
torch.nn.Parameter(packed, requires_grad=False),
|
|
)
|
|
|
|
# Free the original unpacked int8 weight (still registered as "weight")
|
|
# to avoid double-storing both int8 [N, K] and int32 [N, K/8] in memory.
|
|
layer.register_parameter("weight", None)
|
|
|
|
def apply_weights(
|
|
self,
|
|
layer: torch.nn.Module,
|
|
x: torch.Tensor,
|
|
bias: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
reshaped_x = x.reshape(-1, x.shape[-1]) # [M, K]
|
|
from vllm._xpu_ops import xpu_ops as ops
|
|
|
|
# TODO: static and asymmetric quantization case
|
|
# Common code for CompressedTensorsW4A8Int does not read act symmetry data
|
|
quant_x, x_scale, x_zero = ops.dynamic_per_token_int8_quant_ref(
|
|
reshaped_x, True, 8
|
|
)
|
|
|
|
out = torch.ops._xpu_C.int4_gemm_w4a8(
|
|
quant_x,
|
|
x_scale,
|
|
x_zero,
|
|
layer.weight_packed.t(),
|
|
layer.weight_scale,
|
|
layer.weight_zero_point,
|
|
self.config.group_size,
|
|
None, # g_idx not currently supported
|
|
bias,
|
|
)
|
|
|
|
return out.to(x.dtype)
|