Auto-detect fixed kv-cache shape in DefaultKeyValueCache

[akholodnamdcom]

Summary

Some compiled backends (e.g. AMD RyzenAI, and likely future NPU/accelerator targets) emit decoder models where the past_key_values.N.{key,value} ONNX inputs declare a static positive integer in the seq_len dimension instead of the usual symbolic past_sequence_length parameter. Such models require the kv-cache tensor to be allocated to exactly that size and reused as a shared past/present buffer — ORT will reject any cache tensor whose shape doesn't match the model's static dim, and the runtime cannot pick a size from search.max_length after the fact.

This PR makes OGA notice the static dim at cache construction and size the kv-cache accordingly, with no config schema change, no public API change, and no client cooperation required.

How it works

In DefaultKeyValueCache's constructor, right after the existing per-layer head_dim auto-detection (which already reads SessionInfo::GetInputShape), walk each layer's past_key.N input and look at the second-to-last dim:

• All layers report a positive integer for that dim and the values agree → treat the model as fixed-shape; let fixed_kv_seq_len = that value.
• Any layer reports <= 0 (symbolic, typically -1) or values disagree → leave detection off; behavior unchanged.

When detection fires:

• past_present_share_buffer_ is forced true (model demands it; the existing share-buffer code path already does the right thing for a pre-sized buffer reused as both past and present).
• The cache is allocated to fixed_kv_seq_len, not search.max_length. A new local cache_seq_len = fixed_kv_seq_len > 0 ? fixed_kv_seq_len : search.max_length replaces the direct read of search.max_length in the share-buffer branch, so dynamic-shape models are byte-equivalent to before.
• Beam search (num_beams != 1) is rejected with a clear std::runtime_error — beam picking reshuffles past tensors, which is incompatible with a model-mandated fixed shape.
• An info log records the detected size (gated by g_log.enabled, consistent with surrounding logs).
• If search.max_length > fixed_kv_seq_len, a warning is emitted once per process (std::call_once) noting the cache is sized to the model's limit. We do not clamp search.max_length because state_.params_ is std::shared_ptr<const GeneratorParams> and is owned by the caller; clamping would require a public API change.

Files changed

• src/models/kv_cache.cpp — +67/-2. One detection block in DefaultKeyValueCache::DefaultKeyValueCache, one shape_[2] source change in the share-buffer branch, and #include <mutex> for std::call_once.

No header changes. No public API changes. No config schema changes. No changes to Combined/Windowed/LFM2/Cross/ModelManaged cache classes (out of scope for this PR — AMD RyzenAI decoder-only models route through DefaultKeyValueCache).

Reused infrastructure

• SessionInfo::GetInputShape(name) — src/models/model.cpp:533-538. Same call already used at kv_cache.cpp:217 for head_dim auto-detection. Returns the ONNX-declared shape with <= 0 for symbolic dims (confirmed by sibling checks in src/models/multi_modal.cpp and src/models/recurrent_state.cpp).
• Log("warning"|"info", ...) — src/logging.h:82. Same pattern as the existing Log calls in kv_cache.cpp (lines 156, 201, 239).
• The share-buffer data-path — kv_cache.cpp:301-310 (shape), :357-361 (Add: inputs = presents), :366-367 (Update no-op), :409-410 (RewindTo no-op). All already correctly handle a pre-sized buffer.

Verification

Tested end-to-end on Windows x64, RelWithDebInfo, CPU EP, with Llama-3.2-1B-Instruct (Q4F16). Since no existing test model under test/test_models/ declares a static kv-cache seq_len, I produced a patched copy of the Llama model whose past_key/value (and present.key/value) ONNX shapes were rewritten in place — only graph metadata, no weight reload — from symbolic past_sequence_length to a static dim_value = 128. With those static dims in place, ORT will reject any cache tensor that isn't exactly [1, 8, 128, 64]; bound-and-run success is therefore proof that the cache was sized correctly.

#	Model	-ml flag	search.max_length	Detection	Outcome
1	Fixed-shape (patched, static [1,8,128,64])	-1 (use config)	4096 (JSON)	fires	✅ warning printed, model bound + ran 32 tokens against the 128-slot cache. Without detection the cache would be [1,8,4096,64] → bind failure.
2	Fixed-shape	1024 (runtime override)	1024	fires	✅ warning printed, cache stayed at 128, ran 16 tokens. Confirms runtime SetSearchOption overrides of max_length are ignored for cache sizing.
3	Original dynamic-shape Llama	-1 → 131072	131072	does not fire	✅ no warning, no info, normal run. Regression preserved.
4	Fixed-shape, multi-iter (-w 1 -r 3 → 4 cache constructions)	-1	4096	fires once	✅ exactly one warning across the run (std::call_once).

Verbatim warning emitted in test 1:

[warning]  Model has fixed kv-cache seq_len=128 but search.max_length=4096; cache is sized to the model's limit, so generation beyond it will fail.

Backward compatibility

For every model that ships today with symbolic kv-cache dims (i.e. essentially all decoder models built via the OGA model builder or the Microsoft ORT pipelines), seq_dim <= 0 at every layer → detection short-circuits → cache construction is byte-identical to main. No behavior change is exposed to existing users.

[akholodnamdcom]

@baijumeswani, Does this one look like what we agreed on? Thanks!