[Bugfix][CPU] Take an attention group's query head count from its layers (#51852)
Signed-off-by: Ganesh R <Ganesh.R@amd.com> Signed-off-by: R <Ganesh.R@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Ganesh R committed
Aug 17, 2026 at 16:15 UTC
f27ae25473af7520dde3f8b9e0705041940be0c9
3 files changed
+92
-8
.buildkite/hardware_tests/cpu.yaml
+2
@@ -11,6 +11,7 @@ steps:
11
- CMakeLists.txt
12
- vllm/_custom_ops.py
13
- tests/kernels/attention/test_cpu_attn.py
14
+ - tests/v1/attention/test_group_head_counts.py
15
- tests/kernels/moe/test_cpu_fused_moe.py
16
- tests/kernels/moe/test_cpu_quant_fused_moe.py
17
- tests/kernels/test_onednn.py
@@ -27,6 +28,7 @@ steps:
28
- |
29
bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
30
pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
31
+ pytest -x -v -s tests/v1/attention/test_group_head_counts.py
32
pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
33
pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py
34
pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py
tests/v1/attention/test_group_head_counts.py
new
+82
@@ -0,0 +1,82 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Scheduler metadata sizes a scratchpad from the query head count, so it must
4
+come from the builder's own group: the model-wide ``get_num_attention_heads()``
5
+is wrong for models that vary it per layer (e.g. Laguna), and too small a
6
+scratchpad is indexed past its end.
7
+"""
8
+
9
+from types import SimpleNamespace
10
+from unittest.mock import MagicMock, patch
11
+
12
+import pytest
13
+import torch
14
+
15
+from vllm.platforms import current_platform
16
+from vllm.v1.attention.backends.cpu_attn import (
17
+ CPUAttentionBackendImpl,
18
+ CPUAttentionMetadataBuilder,
19
+)
20
+
21
+pytestmark = pytest.mark.skipif(
22
+ not current_platform.is_cpu(), reason="CPU attention backend"
23
+)
24
+
25
+# Laguna's shape: 48 query heads model-wide, 64 on its sliding layers, both
26
+# against 8 KV heads.
27
+MODEL_WIDE_NUM_HEADS = 48
28
+NUM_KV_HEADS = 8
29
+
30
+
31
+def _layers(layer_num_heads: list[int]):
32
+ """Stand-in attention layers, one per head count, as one attention group."""
33
+ return {
34
+ f"layer_{i}": SimpleNamespace(
35
+ impl=MagicMock(
36
+ spec=CPUAttentionBackendImpl,
37
+ num_heads=num_heads,
38
+ sliding_window=None,
39
+ )
40
+ )
41
+ for i, num_heads in enumerate(layer_num_heads)
42
+ }
43
+
44
+
45
+def _build(layer_num_heads: list[int]) -> CPUAttentionMetadataBuilder:
46
+ layers = _layers(layer_num_heads)
47
+ vllm_config = MagicMock()
48
+ vllm_config.model_config.dtype = torch.bfloat16
49
+ vllm_config.model_config.get_num_attention_heads.return_value = MODEL_WIDE_NUM_HEADS
50
+ vllm_config.cache_config.block_size = 16
51
+ vllm_config.cache_config.cache_dtype = "auto"
52
+ kv_cache_spec = SimpleNamespace(num_kv_heads=NUM_KV_HEADS, head_size=64)
53
+
54
+ with (
55
+ patch(
56
+ "vllm.v1.attention.backends.utils.get_layers_from_vllm_config",
57
+ return_value=layers,
58
+ ),
59
+ patch(
60
+ "vllm.v1.attention.backends.cpu_attn.get_layers_from_vllm_config",
61
+ return_value=layers,
62
+ ),
63
+ ):
64
+ return CPUAttentionMetadataBuilder(
65
+ kv_cache_spec=kv_cache_spec,
66
+ layer_names=list(layers),
67
+ vllm_config=vllm_config,
68
+ device=torch.device("cpu"),
69
+ )
70
+
71
+
72
+@pytest.mark.parametrize("group_num_heads", [MODEL_WIDE_NUM_HEADS, 64, 16])
73
+def test_num_heads_comes_from_the_group(group_num_heads):
74
+ """The group's own count wins, even when it is not the model-wide one."""
75
+ builder = _build([group_num_heads, group_num_heads])
76
+ assert builder.num_heads == group_num_heads
77
+
78
+
79
+def test_mixed_head_counts_in_one_group_are_rejected():
80
+ """Grouping guarantees uniformity; a mixed group means that broke."""
81
+ with pytest.raises(AssertionError, match="share num_heads"):
82
+ _build([MODEL_WIDE_NUM_HEADS, 64])
vllm/v1/attention/backends/cpu_attn.py
+8
-8
@@ -31,6 +31,7 @@ from vllm.v1.attention.backend import (
31
)
32
from vllm.v1.attention.backends.utils import (
33
KVCacheLayoutType,
34
+ get_num_attention_heads_from_layers,
35
)
36
from vllm.v1.kv_cache_interface import (
37
AttentionSpec,
@@ -151,13 +152,15 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
152
153
parallel_config = vllm_config.parallel_config
154
self.num_kv_heads = kv_cache_spec.num_kv_heads
154
- self.num_heads = vllm_config.model_config.get_num_attention_heads(
155
- parallel_config
156
- )
155
+ # The scheduler metadata built here sizes a scratchpad from the query
156
+ # head count, so it must come from this group's layers: the model-wide
157
+ # count is wrong for models that vary it per layer (e.g. Laguna).
158
+ self.num_heads = get_num_attention_heads_from_layers(
159
+ vllm_config, layer_names
160
+ ) or vllm_config.model_config.get_num_attention_heads(parallel_config)
161
self.head_dim = kv_cache_spec.head_size
162
self.dtype = vllm_config.model_config.dtype
159
- # Resolved from the layers on the first build(), once they exist.
160
- self.window_size: int | None = None
163
+ self.window_size = self._group_sliding_window()
164
self.block_size = vllm_config.cache_config.block_size
165
self.kv_cache_dtype = vllm_config.cache_config.cache_dtype
166
self.isa = _get_attn_isa(
@@ -198,9 +201,6 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
201
common_attn_metadata: CommonAttentionMetadata,
202
fast_build: bool = False,
203
) -> CPUAttentionMetadata:
201
- if self.window_size is None:
202
- self.window_size = self._group_sliding_window()
203
-
204
num_reqs = common_attn_metadata.num_reqs
205
num_actual_tokens = common_attn_metadata.num_actual_tokens
206
max_query_len = common_attn_metadata.max_query_len