[Bugfix][CPU] Take an attention group's query head count from its layers (vllm-project#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>

Author: 3 people

.buildkite/hardware_tests/cpu.yaml

@@ -11,6 +11,7 @@ steps:
   - CMakeLists.txt
   - vllm/_custom_ops.py
   - tests/kernels/attention/test_cpu_attn.py
+  - tests/v1/attention/test_group_head_counts.py
   - tests/kernels/moe/test_cpu_fused_moe.py
   - tests/kernels/moe/test_cpu_quant_fused_moe.py
   - tests/kernels/test_onednn.py
@@ -27,6 +28,7 @@ steps:
     - |
       bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m "
       pytest -x -v -s tests/kernels/attention/test_cpu_attn.py
+      pytest -x -v -s tests/v1/attention/test_group_head_counts.py
       pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py
       pytest -x -v -s tests/kernels/moe/test_cpu_quant_fused_moe.py
       pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py

tests/v1/attention/test_group_head_counts.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Scheduler metadata sizes a scratchpad from the query head count, so it must
+come from the builder's own group: the model-wide ``get_num_attention_heads()``
+is wrong for models that vary it per layer (e.g. Laguna), and too small a
+scratchpad is indexed past its end.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+from vllm.v1.attention.backends.cpu_attn import (
+    CPUAttentionBackendImpl,
+    CPUAttentionMetadataBuilder,
+)
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cpu(), reason="CPU attention backend"
+)
+
+# Laguna's shape: 48 query heads model-wide, 64 on its sliding layers, both
+# against 8 KV heads.
+MODEL_WIDE_NUM_HEADS = 48
+NUM_KV_HEADS = 8
+
+
+def _layers(layer_num_heads: list[int]):
+    """Stand-in attention layers, one per head count, as one attention group."""
+    return {
+        f"layer_{i}": SimpleNamespace(
+            impl=MagicMock(
+                spec=CPUAttentionBackendImpl,
+                num_heads=num_heads,
+                sliding_window=None,
+            )
+        )
+        for i, num_heads in enumerate(layer_num_heads)
+    }
+
+
+def _build(layer_num_heads: list[int]) -> CPUAttentionMetadataBuilder:
+    layers = _layers(layer_num_heads)
+    vllm_config = MagicMock()
+    vllm_config.model_config.dtype = torch.bfloat16
+    vllm_config.model_config.get_num_attention_heads.return_value = MODEL_WIDE_NUM_HEADS
+    vllm_config.cache_config.block_size = 16
+    vllm_config.cache_config.cache_dtype = "auto"
+    kv_cache_spec = SimpleNamespace(num_kv_heads=NUM_KV_HEADS, head_size=64)
+
+    with (
+        patch(
+            "vllm.v1.attention.backends.utils.get_layers_from_vllm_config",
+            return_value=layers,
+        ),
+        patch(
+            "vllm.v1.attention.backends.cpu_attn.get_layers_from_vllm_config",
+            return_value=layers,
+        ),
+    ):
+        return CPUAttentionMetadataBuilder(
+            kv_cache_spec=kv_cache_spec,
+            layer_names=list(layers),
+            vllm_config=vllm_config,
+            device=torch.device("cpu"),
+        )
+
+
+@pytest.mark.parametrize("group_num_heads", [MODEL_WIDE_NUM_HEADS, 64, 16])
+def test_num_heads_comes_from_the_group(group_num_heads):
+    """The group's own count wins, even when it is not the model-wide one."""
+    builder = _build([group_num_heads, group_num_heads])
+    assert builder.num_heads == group_num_heads
+
+
+def test_mixed_head_counts_in_one_group_are_rejected():
+    """Grouping guarantees uniformity; a mixed group means that broke."""
+    with pytest.raises(AssertionError, match="share num_heads"):
+        _build([MODEL_WIDE_NUM_HEADS, 64])

vllm/v1/attention/backends/cpu_attn.py

@@ -31,6 +31,7 @@
 )
 from vllm.v1.attention.backends.utils import (
     KVCacheLayoutType,
+    get_num_attention_heads_from_layers,
 )
 from vllm.v1.kv_cache_interface import (
     AttentionSpec,
@@ -151,13 +152,15 @@ def __init__(
 
         parallel_config = vllm_config.parallel_config
         self.num_kv_heads = kv_cache_spec.num_kv_heads
-        self.num_heads = vllm_config.model_config.get_num_attention_heads(
-            parallel_config
-        )
+        # The scheduler metadata built here sizes a scratchpad from the query
+        # head count, so it must come from this group's layers: the model-wide
+        # count is wrong for models that vary it per layer (e.g. Laguna).
+        self.num_heads = get_num_attention_heads_from_layers(
+            vllm_config, layer_names
+        ) or vllm_config.model_config.get_num_attention_heads(parallel_config)
         self.head_dim = kv_cache_spec.head_size
         self.dtype = vllm_config.model_config.dtype
-        # Resolved from the layers on the first build(), once they exist.
-        self.window_size: int | None = None
+        self.window_size = self._group_sliding_window()
         self.block_size = vllm_config.cache_config.block_size
         self.kv_cache_dtype = vllm_config.cache_config.cache_dtype
         self.isa = _get_attn_isa(
@@ -198,9 +201,6 @@ def build(
         common_attn_metadata: CommonAttentionMetadata,
         fast_build: bool = False,
     ) -> CPUAttentionMetadata:
-        if self.window_size is None:
-            self.window_size = self._group_sliding_window()
-
         num_reqs = common_attn_metadata.num_reqs
         num_actual_tokens = common_attn_metadata.num_actual_tokens
         max_query_len = common_attn_metadata.max_query_len