| 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:]) |