[Bugfix][CPU] Take an attention group's query head count from its layers

[ganeshr10]

Purpose

CPUAttentionMetadataBuilder sizes the split-KV scratchpad from the model-wide query head count, but the kernel runs with each layer's own count. so models that vary it per layer (e.g. Laguna) index past the end of the scratchpad and
the decode segfaults or hangs. Attention groups are keyed on num_heads_q, so take the count from the group's layers via get_num_attention_heads_from_layers(), as triton_attn and flashinfer already do. Models with a uniform count are unaffected.

Change-Id: I671eadc2b5601f1a3af12f39391657a11df1c9f7

Test Plan

• pytest tests/v1/attention/test_group_head_counts.py — new unit test, mirroring test_group_sliding_window.py: uniform heads, Laguna-style alternating heads, and a default that is not the most common count.
• A tiny Laguna config with per-layer heads [8, 16, 8, 16] against 2 KV heads and random weights (load_format=dummy, so no checkpoint is needed), decoding 512 tokens so split-KV reduction engages. Run both on unpatched main and with this change.
• Laguna-XS-2.1-INT4 end to end on CPU with TP=2: a 300-token decode, plus GSM8K 5-shot over the full 1319 prompts via lm_eval.

Reproducer

import json, os, tempfile
from vllm import LLM, SamplingParams

CONFIG = {
    "architectures": ["LagunaForCausalLM"], "model_type": "laguna",
    "hidden_size": 512, "intermediate_size": 1024, "num_hidden_layers": 4,
    "num_attention_heads": 8, "num_attention_heads_per_layer": [8, 16, 8, 16],
    "num_key_value_heads": 2, "head_dim": 64, "max_position_embeddings": 4096,
    "rms_norm_eps": 1e-6, "rope_theta": 500000.0, "tie_word_embeddings": False,
    "torch_dtype": "bfloat16", "mlp_only_layers": [0, 1, 2, 3],
    "num_experts": 4, "num_experts_per_tok": 2, "moe_intermediate_size": 128,
    "shared_expert_intermediate_size": 128,
}

model_dir = tempfile.mkdtemp()
with open(os.path.join(model_dir, "config.json"), "w") as f:
    json.dump(CONFIG, f)

llm = LLM(model=model_dir, tokenizer="<any tokenizer>", load_format="dummy",
          dtype="bfloat16", max_model_len=2048, enforce_eager=True)
out = llm.generate(["context " * 64],
                   SamplingParams(max_tokens=512, temperature=0.0, ignore_eos=True))
print("SURVIVED", len(out[0].outputs[0].token_ids))

Test Result

• Unit tests: 6 passed (3 new, plus the 3 existing sliding-window cases).

• Reproducer: hangs on current main and is killed at the timeout; passes in 1.5s with this change.

• Laguna-XS 300-token decode: previously segfaulted or hung, now completes in 47s with coherent output.

• GSM8K 5-shot, full 1319 prompts. Before this change the same run died with an execute_model timeout.

  Filter	exact_match	stderr
  flexible-extract	0.9014	±0.0082
  strict-match	0.8908	±0.0086

Instrumenting the scratchpad bounds shows the overrun directly on unpatched main: the region holds 6336 bytes per KV head while the split-KV write path needs 8384, because sizing assumes 6 query heads per KV head (48 heads / 8 KV heads) and the 64-head layers address it with a stride of 8.

[SCRATCHPAD-SIZING]   split_num=2 max_tile=6 q_heads=24 kv_heads=4 q_per_kv=6 per_kv_head=6336
[SCRATCHPAD-OVERFLOW] per_head_need=8384 per_kv_head=6336 q_head_tile_size=8

Models whose layers share one head count keep the existing single-blob path, so there is no extra allocation or lookup for them.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

Hi @ganeshr10, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[mergify]

Hi @ganeshr10, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[bigPYJ1151]

Thanks! LGTM

[bigPYJ1151]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84171 for commit f307f4688309.

[bigPYJ1151]

/ci run

[github-actions]

✅ Triggered Buildkite CI #84175 for commit fcad8f0708a2.