@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e12291d733

[CPU][Zen] Add DA8W4 (W4A8) int4 support for dense and MoE layers (#54024)

Signed-off-by: R <Ganesh.R@amd.com> Signed-off-by: Ganesh R <Ganesh.R@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Ganesh R committed Oct 1, 2026 at 12:48 UTC e12291d7332db897885fdc0e6ed61b28c969d743
9 files changed +1184 -15
.buildkite/hardware_tests/cpu.yaml
+4
@@ -24,6 +24,9 @@ steps:
24 - tests/kernels/test_onednn.py
25 - tests/kernels/test_awq_int4_to_int8.py
26 - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
27 + - tests/kernels/quantization/test_zen_da8w4.py
28 + - vllm/model_executor/kernels/linear/mixed_precision/zentorch.py
29 + - vllm/model_executor/layers/fused_moe/experts/zentorch_moe.py
30 - tests/kernels/mamba/cpu/test_cpu_gdn_ops.py
31 - tests/kernels/mamba/test_cpu_short_conv.py
32 - tests/kernels/mamba/test_causal_conv1d.py
@@ -46,6 +49,7 @@ steps:
49 tests/kernels/test_onednn.py \
50 tests/kernels/test_awq_int4_to_int8.py \
51 tests/kernels/quantization/test_cpu_fp8_scaled_mm.py \
52 + tests/kernels/quantization/test_zen_da8w4.py \
53 tests/kernels/mamba/cpu/test_cpu_gdn_ops.py \
54 tests/kernels/mamba/test_causal_conv1d.py \
55 tests/kernels/mamba/test_mamba_ssm.py \
tests/kernels/quantization/test_zen_da8w4.py new
+661
@@ -0,0 +1,661 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for the zentorch DA8W4 (W4A8) int4 linear and MoE paths on Zen CPUs.
4 +
5 +The zentorch ops are mocked with reference implementations when zentorch is not
6 +installed, so the layout and dispatch contracts are covered in CI.
7 +"""
8 +
9 +import dataclasses
10 +
11 +import pytest
12 +import torch
13 +from compressed_tensors.compressors.pack_quantized.helpers import pack_to_int32
14 +
15 +from tests.kernels.quant_utils import ref_dynamic_per_token_quant
16 +from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (
17 + MPLinearLayerConfig,
18 +)
19 +from vllm.model_executor.kernels.linear.mixed_precision.zentorch import (
20 + ZentorchWNA16LinearKernel,
21 + _import_unpack_from_int32,
22 +)
23 +from vllm.model_executor.layers.fused_moe.oracle import int_wna16
24 +from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import WNA16MoEBackend
25 +from vllm.platforms import current_platform
26 +from vllm.scalar_type import scalar_types
27 +
28 +if not current_platform.is_cpu():
29 + pytest.skip("skipping CPU-only tests", allow_module_level=True)
30 +
31 +GROUP_SIZE = 128
32 +IN_FEATURES = 512
33 +OUT_FEATURES = 128
34 +
35 +
36 +def _unpack_s4(packed: torch.Tensor, in_features: int) -> torch.Tensor:
37 + """int8 [N, K/2] (or int32 [N, K/8]) packed s4 -> float32 [N, K]."""
38 + words = packed.view(torch.int32)
39 + out = torch.zeros(words.shape[0], in_features, dtype=torch.float32)
40 + for i in range(8):
41 + nibble = (words >> (4 * i)) & 0xF
42 + out[:, i::8] = torch.where(nibble > 7, nibble - 16, nibble).float()
43 + return out
44 +
45 +
46 +def _repack_s4(unpacked: torch.Tensor) -> torch.Tensor:
47 + """Reference zentorch_woq_repack_weight: int8 [N, K] -> int32 [N, K/8]."""
48 + n, k = unpacked.shape
49 + values = unpacked.to(torch.int32).reshape(n, k // 8, 8)
50 + out = torch.zeros(n, k // 8, dtype=torch.int32)
51 + for i in range(8):
52 + out |= (values[:, :, i] & 0xF) << (4 * i)
53 + return out
54 +
55 +
56 +def _quantize_per_group(
57 + weight: torch.Tensor, group_size: int
58 +) -> tuple[torch.Tensor, torch.Tensor]:
59 + """Symmetric int4 per-group quantization -> (int8 in [-8, 7], scale [N, G])."""
60 + out_features, in_features = weight.shape
61 + grouped = weight.float().reshape(
62 + out_features, in_features // group_size, group_size
63 + )
64 + scale = grouped.abs().amax(dim=-1) / 8.0
65 + quantized = (grouped / scale.unsqueeze(-1)).round().clamp(-8, 7).to(torch.int8)
66 + return quantized.reshape(out_features, in_features), scale.to(torch.bfloat16)
67 +
68 +
69 +def _dequantize_per_group(
70 + quantized: torch.Tensor, scale: torch.Tensor, group_size: int
71 +) -> torch.Tensor:
72 + """Dequantize with [N, G] scales into a float32 [N, K] weight."""
73 + return quantized.float() * scale.float().repeat_interleave(group_size, dim=1)
74 +
75 +
76 +def _dynamic_quant_matmul(
77 + x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None
78 +) -> torch.Tensor:
79 + """DA8W4 reference: per-token int8 activation quant against a bf16 weight."""
80 + x_q, scale = ref_dynamic_per_token_quant(x.float(), torch.int8)
81 + out = (x_q.float() @ weight.t()) * scale
82 + if bias is not None:
83 + out = out + bias.float()
84 + return out
85 +
86 +
87 +@pytest.fixture
88 +def mock_zentorch_ops():
89 + """Register reference zentorch ops when zentorch is not installed."""
90 + if hasattr(torch.ops.zentorch, "zentorch_dynamic_qlinear"):
91 + yield None
92 + return
93 +
94 + calls: dict[str, tuple] = {}
95 +
96 + def _dynamic_qlinear(input, weight, weight_scales, bias=None, zentorch_op_name=""):
97 + calls["dynamic_qlinear"] = (input, weight, weight_scales, bias)
98 + in_features = weight.view(torch.int32).shape[1] * 8
99 + group_size = in_features // weight_scales.shape[0]
100 + # weight_scales is [G, N]; expand to a [N, K] dequantized weight.
101 + scales = weight_scales.float().t().repeat_interleave(group_size, dim=1)
102 + weight_deq = _unpack_s4(weight, in_features) * scales
103 + out = _dynamic_quant_matmul(input, weight_deq, bias)
104 + return out.to(input.dtype)
105 +
106 + def _fused_moe(
107 + output,
108 + input,
109 + w13,
110 + w2,
111 + w13_bias,
112 + w2_bias,
113 + topk_weights,
114 + topk_id,
115 + skip_weighted,
116 + act,
117 + w13_scales=None,
118 + w2_scales=None,
119 + zentorch_op_name="",
120 + ):
121 + calls["fused_moe"] = (
122 + input,
123 + w13,
124 + w2,
125 + topk_weights,
126 + topk_id,
127 + skip_weighted,
128 + act,
129 + w13_scales,
130 + w2_scales,
131 + )
132 + output.zero_()
133 +
134 + lib_def = torch.library.Library("zentorch", "DEF")
135 + lib_def.define("zentorch_woq_repack_weight(Tensor unpacked_weight) -> Tensor")
136 + lib_def.define(
137 + "zentorch_dynamic_qlinear(Tensor input, Tensor weight, "
138 + "Tensor weight_scales, Tensor? bias=None, *, "
139 + "str zentorch_op_name='zentorch::zentorch_dynamic_qlinear') -> Tensor"
140 + )
141 + lib_def.define(
142 + "zentorch_fused_moe(Tensor(a!) output, Tensor input, Tensor w13, "
143 + "Tensor w2, Tensor? w13_bias, Tensor? w2_bias, Tensor topk_weights, "
144 + "Tensor topk_id, bool skip_weighted, str act, "
145 + "Tensor? w13_scales=None, Tensor? w2_scales=None, *, "
146 + "str zentorch_op_name='zentorch::zentorch_fused_moe') -> ()"
147 + )
148 +
149 + lib_impl = torch.library.Library("zentorch", "IMPL", "CPU")
150 + lib_impl.impl("zentorch_woq_repack_weight", _repack_s4)
151 + lib_impl.impl("zentorch_dynamic_qlinear", _dynamic_qlinear)
152 + lib_impl.impl("zentorch_fused_moe", _fused_moe)
153 +
154 + yield calls
155 +
156 + lib_impl._destroy()
157 + lib_def._destroy()
158 +
159 +
160 +def _make_config(
161 + act_type: torch.dtype = torch.bfloat16,
162 + weight_type=scalar_types.uint4b8,
163 + zero_points: bool = False,
164 + group_size: int = GROUP_SIZE,
165 +) -> MPLinearLayerConfig:
166 + return MPLinearLayerConfig(
167 + full_weight_shape=(IN_FEATURES, OUT_FEATURES),
168 + partition_weight_shape=(IN_FEATURES, OUT_FEATURES),
169 + weight_type=weight_type,
170 + act_type=act_type,
171 + group_size=group_size,
172 + zero_points=zero_points,
173 + )
174 +
175 +
176 +def _make_layer(group_size: int = GROUP_SIZE) -> tuple[torch.nn.Module, torch.Tensor]:
177 + """Build a layer holding a CT-packed int4 weight, plus its dequantized form."""
178 + torch.manual_seed(0)
179 + weight = torch.randn(OUT_FEATURES, IN_FEATURES, dtype=torch.bfloat16)
180 + quantized, scale = _quantize_per_group(weight, group_size)
181 +
182 + # compressed-tensors stores [N, K//8] int32, packed along the input dim.
183 + packed = pack_to_int32(quantized, 4, packed_dim=1)
184 +
185 + layer = torch.nn.Module()
186 + layer.weight_packed = torch.nn.Parameter(packed, requires_grad=False)
187 + layer.weight_packed.input_dim = 1
188 + layer.weight_packed.packed_dim = 1
189 + layer.weight_scale = torch.nn.Parameter(scale, requires_grad=False)
190 + layer.weight_zero_point = None
191 + return layer, _dequantize_per_group(quantized, scale, group_size)
192 +
193 +
194 +def _make_kernel(config: MPLinearLayerConfig) -> ZentorchWNA16LinearKernel:
195 + return ZentorchWNA16LinearKernel(
196 + config, "weight_packed", "weight_scale", "weight_zero_point"
197 + )
198 +
199 +
200 +# ---------------------------------------------------------------------------
201 +# Dense DA8W4
202 +# ---------------------------------------------------------------------------
203 +
204 +
205 +def test_da8w4_eligible_for_symmetric_bf16_layer(mock_zentorch_ops, monkeypatch):
206 + monkeypatch.setattr(
207 + "vllm.model_executor.kernels.linear.mixed_precision.zentorch."
208 + "current_platform.is_zen_cpu",
209 + lambda: True,
210 + )
211 + layer, _ = _make_layer()
212 + kernel = _make_kernel(_make_config())
213 + assert kernel._maybe_process_da8w4_weights(layer)
214 +
215 +
216 +@pytest.mark.parametrize(
217 + "config_kwargs,reason",
218 + [
219 + ({"act_type": torch.float32}, "f32 activations are rejected"),
220 + (
221 + {"weight_type": scalar_types.uint4, "zero_points": True},
222 + "asymmetric int4 is unsupported",
223 + ),
224 + ],
225 +)
226 +def test_da8w4_not_eligible(mock_zentorch_ops, monkeypatch, config_kwargs, reason):
227 + monkeypatch.setattr(
228 + "vllm.model_executor.kernels.linear.mixed_precision.zentorch."
229 + "current_platform.is_zen_cpu",
230 + lambda: True,
231 + )
232 + layer, _ = _make_layer()
233 + kernel = _make_kernel(_make_config(**config_kwargs))
234 + assert not kernel._maybe_process_da8w4_weights(layer), reason
235 +
236 +
237 +def test_da8w4_not_eligible_when_env_disabled(mock_zentorch_ops, monkeypatch):
238 + monkeypatch.setattr(
239 + "vllm.model_executor.kernels.linear.mixed_precision.zentorch."
240 + "current_platform.is_zen_cpu",
241 + lambda: True,
242 + )
243 + monkeypatch.setattr(
244 + "vllm.model_executor.kernels.linear.mixed_precision.zentorch."
245 + "envs.VLLM_CPU_INT4_W4A8",
246 + False,
247 + )
248 + layer, _ = _make_layer()
249 + kernel = _make_kernel(_make_config())
250 + assert not kernel._maybe_process_da8w4_weights(layer)
251 +
252 +
253 +def test_da8w4_process_weights_layout(mock_zentorch_ops, monkeypatch):
254 + monkeypatch.setattr(
255 + "vllm.model_executor.kernels.linear.mixed_precision.zentorch."
256 + "current_platform.is_zen_cpu",
257 + lambda: True,
258 + )
259 + layer, _ = _make_layer()
260 + kernel = _make_kernel(_make_config())
261 + kernel.process_weights_after_loading(layer)
262 +
263 + assert layer._zentorch_da8w4
264 + assert layer._zentorch_kind == "compressed_tensors_w4a8_da8w4"
265 + # Packed s4 holds two nibbles per byte, and scales transpose to {G, N}.
266 + assert layer._zentorch_da8w4_packed.dtype == torch.int8
267 + assert layer._zentorch_da8w4_packed.shape == (OUT_FEATURES, IN_FEATURES // 2)
268 + assert layer._zentorch_da8w4_scale.dtype == torch.bfloat16
269 + assert layer._zentorch_da8w4_scale.shape == (
270 + IN_FEATURES // GROUP_SIZE,
271 + OUT_FEATURES,
272 + )
273 + assert layer._zentorch_da8w4_packed.is_contiguous()
274 + assert layer._zentorch_da8w4_scale.is_contiguous()
275 + # The checkpoint parameters are released once repacked.
276 + assert layer.weight_packed.numel() == 0
277 + assert layer.weight_scale.numel() == 0
278 +
279 +
280 +def test_da8w4_apply_weights_matches_dequantized_reference(
281 + mock_zentorch_ops, monkeypatch
282 +):
283 + monkeypatch.setattr(
284 + "vllm.model_executor.kernels.linear.mixed_precision.zentorch."
285 + "current_platform.is_zen_cpu",
286 + lambda: True,
287 + )
288 + layer, weight_deq = _make_layer()
289 + kernel = _make_kernel(_make_config())
290 + kernel.process_weights_after_loading(layer)
291 +
292 + x = torch.randn(4, IN_FEATURES, dtype=torch.bfloat16)
293 + bias = torch.randn(OUT_FEATURES, dtype=torch.bfloat16)
294 + out = kernel.apply_weights(layer, x, bias)
295 +
296 + expected = _dynamic_quant_matmul(x, weight_deq, bias)
297 + assert out.dtype == torch.bfloat16
298 + torch.testing.assert_close(out.float(), expected, rtol=2e-2, atol=2e-2)
299 +
300 +
301 +def test_da8w4_falls_back_to_w4a16_when_op_missing(monkeypatch):
302 + """Without zentorch_dynamic_qlinear the layer must not take the DA8W4 path."""
303 + monkeypatch.setattr(
304 + "vllm.model_executor.kernels.linear.mixed_precision.zentorch.has_zentorch_op",
305 + lambda ops: "zentorch_dynamic_qlinear" not in ops,
306 + )
307 + layer, _ = _make_layer()
308 + kernel = _make_kernel(_make_config())
309 + assert not kernel._maybe_process_da8w4_weights(layer)
310 +
311 +
312 +def test_can_implement_requires_zen_cpu(monkeypatch):
313 + monkeypatch.setattr(
314 + "vllm.model_executor.kernels.linear.mixed_precision.zentorch."
315 + "current_platform.is_zen_cpu",
316 + lambda: False,
317 + )
318 + ok, reason = ZentorchWNA16LinearKernel.can_implement(_make_config())
319 + assert not ok
320 + assert reason is not None
321 +
322 +
323 +def test_can_implement_accepts_shapes_cpuwna16_rejects(monkeypatch):
324 + """CPUWNA16's N/K % 32 rule is a oneDNN packing constraint zentorch lacks."""
325 + from vllm.model_executor.kernels.linear.mixed_precision.cpu import (
326 + CPUWNA16LinearKernel,
327 + )
328 +
329 + monkeypatch.setattr(
330 + "vllm.model_executor.kernels.linear.mixed_precision.zentorch."
331 + "current_platform.is_zen_cpu",
332 + lambda: True,
333 + )
334 + monkeypatch.setattr(
335 + "vllm.model_executor.kernels.linear.mixed_precision.cpu."
336 + "current_platform.is_cpu",
337 + lambda: True,
338 + )
339 + monkeypatch.setattr(
340 + "vllm.model_executor.kernels.linear.mixed_precision.zentorch.has_zentorch_op",
341 + lambda ops: True,
342 + )
343 +
344 + unaligned = dataclasses.replace(
345 + _make_config(), partition_weight_shape=(IN_FEATURES, OUT_FEATURES - 16)
346 + )
347 + assert unaligned.partition_weight_shape[1] % 32 != 0
348 +
349 + ok, _ = CPUWNA16LinearKernel.can_implement(unaligned)
350 + assert not ok, "the parent is expected to turn this shape away"
351 +
352 + ok, reason = ZentorchWNA16LinearKernel.can_implement(unaligned)
353 + assert ok, reason
354 +
355 +
356 +@pytest.mark.parametrize(
357 + "group_size,expected_ok",
358 + [
359 + (-1, True), # per-channel, one group spanning K
360 + (GROUP_SIZE, True),
361 + (IN_FEATURES, True),
362 + (6, False), # even, so the parent's multiple-of-2 rule would allow it
363 + (0, False),
364 + ],
365 +)
366 +def test_can_implement_group_size_must_divide_k(monkeypatch, group_size, expected_ok):
367 + """ZenDNN WOQ groups along K, so a group that straddles the end is unusable."""
368 + monkeypatch.setattr(
369 + "vllm.model_executor.kernels.linear.mixed_precision.zentorch."
370 + "current_platform.is_zen_cpu",
371 + lambda: True,
372 + )
373 + monkeypatch.setattr(
374 + "vllm.model_executor.kernels.linear.mixed_precision.zentorch.has_zentorch_op",
375 + lambda ops: True,
376 + )
377 +
378 + ok, reason = ZentorchWNA16LinearKernel.can_implement(
379 + _make_config(group_size=group_size)
380 + )
381 + assert ok is expected_ok, reason
382 + if not expected_ok:
383 + assert reason is not None
384 + assert "must divide input size" in reason
385 +
386 +
387 +# ---------------------------------------------------------------------------
388 +# MoE DA8W4
389 +# ---------------------------------------------------------------------------
390 +
391 +NUM_EXPERTS = 4
392 +HIDDEN = 256
393 +INTERMEDIATE = 128
394 +
395 +
396 +def _make_moe_weights(group_size: int = 32):
397 + """Build CT-layout MoE weights: w13 [E, 2I, H//8], w2 [E, H, I//8]."""
398 + torch.manual_seed(0)
399 + w13_q, w13_s, w2_q, w2_s = [], [], [], []
400 + for _ in range(NUM_EXPERTS):
401 + w13 = torch.randn(2 * INTERMEDIATE, HIDDEN, dtype=torch.bfloat16)
402 + q13, s13 = _quantize_per_group(w13, group_size)
403 + # CT is N-first and packs along the input dim: [N, K//8] per expert.
404 + w13_q.append(pack_to_int32(q13, 4, packed_dim=1))
405 + w13_s.append(s13)
406 +
407 + w2 = torch.randn(HIDDEN, INTERMEDIATE, dtype=torch.bfloat16)
408 + q2, s2 = _quantize_per_group(w2, group_size)
409 + w2_q.append(pack_to_int32(q2, 4, packed_dim=1))
410 + w2_s.append(s2)
411 + return (
412 + torch.stack(w13_q),
413 + torch.stack(w2_q),
414 + torch.stack(w13_s),
415 + torch.stack(w2_s),
416 + )
417 +
418 +
419 +def test_zen_cpu_first_in_cpu_backend_priority(monkeypatch):
420 + monkeypatch.setattr(int_wna16.current_platform, "is_cpu", lambda: True)
421 + backends = int_wna16._get_priority_backends()
422 + assert backends[0] == WNA16MoEBackend.ZEN_CPU
423 + # The generic CPU backend stays available as a fallback.
424 + assert WNA16MoEBackend.CPU in backends
425 +
426 +
427 +def test_zen_cpu_process_weights_layout(mock_zentorch_ops):
428 + group_size = 32
429 + w13, w2, w13_scale, w2_scale = _make_moe_weights(group_size)
430 + assert w13.shape == (NUM_EXPERTS, 2 * INTERMEDIATE, HIDDEN // 8)
431 + assert w2.shape == (NUM_EXPERTS, HIDDEN, INTERMEDIATE // 8)
432 +
433 + converted = int_wna16._process_weights_zen_cpu(w13, w2, w13_scale, w2_scale)
434 + w13_out, w2_out, w13_s_out, w2_s_out = converted[:4]
435 +
436 + # zentorch consumes [E, N, K/8] packed s4 with per-group [E, G, N] scales.
437 + assert w13_out.shape == (NUM_EXPERTS, 2 * INTERMEDIATE, HIDDEN // 8)
438 + assert w2_out.shape == (NUM_EXPERTS, HIDDEN, INTERMEDIATE // 8)
439 + assert w13_out.dtype == w2_out.dtype == torch.int32
440 + assert w13_s_out.shape == (NUM_EXPERTS, HIDDEN // group_size, 2 * INTERMEDIATE)
441 + assert w2_s_out.shape == (NUM_EXPERTS, INTERMEDIATE // group_size, HIDDEN)
442 + assert w13_s_out.dtype == w2_s_out.dtype == torch.bfloat16
443 + # Symmetric checkpoints carry no zero points into the kernel.
444 + assert converted[8] is None and converted[9] is None
445 +
446 +
447 +def test_zen_cpu_repack_is_value_exact(mock_zentorch_ops):
448 + """The repacked weight must dequantize back to the checkpoint values."""
449 + group_size = 32
450 + w13, w2, w13_scale, w2_scale = _make_moe_weights(group_size)
451 + w13_out = int_wna16._process_weights_zen_cpu(w13, w2, w13_scale, w2_scale)[0]
452 +
453 + expected = _import_unpack_from_int32()(
454 + w13,
455 + 4,
456 + torch.Size([NUM_EXPERTS, 2 * INTERMEDIATE, HIDDEN]),
457 + packed_dim=1,
458 + )
459 + for expert in range(NUM_EXPERTS):
460 + torch.testing.assert_close(
461 + _unpack_s4(w13_out[expert], HIDDEN),
462 + expected[expert].float(),
463 + )
464 +
465 +
466 +@pytest.mark.parametrize(
467 + "group_size,expected_ok",
468 + [(32, True), (128, True), (-1, False), (2, False)],
469 +)
470 +def test_zen_cpu_group_size_gating(group_size, expected_ok):
471 + quant_config = type("Args", (), {"group_size": group_size})()
472 + reason = int_wna16._backend_incompatibility_reason(
473 + WNA16MoEBackend.ZEN_CPU,
474 + moe_config=None,
475 + quant_config=quant_config,
476 + may_have_zp=False,
477 + may_have_bias=False,
478 + allow_tile_padding=True,
479 + )
480 + assert (reason is None) == expected_ok
481 +
482 +
483 +@pytest.mark.parametrize(
484 + "may_have_zp,may_have_bias,expected_ok",
485 + [(True, False, False), (False, True, True), (True, True, False)],
486 +)
487 +def test_zen_cpu_rejects_zero_points_but_takes_bias(
488 + may_have_zp, may_have_bias, expected_ok
489 +):
490 + quant_config = type("Args", (), {"group_size": 128})()
491 + reason = int_wna16._backend_incompatibility_reason(
492 + WNA16MoEBackend.ZEN_CPU,
493 + moe_config=None,
494 + quant_config=quant_config,
495 + may_have_zp=may_have_zp,
496 + may_have_bias=may_have_bias,
497 + allow_tile_padding=True,
498 + )
499 + assert (reason is None) == expected_ok
500 +
501 +
502 +def test_zen_cpu_rejects_moe_wna16_layout():
503 + from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
504 +
505 + quant_config = MoeWNA16Config(
506 + linear_quant_method="gptq",
507 + weight_bits=4,
508 + group_size=128,
509 + has_zp=False,
510 + lm_head_quantized=False,
511 + modules_to_not_convert=None,
512 + full_config={},
513 + )
514 + reason = int_wna16._backend_incompatibility_reason(
515 + WNA16MoEBackend.ZEN_CPU,
516 + moe_config=None,
517 + quant_config=quant_config,
518 + may_have_zp=False,
519 + may_have_bias=False,
520 + allow_tile_padding=True,
521 + )
522 + assert reason is not None
523 +
524 +
525 +def test_zen_cpu_disabled_by_env(monkeypatch):
526 + monkeypatch.setattr(int_wna16.envs, "VLLM_CPU_INT4_W4A8", False)
527 + quant_config = type("Args", (), {"group_size": 128})()
528 + reason = int_wna16._backend_incompatibility_reason(
529 + WNA16MoEBackend.ZEN_CPU,
530 + moe_config=None,
531 + quant_config=quant_config,
532 + may_have_zp=False,
533 + may_have_bias=False,
534 + allow_tile_padding=True,
535 + )
536 + assert reason is not None
537 +
538 +
539 +def test_zen_cpu_backend_maps_to_experts_class():
540 + from vllm.model_executor.layers.fused_moe.experts.zentorch_moe import (
541 + ZentorchExpertsInt4,
542 + )
543 +
544 + assert int_wna16.backend_to_kernel_cls(WNA16MoEBackend.ZEN_CPU) == [
545 + ZentorchExpertsInt4
546 + ]
547 +
548 +
549 +def test_zen_experts_support_predicates():
550 + from vllm.model_executor.kernels.linear.zentorch_utils import (
551 + _ZENTORCH_MOE_ACTIVATIONS,
552 + )
553 + from vllm.model_executor.layers.fused_moe.activation import MoEActivation
554 + from vllm.model_executor.layers.fused_moe.experts.zentorch_moe import (
555 + ZentorchExpertsInt4,
556 + )
557 + from vllm.model_executor.layers.quantization.utils.quant_utils import (
558 + kInt4Static,
559 + kInt4Static32,
560 + kInt8StaticChannelSym,
561 + )
562 +
563 + assert ZentorchExpertsInt4._supports_quant_scheme(kInt4Static, None)
564 + assert ZentorchExpertsInt4._supports_quant_scheme(kInt4Static32, None)
565 + assert not ZentorchExpertsInt4._supports_quant_scheme(kInt8StaticChannelSym, None)
566 + for act in MoEActivation:
567 + expected = act.value in _ZENTORCH_MOE_ACTIVATIONS
568 + assert ZentorchExpertsInt4._supports_activation(act) == expected
569 + assert ZentorchExpertsInt4.requires_interleaved_w13
570 + assert not ZentorchExpertsInt4._supports_no_act_and_mul()
571 + assert ZentorchExpertsInt4._supports_parallel_config(
572 + type("Par", (), {"use_ep": False})()
573 + )
574 + assert not ZentorchExpertsInt4._supports_parallel_config(
575 + type("Par", (), {"use_ep": True})()
576 + )
577 +
578 +
579 +def test_zen_experts_take_custom_routing():
580 + """Models with their own router (gemma-4) reach select_experts through the
581 + config captured off the layer, since apply() cannot be handed the callable."""
582 + from types import SimpleNamespace
583 +
584 + from tests.kernels.moe.test_cpu_fused_moe import _StubMoELayer
585 + from vllm.model_executor.layers.fused_moe.activation import MoEActivation
586 + from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
587 + from vllm.model_executor.layers.fused_moe.experts.zentorch_moe import (
588 + ZentorchExpertsInt4,
589 + )
590 +
591 + assert ZentorchExpertsInt4._supports_routing_method(
592 + RoutingMethodType.Custom, None, None
593 + )
594 +
595 + def routing_fn(**kwargs):
596 + raise AssertionError("not called")
597 +
598 + layer = _StubMoELayer(
599 + torch.zeros(NUM_EXPERTS, 2 * INTERMEDIATE, HIDDEN),
600 + torch.zeros(NUM_EXPERTS, HIDDEN, INTERMEDIATE),
601 + MoEActivation.SILU,
602 + )
603 + layer.renormalize = True
604 + layer.custom_routing_function = routing_fn
605 +
606 + experts = SimpleNamespace(renormalize=False, custom_routing_function=None)
607 + ZentorchExpertsInt4.process_weights_after_loading(experts, layer)
608 +
609 + assert experts.custom_routing_function is routing_fn
610 + assert experts.renormalize is True
611 +
612 +
613 +def _swigluoai_perm(activation, experts_cls=None):
614 + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
615 + swigluoai_w13_interleave_perm,
616 + )
617 + from vllm.model_executor.layers.fused_moe.experts.zentorch_moe import (
618 + ZentorchExpertsInt4,
619 + )
620 +
621 + return swigluoai_w13_interleave_perm(
622 + experts_cls or ZentorchExpertsInt4,
623 + activation,
624 + 2 * INTERMEDIATE,
625 + torch.device("cpu"),
626 + )
627 +
628 +
629 +def test_swigluoai_perm_only_for_zen_swigluoai():
630 + from vllm.model_executor.layers.fused_moe.activation import MoEActivation
631 +
632 + assert _swigluoai_perm(MoEActivation.SWIGLUOAI) is not None
633 + # Other activations keep the half-split layout the loader leaves.
634 + assert _swigluoai_perm(MoEActivation.SILU) is None
635 + # object stands in for any kernel that never set requires_interleaved_w13.
636 + assert _swigluoai_perm(MoEActivation.SWIGLUOAI, experts_cls=object) is None
637 +
638 +
639 +def test_swigluoai_perm_interleaves_weights_scales_and_bias(mock_zentorch_ops):
640 + """The permutation is a gather on w13's output-channel axis, which is dim 1
641 + for the N-first packed weights and their group scales."""
642 + from vllm.model_executor.layers.fused_moe.activation import MoEActivation
643 +
644 + group_size = 32
645 + w13, w2, w13_scale, w2_scale = _make_moe_weights(group_size)
646 + bias = torch.randn(NUM_EXPERTS, 2 * INTERMEDIATE, dtype=torch.bfloat16)
647 + perm = _swigluoai_perm(MoEActivation.SWIGLUOAI)
648 +
649 + baseline = int_wna16._process_weights_zen_cpu(w13, w2, w13_scale, w2_scale)[0]
650 + permuted = int_wna16._process_weights_zen_cpu(
651 + w13[:, perm].contiguous(), w2, w13_scale[:, perm].contiguous(), w2_scale
652 + )[0]
653 +
654 + for expert in range(NUM_EXPERTS):
655 + before = _unpack_s4(baseline[expert], HIDDEN)
656 + after = _unpack_s4(permuted[expert], HIDDEN)
657 + torch.testing.assert_close(after[0::2], before[:INTERMEDIATE])
658 + torch.testing.assert_close(after[1::2], before[INTERMEDIATE:])
659 +
660 + torch.testing.assert_close(bias[:, perm][:, 0::2], bias[:, :INTERMEDIATE])
661 + torch.testing.assert_close(bias[:, perm][:, 1::2], bias[:, INTERMEDIATE:])
vllm/envs.py
+2 -1
@@ -912,7 +912,8 @@ environment_variables: dict[str, Callable[[], Any]] = {
912 "VLLM_ZENTORCH_WEIGHT_PREPACK": lambda: bool(
913 int(os.getenv("VLLM_ZENTORCH_WEIGHT_PREPACK", "1"))
914 ),
915 - # (CPU backend only) whether to use SGLang INT4 W4A8 kernels for AWQ.
915 + # (CPU backend only) whether to use SGLang INT4 W4A8 kernels for AWQ, and
916 + # on Zen CPUs whether to serve int4 checkpoints as DA8W4 rather than W4A16.
917 "VLLM_CPU_INT4_W4A8": lambda: bool(int(os.getenv("VLLM_CPU_INT4_W4A8", "1"))),
918 # If the env var is set, Ray Compiled Graph uses the specified
919 # channel type to communicate between workers belonging to
vllm/model_executor/kernels/linear/mixed_precision/zentorch.py
+113 -8
@@ -1,6 +1,6 @@
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.
3 +"""Zentorch int4 weight-quantized linear kernels 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
@@ -9,12 +9,13 @@ falls through to the next kernel in ``_POSSIBLE_KERNELS[PlatformEnum.CPU]``.
9
10 import torch
11
12 +from vllm import envs
13 from vllm.logger import init_logger
14 from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op
15 from vllm.platforms import current_platform
16 from vllm.scalar_type import scalar_types
17
17 -from .cpu import CPUWNA16LinearKernel
18 +from .cpu import _CPUWNA16_SUPPORTED_QUANT_TYPES, CPUWNA16LinearKernel
19 from .MPLinearKernel import MPLinearLayerConfig
20
21 logger = init_logger(__name__)
@@ -34,17 +35,35 @@ def _import_unpack_from_int32():
35
36
37 class ZentorchWNA16LinearKernel(CPUWNA16LinearKernel):
37 - """W4A16 GPTQ kernel backed by ``torch.ops.zentorch.zentorch_woq_linear``."""
38 + """Int4 kernel backed by ``zentorch_dynamic_qlinear`` (DA8W4) or
39 + ``zentorch_woq_linear`` (W4A16)."""
40
41 @classmethod
42 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 -
43 if not current_platform.is_zen_cpu():
44 return False, "ZentorchWNA16 requires an AMD Zen CPU."
45
46 + # Not super(): its N/K multiple-of-32 rule is a oneDNN packing
47 + # constraint the zentorch repack does not share.
48 + if c.weight_type not in _CPUWNA16_SUPPORTED_QUANT_TYPES:
49 + return (
50 + False,
51 + f"Quant type ({c.weight_type}) not supported by "
52 + "CPUWNA16, supported types are: "
53 + f"{_CPUWNA16_SUPPORTED_QUANT_TYPES}",
54 + )
55 +
56 + # -1 is per-channel, one group spanning K. Anything else has to divide
57 + # K, which is stricter than the parent's multiple-of-2 rule.
58 + if c.group_size != -1:
59 + in_features = c.partition_weight_shape[0]
60 + if c.group_size <= 0 or in_features % c.group_size != 0:
61 + return (
62 + False,
63 + f"Group size ({c.group_size}) must divide input size "
64 + f"({in_features})",
65 + )
66 +
67 if not has_zentorch_op(["zentorch_woq_repack_weight", "zentorch_woq_linear"]):
68 return (
69 False,
@@ -89,8 +108,81 @@ class ZentorchWNA16LinearKernel(CPUWNA16LinearKernel):
108 num_groups = weight_scale.shape[1]
109 return num_groups > 0 and in_features % num_groups == 0
110
111 + def _maybe_process_da8w4_weights(self, layer: torch.nn.Module) -> bool:
112 + """Repack ``layer`` into the DA8W4 layout, or leave it untouched.
113 +
114 + Returns False without mutating the layer when it is not eligible, so
115 + the caller falls through to the W4A16 paths.
116 + """
117 + if not envs.VLLM_CPU_INT4_W4A8:
118 + return False
119 +
120 + if not has_zentorch_op(["zentorch_dynamic_qlinear"]):
121 + return False
122 +
123 + # DA8W4 is symmetric-only, and the kernel rejects f32 activations.
124 + if self.config.zero_points or self.config.weight_type == scalar_types.uint4:
125 + return False
126 + if self.config.act_type != torch.bfloat16:
127 + return False
128 +
129 + if not self._zentorch_woq_eligible(layer):
130 + return False
131 +
132 + weight_q = getattr(layer, self.w_q_name)
133 + weight_s = getattr(layer, self.w_s_name)
134 + weight_packed = weight_q.data if hasattr(weight_q, "data") else weight_q
135 + weight_scale = weight_s.data if hasattr(weight_s, "data") else weight_s
136 +
137 + bits = self.config.weight_type.mantissa
138 + pack_factor = torch.iinfo(weight_packed.dtype).bits // bits
139 + out_features, num_groups = weight_scale.shape[0], weight_scale.shape[1]
140 + in_features = weight_packed.shape[1] * pack_factor
141 +
142 + # AOCL sym_quant requires K/G to be a multiple of 4; K must be even to
143 + # pack two s4 values per byte.
144 + if (in_features // num_groups) % 4 != 0 or in_features % 2 != 0:
145 + return False
146 +
147 + if self.w_zp_name is not None:
148 + setattr(layer, self.w_zp_name, None)
149 +
150 + unpack_from_int32 = _import_unpack_from_int32()
151 +
152 + weight_unpacked = unpack_from_int32(
153 + weight_packed,
154 + bits,
155 + torch.Size([out_features, in_features]),
156 + packed_dim=weight_q.packed_dim,
157 + )
158 + # The WOQ repack packs 8 int4 per int32, which on little-endian is the
159 + # same byte stream as s4 [N, K/2]; the kernel takes either view.
160 + layer._zentorch_da8w4_packed = (
161 + torch.ops.zentorch.zentorch_woq_repack_weight.default(
162 + weight_unpacked.to(torch.int8).contiguous()
163 + ).view(torch.int8)
164 + )
165 + # CT stores scales as [N, G]; DA8W4 wants per-group {G, N}.
166 + layer._zentorch_da8w4_scale = weight_scale.t().to(torch.bfloat16).contiguous()
167 +
168 + for param_name in (self.w_q_name, self.w_s_name):
169 + param = getattr(layer, param_name, None)
170 + if param is not None and hasattr(param, "data"):
171 + param.data = torch.empty(0)
172 +
173 + layer._zentorch_kind = "compressed_tensors_w4a8_da8w4"
174 + layer._zentorch_da8w4 = True
175 + layer._zentorch_processed_weights = True
176 + logger.info_once(
177 + "[zen_cpu] Using zentorch_dynamic_qlinear for DA8W4 (W4A8) "
178 + "(weight_type=%s, group_size=%d)",
179 + self.config.weight_type,
180 + in_features // num_groups,
181 + )
182 + return True
183 +
184 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
93 - """Repack CT GPTQ weights into the zentorch WOQ layout.
185 + """Repack CT GPTQ weights into the zentorch DA8W4 or WOQ layout.
186
187 Falls back to ``CPUWNA16LinearKernel.process_weights_after_loading``
188 via ``super()`` when the layer doesn't satisfy
@@ -101,6 +193,9 @@ class ZentorchWNA16LinearKernel(CPUWNA16LinearKernel):
193 if getattr(layer, "_zentorch_processed_weights", False):
194 return
195
196 + if self._maybe_process_da8w4_weights(layer):
197 + return
198 +
199 if not self._zentorch_woq_eligible(layer):
200 logger.info_once(
201 "[zen_cpu] ZentorchWNA16 fast path not eligible for this "
@@ -186,6 +281,16 @@ class ZentorchWNA16LinearKernel(CPUWNA16LinearKernel):
281 x: torch.Tensor,
282 bias: torch.Tensor | None = None,
283 ) -> torch.Tensor:
284 + if getattr(layer, "_zentorch_da8w4", False):
285 + # The kernel reads bias through a raw pointer, so it must be
286 + # contiguous; the weight and scales already are.
287 + return torch.ops.zentorch.zentorch_dynamic_qlinear.default(
288 + x if x.dtype == torch.bfloat16 else x.to(torch.bfloat16),
289 + layer._zentorch_da8w4_packed,
290 + layer._zentorch_da8w4_scale,
291 + bias.contiguous() if bias is not None else None,
292 + )
293 +
294 if getattr(layer, "_zentorch_processed_weights", False):
295 return torch.ops.zentorch.zentorch_woq_linear.default(
296 x,
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+26
@@ -51,6 +51,32 @@ from vllm.platforms import CpuArchEnum, current_platform
51 from vllm.utils.math_utils import round_up
52
53 logger = init_logger(__name__)
54 +# ===========================================================================
55 +# Weight layout
56 +# ===========================================================================
57 +
58 +
59 +def swigluoai_w13_interleave_perm(
60 + experts_cls: type[mk.FusedMoEExperts],
61 + activation: MoEActivation,
62 + two_i: int,
63 + device: torch.device,
64 +) -> torch.Tensor | None:
65 + """Index taking a half-split w13 ``[gate | up]`` to the interleaved order
66 + ZenDNN's ``swiglu_oai_mul`` expects, or None when no reorder is needed.
67 + Callers apply it along whichever axis holds their output channels."""
68 + if not getattr(experts_cls, "requires_interleaved_w13", False):
69 + return None
70 + if activation != MoEActivation.SWIGLUOAI:
71 + return None
72 +
73 + i = two_i // 2
74 + return torch.stack(
75 + [torch.arange(0, i, device=device), torch.arange(i, two_i, device=device)],
76 + dim=1,
77 + ).flatten()
78 +
79 +
80 # ===========================================================================
81 # Unquantized (BF16/FP16/FP32) MoE
82 # ===========================================================================
vllm/model_executor/layers/fused_moe/experts/zentorch_moe.py new
+204
@@ -0,0 +1,204 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Zentorch fused MoE experts for AMD Zen CPUs."""
4 +
5 +from collections.abc import Callable
6 +
7 +import torch
8 +
9 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
10 +from vllm.model_executor.kernels.linear.zentorch_utils import (
11 + _ZENTORCH_MOE_ACTIVATIONS,
12 + _moe_activation_to_str,
13 + has_zentorch_op,
14 +)
15 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
16 +from vllm.model_executor.layers.fused_moe.config import (
17 + FusedMoEConfig,
18 + FusedMoEParallelConfig,
19 + FusedMoEQuantConfig,
20 + RoutingMethodType,
21 +)
22 +from vllm.model_executor.layers.fused_moe.router.cpu_router import select_experts
23 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
24 + QuantKey,
25 + kInt4Static,
26 + kInt4Static32,
27 +)
28 +
29 +
30 +class ZentorchExpertsInt4(mk.FusedMoEExpertsMonolithic):
31 + """DA8W4 (W4A8) group-quantized monolithic MoE experts."""
32 +
33 + # swiglu_oai_mul reads gate/up interleaved, not in the half-split order the
34 + # weight loader leaves behind.
35 + requires_interleaved_w13 = True
36 +
37 + def __init__(
38 + self,
39 + moe_config: FusedMoEConfig,
40 + quant_config: FusedMoEQuantConfig,
41 + ):
42 + super().__init__(moe_config, quant_config)
43 + # Defaults hold until process_weights_after_loading sees the layer.
44 + self.renormalize = moe_config.routing_method in (
45 + RoutingMethodType.Renormalize,
46 + RoutingMethodType.RenormalizeNaive,
47 + )
48 + self.custom_routing_function: Callable | None = None
49 +
50 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
51 + """Capture the router config that the monolithic apply() signature
52 + cannot carry; Custom routing has no other source for it."""
53 + self.renormalize = layer.renormalize
54 + self.custom_routing_function = layer.custom_routing_function
55 +
56 + @property
57 + def expects_unquantized_inputs(self) -> bool:
58 + return True
59 +
60 + @staticmethod
61 + def activation_format() -> mk.FusedMoEActivationFormat:
62 + return mk.FusedMoEActivationFormat.Standard
63 +
64 + @staticmethod
65 + def _supports_current_device() -> bool:
66 + return has_zentorch_op(["zentorch_fused_moe", "zentorch_woq_repack_weight"])
67 +
68 + @staticmethod
69 + def is_supported_config(
70 + cls: type[mk.FusedMoEExperts],
71 + moe_config: FusedMoEConfig,
72 + weight_key: QuantKey | None,
73 + activation_key: QuantKey | None,
74 + activation_format: mk.FusedMoEActivationFormat,
75 + ) -> tuple[bool, str | None]:
76 + supported, reason = mk.FusedMoEExperts.is_supported_config(
77 + cls, moe_config, weight_key, activation_key, activation_format
78 + )
79 + if not supported:
80 + return supported, reason
81 + if moe_config.in_dtype != torch.bfloat16:
82 + return False, "kernel requires bfloat16 activations"
83 + # The grouped GEMM needs tokens spread over at least two experts, which
84 + # top_k=1 cannot satisfy once decode drops to a single token.
85 + if moe_config.experts_per_token < 2:
86 + return False, "kernel requires experts_per_token >= 2"
87 + return True, None
88 +
89 + @staticmethod
90 + def _supports_no_act_and_mul() -> bool:
91 + return False
92 +
93 + @staticmethod
94 + def _supports_activation(activation: MoEActivation) -> bool:
95 + return _moe_activation_to_str(activation) in _ZENTORCH_MOE_ACTIVATIONS
96 +
97 + @staticmethod
98 + def _supports_parallel_config(
99 + moe_parallel_config: FusedMoEParallelConfig,
100 + ) -> bool:
101 + # apply() passes select_experts' global ids straight to the kernel, so
102 + # every expert has to be resident on the rank.
103 + return not moe_parallel_config.use_ep
104 +
105 + @staticmethod
106 + def _supports_quant_scheme(
107 + weight_key: QuantKey | None,
108 + activation_key: QuantKey | None,
109 + ) -> bool:
110 + return (weight_key, activation_key) in [
111 + (kInt4Static, None),
112 + (kInt4Static32, None),
113 + ]
114 +
115 + @staticmethod
116 + def _supports_routing_method(
117 + routing_method: RoutingMethodType,
118 + weight_key: QuantKey | None,
119 + activation_key: QuantKey | None,
120 + ) -> bool:
121 + return routing_method in [
122 + RoutingMethodType.Default,
123 + RoutingMethodType.Renormalize,
124 + RoutingMethodType.RenormalizeNaive,
125 + RoutingMethodType.Custom,
126 + ]
127 +
128 + @staticmethod
129 + def _supports_router_logits_dtype(
130 + router_logits_dtype: torch.dtype | None,
131 + routing_method: RoutingMethodType,
132 + ) -> bool:
133 + return True
134 +
135 + def supports_expert_map(self) -> bool:
136 + return False
137 +
138 + def apply(
139 + self,
140 + hidden_states: torch.Tensor,
141 + w1: torch.Tensor,
142 + w2: torch.Tensor,
143 + router_logits: torch.Tensor,
144 + activation: MoEActivation,
145 + global_num_experts: int,
146 + expert_map: torch.Tensor | None,
147 + a1q_scale: torch.Tensor | None,
148 + apply_router_weight_on_input: bool,
149 + # grouped topk + fused topk bias parameters
150 + num_expert_group: int | None = None,
151 + e_score_correction_bias: torch.Tensor | None = None,
152 + routed_scaling_factor: float | None = None,
153 + topk_group: int | None = None,
154 + ) -> torch.Tensor:
155 + if apply_router_weight_on_input:
156 + raise NotImplementedError(
157 + "ZentorchExpertsInt4 does not support "
158 + "apply_router_weight_on_input=True."
159 + )
160 +
161 + if (
162 + self.moe_config.routing_method == RoutingMethodType.Custom
163 + and self.custom_routing_function is None
164 + ):
165 + raise RuntimeError(
166 + "ZentorchExpertsInt4 needs the model's custom_routing_function "
167 + "for RoutingMethodType.Custom."
168 + )
169 +
170 + topk_weights, topk_ids = select_experts(
171 + hidden_states=hidden_states,
172 + router_logits=router_logits,
173 + use_grouped_topk=num_expert_group is not None,
174 + top_k=self.moe_config.experts_per_token,
175 + renormalize=self.renormalize,
176 + topk_group=topk_group,
177 + num_expert_group=num_expert_group,
178 + custom_routing_function=self.custom_routing_function,
179 + scoring_func="softmax",
180 + routed_scaling_factor=(
181 + routed_scaling_factor if routed_scaling_factor is not None else 1.0
182 + ),
183 + e_score_correction_bias=e_score_correction_bias,
184 + )
185 +
186 + output = torch.empty_like(hidden_states)
187 + torch.ops.zentorch.zentorch_fused_moe(
188 + output,
189 + hidden_states,
190 + w1,
191 + w2,
192 + self.w1_bias,
193 + self.w2_bias,
194 + topk_weights,
195 + topk_ids,
196 + False, # skip_weighted
197 + _moe_activation_to_str(activation),
198 + self.w1_scale,
199 + self.w2_scale,
200 + )
201 + return output
202 +
203 +
204 +__all__ = ["ZentorchExpertsInt4"]
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+104 -3
@@ -12,6 +12,7 @@ from compressed_tensors.quantization import (
12
13 import vllm._custom_ops as ops
14 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
15 +from vllm import envs
16 from vllm.config.kernel import MoEBackend
17 from vllm.logger import init_logger
18 from vllm.model_executor.layers.fused_moe.config import (
@@ -55,6 +56,7 @@ class WNA16MoEBackend(Enum):
56 BATCHED_MARLIN = "BATCHED_MARLIN"
57 HUMMING = "HUMMING"
58 CPU = "CPU"
59 + ZEN_CPU = "ZEN_CPU"
60 FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM"
61 TRITON = "TRITON"
62 XPU = "XPU"
@@ -98,6 +100,12 @@ def backend_to_kernel_cls(
100 )
101
102 return [CPUExpertsInt4]
103 + elif backend == WNA16MoEBackend.ZEN_CPU:
104 + from vllm.model_executor.layers.fused_moe.experts.zentorch_moe import (
105 + ZentorchExpertsInt4,
106 + )
107 +
108 + return [ZentorchExpertsInt4]
109 elif backend == WNA16MoEBackend.EMULATION:
110 from vllm.model_executor.layers.fused_moe.experts.int4_emulation_moe import (
111 Int4EmulationTritonExperts,
@@ -117,7 +125,7 @@ def backend_to_kernel_cls(
125 def _get_priority_backends() -> list[WNA16MoEBackend]:
126 """Get available backends in priority order based on platform and config."""
127 if current_platform.is_cpu():
120 - return [WNA16MoEBackend.CPU]
128 + return [WNA16MoEBackend.ZEN_CPU, WNA16MoEBackend.CPU]
129 if current_platform.is_xpu():
130 return [WNA16MoEBackend.XPU]
131
@@ -160,6 +168,27 @@ def _backend_incompatibility_reason(
168 from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
169 from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
170
171 + if backend == WNA16MoEBackend.ZEN_CPU:
172 + if not envs.VLLM_CPU_INT4_W4A8:
173 + return "VLLM_CPU_INT4_W4A8=0 disables the DA8W4 path"
174 + if may_have_zp:
175 + return "zero points are not supported"
176 + if isinstance(quant_config, MoeWNA16Config):
177 + return "the MoeWNA16 weight layout is not supported"
178 + if isinstance(quant_config, AutoGPTQConfig) and quant_config.desc_act:
179 + return "GPTQ activation ordering is not supported"
180 + if (
181 + isinstance(quant_config, QuantizationArgs)
182 + and quant_config.actorder == "group"
183 + ):
184 + return "group activation ordering is not supported"
185 + group_size = getattr(quant_config, "group_size", None)
186 + if group_size is None or group_size <= 0:
187 + return "DA8W4 requires group-quantized weights"
188 + # AOCL sym_quant requires the group size to be a multiple of 4.
189 + if group_size % 4 != 0:
190 + return f"group size {group_size} is not a multiple of 4"
191 +
192 if backend == WNA16MoEBackend.TRITON:
193 if may_have_bias:
194 return "expert bias is not supported"
@@ -403,10 +432,14 @@ def make_wna16_moe_kernel(
432 from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
433 XPUExpertsWNA16,
434 )
435 + from vllm.model_executor.layers.fused_moe.experts.zentorch_moe import (
436 + ZentorchExpertsInt4,
437 + )
438
439 # Currently, we only support TrtLlmMxint4ExpertsMonolithic, MarlinExperts,
408 - # BatchedMarlinExperts, XPUExpertsWNA16, CPUExpertsInt4, the Humming
409 - # grouped/indexed experts, and Int4EmulationTritonExperts
440 + # BatchedMarlinExperts, XPUExpertsWNA16, CPUExpertsInt4,
441 + # ZentorchExpertsInt4, the Humming grouped/indexed experts, and
442 + # Int4EmulationTritonExperts
443 allowed_experts: tuple[type[mk.FusedMoEExperts], ...] = (
444 MarlinExperts,
445 BatchedMarlinExperts,
@@ -414,6 +447,7 @@ def make_wna16_moe_kernel(
447 TrtLlmMxint4ExpertsMonolithic,
448 XPUExpertsWNA16,
449 CPUExpertsInt4,
450 + ZentorchExpertsInt4,
451 Int4EmulationTritonExperts,
452 Rdna3WNA16Experts,
453 )
@@ -1034,6 +1068,64 @@ def _process_weights_cpu(
1068 )
1069
1070
1071 +def _zen_repack_s4(packed: torch.Tensor, in_features: int) -> torch.Tensor:
1072 + """Repack CT int4 ``[E, N, K//8]`` into zentorch s4 ``[E, N, K//8]``."""
1073 + from vllm.model_executor.kernels.linear.mixed_precision.zentorch import (
1074 + _import_unpack_from_int32,
1075 + )
1076 +
1077 + num_experts, out_features, _ = packed.shape
1078 + # packed_dim indexes the 2D expert slice, which is N-first, so K is dim 1.
1079 + unpacked = _import_unpack_from_int32()(
1080 + packed,
1081 + 4,
1082 + torch.Size([num_experts, out_features, in_features]),
1083 + packed_dim=1,
1084 + )
1085 + # CT is already N-first, which is what the repack consumes.
1086 + repack_op = torch.ops.zentorch.zentorch_woq_repack_weight.default
1087 + return torch.stack([repack_op(w) for w in unpacked])
1088 +
1089 +
1090 +def _process_weights_zen_cpu(
1091 + w13: torch.Tensor,
1092 + w2: torch.Tensor,
1093 + w13_scale: torch.Tensor,
1094 + w2_scale: torch.Tensor,
1095 + w13_bias: torch.Tensor | None = None,
1096 + w2_bias: torch.Tensor | None = None,
1097 +) -> tuple[
1098 + torch.Tensor, # w13_qweight
1099 + torch.Tensor, # w2_qweight
1100 + torch.Tensor, # w13_scales
1101 + torch.Tensor, # w2_scales
1102 + torch.Tensor | None, # w13_qzeros
1103 + torch.Tensor | None, # w2_qzeros
1104 + torch.Tensor | None, # w13_input_global_scale
1105 + torch.Tensor | None, # w2_input_global_scale
1106 + torch.Tensor | None, # w13_bias
1107 + torch.Tensor | None, # w2_bias
1108 +]:
1109 + """Zen CPU INT4 DA8W4 (W4A8) weight post-processing."""
1110 + # CT WNA16 buffers are N-first, so the scales are [E, N, G].
1111 + hidden_size = w2_scale.shape[1]
1112 + intermediate_size = w13_scale.shape[1] // 2
1113 +
1114 + # The kernel reads per-group scales as [E, G, N], so N-first needs a swap.
1115 + return (
1116 + _zen_repack_s4(w13.data, hidden_size),
1117 + _zen_repack_s4(w2.data, intermediate_size),
1118 + w13_scale.data.transpose(1, 2).to(torch.bfloat16).contiguous(),
1119 + w2_scale.data.transpose(1, 2).to(torch.bfloat16).contiguous(),
1120 + None, # w13_qzeros (symmetric)
1121 + None, # w2_qzeros (symmetric)
1122 + None, # w13_input_global_scale
1123 + None, # w2_input_global_scale
1124 + w13_bias,
1125 + w2_bias,
1126 + )
1127 +
1128 +
1129 def _process_weights_xpu(
1130 layer: torch.nn.Module,
1131 quant_config: QuantizationConfig,
@@ -1584,6 +1676,15 @@ def convert_to_wna16_moe_kernel_format(
1676 w13_bias,
1677 w2_bias,
1678 )
1679 + elif backend == WNA16MoEBackend.ZEN_CPU:
1680 + return _process_weights_zen_cpu(
1681 + w13,
1682 + w2,
1683 + w13_scale,
1684 + w2_scale,
1685 + w13_bias,
1686 + w2_bias,
1687 + )
1688 elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM:
1689 return _process_weights_flashinfer(
1690 w13,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py
+66 -2
@@ -210,6 +210,27 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
210 layer.register_parameter("w2_weight_packed", w2_weight)
211 set_weight_attrs(w2_weight, extra_weight_attrs)
212
213 + # Per-expert biases, as gpt-oss carries.
214 + if self.moe.has_bias:
215 + w13_num_shards = 2 if self.moe.is_act_and_mul else 1
216 + w13_bias = torch.nn.Parameter(
217 + torch.zeros(
218 + num_experts,
219 + w13_num_shards * intermediate_size_per_partition,
220 + dtype=params_dtype,
221 + ),
222 + requires_grad=False,
223 + )
224 + layer.register_parameter("w13_bias", w13_bias)
225 + set_weight_attrs(w13_bias, extra_weight_attrs)
226 +
227 + w2_bias = torch.nn.Parameter(
228 + torch.zeros(num_experts, hidden_size, dtype=params_dtype),
229 + requires_grad=False,
230 + )
231 + layer.register_parameter("w2_bias", w2_bias)
232 + set_weight_attrs(w2_bias, extra_weight_attrs)
233 +
234 if self.strategy == "channel":
235 num_groups_w2 = num_groups_w13 = 1
236 self.group_size = -1
@@ -350,7 +371,41 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
371 routing_tables=layer._expert_routing_tables(),
372 )
373
374 + if self.wna16_backend == WNA16MoEBackend.ZEN_CPU:
375 + # Monolithic experts route inside apply(), so like the unquantized
376 + # CPU path they need the router config that only the layer carries.
377 + self.moe_kernel.fused_experts.process_weights_after_loading(layer)
378 +
379 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
380 + from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
381 + swigluoai_w13_interleave_perm,
382 + )
383 +
384 + # Reorder before the conversion below, which repacks w13 for the kernel.
385 + # The CT WNA16 buffers are N-first, so 2I is dim 1 of the packed weight
386 + # [E, 2I, K/pack] and of its group scales [E, 2I, G], as for the bias.
387 + perm = swigluoai_w13_interleave_perm(
388 + self.experts_cls,
389 + self.moe.activation,
390 + layer.w13_weight_packed.size(1),
391 + layer.w13_weight_packed.device,
392 + )
393 + if perm is not None:
394 + replace_parameter(
395 + layer,
396 + "w13_weight_packed",
397 + layer.w13_weight_packed.data[:, perm].contiguous(),
398 + )
399 + replace_parameter(
400 + layer,
401 + "w13_weight_scale",
402 + layer.w13_weight_scale.data[:, perm].contiguous(),
403 + )
404 + if self.moe.has_bias:
405 + replace_parameter(
406 + layer, "w13_bias", layer.w13_bias.data[:, perm].contiguous()
407 + )
408 +
409 # Process weights using the shared oracle infrastructure
410 converted = convert_to_wna16_moe_kernel_format(
411 backend=self.wna16_backend,
@@ -363,6 +418,8 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
418 w2_scale=layer.w2_weight_scale,
419 w13_qzeros=getattr(layer, "w13_weight_zero_point", None),
420 w2_qzeros=getattr(layer, "w2_weight_zero_point", None),
421 + w13_bias=getattr(layer, "w13_bias", None),
422 + w2_bias=getattr(layer, "w2_bias", None),
423 )
424
425 if converted is None:
@@ -378,8 +435,8 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
435 w2_qzeros,
436 w13_input_global_scale,
437 w2_input_global_scale,
381 - _, # w13_bias
382 - _, # w2_bias
438 + w13_bias,
439 + w2_bias,
440 ) = converted
441
442 # Replace common parameters
@@ -388,6 +445,11 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
445 replace_parameter(layer, "w13_weight_scale", w13_scales)
446 replace_parameter(layer, "w2_weight_scale", w2_scales)
447
448 + if w13_bias is not None:
449 + replace_parameter(layer, "w13_bias", w13_bias)
450 + if w2_bias is not None:
451 + replace_parameter(layer, "w2_bias", w2_bias)
452 +
453 # CPU fused_experts_cpu requires zero points even for symmetric quant.
454 # EMULATION bakes ZP into the dequantized bf16 weights — ZP is None.
455 if (
@@ -437,6 +499,8 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
499 num_bits=self.num_bits,
500 w1_zp=getattr(layer, "w13_weight_zero_point", None),
501 w2_zp=getattr(layer, "w2_weight_zero_point", None),
502 + w1_bias=getattr(layer, "w13_bias", None),
503 + w2_bias=getattr(layer, "w2_bias", None),
504 gemm1_clamp_limit=getattr(layer, "swiglu_limit", None),
505 gemm1_alpha=getattr(layer, "swiglu_alpha", None),
506 gemm1_beta=getattr(layer, "swiglu_beta", None),
vllm/model_executor/models/gpt_oss.py
+4 -1
@@ -76,6 +76,8 @@ _GPT_OSS_STREAMED_EXPERT_SUFFIX_TO_SHARD = {
76 "w2_bias": "gpt_oss_w2",
77 "w13_weight_scale": "gpt_oss_w13",
78 "w2_weight_scale": "gpt_oss_w2",
79 + "w13_weight_packed": "gpt_oss_w13",
80 + "w2_weight_packed": "gpt_oss_w2",
81 }
82
83
@@ -263,7 +265,7 @@ class GptOssRoutedExperts(RoutedExperts):
265 ) -> None:
266 tp_rank = self.moe_config.moe_parallel_config.tp_rank
267 is_w13 = shard_id == "gpt_oss_w13"
266 - is_weight = weight_name.endswith("_weight")
268 + is_weight = weight_name.endswith(("_weight", "_weight_packed"))
269 is_partitioned_scale = weight_name.endswith("_weight_scale")
270
271 if is_weight:
@@ -347,6 +349,7 @@ class GptOssRoutedExperts(RoutedExperts):
349 ) or quant_method_name in (
350 "CompressedTensorsW4A4Nvfp4MoEMethod",
351 "Nvfp4OnlineMoEMethod",
352 + "CompressedTensorsWNA16MoEMethod",
353 ):
354 self._load_packed_expert(expert_data, loaded_weight, weight_name, shard_id)
355 else: