Fix DefaultKeyValueCache: per-layer num_kv_heads for Gemma 4 dual/MQA attention

[shreyshah-microsoft]

Problem

Loading and running Gemma 4 (gemma4 / gemma4_unified, e.g. the 12B) on the CPU EP fails at prefill:

RuntimeError: Got invalid dimensions for input: past_key_values.5.key
 index: 1 Got: 8 Expected: 1

Root cause

DefaultKeyValueCache already auto-detects a per-layer head_dim (shape index 3) from each
past_key_values.* input shape (added with Gemma 4 support for its dual head_dim). But it keeps a
single uniform num_key_value_heads (shape index 1) for every layer. Gemma 4 uses a dual attention
pattern:

• sliding-window layers: GQA — num_key_value_heads = 8, head_dim = 256
• global/full-attention layers: MQA — num_key_value_heads = 1, head_dim = 512

So layer_shapes_ ends up with the right per-layer head_dim but the uniform KV-head count. The
per-layer empty/present KV tensors for the global layers are then allocated as [B, 8, 0, 512] while
the decoder graph expects [B, 1, 0, 512] — the index-1 mismatch above.

Fix

Extend the existing per-layer detection to also read num_kv_heads (input shape index 1), and trigger
per-layer allocation when either num_kv_heads or head_dim varies (set layer_shapes_[i][1]).
The change propagates through empty_pasts_, presents_, Add() and Update(). Uniform-KV models
are unaffected — with no variation, has_varying_shape stays false and the original single-shape path
is used. Guarded with input_shape.size() >= 4.

Testing

• Gemma 4 12B (gemma4_unified) now loads and decodes on the CPU EP (previously failed at prefill).
• An fp32 ONNX export of the model matches HuggingFace transformers exactly on a 3-prompt suite:
  last-position logit cosine 1.0000, top-1 match, greedy 12/12 identical tokens.
• Uniform-KV architectures (Llama / Gemma 3 / etc.) are unchanged — the new branch only activates when
  per-layer KV shapes actually differ.

Note: a regression pass over the RewindTo / past_present_share_buffer / sampling paths is
recommended before merge; this change was validated on the standard prefill + greedy-decode path.

[Copilot]

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Updates KV-cache initialization to detect and support per-layer variation in KV tensor shapes (both num_kv_heads and head_dim) based on ONNX session input shapes, fixing model variants where global-attention layers use MQA (single KV head).

Changes:

• Detect per-layer num_kv_heads and head_dim from ONNX input shapes instead of only head_dim.
• Populate layer_shapes_ with per-layer KV shape overrides and create per-layer empty past tensors accordingly.
• Improve logging and inline documentation to reflect the expanded shape handling.

Show a summary per file

File	Description
src/models/kv_cache.cpp	Extends per-layer KV cache shape detection to include both KV head count and head dimension.

Copilot's findings

• Files reviewed: 1/1 changed files
• Comments generated: 2

[kunal-vaishnavi]

/azp run Integration Tests

[azure-pipelines]

Azure Pipelines successfully started running 1 pipeline(s).

[shreyshah-microsoft]

shreyshah-microsoft please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.

@microsoft-github-policy-service agree [company="{your company}"]

Options:

• (default - no company specified) I have sole ownership of intellectual property rights to my Submissions and I am not making Submissions in the course of work for my employer.

@microsoft-github-policy-service agree

• (when company given) I am making Submissions in the course of work for my employer (or my employer has intellectual property rights in my Submissions by contract or applicable law). I have permission from my employer to make Submissions and enter into this Agreement on behalf of my employer. By signing below, the defined term “You” includes me and my employer.

@microsoft-github-policy-service agree company="Microsoft"

Contributor License Agreement

@microsoft-github-policy-service agree company="Microsoft"