Fix LinearAttention output shape inference for standard GQA (q > kv) (#29892)
Fix LinearAttention output shape inference for standard GQA (q_num_heads
> kv_num_heads)
Problem
LinearAttention's shape-inference function disagreed with its own
Compute() for standard grouped-query attention (q_num_heads >
kv_num_heads):
Compute() allocates and writes an output of width max(q_num_heads,
kv_num_heads) * d_v, emitting one output per query head
(linear_attention.cc, output sized as std::max(q_num_heads_,
kv_num_heads_) * d_v, readout indexed by h_q).
Shape inference (bert_defs.cc) declared the output width as kv_num_heads
* d_v.
For standard GQA these differ (q*d_v vs kv*d_v). Any graph where the
output dim is validated against the inferred shape — a declared graph
output, a shape-sensitive consumer, or an optimizer — fails at model
load with:
[ShapeInferenceError] Can't merge shape info. Both inferred and declared
dimension have values but they differ. Inferred=128 Declared=256
The bug was latent because it only manifests when q_num_heads >
kv_num_heads. For MHA (q == kv) and inverse GQA (q < kv), max(q,kv) ==
kv, so the inferred value happened to match and the inconsistency stayed
hidden.
Fix
One line in the shape-inference function: compute the output hidden dim
as max(q_num_heads, kv_num_heads) * d_v, matching what Compute() already
produces. No kernel/behavioral change — the compute path was already
correct; only the advertised shape was wrong.
Tests
Added standard-GQA regression tests (*_StandardGQA_*) covering q>kv
across head ratios (N2/N4/N8/N16), all update rules (linear, gated,
delta, gated_delta), and both d_k=32/d_v=64 and d_k=d_v=128 shapes.
These reproduce the shape-inference failure before the fix and pass
after it. Also added multi-chunk (T > 64) coverage for long-sequence
correctness. Full LinearAttention suite: 37/37 passing.
Scope
This is a self-contained correctness fix (schema + regression tests) and
is independent of any performance work on the operator.