onnxruntime
b725f3d3 - Fix LinearAttention output shape inference for standard GQA (q > kv) (#29892)

Commit
5 days ago
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.
Author
Parents
Loading