@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3fd9d2d357
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 08:30 UTC

[CPU][Zen] Route W8A8 and W4A16 linear inference through zentorch on AMD Zen CPUs (#41813)

Signed-off-by: R <Ganesh.R@amd.com> Signed-off-by: Harshal Adhav <harshal.adhav@amd.com> Signed-off-by: Aakar Dwivedi <aadwived@amd.com> Co-authored-by: R <Ganesh.R@amd.com> Co-authored-by: Harshal Adhav <harshal.adhav@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>

Aakar Dwivedi committed May 31, 2026 at 00:47 UTC 3fd9d2d35714e80b4cb3fcd3c408a0398fa2525f
7 files changed +351 -4
setup.py
+1 -3
@@ -1165,9 +1165,7 @@ setup(
1165 install_requires=get_requirements(),
1166 extras_require={
1167 # AMD Zen CPU optimizations via zentorch
1168 - "zen": [
1169 - "zentorch-weekly==5.2.1.dev20260408"
1170 - ], # Zentorch has weekly releases. This pulls the known-good version.
1168 + "zen": ["zentorch==2.11.0.0"],
1169 "bench": ["pandas", "matplotlib", "seaborn", "datasets", "scipy", "plotly"],
1170 "tensorizer": ["tensorizer==2.10.1"],
1171 "fastsafetensors": ["fastsafetensors >= 0.2.2"],
vllm/model_executor/kernels/linear/__init__.py
+10 -1
@@ -61,6 +61,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.xpu import (
61 XPUW4A8IntLinearKernel,
62 XPUwNa16LinearKernel,
63 )
64 +from vllm.model_executor.kernels.linear.mixed_precision.zentorch import (
65 + ZentorchWNA16LinearKernel,
66 +)
67 from vllm.model_executor.kernels.linear.mxfp4 import (
68 MxFp4LinearKernel,
69 MxFp4LinearLayerConfig,
@@ -160,6 +163,9 @@ from vllm.model_executor.kernels.linear.scaled_mm.triton import (
163 from vllm.model_executor.kernels.linear.scaled_mm.xpu import (
164 XPUFP8ScaledMMLinearKernel,
165 )
166 +from vllm.model_executor.kernels.linear.scaled_mm.zentorch import (
167 + ZentorchInt8ScaledMMLinearKernel,
168 +)
169 from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
170 from vllm.platforms import PlatformEnum, current_platform
171
@@ -257,7 +263,7 @@ def _filter_kernels_by_backend(
263
264 # in priority/performance order (when available)
265 _POSSIBLE_INT8_KERNELS: dict[PlatformEnum, list[type[Int8ScaledMMLinearKernel]]] = {
260 - PlatformEnum.CPU: [CPUInt8ScaledMMLinearKernel],
266 + PlatformEnum.CPU: [ZentorchInt8ScaledMMLinearKernel, CPUInt8ScaledMMLinearKernel],
267 PlatformEnum.CUDA: [
268 CutlassInt8ScaledMMLinearKernel,
269 TritonInt8ScaledMMLinearKernel,
@@ -353,6 +359,7 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
359 ],
360 PlatformEnum.CPU: [
361 Dynamic4bitLinearKernel,
362 + ZentorchWNA16LinearKernel,
363 CPUWNA16LinearKernel,
364 ],
365 }
@@ -1023,6 +1030,8 @@ __all__ = [
1030 "RowWiseTorchFP8ScaledMMLinearKernel",
1031 "ROCmFP8ScaledMMLinearKernel",
1032 "TritonInt8ScaledMMLinearKernel",
1033 + "ZentorchInt8ScaledMMLinearKernel",
1034 + "ZentorchWNA16LinearKernel",
1035 "MPLinearKernel",
1036 "MPLinearLayerConfig",
1037 "AllSparkLinearKernel",
vllm/model_executor/kernels/linear/mixed_precision/__init__.py
+4
@@ -39,6 +39,9 @@ from vllm.model_executor.kernels.linear.mixed_precision.xpu import (
39 XPUW4A8IntLinearKernel,
40 XPUwNa16LinearKernel,
41 )
42 +from vllm.model_executor.kernels.linear.mixed_precision.zentorch import (
43 + ZentorchWNA16LinearKernel,
44 +)
45
46 __all__ = [
47 "MPLinearKernel",
@@ -55,4 +58,5 @@ __all__ = [
58 "TritonW4A16LinearKernel",
59 "XPUW4A8IntLinearKernel",
60 "XPUwNa16LinearKernel",
61 + "ZentorchWNA16LinearKernel",
62 ]
vllm/model_executor/kernels/linear/mixed_precision/zentorch.py new
+211
@@ -0,0 +1,211 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Zentorch W4A16 GPTQ weight-only-quantized linear kernel for AMD Zen CPUs.
4 +
5 +Selected by ``choose_mp_linear_kernel`` ahead of the generic oneDNN-backed
6 +``CPUWNA16LinearKernel``. When ``can_implement`` rejects a layer, the selector
7 +falls through to the next kernel in ``_POSSIBLE_KERNELS[PlatformEnum.CPU]``.
8 +"""
9 +
10 +import torch
11 +
12 +from vllm.logger import init_logger
13 +from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op
14 +from vllm.platforms import current_platform
15 +from vllm.scalar_type import scalar_types
16 +
17 +from .cpu import CPUWNA16LinearKernel
18 +from .MPLinearKernel import MPLinearLayerConfig
19 +
20 +logger = init_logger(__name__)
21 +
22 +
23 +def _import_unpack_from_int32():
24 + """Import compressed-tensors' ``unpack_from_int32`` across versions."""
25 + try:
26 + from compressed_tensors.compressors.pack_quantized.helpers import (
27 + unpack_from_int32,
28 + )
29 + except ImportError:
30 + from compressed_tensors.compressors.quantized_compressors.pack_quantized import ( # type: ignore[import-not-found] # noqa: E501
31 + unpack_from_int32,
32 + )
33 + return unpack_from_int32
34 +
35 +
36 +class ZentorchWNA16LinearKernel(CPUWNA16LinearKernel):
37 + """W4A16 GPTQ kernel backed by ``torch.ops.zentorch.zentorch_woq_linear``."""
38 +
39 + @classmethod
40 + def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
41 + ok, reason = super().can_implement(c)
42 + if not ok:
43 + return ok, reason
44 +
45 + if not current_platform.is_zen_cpu():
46 + return False, "ZentorchWNA16 requires an AMD Zen CPU."
47 +
48 + if not has_zentorch_op(["zentorch_woq_repack_weight", "zentorch_woq_linear"]):
49 + return (
50 + False,
51 + "torch.ops.zentorch.{zentorch_woq_repack_weight, "
52 + "zentorch_woq_linear} are not registered.",
53 + )
54 +
55 + if c.has_g_idx:
56 + return False, "ZentorchWNA16 does not support activation re-ordering."
57 + return True, None
58 +
59 + def _zentorch_woq_eligible(self, layer: torch.nn.Module) -> bool:
60 + """Eligibility predicate for the zentorch W4A16 GPTQ fast path.
61 +
62 + Constraints (any failure -> ``cpu_gemm_wna16`` path via ``super()``
63 + with ``layer`` untouched).
64 + """
65 + if (
66 + self.w_gidx_name is not None
67 + and getattr(layer, self.w_gidx_name, None) is not None
68 + ) or (getattr(self.config, "has_g_idx", False)):
69 + return False
70 +
71 + weight_packed = getattr(layer, self.w_q_name, None)
72 + weight_scale = getattr(layer, self.w_s_name, None)
73 + if weight_packed is None or weight_scale is None:
74 + return False
75 +
76 + bits = self.config.weight_type.mantissa
77 + pack_factor = torch.iinfo(weight_packed.dtype).bits // bits
78 + # 4-bit -> 8 values per int32;
79 + if pack_factor != 8:
80 + return False
81 +
82 + # GPTQ-only. AWQ packs along the output dim instead.
83 + in_dim = getattr(weight_packed, "input_dim", None)
84 + pk_dim = getattr(weight_packed, "packed_dim", None)
85 + if in_dim is None or pk_dim is None or in_dim != pk_dim:
86 + return False
87 +
88 + is_ct_format = in_dim == pk_dim == 1
89 + if not is_ct_format:
90 + return False
91 +
92 + if weight_packed.dim() != 2 or weight_scale.dim() != 2:
93 + return False
94 +
95 + # 4-bit -> 8 values per int32; in_features must be divisible by num_groups.
96 + in_features = weight_packed.shape[1] * 8
97 + num_groups = weight_scale.shape[1]
98 + return num_groups > 0 and in_features % num_groups == 0
99 +
100 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
101 + """Repack CT GPTQ weights into the zentorch WOQ layout.
102 +
103 + Falls back to ``CPUWNA16LinearKernel.process_weights_after_loading``
104 + via ``super()`` when the layer doesn't satisfy
105 + ``_zentorch_woq_eligible``.
106 +
107 + On success, ``layer._zentorch_processed_weights`` is set to ``True``
108 + """
109 + if getattr(layer, "_zentorch_processed_weights", False):
110 + return
111 +
112 + if not self._zentorch_woq_eligible(layer):
113 + logger.info_once(
114 + "[zen_cpu] ZentorchWNA16 fast path not eligible for this "
115 + "layer (AWQ pack layout, g_idx, or non-int32 storage); "
116 + "falling back to CPUWNA16LinearKernel (cpu_gemm_wna16)."
117 + )
118 + super().process_weights_after_loading(layer)
119 + return
120 +
121 + if (not self.config.zero_points) and (self.w_zp_name is not None):
122 + setattr(layer, self.w_zp_name, None)
123 +
124 + if (not self.config.has_g_idx) and (self.w_gidx_name is not None):
125 + setattr(layer, self.w_gidx_name, None)
126 +
127 + weight_q = getattr(layer, self.w_q_name)
128 + weight_s = getattr(layer, self.w_s_name)
129 + weight_packed = weight_q.data if hasattr(weight_q, "data") else weight_q
130 + weight_scale = weight_s.data if hasattr(weight_s, "data") else weight_s
131 +
132 + bits = self.config.weight_type.mantissa
133 + pack_factor = torch.iinfo(weight_packed.dtype).bits // bits
134 + out_features, num_groups = weight_scale.shape[0], weight_scale.shape[1]
135 + in_features = weight_packed.shape[1] * pack_factor
136 + original_shape = torch.Size([out_features, in_features])
137 + unpack_from_int32 = _import_unpack_from_int32()
138 + repack_op = torch.ops.zentorch.zentorch_woq_repack_weight.default
139 +
140 + weight_unpacked = unpack_from_int32(
141 + weight_packed,
142 + bits,
143 + original_shape,
144 + packed_dim=weight_q.packed_dim,
145 + )
146 +
147 + zp_param = (
148 + getattr(layer, self.w_zp_name, None) if self.w_zp_name is not None else None
149 + )
150 + needs_unsigned_offset = self.config.weight_type == scalar_types.uint4
151 +
152 + if needs_unsigned_offset:
153 + weight_unpacked = (weight_unpacked.to(torch.int32) + 8).clamp(0, 15)
154 + repacked = repack_op(weight_unpacked.to(torch.int8).contiguous())
155 +
156 + if zp_param is None:
157 + zp_tc = None
158 + else:
159 + zp_tensor = zp_param.data if hasattr(zp_param, "data") else zp_param
160 + zp = unpack_from_int32(
161 + zp_tensor,
162 + bits,
163 + (out_features, num_groups),
164 + packed_dim=zp_param.packed_dim,
165 + )
166 + if needs_unsigned_offset:
167 + zp = (zp.to(torch.int32) + 8).clamp(0, 15)
168 + zp_tc = zp.to(torch.int8).t().contiguous()
169 +
170 + layer._zentorch_woq_packed = repacked.t()
171 + layer._zentorch_woq_scale = weight_scale.t().contiguous()
172 + layer._zentorch_woq_zero_point = zp_tc
173 +
174 + for param_name in (self.w_q_name, self.w_s_name, self.w_zp_name):
175 + if param_name is None:
176 + continue
177 + param = getattr(layer, param_name, None)
178 + if param is None:
179 + continue
180 + if hasattr(param, "data"):
181 + param.data = torch.empty(0)
182 + else:
183 + setattr(layer, param_name, torch.empty(0))
184 +
185 + layer._zentorch_kind = "compressed_tensors_w4a16_gptq"
186 + layer._zentorch_processed_weights = True
187 + logger.info_once(
188 + "[zen_cpu] Using zentorch_woq_linear for W4A16 GPTQ "
189 + "(weight_type=%s, has_zp=%s)",
190 + self.config.weight_type,
191 + zp_tc is not None,
192 + )
193 +
194 + def apply_weights(
195 + self,
196 + layer: torch.nn.Module,
197 + x: torch.Tensor,
198 + bias: torch.Tensor | None = None,
199 + ) -> torch.Tensor:
200 + if getattr(layer, "_zentorch_processed_weights", False):
201 + return torch.ops.zentorch.zentorch_woq_linear.default(
202 + x,
203 + layer._zentorch_woq_packed,
204 + layer._zentorch_woq_scale,
205 + layer._zentorch_woq_zero_point,
206 + bias,
207 + )
208 + return super().apply_weights(layer, x, bias)
209 +
210 +
211 +__all__ = ["ZentorchWNA16LinearKernel"]
vllm/model_executor/kernels/linear/scaled_mm/__init__.py
+4
@@ -39,6 +39,9 @@ from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
39 from vllm.model_executor.kernels.linear.scaled_mm.triton import (
40 TritonInt8ScaledMMLinearKernel,
41 )
42 +from vllm.model_executor.kernels.linear.scaled_mm.zentorch import (
43 + ZentorchInt8ScaledMMLinearKernel,
44 +)
45
46 __all__ = [
47 "FP8ScaledMMLinearKernel",
@@ -58,6 +61,7 @@ __all__ = [
61 "RowWiseTorchFP8ScaledMMLinearKernel",
62 "ROCmFP8ScaledMMLinearKernel",
63 "TritonInt8ScaledMMLinearKernel",
64 + "ZentorchInt8ScaledMMLinearKernel",
65 "Fp8BlockScaledMMLinearKernel",
66 "CPUFp8BlockScaledMMKernel",
67 ]
vllm/model_executor/kernels/linear/scaled_mm/zentorch.py new
+98
@@ -0,0 +1,98 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Zentorch dynamic-symmetric W8A8 int8 linear kernel for AMD Zen CPUs.
4 +
5 +Selected by ``choose_scaled_mm_linear_kernel`` ahead of the generic
6 +oneDNN-backed ``CPUInt8ScaledMMLinearKernel``. When ``is_supported`` or
7 +``can_implement`` rejects a layer, the selector falls through to the next
8 +kernel in ``_POSSIBLE_INT8_KERNELS[PlatformEnum.CPU]``.
9 +"""
10 +
11 +import torch
12 +
13 +from vllm.logger import init_logger
14 +from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op
15 +from vllm.model_executor.layers.quantization.utils import replace_parameter
16 +from vllm.platforms import current_platform
17 +
18 +from .ScaledMMLinearKernel import (
19 + Int8ScaledMMLinearKernel,
20 + Int8ScaledMMLinearLayerConfig,
21 +)
22 +
23 +logger = init_logger(__name__)
24 +
25 +
26 +class ZentorchInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
27 + @classmethod
28 + def is_supported(
29 + cls, compute_capability: int | None = None
30 + ) -> tuple[bool, str | None]:
31 + if not current_platform.is_cpu():
32 + return False, "requires CPU."
33 + if not current_platform.is_zen_cpu():
34 + return False, "requires AMD Zen CPU."
35 + if not has_zentorch_op(["zentorch_dynamic_qlinear"]):
36 + return (
37 + False,
38 + "torch.ops.zentorch.zentorch_dynamic_qlinear is not registered.",
39 + )
40 + return True, None
41 +
42 + @classmethod
43 + def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
44 + if c.is_static_input_scheme:
45 + return False, "requires dynamic activation quantization."
46 + if not c.input_symmetric:
47 + return False, "requires symmetric activation quantization."
48 + if not c.is_channelwise:
49 + return False, "requires per-channel weight quantization."
50 + return True, None
51 +
52 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
53 + """Prepare weights for ``zentorch_dynamic_qlinear``.
54 +
55 + Keeps weight in [N, K] layout (int8, contiguous) and converts the
56 + per-channel weight scale to bf16 with shape ``(N,)``.
57 + """
58 + w_q_name, w_s_name, _, _, _ = self.layer_param_names
59 + weight = getattr(layer, w_q_name)
60 + n = weight.shape[0]
61 + replace_parameter(
62 + layer,
63 + w_q_name,
64 + torch.nn.Parameter(weight.data.contiguous(), requires_grad=False),
65 + )
66 +
67 + weight_scale = getattr(layer, w_s_name)
68 + ws = weight_scale.data
69 + if ws.dim() == 2 and ws.shape[-1] == 1:
70 + ws = ws.squeeze(-1)
71 + ws = ws.to(torch.bfloat16).contiguous()
72 + assert ws.shape == (n,), (
73 + f"[zen_cpu] expected weight scale shape ({n},), got {tuple(ws.shape)}"
74 + )
75 +
76 + replace_parameter(
77 + layer,
78 + w_s_name,
79 + torch.nn.Parameter(ws, requires_grad=False),
80 + )
81 + logger.info_once(
82 + "[zen_cpu] Using zentorch_dynamic_qlinear for W8A8 (dynamic-symmetric)"
83 + )
84 +
85 + def apply_weights(
86 + self,
87 + layer: torch.nn.Module,
88 + x: torch.Tensor,
89 + bias: torch.Tensor | None = None,
90 + ) -> torch.Tensor:
91 + w_q_name, w_s_name, _, _, _ = self.layer_param_names
92 + return torch.ops.zentorch.zentorch_dynamic_qlinear(
93 + x,
94 + getattr(layer, w_q_name),
95 + getattr(layer, w_s_name),
96 + bias,
97 + zentorch_op_name="zentorch::zentorch_dynamic_qlinear",
98 + )
vllm/model_executor/kernels/linear/zentorch_utils.py new
+23
@@ -0,0 +1,23 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Gates zentorch CPU linear dispatch on platform/op availability."""
4 +
5 +from __future__ import annotations
6 +
7 +import torch
8 +
9 +from vllm.platforms import current_platform
10 +
11 +__all__ = ["has_zentorch_op"]
12 +
13 +
14 +def has_zentorch_op(op_names: list[str]) -> bool:
15 + """Return ``True`` when running on Zen CPU with all named ops registered."""
16 + if not op_names:
17 + raise ValueError("has_zentorch_op requires at least one op name")
18 + if not current_platform.is_zen_cpu():
19 + return False
20 + ns = getattr(torch.ops, "zentorch", None)
21 + if ns is None:
22 + return False
23 + return all(hasattr(ns, op_name) for op_name in op_names)