From bc1c8506d99b1c86e0671973949bb4143dd0e928 Mon Sep 17 00:00:00 2001 From: Oleksandr Kholodnyi Date: Mon, 18 May 2026 13:57:13 -0700 Subject: [PATCH 1/4] Auto-detect fixed kv-cache shape in DefaultKeyValueCache Some compiled backends (e.g. AMD RyzenAI) emit decoder models where the past_key/past_value ONNX inputs declare a static positive integer in the seq_len dimension instead of a symbolic dim. Such models require the kv-cache to be allocated to that exact size and reused as a shared past/present buffer; max_length cannot drive the size because ORT rejects any tensor whose shape doesn't match the model's static dim. Inspect past_key shapes per layer in DefaultKeyValueCache's constructor. When all layers agree on a positive seq_len, treat the model as fixed-shape: force past_present_share_buffer_=true, allocate the cache to the detected size, reject beam search, and warn if search.max_length exceeds the cache capacity. Behavior is unchanged for models with symbolic kv-cache dims (detection does not fire). Co-Authored-By: Claude Opus 4 --- src/models/kv_cache.cpp | 69 +++++++++++++++++++++++++++++++++++++++-- 1 file changed, 67 insertions(+), 2 deletions(-) diff --git a/src/models/kv_cache.cpp b/src/models/kv_cache.cpp index 01aaf0c2a3..3869439049 100644 --- a/src/models/kv_cache.cpp +++ b/src/models/kv_cache.cpp @@ -8,6 +8,7 @@ #include "../openvino/interface.h" #include "../qnn/interface.h" #include +#include namespace Generators { @@ -250,6 +251,65 @@ DefaultKeyValueCache::DefaultKeyValueCache(State& state) } } + // Auto-detect a fixed kv-cache shape from the model's past_key input shapes. + // Some compiled backends (e.g. AMD RyzenAI) emit models where the kv-cache + // seq_len dimension is a static positive integer instead of a symbolic dim. + // When that's the case the cache must be allocated to exactly that size and + // reused as a shared past/present buffer; max_length cannot drive the size + // because ORT rejects any tensor that doesn't match the model's static dim. + int64_t fixed_kv_seq_len = 0; + { + bool all_fixed_and_uniform = layer_count_ > 0; + int64_t common_seq_len = 0; + for (int i = 0; i < layer_count_; ++i) { + auto input_shape = model_.session_info_.GetInputShape(input_name_strings_[i * 2]); + if (input_shape.size() < 2) { + all_fixed_and_uniform = false; + break; + } + const int64_t seq_dim = input_shape[input_shape.size() - 2]; + if (seq_dim <= 0) { // symbolic/dynamic dim (typically -1) + all_fixed_and_uniform = false; + break; + } + if (common_seq_len == 0) { + common_seq_len = seq_dim; + } else if (common_seq_len != seq_dim) { + all_fixed_and_uniform = false; + break; + } + } + if (all_fixed_and_uniform && common_seq_len > 0) { + fixed_kv_seq_len = common_seq_len; + } + } + + if (fixed_kv_seq_len > 0) { + if (state_.params_->search.num_beams != 1) { + throw std::runtime_error( + "Beam search (num_beams > 1) is not supported for models with a fixed kv-cache " + "shape (model expects seq_len=" + + std::to_string(fixed_kv_seq_len) + ")."); + } + past_present_share_buffer_ = true; + if (g_log.enabled) { + Log("info", "DefaultKeyValueCache: auto-detected fixed kv-cache seq_len=" + + std::to_string(fixed_kv_seq_len) + + "; allocating shared past/present buffer to that size."); + } + if (state_.params_->search.max_length > static_cast(fixed_kv_seq_len)) { + // De-duplicate across repeated Generator constructions (e.g. benchmark loops). + static std::once_flag warned_once; + std::call_once(warned_once, [&] { + Log("warning", "Model has fixed kv-cache seq_len=" + + std::to_string(fixed_kv_seq_len) + + " but search.max_length=" + + std::to_string(state_.params_->search.max_length) + + "; cache is sized to the model's limit, so generation beyond it will fail."); + }); + } + } + if (state_.params_->use_graph_capture && !past_present_share_buffer_) { // share buffer is a precondition for graph capture throw std::runtime_error("Graph capture is not supported with past_present_share_buffer set to false."); @@ -299,12 +359,17 @@ DefaultKeyValueCache::DefaultKeyValueCache(State& state) shape_[2] = std::min(max_length, sliding_window_size); } } else if (past_present_share_buffer_) { - shape_[2] = state_.params_->search.max_length; + // For fixed kv-cache models the cache size comes from the model graph, + // not from max_length — see the auto-detection block earlier in this ctor. + const int64_t cache_seq_len = fixed_kv_seq_len > 0 + ? fixed_kv_seq_len + : static_cast(state_.params_->search.max_length); + shape_[2] = cache_seq_len; // If per-layer shapes exist (from head_dim auto-detection), update their sequence dim too if (!layer_shapes_.empty()) { for (int i = 0; i < layer_count_; ++i) { - layer_shapes_[i][2] = state_.params_->search.max_length; + layer_shapes_[i][2] = cache_seq_len; } } } From 877022ad43222fe4d899efa7467b140318568aac Mon Sep 17 00:00:00 2001 From: Oleksandr Kholodnyi Date: Tue, 19 May 2026 10:32:09 -0700 Subject: [PATCH 2/4] Document per-layer-uniform limitation in fixed kv-cache detection Address review feedback: the auto-detection only recognises models where every past_key layer declares the same fixed seq_len. Add a comment noting the restriction and outlining how to extend layer_shapes_ to support per-layer static seq_lens when a model in the wild needs it. Co-Authored-By: Claude Opus 4 --- src/models/kv_cache.cpp | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/models/kv_cache.cpp b/src/models/kv_cache.cpp index 3869439049..dfb3c459bf 100644 --- a/src/models/kv_cache.cpp +++ b/src/models/kv_cache.cpp @@ -257,6 +257,16 @@ DefaultKeyValueCache::DefaultKeyValueCache(State& state) // When that's the case the cache must be allocated to exactly that size and // reused as a shared past/present buffer; max_length cannot drive the size // because ORT rejects any tensor that doesn't match the model's static dim. + // + // Limitation: only uniform per-layer static sizes are recognised here — + // every past_key layer must declare the same fixed seq_len. Models that + // declare *different* static seq_lens per layer (e.g. a mix of full-attention + // and sliding-window layers with distinct static caps) fall through to + // dynamic handling. Lifting this restriction would extend the existing + // layer_shapes_ infrastructure used for per-layer head_dim detection above: + // store the per-layer detected seq_len into layer_shapes_[i][2] instead of + // a single scalar, and let the share-buffer branch's per-layer loop do the + // rest. Deferred until a model in the wild actually needs it. int64_t fixed_kv_seq_len = 0; { bool all_fixed_and_uniform = layer_count_ > 0; From 52cc2ba47b2b4117a3bb32e7240f55118b2e9c7e Mon Sep 17 00:00:00 2001 From: Oleksandr Kholodnyi Date: Tue, 19 May 2026 10:39:29 -0700 Subject: [PATCH 3/4] Extract fixed kv-cache detection into helper function Address review feedback: pull the auto-detection and its consequence handling (beam-search reject, force share-buffer, info log, warn-once on max_length mismatch) out of DefaultKeyValueCache's constructor into a file-local helper DetectAndConfigureFixedKvShape in an anonymous namespace. The constructor now calls the helper once and uses its return value (the detected static seq_len, or 0) in the share-buffer branch. No behavior change: pure refactor verified against both the fixed-shape patched model (warning fires once, generation succeeds against the 128-slot cache) and the dynamic-shape Llama (detection short-circuits, no warning, normal run). Co-Authored-By: Claude Opus 4 --- src/models/kv_cache.cpp | 148 ++++++++++++++++++++++------------------ 1 file changed, 80 insertions(+), 68 deletions(-) diff --git a/src/models/kv_cache.cpp b/src/models/kv_cache.cpp index dfb3c459bf..14f9c88be2 100644 --- a/src/models/kv_cache.cpp +++ b/src/models/kv_cache.cpp @@ -149,6 +149,83 @@ void CombinedKeyValueCache::PickPastState(DeviceSpan beam_indices, int } } +namespace { + +// Auto-detect a fixed kv-cache shape from the model's past_key input shapes, +// and, when detected, apply the implied configuration: +// - reject beam search (num_beams != 1), +// - force past_present_share_buffer on, +// - log info on the detected size, +// - warn once per process if search.max_length exceeds the detected size. +// Returns the detected static seq_len, or 0 if the model has symbolic +// kv-cache dims or per-layer static sizes disagree. +// +// Background: some compiled backends (e.g. AMD RyzenAI) emit models where the +// kv-cache seq_len dimension is a static positive integer instead of a +// symbolic dim. In that case the cache must be allocated to exactly that size +// and reused as a shared past/present buffer; max_length cannot drive the +// size because ORT rejects any tensor that doesn't match the model's static +// dim. +// +// Limitation: only uniform per-layer static sizes are recognised — every +// past_key layer must declare the same fixed seq_len. Models that declare +// different static seq_lens per layer (e.g. a mix of full-attention and +// sliding-window layers with distinct static caps) fall through to dynamic +// handling. Lifting this restriction would extend the existing layer_shapes_ +// infrastructure used for per-layer head_dim detection in +// DefaultKeyValueCache: store the per-layer detected seq_len into +// layer_shapes_[i][2] instead of a single scalar, and let the share-buffer +// branch's per-layer loop do the rest. Deferred until a model in the wild +// actually needs it. +int64_t DetectAndConfigureFixedKvShape(const SessionInfo& session_info, + const std::vector& input_name_strings, + int layer_count, + const Config::Search& search, + bool& past_present_share_buffer) { + if (layer_count <= 0) return 0; + + // input_name_strings stores [past_key.0, past_value.0, past_key.1, past_value.1, ...]. + int64_t common_seq_len = 0; + for (int i = 0; i < layer_count; ++i) { + auto input_shape = session_info.GetInputShape(input_name_strings[i * 2]); + if (input_shape.size() < 2) return 0; + const int64_t seq_dim = input_shape[input_shape.size() - 2]; + if (seq_dim <= 0) return 0; // symbolic/dynamic dim (typically -1) + if (common_seq_len == 0) { + common_seq_len = seq_dim; + } else if (common_seq_len != seq_dim) { + return 0; + } + } + + if (search.num_beams != 1) { + throw std::runtime_error( + "Beam search (num_beams > 1) is not supported for models with a fixed kv-cache " + "shape (model expects seq_len=" + + std::to_string(common_seq_len) + ")."); + } + past_present_share_buffer = true; + if (g_log.enabled) { + Log("info", "DefaultKeyValueCache: auto-detected fixed kv-cache seq_len=" + + std::to_string(common_seq_len) + + "; allocating shared past/present buffer to that size."); + } + if (search.max_length > static_cast(common_seq_len)) { + // De-duplicate across repeated Generator constructions (e.g. benchmark loops). + static std::once_flag warned_once; + std::call_once(warned_once, [&] { + Log("warning", "Model has fixed kv-cache seq_len=" + + std::to_string(common_seq_len) + + " but search.max_length=" + + std::to_string(search.max_length) + + "; cache is sized to the model's limit, so generation beyond it will fail."); + }); + } + return common_seq_len; +} + +} // namespace + DefaultKeyValueCache::DefaultKeyValueCache(State& state) : state_{state}, layer_count_{model_.config_->model.decoder.num_hidden_layers}, @@ -251,74 +328,9 @@ DefaultKeyValueCache::DefaultKeyValueCache(State& state) } } - // Auto-detect a fixed kv-cache shape from the model's past_key input shapes. - // Some compiled backends (e.g. AMD RyzenAI) emit models where the kv-cache - // seq_len dimension is a static positive integer instead of a symbolic dim. - // When that's the case the cache must be allocated to exactly that size and - // reused as a shared past/present buffer; max_length cannot drive the size - // because ORT rejects any tensor that doesn't match the model's static dim. - // - // Limitation: only uniform per-layer static sizes are recognised here — - // every past_key layer must declare the same fixed seq_len. Models that - // declare *different* static seq_lens per layer (e.g. a mix of full-attention - // and sliding-window layers with distinct static caps) fall through to - // dynamic handling. Lifting this restriction would extend the existing - // layer_shapes_ infrastructure used for per-layer head_dim detection above: - // store the per-layer detected seq_len into layer_shapes_[i][2] instead of - // a single scalar, and let the share-buffer branch's per-layer loop do the - // rest. Deferred until a model in the wild actually needs it. - int64_t fixed_kv_seq_len = 0; - { - bool all_fixed_and_uniform = layer_count_ > 0; - int64_t common_seq_len = 0; - for (int i = 0; i < layer_count_; ++i) { - auto input_shape = model_.session_info_.GetInputShape(input_name_strings_[i * 2]); - if (input_shape.size() < 2) { - all_fixed_and_uniform = false; - break; - } - const int64_t seq_dim = input_shape[input_shape.size() - 2]; - if (seq_dim <= 0) { // symbolic/dynamic dim (typically -1) - all_fixed_and_uniform = false; - break; - } - if (common_seq_len == 0) { - common_seq_len = seq_dim; - } else if (common_seq_len != seq_dim) { - all_fixed_and_uniform = false; - break; - } - } - if (all_fixed_and_uniform && common_seq_len > 0) { - fixed_kv_seq_len = common_seq_len; - } - } - - if (fixed_kv_seq_len > 0) { - if (state_.params_->search.num_beams != 1) { - throw std::runtime_error( - "Beam search (num_beams > 1) is not supported for models with a fixed kv-cache " - "shape (model expects seq_len=" + - std::to_string(fixed_kv_seq_len) + ")."); - } - past_present_share_buffer_ = true; - if (g_log.enabled) { - Log("info", "DefaultKeyValueCache: auto-detected fixed kv-cache seq_len=" + - std::to_string(fixed_kv_seq_len) + - "; allocating shared past/present buffer to that size."); - } - if (state_.params_->search.max_length > static_cast(fixed_kv_seq_len)) { - // De-duplicate across repeated Generator constructions (e.g. benchmark loops). - static std::once_flag warned_once; - std::call_once(warned_once, [&] { - Log("warning", "Model has fixed kv-cache seq_len=" + - std::to_string(fixed_kv_seq_len) + - " but search.max_length=" + - std::to_string(state_.params_->search.max_length) + - "; cache is sized to the model's limit, so generation beyond it will fail."); - }); - } - } + const int64_t fixed_kv_seq_len = DetectAndConfigureFixedKvShape( + model_.session_info_, input_name_strings_, layer_count_, + state_.params_->search, past_present_share_buffer_); if (state_.params_->use_graph_capture && !past_present_share_buffer_) { // share buffer is a precondition for graph capture From 3c9481e63d5377de9f19801bf3dee877d482f2cd Mon Sep 17 00:00:00 2001 From: Oleksandr Kholodnyi Date: Wed, 20 May 2026 09:44:07 -0700 Subject: [PATCH 4/4] Guard fixed kv-cache warning on g_log.enabled and drop call_once dedup Address review feedback: - Add the standard 'g_log.enabled && g_log.warning' guard around the warning Log call, matching the pattern already used at the past_present_share_buffer warning a few lines above. The bare Log() asserts on g_log.enabled in debug builds. - Drop std::once_flag/std::call_once (and the now-unused include). Per-process dedup is wrong for multi-model hosts (e.g. Foundry Local), where every distinct model loaded into the process should be able to emit its own warning. Callers that re-run generation against the same model should use Generator::RewindTo() rather than constructing a new Generator (and model_benchmark exposes that as --reuse_generator), so dedup machinery for repeated constructions isn't justified. Co-Authored-By: Claude Opus 4 --- src/models/kv_cache.cpp | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/src/models/kv_cache.cpp b/src/models/kv_cache.cpp index 14f9c88be2..02a9fa6b6c 100644 --- a/src/models/kv_cache.cpp +++ b/src/models/kv_cache.cpp @@ -8,7 +8,6 @@ #include "../openvino/interface.h" #include "../qnn/interface.h" #include -#include namespace Generators { @@ -156,7 +155,7 @@ namespace { // - reject beam search (num_beams != 1), // - force past_present_share_buffer on, // - log info on the detected size, -// - warn once per process if search.max_length exceeds the detected size. +// - warn if search.max_length exceeds the detected size. // Returns the detected static seq_len, or 0 if the model has symbolic // kv-cache dims or per-layer static sizes disagree. // @@ -210,16 +209,13 @@ int64_t DetectAndConfigureFixedKvShape(const SessionInfo& session_info, std::to_string(common_seq_len) + "; allocating shared past/present buffer to that size."); } - if (search.max_length > static_cast(common_seq_len)) { - // De-duplicate across repeated Generator constructions (e.g. benchmark loops). - static std::once_flag warned_once; - std::call_once(warned_once, [&] { - Log("warning", "Model has fixed kv-cache seq_len=" + - std::to_string(common_seq_len) + - " but search.max_length=" + - std::to_string(search.max_length) + - "; cache is sized to the model's limit, so generation beyond it will fail."); - }); + if (search.max_length > static_cast(common_seq_len) && + g_log.enabled && g_log.warning) { + Log("warning", "Model has fixed kv-cache seq_len=" + + std::to_string(common_seq_len) + + " but search.max_length=" + + std::to_string(search.max_length) + + "; cache is sized to the model's limit, so generation beyond it will fail."); } return common_seq_len; }