[CPU][Whisper] Support W4A16 quantized Whisper on the CPU WNA16 kernel (#58268)
Signed-off-by: Harshal Adhav <harshal.adhav@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Harshal Adhav committed
Sep 30, 2026 at 11:22 UTC
5faf81a4297921c699f9aa58f9c6395e95ede718
5 files changed
+144
-15
tests/kernels/test_awq_int4_to_int8.py
+5
-4
@@ -203,7 +203,7 @@ class TestInt4ScaledMmCpu:
203
print(f" [PASS] INT4 GEMM with bias: M={M}")
204
205
def test_gemm_3d_input(self):
206
- """apply() reshapes 3D input [B, S, K] -> [B*S, K] -> back to 3D."""
206
+ """int4_scaled_mm_cpu flattens 3D [B, S, K] to 2D and reshapes back."""
207
K, N, group_size = 256, 128, 128
208
(packed_qweight, packed_qzeros, scales, float_ref, _, _) = (
209
make_awq_checkpoint_data(K, N, group_size)
@@ -220,12 +220,13 @@ class TestInt4ScaledMmCpu:
220
x_3d = torch.randn(B, S, K, dtype=torch.bfloat16)
221
x_2d = x_3d.reshape(-1, K)
222
223
+ out_3d = int4_scaled_mm_cpu(x_3d, blocked_w, blocked_zp, blocked_s, None)
224
out_2d = int4_scaled_mm_cpu(x_2d, blocked_w, blocked_zp, blocked_s, None)
224
- out_3d = out_2d.reshape(B, S, N)
225
-
226
- ref_out = torch.mm(x_2d.float(), float_ref).reshape(B, S, N)
225
226
assert out_3d.shape == (B, S, N)
227
+ torch.testing.assert_close(out_3d.reshape(-1, N), out_2d)
228
+
229
+ ref_out = torch.mm(x_2d.float(), float_ref).reshape(B, S, N)
230
abs_diff = (out_3d.float() - ref_out).abs()
231
mean_abs = abs_diff.mean().item()
232
ref_mag = ref_out.abs().mean().item() + 1e-6
tests/models/multimodal/generation/test_whisper.py
+44
@@ -5,6 +5,7 @@ from collections.abc import Sequence
5
from typing import Any
6
7
import pytest
8
+import torch
9
from transformers import AutoModelForSpeechSeq2Seq
10
11
from vllm.assets.audio import AudioAsset
@@ -264,6 +265,49 @@ def test_parse_language_detection_output():
265
cls.parse_language_detection_output([], make_tokenizer("anything"))
266
267
268
+def test_create_fake_bias_for_k_proj_injects_out_features():
269
+ """Packed k_proj weights get a zeros bias of length ``out_features``."""
270
+ from vllm.model_executor.models.whisper import _create_fake_bias_for_k_proj
271
+
272
+ packed = torch.ones(2, 16)
273
+ result = dict(
274
+ _create_fake_bias_for_k_proj(
275
+ [("model.encoder.layers.0.self_attn.k_proj.weight_packed", packed)],
276
+ ".k_proj.weight_packed",
277
+ out_features=8,
278
+ )
279
+ )
280
+ bias_name = "model.encoder.layers.0.self_attn.k_proj.bias"
281
+ assert bias_name in result
282
+ assert "model.encoder.layers.0.self_attn.k_proj.bias_packed" not in result
283
+ assert tuple(result[bias_name].shape) == (8,)
284
+ assert torch.equal(result[bias_name], torch.zeros(8))
285
+
286
+
287
+def test_create_fake_bias_for_k_proj_forwards_real_bias():
288
+ """A checkpoint bias is forwarded and no extra fake bias is emitted."""
289
+ from vllm.model_executor.models.whisper import _create_fake_bias_for_k_proj
290
+
291
+ weight = torch.ones(8, 16)
292
+ real_bias = torch.ones(8)
293
+ items = list(
294
+ _create_fake_bias_for_k_proj(
295
+ [
296
+ ("layers.0.self_attn.k_proj.weight", weight),
297
+ ("layers.0.self_attn.k_proj.bias", real_bias),
298
+ ],
299
+ ".k_proj.weight",
300
+ out_features=8,
301
+ )
302
+ )
303
+ names = [n for n, _ in items]
304
+ assert names == [
305
+ "layers.0.self_attn.k_proj.weight",
306
+ "layers.0.self_attn.k_proj.bias",
307
+ ]
308
+ assert items[1][1] is real_bias
309
+
310
+
311
@pytest.mark.core_model
312
@pytest.mark.cpu_model
313
@pytest.mark.parametrize("model", ["openai/whisper-large-v3-turbo"])
tests/quantization/test_cpu_wna16.py
+45
@@ -1,12 +1,57 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import pytest
4
+import torch
5
6
from vllm.platforms import current_platform
7
8
if not current_platform.is_cpu():
9
pytest.skip("skipping CPU-only tests", allow_module_level=True)
10
11
+from vllm._custom_ops import cpu_gemm_wna16 # noqa: E402
12
+from vllm.model_executor.kernels.linear.mixed_precision.cpu import ( # noqa: E402
13
+ _get_isa_hint,
14
+)
15
+
16
+# Isolated runs of cpu_gemm_wna16 can hang unless AMX tiles are initialized.
17
+# Guard so non-AMX hosts still take the vec/rvv path.
18
+if torch.cpu._is_amx_tile_supported():
19
+ torch.cpu._init_amx()
20
+
21
+
22
+@pytest.mark.cpu_test
23
+def test_cpu_gemm_wna16_3d_input():
24
+ """cpu_gemm_wna16 flattens 3D [B, T, K] to 2D and reshapes the output back."""
25
+ if not hasattr(torch.ops._C, "cpu_gemm_wna16"):
26
+ pytest.skip("cpu_gemm_wna16 op is not available")
27
+
28
+ B, T, K, N = 2, 5, 64, 64
29
+ pack_factor = 8
30
+ group_size = 32
31
+ dtype = torch.bfloat16
32
+ q_weight = torch.zeros(N // 16, K * 16 // pack_factor, dtype=torch.int32)
33
+ scales = torch.ones(K // group_size, N, dtype=dtype)
34
+
35
+ x_3d = torch.randn(B, T, K, dtype=dtype)
36
+ x_2d = x_3d.reshape(-1, K)
37
+ isa_hint = _get_isa_hint(dtype)
38
+ kwargs = dict(
39
+ q_weight=q_weight,
40
+ scales=scales,
41
+ zeros=None,
42
+ bias=None,
43
+ pack_factor=pack_factor,
44
+ isa_hint=isa_hint,
45
+ )
46
+
47
+ out_3d = cpu_gemm_wna16(input=x_3d, **kwargs)
48
+ out_2d = cpu_gemm_wna16(input=x_2d, **kwargs)
49
+
50
+ assert out_3d.shape == (B, T, N)
51
+ assert out_2d.shape == (B * T, N)
52
+ torch.testing.assert_close(out_3d.reshape(-1, N), out_2d)
53
+
54
+
55
MODELS = [
56
"TheBloke/TinyLlama-1.1B-Chat-v1.0-AWQ",
57
"Qwen/Qwen1.5-0.5B-Chat-GPTQ-Int4",
vllm/_custom_ops.py
+9
-4
@@ -4434,18 +4434,23 @@ def cpu_gemm_wna16(
4434
pack_factor: int,
4435
isa_hint: str,
4436
) -> torch.Tensor:
4437
- output = torch.empty((input.size(0), scales.size(1)), dtype=input.dtype)
4437
+ # Match int4_scaled_mm_cpu: flatten >2-D activations to [M, K] for the
4438
+ # C++ kernel, then restore the original leading dims on the output.
4439
+ x_shape = input.shape
4440
+ x_2d = input.reshape(-1, x_shape[-1]) if len(x_shape) > 2 else input
4441
+ out = torch.empty((x_2d.size(0), scales.size(1)), dtype=input.dtype)
4442
torch.ops._C.cpu_gemm_wna16(
4439
- input,
4443
+ x_2d,
4444
q_weight,
4441
- output,
4445
+ out,
4446
scales,
4447
zeros,
4448
bias,
4449
pack_factor,
4450
isa_hint,
4451
)
4448
- return output
4452
+ out = out.reshape(x_shape[:-1] + (out.size(-1),)) if len(x_shape) > 2 else out
4453
+ return out
4454
4455
4456
def cpu_activation_lut_bf16(input: torch.Tensor, activation: str) -> torch.Tensor:
vllm/model_executor/models/whisper.py
+41
-7
@@ -1039,20 +1039,54 @@ class WhisperForConditionalGeneration(
1039
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1040
loader = AutoWeightsLoader(self)
1041
1042
- # add fake zeros bias for k_proj to state_dict
1043
- weights = _create_fake_bias_for_k_proj(weights, ".k_proj.weight")
1042
+ # add fake zeros bias for k_proj to state_dict. Covers the
1043
+ # ``.k_proj.weight`` and the compressed-tensors packed
1044
+ # ``.k_proj.weight_packed`` checkpoints.
1045
+ k_proj_weight_suffixes = (".k_proj.weight", ".k_proj.weight_packed")
1046
+ for suffix in k_proj_weight_suffixes:
1047
+ weights = _create_fake_bias_for_k_proj(
1048
+ weights, suffix, out_features=self.config.d_model
1049
+ )
1050
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1051
1052
1053
def _create_fake_bias_for_k_proj(
1048
- weights: Iterable[tuple[str, torch.Tensor]], fake_bias_key_name: str
1054
+ weights: Iterable[tuple[str, torch.Tensor]],
1055
+ fake_bias_key_name: str,
1056
+ out_features: int | None = None,
1057
) -> Iterable[tuple[str, torch.Tensor]]:
1058
"""Create full zeros bias for k_proj weight in self-attn and x-attn layers.
1051
- So that the bias for k_proj in qkv_proj can be initialized with zeros.
1059
+ So that the bias for k_proj in qkv_proj or kv_proj can be initialized with
1060
+ zeros.
1061
+
1062
+ If the checkpoint already provides a real ``.bias`` entry for the given
1063
+ weight, it is forwarded as-is and no fake bias is injected for that layer.
1064
"""
1065
+ # Map the weight-name suffix to the corresponding bias-name suffix, e.g.
1066
+ # ".k_proj.weight_packed" / ".k_proj.weight" -> ".k_proj.bias",
1067
+ # ".wk.weight" -> ".wk.bias". ``.weight_packed`` is stripped first so the
1068
+ # longer suffix wins.
1069
+ bias_key_name = (
1070
+ fake_bias_key_name.removesuffix(".weight_packed").removesuffix(".weight")
1071
+ + ".bias"
1072
+ )
1073
+
1074
+ real_bias_names: set[str] = set()
1075
+ pending: dict[str, torch.Tensor] = {}
1076
for name, weight in weights:
1077
yield name, weight
1078
+
1079
+ if name.endswith(bias_key_name):
1080
+ real_bias_names.add(name)
1081
+ pending.pop(name, None)
1082
+ continue
1083
+
1084
if name.endswith(fake_bias_key_name):
1056
- bias = torch.zeros(weight.size(0))
1057
- bias_name = name.replace("weight", "bias")
1058
- yield bias_name, bias
1085
+ bias_name = name[: -len(fake_bias_key_name)] + bias_key_name
1086
+ if bias_name not in real_bias_names:
1087
+ pending[bias_name] = torch.zeros(
1088
+ out_features if out_features is not None else weight.size(0)
1089
+ )
1090
+
1091
+ for bias_name, bias in pending.items():
1092
+ yield bias_name, bias