[CPU][Zen] Route Int8 MoE inference through zentorch on AMD (#44834)
Signed-off-by: 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
Sep 1, 2026 at 15:01 UTC
8f03625b3d14969e59b2adb06097b2f20be5ce8c
6 files changed
+499
-2
.buildkite/hardware_tests/cpu.yaml
+3
@@ -18,6 +18,8 @@ steps:
18
- vllm/model_executor/layers/attention/mla_attention.py
19
- tests/kernels/moe/test_cpu_fused_moe.py
20
- tests/kernels/moe/test_cpu_quant_fused_moe.py
21
+ - tests/kernels/moe/test_zen_cpu_int8_moe.py
22
+ - vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
23
- tests/kernels/test_onednn.py
24
- tests/kernels/test_awq_int4_to_int8.py
25
- tests/kernels/quantization/test_cpu_fp8_scaled_mm.py
@@ -36,6 +38,7 @@ steps:
38
pytest -x -v -s tests/kernels/attention/test_amx_mla.py
39
pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
40
pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py
41
+ pytest -x -v -s tests/kernels/moe/test_zen_cpu_int8_moe.py
42
pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
43
pytest -x -v -s tests/kernels/test_onednn.py
44
pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py
tests/kernels/moe/test_zen_cpu_int8_moe.py
new
+253
@@ -0,0 +1,253 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for the Zen CPU INT8 MoE experts (ZenCPUExpertsInt8)."""
4
+
5
+import pytest
6
+import torch
7
+
8
+from tests.kernels.moe.test_cpu_fused_moe import (
9
+ _make_moe_config,
10
+ quantize_per_channel,
11
+ ref_fused_moe_int8,
12
+)
13
+from vllm.model_executor.kernels.linear.zentorch_utils import (
14
+ _ZENTORCH_MOE_ACTIVATIONS,
15
+)
16
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
17
+from vllm.model_executor.layers.fused_moe.config import (
18
+ FusedMoEParallelConfig,
19
+ int8_w8a8_moe_quant_config,
20
+)
21
+from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
22
+ ZenCPUExpertsInt8,
23
+ select_experts,
24
+)
25
+from vllm.model_executor.layers.fused_moe.oracle.int8 import (
26
+ Int8MoeBackend,
27
+ backend_to_kernel_cls,
28
+)
29
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
30
+ kInt8DynamicTokenSym,
31
+ kInt8StaticChannelSym,
32
+)
33
+from vllm.platforms import current_platform
34
+from vllm.utils.torch_utils import set_random_seed
35
+
36
+if not current_platform.is_cpu():
37
+ pytest.skip("skipping CPU-only tests", allow_module_level=True)
38
+
39
+EXPERT_NUM = 8
40
+HIDDEN_SIZE = 128
41
+INTERMEDIATE_SIZE = 128
42
+BATCH_SIZE = 32
43
+# zentorch names its activations as strings; the experts API is enum-typed.
44
+ZENTORCH_ACT = [act for act in MoEActivation if act.value in _ZENTORCH_MOE_ACTIVATIONS]
45
+
46
+
47
+@pytest.fixture
48
+def mock_zentorch_fused_moe():
49
+ """Reference zentorch_fused_moe, registered when zentorch is absent."""
50
+ if hasattr(torch.ops.zentorch, "zentorch_fused_moe"):
51
+ yield None
52
+ return
53
+
54
+ calls: dict[str, dict] = {}
55
+
56
+ def _fused_moe(
57
+ output,
58
+ input,
59
+ w13,
60
+ w2,
61
+ w13_bias,
62
+ w2_bias,
63
+ topk_weights,
64
+ topk_id,
65
+ skip_weighted,
66
+ act,
67
+ w13_scales=None,
68
+ w2_scales=None,
69
+ zentorch_op_name="",
70
+ ):
71
+ calls["fused_moe"] = {
72
+ "w13": w13,
73
+ "w2": w2,
74
+ "w13_bias": w13_bias,
75
+ "w2_bias": w2_bias,
76
+ "topk_weights": topk_weights,
77
+ "topk_id": topk_id,
78
+ "skip_weighted": skip_weighted,
79
+ "act": act,
80
+ "w13_scales": w13_scales,
81
+ "w2_scales": w2_scales,
82
+ }
83
+ output.copy_(
84
+ ref_fused_moe_int8(
85
+ input,
86
+ w13,
87
+ w2,
88
+ w13_scales.float(),
89
+ w2_scales.float(),
90
+ None if w13_bias is None else w13_bias.float(),
91
+ None if w2_bias is None else w2_bias.float(),
92
+ topk_weights,
93
+ topk_id,
94
+ MoEActivation(act),
95
+ )
96
+ )
97
+
98
+ lib_def = torch.library.Library("zentorch", "DEF")
99
+ lib_def.define(
100
+ "zentorch_fused_moe(Tensor(a!) output, Tensor input, Tensor w13, "
101
+ "Tensor w2, Tensor? w13_bias, Tensor? w2_bias, Tensor topk_weights, "
102
+ "Tensor topk_id, bool skip_weighted, str act, "
103
+ "Tensor? w13_scales=None, Tensor? w2_scales=None, *, "
104
+ "str zentorch_op_name='zentorch::zentorch_fused_moe') -> ()"
105
+ )
106
+ lib_impl = torch.library.Library("zentorch", "IMPL", "CPU")
107
+ lib_impl.impl("zentorch_fused_moe", _fused_moe)
108
+
109
+ yield calls
110
+
111
+ lib_impl._destroy()
112
+ lib_def._destroy()
113
+
114
+
115
+def _make_int8_moe_weights(use_bias: bool, dtype: torch.dtype = torch.bfloat16):
116
+ up_dim = 2 * INTERMEDIATE_SIZE
117
+ w13_fp = torch.randn((EXPERT_NUM, up_dim, HIDDEN_SIZE), dtype=dtype) / (
118
+ 0.5 * HIDDEN_SIZE**0.5
119
+ )
120
+ w2_fp = torch.randn((EXPERT_NUM, HIDDEN_SIZE, INTERMEDIATE_SIZE), dtype=dtype) / (
121
+ 0.5 * INTERMEDIATE_SIZE**0.5
122
+ )
123
+ w13, w13_scale = quantize_per_channel(w13_fp)
124
+ w2, w2_scale = quantize_per_channel(w2_fp)
125
+ w13_scale = w13_scale.to(torch.bfloat16).float()
126
+ w2_scale = w2_scale.to(torch.bfloat16).float()
127
+
128
+ w13_bias = w2_bias = None
129
+ if use_bias:
130
+ w13_bias = torch.randn((EXPERT_NUM, up_dim), dtype=dtype) / (0.5 * up_dim**0.5)
131
+ w2_bias = torch.randn((EXPERT_NUM, HIDDEN_SIZE), dtype=dtype) / (
132
+ 0.5 * HIDDEN_SIZE**0.5
133
+ )
134
+ return w13, w2, w13_scale, w2_scale, w13_bias, w2_bias
135
+
136
+
137
+def test_zen_int8_first_in_cpu_backend_priority():
138
+ """The Zen experts are preferred over the generic CPU int8 experts."""
139
+ assert backend_to_kernel_cls(Int8MoeBackend.CPU)[0] is ZenCPUExpertsInt8
140
+
141
+
142
+def test_zen_int8_support_predicates(mock_zentorch_fused_moe, monkeypatch):
143
+ monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: False)
144
+ assert not ZenCPUExpertsInt8._supports_current_device()
145
+ monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
146
+ assert ZenCPUExpertsInt8._supports_current_device()
147
+ assert ZenCPUExpertsInt8._supports_quant_scheme(
148
+ kInt8StaticChannelSym, kInt8DynamicTokenSym
149
+ )
150
+ assert not ZenCPUExpertsInt8._supports_quant_scheme(kInt8StaticChannelSym, None)
151
+ for act in ZENTORCH_ACT:
152
+ assert ZenCPUExpertsInt8._supports_activation(act)
153
+ assert not ZenCPUExpertsInt8._supports_activation(MoEActivation.RELU2)
154
+ assert not ZenCPUExpertsInt8._supports_activation(
155
+ MoEActivation.SWIGLUOAI_UNINTERLEAVE
156
+ )
157
+ assert not ZenCPUExpertsInt8._supports_no_act_and_mul()
158
+ assert ZenCPUExpertsInt8._supports_parallel_config(
159
+ FusedMoEParallelConfig.make_no_parallel()
160
+ )
161
+
162
+
163
+@pytest.mark.parametrize("act", ZENTORCH_ACT)
164
+@pytest.mark.parametrize("use_bias", [False, True])
165
+def test_zen_int8_dispatch_contract(
166
+ default_vllm_config,
167
+ mock_zentorch_fused_moe,
168
+ act: MoEActivation,
169
+ use_bias: bool,
170
+):
171
+ """apply() hands zentorch the weights, scales and biases it expects."""
172
+ if mock_zentorch_fused_moe is None:
173
+ pytest.skip("real zentorch is installed; the recorded call is unavailable")
174
+ set_random_seed(0)
175
+
176
+ topk_num = EXPERT_NUM // 2
177
+ w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = _make_int8_moe_weights(use_bias)
178
+ moe_config = _make_moe_config(
179
+ EXPERT_NUM,
180
+ HIDDEN_SIZE,
181
+ INTERMEDIATE_SIZE,
182
+ topk_num,
183
+ torch.bfloat16,
184
+ act,
185
+ use_bias,
186
+ )
187
+ # A layer carries per-channel scales as [E, N, 1].
188
+ quant_config = int8_w8a8_moe_quant_config(
189
+ w1_scale=w13_scale.unsqueeze(-1),
190
+ w2_scale=w2_scale.unsqueeze(-1),
191
+ a1_scale=None,
192
+ a2_scale=None,
193
+ w1_bias=w13_bias,
194
+ w2_bias=w2_bias,
195
+ per_act_token_quant=True,
196
+ )
197
+
198
+ input = torch.randn((BATCH_SIZE, HIDDEN_SIZE), dtype=torch.bfloat16) / (
199
+ 0.5 * HIDDEN_SIZE**0.5
200
+ )
201
+ router_logits = torch.randn((BATCH_SIZE, EXPERT_NUM), dtype=torch.bfloat16)
202
+
203
+ experts = ZenCPUExpertsInt8(moe_config, quant_config)
204
+ # zentorch_fused_moe takes no expert_map.
205
+ assert not experts.supports_expert_map()
206
+ output = experts.apply(
207
+ hidden_states=input,
208
+ w1=w13,
209
+ w2=w2,
210
+ router_logits=router_logits,
211
+ activation=act,
212
+ global_num_experts=EXPERT_NUM,
213
+ expert_map=None,
214
+ a1q_scale=None,
215
+ apply_router_weight_on_input=False,
216
+ )
217
+
218
+ call = mock_zentorch_fused_moe["fused_moe"]
219
+ assert call["w13_scales"].shape == (EXPERT_NUM, 2 * INTERMEDIATE_SIZE)
220
+ assert call["w2_scales"].shape == (EXPERT_NUM, HIDDEN_SIZE)
221
+ assert call["w13_scales"].dtype == torch.bfloat16
222
+ assert call["w2_scales"].dtype == torch.bfloat16
223
+ assert call["act"] == act.value.lower()
224
+ assert call["skip_weighted"] is False
225
+ assert call["topk_weights"].dtype == torch.float32
226
+ assert call["topk_id"].dtype == torch.int32
227
+ if use_bias:
228
+ assert call["w13_bias"].shape == (EXPERT_NUM, 2 * INTERMEDIATE_SIZE)
229
+ assert call["w2_bias"].shape == (EXPERT_NUM, HIDDEN_SIZE)
230
+ else:
231
+ assert call["w13_bias"] is None and call["w2_bias"] is None
232
+
233
+ topk_weights, topk_ids = select_experts(
234
+ hidden_states=input,
235
+ router_logits=router_logits,
236
+ use_grouped_topk=False,
237
+ top_k=topk_num,
238
+ renormalize=False,
239
+ scoring_func="softmax",
240
+ )
241
+ ref = ref_fused_moe_int8(
242
+ input,
243
+ w13,
244
+ w2,
245
+ w13_scale,
246
+ w2_scale,
247
+ None if w13_bias is None else w13_bias.float(),
248
+ None if w2_bias is None else w2_bias.float(),
249
+ topk_weights.to(torch.float32),
250
+ topk_ids.to(torch.int32),
251
+ act,
252
+ )
253
+ torch.testing.assert_close(output, ref)
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+162
@@ -20,7 +20,10 @@ from vllm._custom_ops import (
20
cpu_prepack_moe_weight_int8,
21
fused_experts_cpu,
22
)
23
+from vllm.logger import init_logger
24
from vllm.model_executor.kernels.linear.zentorch_utils import (
25
+ _ZENTORCH_MOE_ACTIVATIONS,
26
+ has_zentorch_op,
27
is_zentorch_moe_config_supported,
28
is_zentorch_moe_supported,
29
)
@@ -44,6 +47,7 @@ from vllm.model_executor.utils import replace_parameter
47
from vllm.platforms import CpuArchEnum, current_platform
48
from vllm.utils.math_utils import round_up
49
50
+logger = init_logger(__name__)
51
# ===========================================================================
52
# Routing
53
# ===========================================================================
@@ -1382,3 +1386,161 @@ class ArmCPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
1386
"neon",
1387
skip_weighted=apply_router_weight_on_input,
1388
)
1389
+
1390
+
1391
+class ZenCPUExpertsInt8(mk.FusedMoEExpertsMonolithic):
1392
+ """AMD Zen INT8 MoE with per-token activation and channelwise weight
1393
+ quantization, dispatched through zentorch."""
1394
+
1395
+ def __init__(
1396
+ self,
1397
+ moe_config: FusedMoEConfig,
1398
+ quant_config: FusedMoEQuantConfig,
1399
+ ):
1400
+ super().__init__(moe_config, quant_config)
1401
+ assert self.w1_scale is not None and self.w2_scale is not None, (
1402
+ "ZenCPUExpertsInt8 requires per-channel weight scales on the layer."
1403
+ )
1404
+ num_experts = self.w1_scale.shape[0]
1405
+ self._w1_scale_bf16 = (
1406
+ self.w1_scale.detach()
1407
+ .to(torch.bfloat16)
1408
+ .reshape(num_experts, -1)
1409
+ .contiguous()
1410
+ )
1411
+ self._w2_scale_bf16 = (
1412
+ self.w2_scale.detach()
1413
+ .to(torch.bfloat16)
1414
+ .reshape(num_experts, -1)
1415
+ .contiguous()
1416
+ )
1417
+ self._w1_bias_bf16 = (
1418
+ None
1419
+ if self.w1_bias is None
1420
+ else self.w1_bias.detach().to(torch.bfloat16).contiguous()
1421
+ )
1422
+ self._w2_bias_bf16 = (
1423
+ None
1424
+ if self.w2_bias is None
1425
+ else self.w2_bias.detach().to(torch.bfloat16).contiguous()
1426
+ )
1427
+ logger.info_once(
1428
+ "[zen_cpu] Using zentorch_fused_moe for W8A8 INT8 MoE (monolithic experts)"
1429
+ )
1430
+
1431
+ @property
1432
+ def expects_unquantized_inputs(self) -> bool:
1433
+ # zentorch_fused_moe quantizes activations itself.
1434
+ return True
1435
+
1436
+ @staticmethod
1437
+ def activation_format() -> mk.FusedMoEActivationFormat:
1438
+ return mk.FusedMoEActivationFormat.Standard
1439
+
1440
+ @staticmethod
1441
+ def _supports_current_device() -> bool:
1442
+ return has_zentorch_op(["zentorch_fused_moe"])
1443
+
1444
+ @staticmethod
1445
+ def _supports_no_act_and_mul() -> bool:
1446
+ return False
1447
+
1448
+ @staticmethod
1449
+ def _supports_activation(activation: MoEActivation) -> bool:
1450
+ return str(activation.value).lower() in _ZENTORCH_MOE_ACTIVATIONS
1451
+
1452
+ @staticmethod
1453
+ def _supports_parallel_config(
1454
+ moe_parallel_config: FusedMoEParallelConfig,
1455
+ ) -> bool:
1456
+ return not moe_parallel_config.use_ep
1457
+
1458
+ @staticmethod
1459
+ def _supports_quant_scheme(
1460
+ weight_key: QuantKey | None,
1461
+ activation_key: QuantKey | None,
1462
+ ) -> bool:
1463
+ return (weight_key, activation_key) == (
1464
+ kInt8StaticChannelSym,
1465
+ kInt8DynamicTokenSym,
1466
+ )
1467
+
1468
+ @staticmethod
1469
+ def _supports_routing_method(
1470
+ routing_method: RoutingMethodType,
1471
+ weight_key: QuantKey | None,
1472
+ activation_key: QuantKey | None,
1473
+ ) -> bool:
1474
+ return routing_method in [
1475
+ RoutingMethodType.Default,
1476
+ RoutingMethodType.Renormalize,
1477
+ RoutingMethodType.RenormalizeNaive,
1478
+ ]
1479
+
1480
+ @staticmethod
1481
+ def _supports_router_logits_dtype(
1482
+ router_logits_dtype: torch.dtype | None,
1483
+ routing_method: RoutingMethodType,
1484
+ ) -> bool:
1485
+ return True
1486
+
1487
+ def supports_expert_map(self) -> bool:
1488
+ return False
1489
+
1490
+ def apply(
1491
+ self,
1492
+ hidden_states: torch.Tensor,
1493
+ w1: torch.Tensor,
1494
+ w2: torch.Tensor,
1495
+ router_logits: torch.Tensor,
1496
+ activation: MoEActivation,
1497
+ global_num_experts: int,
1498
+ expert_map: torch.Tensor | None,
1499
+ a1q_scale: torch.Tensor | None,
1500
+ apply_router_weight_on_input: bool,
1501
+ num_expert_group: int | None = None,
1502
+ e_score_correction_bias: torch.Tensor | None = None,
1503
+ routed_scaling_factor: float | None = None,
1504
+ topk_group: int | None = None,
1505
+ ) -> torch.Tensor:
1506
+ topk_weights, topk_ids = select_experts(
1507
+ hidden_states=hidden_states,
1508
+ router_logits=router_logits,
1509
+ use_grouped_topk=num_expert_group is not None,
1510
+ top_k=self.moe_config.experts_per_token,
1511
+ renormalize=self.moe_config.routing_method
1512
+ in (
1513
+ RoutingMethodType.Renormalize,
1514
+ RoutingMethodType.RenormalizeNaive,
1515
+ ),
1516
+ topk_group=topk_group,
1517
+ num_expert_group=num_expert_group,
1518
+ scoring_func="softmax",
1519
+ routed_scaling_factor=(
1520
+ routed_scaling_factor if routed_scaling_factor is not None else 1.0
1521
+ ),
1522
+ e_score_correction_bias=e_score_correction_bias,
1523
+ )
1524
+
1525
+ if apply_router_weight_on_input:
1526
+ assert topk_ids.size(1) == 1, (
1527
+ "apply_router_weight_on_input is only implemented for topk=1"
1528
+ )
1529
+ hidden_states = hidden_states.mul(topk_weights.to(hidden_states.dtype))
1530
+
1531
+ output = torch.empty_like(hidden_states)
1532
+ torch.ops.zentorch.zentorch_fused_moe(
1533
+ output,
1534
+ hidden_states,
1535
+ w1,
1536
+ w2,
1537
+ self._w1_bias_bf16,
1538
+ self._w2_bias_bf16,
1539
+ topk_weights.to(torch.float32).contiguous(),
1540
+ topk_ids.to(torch.int32).contiguous(),
1541
+ apply_router_weight_on_input, # skip_weighted
1542
+ str(activation.value).lower(),
1543
+ self._w1_scale_bf16,
1544
+ self._w2_scale_bf16,
1545
+ )
1546
+ return output
vllm/model_executor/layers/fused_moe/oracle/int8.py
+2
-1
@@ -73,9 +73,10 @@ def backend_to_kernel_cls(
73
from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
74
ArmCPUExpertsInt8,
75
CPUExpertsInt8,
76
+ ZenCPUExpertsInt8,
77
)
78
78
- return [ArmCPUExpertsInt8, CPUExpertsInt8]
79
+ return [ZenCPUExpertsInt8, ArmCPUExpertsInt8, CPUExpertsInt8]
80
else:
81
raise ValueError(f"Unknown Int8 MoE backend: {backend.value}")
82
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_int8.py
+71
@@ -83,6 +83,8 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
83
params_dtype: torch.dtype,
84
**extra_weight_attrs,
85
):
86
+ # Capture bias dtype before params_dtype is overridden to int8 below.
87
+ bias_dtype = params_dtype
88
params_dtype = torch.int8
89
w13_num_shards = 2 if self.moe.is_act_and_mul else 1
90
@@ -111,6 +113,26 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
113
layer.register_parameter("w2_weight", w2_weight)
114
set_weight_attrs(w2_weight, extra_weight_attrs)
115
116
+ # Per-expert biases (gpt-oss), allocated only when has_bias.
117
+ if self.moe.has_bias:
118
+ w13_bias = torch.nn.Parameter(
119
+ torch.zeros(
120
+ num_experts,
121
+ w13_num_shards * intermediate_size_per_partition,
122
+ dtype=bias_dtype,
123
+ ),
124
+ requires_grad=False,
125
+ )
126
+ layer.register_parameter("w13_bias", w13_bias)
127
+ set_weight_attrs(w13_bias, extra_weight_attrs)
128
+
129
+ w2_bias = torch.nn.Parameter(
130
+ torch.zeros(num_experts, hidden_size, dtype=bias_dtype),
131
+ requires_grad=False,
132
+ )
133
+ layer.register_parameter("w2_bias", w2_bias)
134
+ set_weight_attrs(w2_bias, extra_weight_attrs)
135
+
136
# WEIGHT_SCALES
137
assert self.weight_quant.strategy == QuantizationStrategy.CHANNEL
138
w13_weight_scale = torch.nn.Parameter(
@@ -141,6 +163,53 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
163
layer.w2_input_scale = None
164
165
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
166
+ # ZenDNN's swiglu_oai_mul reads gate/up interleaved (stride 2), not the
167
+ # half-split layout _load_w13 leaves; reorder before the conversion below.
168
+ from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
169
+ ZenCPUExpertsInt8,
170
+ )
171
+
172
+ _act = getattr(layer, "activation", None)
173
+ _act_str = getattr(_act, "value", _act)
174
+ if (
175
+ self.experts_cls is ZenCPUExpertsInt8
176
+ and isinstance(_act_str, str)
177
+ and _act_str.lower() == "swigluoai"
178
+ ):
179
+ two_i = layer.w13_weight.size(1)
180
+ i = two_i // 2
181
+ device = layer.w13_weight.device
182
+ perm = torch.stack(
183
+ [
184
+ torch.arange(0, i, device=device),
185
+ torch.arange(i, two_i, device=device),
186
+ ],
187
+ dim=1,
188
+ ).flatten()
189
+ has_w13_bias = getattr(layer, "w13_bias", None) is not None
190
+ logger.info_once(
191
+ "[zen_cpu][swigluoai-permute] Reordering w13 from "
192
+ "half-split to interleaved for ZenDNN swiglu_oai_mul "
193
+ "kernel: E=%d, 2I=%d, I=%d, has_w13_bias=%s",
194
+ layer.w13_weight.size(0),
195
+ two_i,
196
+ i,
197
+ has_w13_bias,
198
+ )
199
+ layer.w13_weight = torch.nn.Parameter(
200
+ layer.w13_weight.data[:, perm, :].contiguous(),
201
+ requires_grad=False,
202
+ )
203
+ layer.w13_weight_scale = torch.nn.Parameter(
204
+ layer.w13_weight_scale.data[:, perm].contiguous(),
205
+ requires_grad=False,
206
+ )
207
+ if has_w13_bias:
208
+ layer.w13_bias = torch.nn.Parameter(
209
+ layer.w13_bias.data[:, perm].contiguous(),
210
+ requires_grad=False,
211
+ )
212
+
213
w13, w2 = convert_to_int8_moe_kernel_format(
214
int8_backend=self.int8_backend,
215
w13=layer.w13_weight,
@@ -169,6 +238,8 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
238
w2_scale=layer.w2_weight_scale,
239
a1_scale=layer.w13_input_scale,
240
a2_scale=layer.w2_input_scale,
241
+ w1_bias=getattr(layer, "w13_bias", None),
242
+ w2_bias=getattr(layer, "w2_bias", None),
243
per_act_token_quant=True,
244
layer=layer,
245
)
vllm/model_executor/models/mixtral.py
+8
-1
@@ -288,7 +288,14 @@ class MixtralModel(nn.Module):
288
".q_proj": (".qkv_proj", "q"),
289
".k_proj": (".qkv_proj", "k"),
290
".v_proj": (".qkv_proj", "v"),
291
- }
291
+ },
292
+ orig_to_new_substr={
293
+ # W8A8 compressed-tensors checkpoints name experts gate/up/down;
294
+ # map to the w1/w3/w2 the FusedMoE loader expects (no-op if native).
295
+ ".gate_proj.": ".w1.",
296
+ ".up_proj.": ".w3.",
297
+ ".down_proj.": ".w2.",
298
+ },
299
)
300
301
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):