diff --git a/.gitignore b/.gitignore index d8a1b8691a..8d1dba5ae8 100644 --- a/.gitignore +++ b/.gitignore @@ -39,6 +39,7 @@ examples/csharp/ModelChat/models !test/models/qwen2-5-vl/* !test/models/qwen3-5/* !test/models/qwen3-vl/* +!test/models/static-scatter-bias-decoder/* !test/models/whisper/* .ipynb_checkpoints/ diff --git a/src/config.cpp b/src/config.cpp index a32a986c19..605bd13908 100644 --- a/src/config.cpp +++ b/src/config.cpp @@ -321,6 +321,10 @@ struct DecoderInputs_Element : JSON::Element { v_.current_sequence_length = JSON::Get(value); } else if (name == "total_sequence_length") { v_.total_sequence_length = JSON::Get(value); + } else if (name == "write_indices") { + v_.write_indices = JSON::Get(value); + } else if (name == "nonpad_kv_seqlen") { + v_.nonpad_kv_seqlen = JSON::Get(value); } else if (name == "encoder_hidden_states") { v_.encoder_hidden_states = JSON::Get(value); } else if (name == "encoder_attention_mask") { diff --git a/src/config.h b/src/config.h index 8be9de3ecb..599f1cf2a8 100644 --- a/src/config.h +++ b/src/config.h @@ -54,6 +54,8 @@ struct Config { static constexpr std::string_view PastSequenceLengthName = "past_sequence_length"; static constexpr std::string_view CurrentSequenceLengthName = "current_sequence_length"; static constexpr std::string_view TotalSequenceLengthName = "total_sequence_length"; + static constexpr std::string_view WriteIndicesName = "write_indices"; + static constexpr std::string_view NonpadKvSeqlenName = "nonpad_kv_seqlen"; static constexpr std::string_view CacheIndirectionName = "cache_indirection"; static constexpr std::string_view AlignmentHeadsName = "alignment_heads"; static constexpr std::string_view TokenTypeIdsName = "token_type_ids"; @@ -344,6 +346,12 @@ struct Config { std::string past_sequence_length{Defaults::PastSequenceLengthName}; std::string current_sequence_length{Defaults::CurrentSequenceLengthName}; std::string total_sequence_length{Defaults::TotalSequenceLengthName}; + // Static-scatter (TensorScatter) KV cache driver inputs ([batch] int64). + // Both write_indices AND nonpad_kv_seqlen must be present to select + // StaticScatterKeyValueCache: IsStaticScatterCache() (the factory + // predicate) and the DefaultInputIDs producer both require both inputs. + std::string write_indices{Defaults::WriteIndicesName}; + std::string nonpad_kv_seqlen{Defaults::NonpadKvSeqlenName}; std::string cache_indirection{Defaults::CacheIndirectionName}; std::string encoder_hidden_states{Defaults::EncoderHiddenStatesName}; std::string rnn_prev_states{Defaults::RnnStatesPrevName}; diff --git a/src/models/input_ids.cpp b/src/models/input_ids.cpp index 9adc99fe05..4bbc2b7dc6 100644 --- a/src/models/input_ids.cpp +++ b/src/models/input_ids.cpp @@ -29,6 +29,22 @@ DefaultInputIDs::DefaultInputIDs(State& state) *past_sequence_length_->GetTensorMutableData() = -1; } + if (model_.session_info_.HasInput(model_.config_->model.decoder.inputs.write_indices) && + model_.session_info_.HasInput(model_.config_->model.decoder.inputs.nonpad_kv_seqlen)) { + if (state_.params_->BatchBeamSize() != 1) { + throw std::runtime_error("Batch beam size (batch_size * num_beams) must be 1 for write_indices and nonpad_kv_seqlen inputs"); + } + if (model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.write_indices) != Ort::TypeToTensorType || + model_.session_info_.GetInputDataType(model_.config_->model.decoder.inputs.nonpad_kv_seqlen) != Ort::TypeToTensorType) + throw std::runtime_error("write_indices and nonpad_kv_seqlen must be int64"); + + const std::array static_scatter_shape{1}; + write_indices_ = OrtValue::CreateTensor(model_.allocator_cpu_, static_scatter_shape, Ort::TypeToTensorType); + nonpad_kv_seqlen_ = OrtValue::CreateTensor(model_.allocator_cpu_, static_scatter_shape, Ort::TypeToTensorType); + *write_indices_->GetTensorMutableData() = 0; + *nonpad_kv_seqlen_->GetTensorMutableData() = 0; + } + value_ = std::make_unique(model_.p_device_inputs_, Ort::TypeToTensorType); cast_value_ = std::make_unique(model_.p_device_inputs_, Ort::TypeToTensorType); } @@ -45,6 +61,13 @@ void DefaultInputIDs::Add() { state_.input_names_.push_back(model_.config_->model.decoder.inputs.past_sequence_length.c_str()); state_.inputs_.push_back(past_sequence_length_.get()); } + + if (write_indices_ && nonpad_kv_seqlen_) { + state_.input_names_.push_back(model_.config_->model.decoder.inputs.write_indices.c_str()); + state_.inputs_.push_back(write_indices_.get()); + state_.input_names_.push_back(model_.config_->model.decoder.inputs.nonpad_kv_seqlen.c_str()); + state_.inputs_.push_back(nonpad_kv_seqlen_.get()); + } } void DefaultInputIDs::Update(DeviceSpan new_tokens) { @@ -67,6 +90,16 @@ void DefaultInputIDs::Update(DeviceSpan new_tokens) { *past_sequence_length_->GetTensorMutableData() += new_sequence_length; } + if (write_indices_ && nonpad_kv_seqlen_) { + if (state_.params_->BatchBeamSize() != 1) { + throw std::runtime_error("Batch beam size (batch_size * num_beams) must be 1 for write_indices and nonpad_kv_seqlen inputs"); + } + auto new_sequence_length = get_unpadded_sequence_length(new_tokens_cpu, model_.config_->model.pad_token_id); + const StaticScatterIndices indices = static_scatter_indices_.Advance(new_sequence_length); + *write_indices_->GetTensorMutableData() = indices.write_index; + *nonpad_kv_seqlen_->GetTensorMutableData() = indices.nonpad_kv_seqlen; + } + // For beam search, resize input_ids shape based on new_tokens size_t sequence_length = static_cast(new_tokens.size()) / state_.params_->BatchBeamSize(); if (is_prompt_ && state_.params_->search.num_beams > 1) diff --git a/src/models/input_ids.h b/src/models/input_ids.h index 94bed23402..fe936d768a 100644 --- a/src/models/input_ids.h +++ b/src/models/input_ids.h @@ -1,5 +1,7 @@ #pragma once +#include "static_scatter_indices.h" + namespace Generators { struct InputIDs { @@ -41,6 +43,13 @@ struct DefaultInputIDs : InputIDs { std::unique_ptr current_sequence_length_; std::unique_ptr past_sequence_length_; + + // Static-scatter (TensorScatter) KV cache driver inputs, created only when the + // model declares write_indices + nonpad_kv_seqlen. Both are [batch] int64 CPU + // tensors; their per-step values come from static_scatter_indices_. + std::unique_ptr write_indices_; + std::unique_ptr nonpad_kv_seqlen_; + StaticScatterIndexTracker static_scatter_indices_; }; // Certain models can only process a fixed number of tokens at a time. diff --git a/src/models/kv_cache.cpp b/src/models/kv_cache.cpp index 04111afb6a..32743c361c 100644 --- a/src/models/kv_cache.cpp +++ b/src/models/kv_cache.cpp @@ -4,6 +4,7 @@ #include "../generators.h" #include "model.h" #include "kv_cache.h" +#include "static_scatter_indices.h" #include "windowed_kv_cache.h" #include "../openvino/interface.h" #include "../qnn/interface.h" @@ -708,6 +709,130 @@ void ModelManagedKeyValueCache::RewindTo(size_t index) { state_.ep_dynamic_options_next_run_.push_back({"kvcache_rewind", std::to_string(index)}); } +bool StaticScatterKeyValueCache::IsStaticScatterCache(const Model& model) { + // Both driver inputs must be present. input_ids.cpp only produces the indices + // when it sees write_indices AND nonpad_kv_seqlen, so requiring just one here + // would create the cache for a model whose indices never get bound, surfacing + // as an obscure unbound-input error at Run. Keep this predicate in lockstep + // with the producer gate in input_ids.cpp. + // + // NAME-COLLISION ASSUMPTION (rank-3 tightening = TRACKED FOLLOW-UP, not done + // here): selection is purely on the presence of both driver inputs, BEFORE the + // rank-3 static layout is checked. A model that declares both write_indices and + // nonpad_kv_seqlen but uses a 4D KV layout would be routed here and then hit the + // ctor's `rank != 3` throw, rather than falling back to Default/Windowed. + // + // A proper fix would factor KV-input discovery into a shared helper (the index + // parse already lives in DiscoverKvLayerIndices(); the missing piece is locating + // a KV input name and reading its rank at selection time) so that this predicate + // returns false on a rank-4 layout and the factory falls back to Default/ + // Windowed. That is DEFERRED for this PR: it touches the sensitive factory- + // selection path and risks layer-discovery drift for sparse/hybrid models that + // may lack a layer-0 KV input. The collision is considered unlikely (the two + // driver inputs are specific to this layout), so we accept the ctor throw over a + // risky attempt-and-fall-back restructuring until the follow-up lands. + return model.session_info_.HasInput(model.config_->model.decoder.inputs.write_indices) && + model.session_info_.HasInput(model.config_->model.decoder.inputs.nonpad_kv_seqlen); +} + +StaticScatterKeyValueCache::StaticScatterKeyValueCache(State& state) + : state_{state}, + layer_count_{model_.config_->model.decoder.num_hidden_layers} { + if (state_.params_->search.num_beams != 1) { + throw std::runtime_error("Beam search (num_beams > 1) is not supported by the static-scatter KV cache."); + } + // The index tracker advances a single slot per step and cannot be forked + // across batch or beam dimensions; keep this aligned with the DefaultInputIDs + // producer gate, which also requires BatchBeamSize()==1. + if (state_.params_->BatchBeamSize() != 1) { + throw std::runtime_error("The static-scatter KV cache requires batch beam size (batch_size * num_beams) == 1."); + } + + // Auto-discover which layer indices have KV cache inputs (mirrors + // DefaultKeyValueCache so sparse/hybrid layouts work the same way). The strict + // parse / dedup lives in DiscoverKvLayerIndices (static_scatter_indices.h) so + // it can be unit-tested without standing up a Model. + { + const auto& key_template = model_.config_->model.decoder.inputs.past_key_names; + auto prefix = key_template.substr(0, key_template.find('%')); + auto suffix = key_template.substr(key_template.find('%') + 2); + kv_layer_indices_ = DiscoverKvLayerIndices(model_.session_info_.GetInputNames(), prefix, suffix); + } + + if (!kv_layer_indices_.empty()) { + layer_count_ = static_cast(kv_layer_indices_.size()); + } + + for (int i = 0; i < layer_count_; ++i) { + int layer_idx = kv_layer_indices_.empty() ? i : kv_layer_indices_[i]; + input_name_strings_.emplace_back(ComposeKeyValueName(model_.config_->model.decoder.inputs.past_key_names, layer_idx)); + input_name_strings_.emplace_back(ComposeKeyValueName(model_.config_->model.decoder.inputs.past_value_names, layer_idx)); + output_name_strings_.emplace_back(ComposeKeyValueName(model_.config_->model.decoder.outputs.present_key_names, layer_idx)); + output_name_strings_.emplace_back(ComposeKeyValueName(model_.config_->model.decoder.outputs.present_value_names, layer_idx)); + } + + type_ = model_.session_info_.GetInputDataType(input_name_strings_[0]); + + // Each KV input declares shape [batch, max_seq_len, kv_hidden]. The batch dim + // is a runtime property (symbolic in the graph), so take it from the params + // like DefaultKeyValueCache; only max_seq_len and kv_hidden must be static. + // kv_hidden = num_kv_heads * head_dim and may vary per layer, so read each + // layer's own declared shape rather than assuming a uniform value. + const int64_t batch_size = state_.params_->BatchBeamSize(); + caches_.reserve(layer_count_ * 2); + for (int i = 0; i < layer_count_ * 2; ++i) { + const auto input_shape = model_.session_info_.GetInputShape(input_name_strings_[i]); + if (input_shape.size() != 3) { + throw std::runtime_error( + "StaticScatterKeyValueCache expects 3D [batch, max_seq_len, kv_hidden] KV inputs, but '" + + input_name_strings_[i] + "' has rank " + std::to_string(input_shape.size()) + "."); + } + // max_seq_len (axis 1) and kv_hidden (axis 2) size the fixed buffer and must + // be concrete; the batch dim (axis 0) is allowed to be symbolic. + for (size_t axis = 1; axis < 3; ++axis) { + if (input_shape[axis] <= 0) { + throw std::runtime_error( + "StaticScatterKeyValueCache requires a static max_seq_len and kv_hidden, but '" + + input_name_strings_[i] + "' has a non-concrete dim at axis " + std::to_string(axis) + "."); + } + } + std::array tensor_shape{batch_size, input_shape[1], input_shape[2]}; + + caches_.push_back(OrtValue::CreateTensor(Allocator(), tensor_shape, type_)); + if (Device().GetType() != DeviceType::WEBGPU) { + ByteWrapTensor(Device(), *caches_.back()).Zero(); + } + } +} + +void StaticScatterKeyValueCache::Add() { + // Past and present share one buffer: TensorScatter writes new rows in place, + // so key_cache.{i} (input) and updated_key_cache.{i} (output) point at the + // same OrtValue and never need rebinding between steps. + for (int i = 0; i < layer_count_ * 2; ++i) { + state_.inputs_.push_back(caches_[i].get()); + state_.input_names_.push_back(input_name_strings_[i].c_str()); + state_.outputs_.push_back(caches_[i].get()); + state_.output_names_.push_back(output_name_strings_[i].c_str()); + } +} + +void StaticScatterKeyValueCache::Update(DeviceSpan /*beam_indices*/, int /*total_length*/) { + // No-op: the shared buffer is updated in place by the graph's TensorScatter, + // and the write offset / valid length are carried by the write_indices / + // nonpad_kv_seqlen inputs (see input_ids.cpp), not by rebinding tensors here. +} + +void StaticScatterKeyValueCache::RewindTo(size_t /*index*/) { + // Fail loud: rewind is NOT wired for the static-scatter cache. The + // write_indices/nonpad_kv_seqlen stream lives in InputIDs and has no RewindTo + // hook, so a silent no-op here would leave the index tracker stale -> wrong + // scatter slots and an over-reported nonpad => silently wrong logits with no + // error. Throw until rewind is properly wired, matching the LFM2Cache and + // WindowedKeyValueCache siblings. + throw std::runtime_error("StaticScatterKeyValueCache does not support RewindTo."); +} + LFM2Cache::LFM2Cache(State& state) : state_{state}, layer_types_{model_.config_->model.decoder.layer_types}, @@ -947,6 +1072,15 @@ std::unique_ptr CreateKeyValueCache(State& state) { return nullptr; } + // mobius static-cache decoders drive an in-place TensorScatter KV buffer via + // the write_indices input; auto-detect that (no user-visible search flag, + // mirroring DetectAndConfigureFixedKvShape) before the default fallback. + if (StaticScatterKeyValueCache::IsStaticScatterCache(state.model_)) { + if (g_log.enabled) + Log("info", "CreateKeyValueCache: Creating StaticScatterKeyValueCache"); + return std::make_unique(state); + } + if (state.model_.p_device_->GetType() != DeviceType::NvTensorRtRtx && state.model_.config_->model.decoder.sliding_window && state.model_.config_->model.decoder.sliding_window->slide_key_value_cache) { diff --git a/src/models/kv_cache.h b/src/models/kv_cache.h index e20e40e1d6..88d4bf8dee 100644 --- a/src/models/kv_cache.h +++ b/src/models/kv_cache.h @@ -194,6 +194,54 @@ struct LFM2Cache : KeyValueCache { std::string ComposeKeyValueName(const std::string& template_string, int index); +// A static-scatter KV cache for mobius-exported static-cache decoders. +// +// The model pre-allocates each layer's KV as a fixed 3D buffer +// [batch, max_seq_len, kv_hidden] and writes new rows in place via opset-24 +// TensorScatter (driven by the write_indices / nonpad_kv_seqlen inputs produced +// in input_ids.cpp), reading them back through Attention. Because the scatter is +// in place, past and present share one buffer: Add() binds key_cache.{i} and +// updated_key_cache.{i} to the same OrtValue, and Update() is a no-op (mirroring +// DefaultKeyValueCache's past_present_share_buffer path). RewindTo() is NOT +// supported and throws: the write_indices/nonpad_kv_seqlen index stream lives in +// InputIDs with no rewind hook, so rewinding would silently desynchronize it. +// +// Distinct from DefaultKeyValueCache in two ways. (1) Layout: it consumes +// mobius's 3D emission directly (vs the 4D [batch, num_kv_heads, seq, head_dim] +// layout), and kv_hidden may vary per layer (e.g. Gemma-4 sliding GQA 8*256 vs +// global MQA 1*512), read from each layer's own declared input shape. +// (2) RewindTo: DefaultKeyValueCache rewinds by reshaping its buffers, whereas +// this cache THROWS (rewind is unsupported), because the write_indices / +// nonpad_kv_seqlen index stream lives in InputIDs with no rewind hook and a +// silent no-op would desynchronize it. +struct StaticScatterKeyValueCache : KeyValueCache { + StaticScatterKeyValueCache(State& state); + + // True if the model declares BOTH static-scatter driver inputs (write_indices + // and nonpad_kv_seqlen); kept in lockstep with the producer gate in input_ids. + static bool IsStaticScatterCache(const Model& model); + + void Add() override; + void Update(DeviceSpan beam_indices, int total_length) override; + void RewindTo(size_t index) override; + + private: + DeviceInterface& Device() { return *model_.p_device_kvcache_; } + Ort::Allocator& Allocator() { return model_.p_device_kvcache_->GetAllocator(); } + + State& state_; + const Model& model_{state_.model_}; + int layer_count_; + + // Auto-discovered KV layer indices (sparse for hybrid models). + std::vector kv_layer_indices_; + ONNXTensorElementDataType type_; + + // One shared past/present buffer per key and per value tensor (2 per layer). + std::vector> caches_; + std::vector input_name_strings_, output_name_strings_; +}; + std::unique_ptr CreateKeyValueCache(State& state); } // namespace Generators diff --git a/src/models/static_scatter_indices.h b/src/models/static_scatter_indices.h new file mode 100644 index 0000000000..b30565d1ec --- /dev/null +++ b/src/models/static_scatter_indices.h @@ -0,0 +1,121 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace Generators { + +// The per-step index pair a static-scatter (TensorScatter) KV cache needs. +// +// A mobius-exported static-cache decoder consumes a pre-allocated KV buffer of +// shape [batch, max_seq_len, kv_hidden] and writes each step's new key/value +// rows into it in place via TensorScatter, then reads them back through +// Attention. Two int64 [batch] inputs drive that: +// * write_indices - the cache row offset TensorScatter writes this step's +// rows at (i.e. how many valid tokens are already cached +// BEFORE this step). +// * nonpad_kv_seqlen - the number of valid cached tokens AFTER this step, +// which Attention reads as the per-batch seqlens_k. +struct StaticScatterIndices { + int64_t write_index; // valid cache tokens before this step (scatter offset) + int64_t nonpad_kv_seqlen; // valid cache tokens after this step (Attention seqlens_k) +}; + +// Tracks the running static-scatter cache indices for a single (batch==1) +// generation stream. Kept as a standalone, dependency-free helper so the +// off-by-one / init behaviour can be unit-tested without standing up a Model. +// +// Sequencing contract (the crux mobius and genai must agree on): +// * The very first step (prefill) writes at row 0 and reports nonpad equal to +// the number of prefill tokens. +// * Each subsequent step's write_index is the PREVIOUS step's nonpad_kv_seqlen, +// so rows are appended contiguously with no gap or overlap. +// This deliberately does NOT reuse genai's existing past_sequence_length scalar +// (which inits to -1 and is consumed differently); mixing the two would yield +// nonpad = 2N-1 after a length-N prefill instead of N. +class StaticScatterIndexTracker { + public: + // Advance one generation step that appended new_unpadded_tokens valid tokens + // (the prompt length on prefill, normally 1 per decode step). Returns the + // index pair to bind for THIS step, then folds the new tokens into the + // running total for the next step. + // + // pad_token aliasing assumption: a decode step whose generated token equals + // pad_token_id yields new_unpadded_tokens == 0 (the unpadded-length probe in + // input_ids.cpp counts it as padding), so write_index / valid_tokens_ do NOT + // advance and the NEXT real token would scatter onto the same cache slot. + // The CONDITION that makes this safe: for the targeted models pad_token_id == + // eos_token_id, so the very step that produces a pad token also ENDS the + // sequence -- generation terminates before any later step could read the + // stalled slot. Given that, the aliasing is benign (and mirrors genai's + // existing current_sequence_length logic). It would only break if a model made + // pad_token_id a legal mid-stream generated token, which the targeted models + // do not. + StaticScatterIndices Advance(int64_t new_unpadded_tokens) { + const int64_t write_index = valid_tokens_; + valid_tokens_ += new_unpadded_tokens; + return {write_index, valid_tokens_}; + } + + // Valid cached tokens before the next step. Zero before any Advance(). + int64_t valid_tokens() const { return valid_tokens_; } + + private: + int64_t valid_tokens_{0}; +}; + +// Discover which decoder layers expose a KV cache input, parsing the layer index +// out of each matching input name. `prefix` / `suffix` bracket the numeric index +// in the past-key name template (e.g. "past_key_values.%d.key" -> prefix +// "past_key_values.", suffix ".key"). +// +// The parse is STRICT: the index segment must be a COMPLETE, non-negative +// integer. std::stoi would accept trailing junk (e.g. "past_key_values.0.bad.key" +// -> 0) and silently mis-map a layer, so std::from_chars must consume the whole +// segment with no leftover characters. Duplicate indices are rejected, because +// two inputs mapping to one layer would double-count layer_count_ and bind the +// same cache slot twice. Throws std::runtime_error on a malformed or duplicate +// index. Returned indices are sorted ascending. +// +// Kept here as a standalone, model-free helper (alongside StaticScatterIndex- +// Tracker) so the strict-parse / dedup behaviour can be unit-tested directly, +// without standing up a Model. +inline std::vector DiscoverKvLayerIndices(const std::vector& input_names, + const std::string& prefix, + const std::string& suffix) { + std::vector indices; + for (const auto& name : input_names) { + if (name.size() > prefix.size() + suffix.size() && + name.compare(0, prefix.size(), prefix) == 0 && + name.compare(name.size() - suffix.size(), suffix.size(), suffix) == 0) { + const auto idx_str = name.substr(prefix.size(), name.size() - prefix.size() - suffix.size()); + int layer_idx = 0; + const char* begin = idx_str.data(); + const char* end = begin + idx_str.size(); + auto [parse_end, ec] = std::from_chars(begin, end, layer_idx); + if (ec != std::errc{} || parse_end != end || layer_idx < 0) { + throw std::runtime_error( + "StaticScatterKeyValueCache: input '" + name + + "' has a malformed KV layer index '" + idx_str + + "' (expected a non-negative integer)."); + } + if (std::find(indices.begin(), indices.end(), layer_idx) != indices.end()) { + throw std::runtime_error( + "StaticScatterKeyValueCache: duplicate KV layer index '" + + std::to_string(layer_idx) + "' (from input '" + name + "')."); + } + indices.push_back(layer_idx); + } + } + std::sort(indices.begin(), indices.end()); + return indices; +} + +} // namespace Generators diff --git a/test/models/static-scatter-bias-decoder/genai_config.json b/test/models/static-scatter-bias-decoder/genai_config.json new file mode 100644 index 0000000000..ef0e2ec6a2 --- /dev/null +++ b/test/models/static-scatter-bias-decoder/genai_config.json @@ -0,0 +1,37 @@ +{ + "model": { + "type": "mistral", + "pad_token_id": 0, + "bos_token_id": 1, + "eos_token_id": 2, + "vocab_size": 256, + "context_length": 16, + "decoder": { + "session_options": { + "provider_options": [] + }, + "filename": "model.onnx", + "num_key_value_heads": 2, + "head_size": 16, + "num_attention_heads": 4, + "num_hidden_layers": 2, + "inputs": { + "input_ids": "input_ids", + "position_ids": "position_ids", + "past_key_names": "key_cache.%d", + "past_value_names": "value_cache.%d", + "write_indices": "write_indices", + "nonpad_kv_seqlen": "nonpad_kv_seqlen" + }, + "outputs": { + "logits": "logits", + "present_key_names": "updated_key_cache.%d", + "present_value_names": "updated_value_cache.%d" + } + } + }, + "search": { + "max_length": 16, + "do_sample": false + } +} diff --git a/test/models/static-scatter-bias-decoder/golden_io.npz b/test/models/static-scatter-bias-decoder/golden_io.npz new file mode 100644 index 0000000000..ae13f05dca Binary files /dev/null and b/test/models/static-scatter-bias-decoder/golden_io.npz differ diff --git a/test/models/static-scatter-bias-decoder/model.onnx b/test/models/static-scatter-bias-decoder/model.onnx new file mode 100644 index 0000000000..0e203c4767 Binary files /dev/null and b/test/models/static-scatter-bias-decoder/model.onnx differ diff --git a/test/static_scatter_golden.h b/test/static_scatter_golden.h new file mode 100644 index 0000000000..e9bc71e0ab --- /dev/null +++ b/test/static_scatter_golden.h @@ -0,0 +1,297 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +// +// ============================================================================= +// GENERATED FILE -- DO NOT EDIT BY HAND. +// ============================================================================= +// +// Full-tensor golden outputs for the static-scatter slice-A fixture +// (test/models/static-scatter-bias-decoder), consumed by EndToEndFixtureParity +// in static_scatter_kv_cache_test.cpp to assert element-wise parity (a sum/argmax +// check alone passes wrong-but-sum-preserving outputs). +// +// AUTHORITATIVE SOURCE +// test/models/static-scatter-bias-decoder/golden_io.npz (committed alongside +// the fixture model.onnx) -- a recorded ORT CPU MEA reference run (opset 24) of +// the fixture model above, the SAME independent reference the scalar argmax/sum +// goldens in static_scatter_kv_cache_test.cpp come from. The fixture's +// genai_config declares no execution provider, so it always runs the CPU EP; +// that is why these goldens match a C++ run to within kParityAtol (1e-3). +// +// HOW golden_io.npz WAS PRODUCED (already committed; recipe kept for refresh) +// Feed the fixture model the documented fixed inputs through ORT CPU and record +// every output. Concretely, with onnxruntime in Python: +// +// import numpy as np, onnxruntime as ort +// M = "test/models/static-scatter-bias-decoder/model.onnx" +// sess = ort.InferenceSession(M, providers=["CPUExecutionProvider"]) +// B, MAXS, KVH = 1, 16, 32 # batch, max_seq_len, kv_hidden (per fixture) +// z = lambda: np.zeros((B, MAXS, KVH), np.float32) +// # --- prefill: 4 tokens at row 0, caches zero-initialized --- +// pf = { +// "input_ids": np.array([[1, 2, 3, 4]], np.int64), +// "position_ids": np.array([[0, 1, 2, 3]], np.int64), +// "write_indices": np.array([0], np.int64), # shape [B] +// "nonpad_kv_seqlen": np.array([4], np.int64), # shape [B] +// "key_cache.0": z(), "value_cache.0": z(), +// "key_cache.1": z(), "value_cache.1": z(), +// } +// pout = sess.run(None, pf) +// onames = [o.name for o in sess.get_outputs()] +// pf_out = dict(zip(onames, pout)) +// # --- decode: 1 token at row 4, reusing prefill-updated caches --- +// dec = { +// "input_ids": np.array([[5]], np.int64), +// "position_ids": np.array([[4]], np.int64), +// "write_indices": np.array([4], np.int64), +// "nonpad_kv_seqlen": np.array([5], np.int64), +// "key_cache.0": pf_out["updated_key_cache.0"], +// "value_cache.0": pf_out["updated_value_cache.0"], +// "key_cache.1": pf_out["updated_key_cache.1"], +// "value_cache.1": pf_out["updated_value_cache.1"], +// } +// dout = dict(zip(onames, sess.run(None, dec))) +// np.savez("golden_io.npz", +// **{f"prefill_out.{k}": v for k, v in pf_out.items()}, +// **{f"decode_out.{k}": v for k, v in dout.items()}, +// **{f"prefill_in.{k}": v for k, v in pf.items()}, +// **{f"decode_in.{k}": v for k, v in dec.items()}) +// (Input names/dtypes/values above are the exact bindings the genai generator +// uses for this fixture; verify against the model's get_inputs() if it changes.) +// +// HOW TO REGENERATE THIS HEADER FROM golden_io.npz +// Run the following Python (numpy) snippet, pointing NPZ at the committed +// golden_io.npz above. It slices the per-array source below and emits valid C++ +// float literals (%.9g, with a '.0' appended when the formatted value has no +// '.'/exponent, so e.g. 0 -> "0.0f" rather than the invalid "0f"): +// +// import numpy as np +// d = np.load(NPZ) # test/models/static-scatter-bias-decoder/golden_io.npz +// def lit(x): +// x = float(np.float32(x)); s = "{:.9g}".format(x) +// return (s + (".0" if not any(c in s for c in ".eEnN") else "")) + "f" +// # kPrefillLastTokLogitsGolden = d['prefill_out.logits'][0, -1, :] # [256] +// # kDecodeLastTokLogitsGolden = d['decode_out.logits'][0, 0, :] # [256] +// # kPrefillUpdatedKeyCache0Golden = d['prefill_out.updated_key_cache.0'].ravel() # [512] +// # kDecodeUpdatedKeyCache0Golden = d['decode_out.updated_key_cache.0'].ravel() # [512] +// # emit each as "inline constexpr float [] = {", then the values 8 +// # per line, and place the closing "};" on the SAME line as the FINAL value +// # (no line break before "}"). With the repo .clang-format (ColumnLimit: 0) +// # this output is already clang-format-clean, so the regen is a single step +// # and "clang-format -i" is a verified no-op (no digits/whitespace change). +// +// Last-token logits are [vocab=256]; the updated_key_cache.0 buffer is +// [batch=1, max_seq_len=16, kv_hidden=32] = 512 elements. +#pragma once + +namespace Generators::test { + +// d['prefill_out.logits'][0, -1, :] -- prefill last-token logits [256] +// (input_ids [1,2,3,4] @ write_index 0, nonpad 4). +inline constexpr float kPrefillLastTokLogitsGolden[] = { + 7.32147932f, 2.07709384f, -0.557281017f, -6.75512075f, 5.06091118f, 0.437584609f, 6.63706923f, -5.08704519f, + -5.14650249f, -12.3084555f, -3.72699022f, 18.6615105f, 5.6121521f, -10.7090082f, -12.3951082f, -2.01596904f, + -10.7374859f, -0.977514625f, 6.50934172f, 9.12832832f, 4.89916182f, -3.62374616f, 0.852583587f, -4.97442436f, + 5.94551849f, 16.0105114f, 4.52778435f, 1.25030994f, 9.71587372f, 6.40192604f, -0.496018231f, -1.74474549f, + 1.72573113f, -2.23981524f, 5.88913441f, 8.15920162f, 12.7738361f, 0.313029617f, 10.7856092f, 6.67838192f, + -4.10607386f, -0.0159605611f, -7.4326148f, 4.09218597f, -4.13840342f, 0.616579533f, -4.86502552f, 0.116430476f, + 3.47575617f, 0.261647195f, 1.42642105f, 3.40532207f, -6.75533295f, 5.10773849f, -6.63309193f, -6.44912338f, + -6.58160067f, 4.76952076f, 7.09192753f, -3.52454209f, -4.45940733f, 12.9240313f, -2.7946744f, 10.4820518f, + 0.139866471f, 5.75305939f, -5.44895744f, 3.75838327f, -2.66457272f, 0.108811028f, -5.38652706f, -0.48316288f, + 0.760434091f, -1.30405056f, 6.36170149f, 8.11850929f, -2.83244157f, -7.65983248f, -8.02928352f, -1.91442454f, + -4.48358154f, -0.943943799f, -9.61816883f, 1.69027245f, 2.26689196f, 2.30658293f, 2.03520989f, -10.2187567f, + -12.1563702f, 0.973694026f, 0.303455114f, 1.21975017f, 4.31499672f, 12.3120317f, -11.8075132f, -3.50527024f, + 5.39970636f, 2.29451704f, 10.0618114f, 0.289482832f, 7.66823769f, 12.5025358f, 13.9188004f, 2.21803069f, + -5.29775953f, 1.8354311f, -4.70493698f, 8.22589874f, -8.9475193f, 15.289547f, -2.45194435f, 2.0209589f, + 3.06274104f, -4.78041935f, -4.47679949f, 7.05563784f, 4.40434122f, 2.24652028f, 1.37222958f, -8.85774612f, + -10.3043776f, 14.1776772f, -6.67577696f, -9.35319233f, 4.03369474f, -6.4065094f, -4.33496237f, -10.1195745f, + 0.0400873311f, 6.02156305f, -6.16679001f, 6.47703648f, 7.59334469f, -1.25835943f, -11.4756517f, -11.8567305f, + -7.69087791f, -0.49272716f, -8.2004776f, -8.14299488f, 8.56215286f, -3.50410986f, -10.3120861f, -4.78805876f, + 4.05592108f, 4.54013586f, -3.73437977f, -2.96257687f, -0.0869583115f, 10.264884f, 6.63542461f, -19.043808f, + -6.53306341f, 2.65079665f, -7.06935453f, 1.93224776f, 8.81772614f, -8.46978855f, 5.66759729f, 3.8439517f, + -9.52455711f, 11.7744236f, -7.27653837f, 5.26825285f, -7.46642017f, 5.20434332f, -3.09576082f, 2.28935027f, + 6.33351612f, 10.0633717f, -0.762654662f, 3.52562618f, 2.35045457f, -6.52551126f, -4.02287865f, -9.03767872f, + 8.09876728f, -3.16576624f, -1.81629372f, 2.13947153f, -9.1076746f, -3.20141411f, -9.61193943f, 8.21758938f, + 5.57498169f, -2.08062291f, -8.64711761f, -2.88570189f, -13.2547789f, -12.0987511f, 0.587717354f, -2.39974165f, + -6.61053801f, -8.28171825f, 1.22169876f, 1.71424627f, -11.3196888f, 7.09847641f, -3.95212436f, -9.92163849f, + 7.80496931f, -10.8219032f, 1.61763167f, 12.5319614f, 0.689085901f, -2.76638532f, 3.94318318f, 4.52733612f, + 6.38785028f, -3.88717103f, 3.66903424f, 16.5564899f, -1.41133749f, 0.627034485f, 2.73648953f, 6.34655619f, + -7.72259998f, -1.10194695f, 0.417099744f, -3.37768626f, 5.9978013f, 4.84983873f, -7.89546442f, -1.04398966f, + -2.90231109f, 4.62135983f, -3.61740851f, -1.73753154f, -12.2318411f, 5.68573856f, 5.11865711f, 2.92729402f, + -2.89803171f, 10.0524645f, -10.3232079f, 5.24708462f, -8.62805653f, 2.74606276f, -6.01792622f, 4.47605181f, + 3.9671011f, -4.22735834f, -10.7976007f, -1.5516504f, 3.69843888f, 4.66894436f, -9.34679031f, 2.2463336f, + 3.29976392f, 11.0057621f, 4.72564793f, -2.26639509f, 4.78171635f, 6.97382593f, 4.53631353f, -8.65044498f}; + +// d['decode_out.logits'][0, 0, :] -- decode last-token logits [256] +// (input_ids [5] @ write_index 4, nonpad 5). +inline constexpr float kDecodeLastTokLogitsGolden[] = { + 3.04830623f, -2.1821003f, 0.509211361f, 8.00238228f, 14.1502399f, 2.56521225f, 11.3899832f, -8.03155613f, + -10.2850733f, 5.46324444f, -1.4492147f, 13.1968269f, 0.144806564f, -11.2526817f, 5.22660351f, 0.672616303f, + 7.9224577f, 2.75280857f, 7.81075382f, -4.61093569f, 6.89437056f, 9.3377943f, 1.89660597f, 0.31125471f, + 5.129951f, -3.38609195f, -5.59743071f, -5.63517857f, 0.40092662f, -0.951532066f, 6.53623533f, 0.385277867f, + 9.96651554f, -2.30869317f, 11.2010689f, 6.41328192f, 20.6644154f, -6.04290247f, 2.68440533f, 5.99150038f, + -3.32973623f, 5.83996677f, 7.57308006f, 5.85156679f, -2.90182567f, 1.25172186f, -1.55216146f, 2.48068643f, + -8.98115158f, 2.26843834f, -6.25523996f, -3.76567054f, -1.55692065f, -1.31021917f, -17.2543621f, -4.73680115f, + 4.79129696f, 4.43301868f, -5.04191256f, -4.45821571f, 1.11839592f, 10.7114334f, 10.3363142f, 4.5869875f, + -4.24611187f, 8.66466141f, -15.1922913f, 3.81752467f, 1.27213383f, -0.277198255f, 4.25113297f, 1.98459947f, + 2.50251341f, -8.2445097f, -0.570915282f, 11.4630165f, -0.121401325f, -2.09976959f, -2.0684073f, 1.75497413f, + -5.32069016f, -1.1862613f, -3.7123301f, -2.49625731f, 9.96527576f, 7.61265469f, -4.80026245f, -4.73949194f, + 0.0539198145f, 5.81623268f, 1.23488975f, -0.108031996f, 2.0823586f, 8.60173607f, 3.52535558f, -4.20718479f, + -8.58668804f, -5.12117004f, -3.36042643f, -3.61958051f, 10.2176113f, 11.5389862f, 13.9725113f, 2.00129437f, + 13.7092819f, 6.71030903f, 4.75589991f, 1.31852889f, 4.24631548f, 8.67805958f, 0.153570324f, 0.41846031f, + 4.73776436f, -4.99446487f, -7.07331419f, 2.22066903f, 4.53326416f, 3.46435976f, 0.557397425f, -12.1032677f, + 0.261833549f, 0.410538346f, -23.9129353f, 1.84614754f, -0.651473463f, -0.852924824f, 0.844337463f, -9.90745831f, + 11.5527029f, -7.32661343f, -2.80231261f, 3.81387115f, 14.4366531f, -9.4370327f, -0.571381092f, -6.06561565f, + -9.75665569f, 5.98297119f, -6.6226244f, -0.173047945f, -2.20043612f, -8.93277645f, -5.04397392f, -9.34880257f, + 8.45656776f, 8.29689789f, -3.09127259f, 2.93323541f, -6.10184717f, 7.07846689f, 6.14048195f, -5.34002638f, + -5.54564571f, -6.00822163f, -6.23473883f, 6.58850336f, 1.89152205f, -5.75556183f, -10.6014547f, -2.89770699f, + 1.02006948f, 6.62004471f, -5.33858013f, 13.269515f, -6.93260193f, 13.1594839f, -1.48868668f, -0.605923474f, + 4.36255455f, -2.79251456f, 6.51062775f, 0.959224105f, -8.34952641f, -3.33682346f, -0.863444448f, -6.24669647f, + -1.75509024f, 10.5798025f, -0.465528697f, -10.0351458f, 5.64303446f, -14.4136744f, -5.64676809f, 8.43798447f, + 0.449698329f, 0.899144351f, -8.04611015f, -1.75362051f, -10.9259634f, 0.192091823f, -3.76766443f, 5.7393651f, + 9.85207748f, -7.27534533f, 4.33441925f, 10.1686993f, -5.63948059f, -5.6496172f, -11.7046385f, -2.26661181f, + -0.702790499f, -7.61027718f, 0.12830475f, -1.4042449f, -0.0756285116f, -2.6217463f, 12.5725212f, -0.217608988f, + 0.212240249f, -3.66012502f, -4.30753946f, 1.07191825f, 0.3540591f, 2.14598584f, 0.651568592f, -3.76242805f, + -0.243642583f, -5.23432207f, -8.71735477f, 8.90773106f, 2.4235363f, 0.725729346f, -14.0902624f, -6.71332598f, + 1.68114519f, 8.91252518f, -8.34446812f, 6.05059433f, -6.9353776f, 10.9396544f, 11.5374079f, -0.411692858f, + -9.3254776f, 2.73739123f, -20.5660419f, 6.62059069f, -5.05113792f, 8.66979694f, 0.142779008f, -0.524229169f, + 13.7954569f, -15.1851072f, 8.67184448f, 8.80776119f, 8.57041073f, -7.29395485f, 6.60269308f, -2.3672061f, + 7.83773088f, 2.50839329f, 3.15586114f, -3.38412929f, -0.403188199f, -1.8428973f, 2.26576328f, -3.61321878f}; + +// d['prefill_out.updated_key_cache.0'].ravel() -- [1,16,32]=512 +// (rows 0..3 written by TensorScatter, remaining rows zero). +inline constexpr float kPrefillUpdatedKeyCache0Golden[] = { + -9.29390717f, 10.8017168f, 25.0232124f, 0.842968822f, -17.7853756f, 4.99596262f, -0.980121076f, -2.60616899f, + -8.24257565f, -11.0636349f, 3.0645926f, 7.87195587f, 1.76100171f, -1.47335982f, 4.82979012f, 5.88793516f, + -3.26263547f, 0.770600617f, -2.29790735f, -2.36862731f, 2.49676466f, -1.93307781f, 6.707623f, 1.41513669f, + -4.83117056f, -19.7080727f, 2.33562326f, -1.8799938f, -10.6249771f, 7.53300285f, 2.44383907f, 1.81175351f, + -8.7004118f, 0.885841012f, 1.50503278f, 9.06685543f, -5.85096979f, -7.78261471f, -0.546558559f, 4.02744627f, + -1.18524015f, 3.05206752f, 4.7859354f, -4.7297349f, -6.09195471f, -0.309818447f, 3.91187596f, -3.49300504f, + -9.38397026f, 4.05882692f, -1.48766255f, 9.57935619f, 7.66612053f, -3.9526403f, 7.6680336f, -0.246113881f, + 0.664837956f, -2.76427627f, -2.89326382f, -1.66666329f, -16.4167976f, -10.2596884f, 0.906531572f, -6.18131065f, + 12.3709888f, -5.29498053f, 4.97050238f, -6.07616234f, 0.585708082f, -1.56864667f, -7.46582794f, -5.26050997f, + -2.03424692f, -0.219650656f, 12.2274323f, 0.0804507658f, -7.1458087f, 4.13784742f, 2.58180141f, -4.6287384f, + -0.0828146935f, 2.93192792f, 4.34070587f, -6.41805553f, -0.389455259f, 5.78705692f, -2.35613275f, 0.189541504f, + -5.01314926f, 0.296766788f, -4.11742115f, 8.73949528f, 9.60200691f, -0.927558661f, 7.58941412f, -4.59501553f, + 5.08014679f, -0.397514224f, -5.88721609f, 11.1417484f, 11.4109135f, -3.77362061f, -0.112290829f, 4.59168434f, + 1.40614533f, -1.46569026f, -5.63041496f, 15.0410271f, -0.342983961f, 1.49602675f, 1.59787953f, 2.79422665f, + -2.9813602f, 9.32386112f, 3.4557066f, 3.120193f, -4.04099369f, -4.30661869f, 3.79902935f, -1.27486908f, + -1.02755749f, 2.87558699f, -2.4337163f, -0.757302046f, -0.556162f, -2.3841033f, -3.26099253f, -1.10462368f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + +// d['decode_out.updated_key_cache.0'].ravel() -- [1,16,32]=512 +// (row 4 additionally written on the decode step). +inline constexpr float kDecodeUpdatedKeyCache0Golden[] = { + -9.29390717f, 10.8017168f, 25.0232124f, 0.842968822f, -17.7853756f, 4.99596262f, -0.980121076f, -2.60616899f, + -8.24257565f, -11.0636349f, 3.0645926f, 7.87195587f, 1.76100171f, -1.47335982f, 4.82979012f, 5.88793516f, + -3.26263547f, 0.770600617f, -2.29790735f, -2.36862731f, 2.49676466f, -1.93307781f, 6.707623f, 1.41513669f, + -4.83117056f, -19.7080727f, 2.33562326f, -1.8799938f, -10.6249771f, 7.53300285f, 2.44383907f, 1.81175351f, + -8.7004118f, 0.885841012f, 1.50503278f, 9.06685543f, -5.85096979f, -7.78261471f, -0.546558559f, 4.02744627f, + -1.18524015f, 3.05206752f, 4.7859354f, -4.7297349f, -6.09195471f, -0.309818447f, 3.91187596f, -3.49300504f, + -9.38397026f, 4.05882692f, -1.48766255f, 9.57935619f, 7.66612053f, -3.9526403f, 7.6680336f, -0.246113881f, + 0.664837956f, -2.76427627f, -2.89326382f, -1.66666329f, -16.4167976f, -10.2596884f, 0.906531572f, -6.18131065f, + 12.3709888f, -5.29498053f, 4.97050238f, -6.07616234f, 0.585708082f, -1.56864667f, -7.46582794f, -5.26050997f, + -2.03424692f, -0.219650656f, 12.2274323f, 0.0804507658f, -7.1458087f, 4.13784742f, 2.58180141f, -4.6287384f, + -0.0828146935f, 2.93192792f, 4.34070587f, -6.41805553f, -0.389455259f, 5.78705692f, -2.35613275f, 0.189541504f, + -5.01314926f, 0.296766788f, -4.11742115f, 8.73949528f, 9.60200691f, -0.927558661f, 7.58941412f, -4.59501553f, + 5.08014679f, -0.397514224f, -5.88721609f, 11.1417484f, 11.4109135f, -3.77362061f, -0.112290829f, 4.59168434f, + 1.40614533f, -1.46569026f, -5.63041496f, 15.0410271f, -0.342983961f, 1.49602675f, 1.59787953f, 2.79422665f, + -2.9813602f, 9.32386112f, 3.4557066f, 3.120193f, -4.04099369f, -4.30661869f, 3.79902935f, -1.27486908f, + -1.02755749f, 2.87558699f, -2.4337163f, -0.757302046f, -0.556162f, -2.3841033f, -3.26099253f, -1.10462368f, + -5.66989565f, -5.78240442f, 1.98012221f, -6.78923988f, 12.5755196f, 10.0883236f, 4.15663481f, -5.95230722f, + -5.14942312f, -6.5461688f, -2.30390453f, 2.63517213f, 3.67549992f, -5.18695545f, -11.2737637f, 2.47666097f, + -2.95048976f, -3.02296424f, -3.18152666f, 14.9827433f, -11.2399702f, -3.75876999f, -1.66768241f, 7.25395632f, + -5.13640833f, -1.79465723f, -10.2185383f, -4.30390215f, 6.81118441f, 0.0244651996f, 4.45140266f, -10.8126154f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + +} // namespace Generators::test diff --git a/test/static_scatter_kv_cache_test.cpp b/test/static_scatter_kv_cache_test.cpp new file mode 100644 index 0000000000..fff68c91d8 --- /dev/null +++ b/test/static_scatter_kv_cache_test.cpp @@ -0,0 +1,298 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "span.h" +#define OGA_USE_SPAN 1 +#include +#include + +#include "models/static_scatter_indices.h" +#include "static_scatter_golden.h" +#include "test_utils.h" + +namespace Generators::test { + +// The StaticScatterIndexTracker drives the write_indices / nonpad_kv_seqlen +// inputs of a mobius static-scatter (TensorScatter) KV cache. These tests pin +// the off-by-one / init contract: the first (prefill) step writes at row 0 and +// reports nonpad == prefill length, and each later step appends contiguously. + +TEST(StaticScatterIndexTracker, StartsEmpty) { + StaticScatterIndexTracker tracker; + EXPECT_EQ(tracker.valid_tokens(), 0); +} + +TEST(StaticScatterIndexTracker, FirstPrefillWritesAtRowZero) { + StaticScatterIndexTracker tracker; + // Prompt of 5 tokens: write at row 0, 5 valid tokens after. + const StaticScatterIndices prefill = tracker.Advance(5); + EXPECT_EQ(prefill.write_index, 0); // NOT -1: the crux off-by-one + EXPECT_EQ(prefill.nonpad_kv_seqlen, 5); + EXPECT_EQ(tracker.valid_tokens(), 5); +} + +TEST(StaticScatterIndexTracker, DecodeStepsAppendContiguously) { + StaticScatterIndexTracker tracker; + tracker.Advance(5); // prefill + + const StaticScatterIndices decode1 = tracker.Advance(1); + EXPECT_EQ(decode1.write_index, 5); // this step's offset == previous nonpad + EXPECT_EQ(decode1.nonpad_kv_seqlen, 6); + + const StaticScatterIndices decode2 = tracker.Advance(1); + EXPECT_EQ(decode2.write_index, 6); + EXPECT_EQ(decode2.nonpad_kv_seqlen, 7); +} + +TEST(StaticScatterIndexTracker, SingleTokenPrefillThenDecode) { + // Degenerate prompt length 1 must still write at row 0, then decode at row 1. + StaticScatterIndexTracker tracker; + const StaticScatterIndices prefill = tracker.Advance(1); + EXPECT_EQ(prefill.write_index, 0); + EXPECT_EQ(prefill.nonpad_kv_seqlen, 1); + + const StaticScatterIndices decode = tracker.Advance(1); + EXPECT_EQ(decode.write_index, 1); + EXPECT_EQ(decode.nonpad_kv_seqlen, 2); +} + +TEST(StaticScatterIndexTracker, MatchesSliceAFixtureGolden) { + // Exact values from the #366 slice-A fixture golden_io.npz: a 4-token prefill + // then a 1-token decode. Keeps this producer locked to the frozen contract. + StaticScatterIndexTracker tracker; + const StaticScatterIndices prefill = tracker.Advance(4); + EXPECT_EQ(prefill.write_index, 0); + EXPECT_EQ(prefill.nonpad_kv_seqlen, 4); + + const StaticScatterIndices decode = tracker.Advance(1); + EXPECT_EQ(decode.write_index, 4); + EXPECT_EQ(decode.nonpad_kv_seqlen, 5); +} + +// --- DiscoverKvLayerIndices: the strict layer-index parse the StaticScatter +// constructor uses to map past_key_values.N.key inputs to layer indices. --- + +TEST(DiscoverKvLayerIndices, ParsesAndSortsWellFormedNames) { + // Out-of-order, sparse (missing layer 1) -> parsed and sorted ascending. + const std::vector names{ + "past_key_values.2.key", "past_key_values.0.key", "other_input", + "past_key_values.5.key"}; + const auto indices = DiscoverKvLayerIndices(names, "past_key_values.", ".key"); + EXPECT_EQ(indices, (std::vector{0, 2, 5})); +} + +TEST(DiscoverKvLayerIndices, ThrowsOnMalformedIndex) { + // "0.bad" must be rejected: std::stoi would have silently accepted it as 0 and + // mis-mapped the layer. from_chars rejects the leftover ".bad". + const std::vector names{"past_key_values.0.bad.key"}; + EXPECT_THROW(DiscoverKvLayerIndices(names, "past_key_values.", ".key"), + std::runtime_error); +} + +TEST(DiscoverKvLayerIndices, ThrowsOnDuplicateIndex) { + // Two inputs mapping to the same layer index would double-count and bind the + // same cache slot twice -> rejected. + const std::vector names{"past_key_values.1.key", + "past_key_values.1.key"}; + EXPECT_THROW(DiscoverKvLayerIndices(names, "past_key_values.", ".key"), + std::runtime_error); +} + +// Shared helpers + frozen golden for the fixture-backed cache tests below. The +// #366 slice-A fixture (test/models/static-scatter-bias-decoder) is a bias-aware +// external-KV static-cache decoder: mistral backbone, 2 layers, +// key_cache.{i}/value_cache.{i} [batch,16,32] FLOAT, write_indices / +// nonpad_kv_seqlen [batch] int64. It runs on the CPU EP (empty provider_options) +// so it matches the CPU MEA (opset 24) golden captured in golden_io.npz. +// +// These fixture-backed tests (and their helpers/golden) drive the real C++ cache +// path end-to-end through the Oga generator, which executes the fixture's +// TensorScatter(24) node. TensorScatter has CPU and CUDA kernels but no DirectML +// implementation, so under a USE_DML build (which routes these models onto the +// DML EP) session.run throws "Could not find an implementation for +// TensorScatter(24)". The static-cache Flash feature is CPU/CUDA-targeted, so +// these tests are excluded only on DML builds (the StaticScatterIndexTracker +// unit tests above are pure C++ and run on every EP). +#if !USE_DML +namespace { + +// Last-token logits summary from golden_io.npz (CPU MEA, opset 24). +// prefill: input_ids [1,2,3,4] @ write_index 0, nonpad 4 -> argmax(last)=11. +// decode: input_ids [5] @ write_index 4, nonpad 5 -> argmax=36. +constexpr int kPrefillArgmax = 11; +constexpr float kPrefillLastTokLogitsSum = -9.80959f; +constexpr int kDecodeArgmax = 36; +constexpr float kDecodeLogitsSum = 78.96353f; // matches manifest logits_sum. + +// updated_key_cache.0 element-sums from golden (proves the in-place scatter +// actually wrote the cache, not just that logits are right). +constexpr float kPrefillUpdatedKeyCache0Sum = 4.31156f; +constexpr float kDecodeUpdatedKeyCache0Sum = -37.31835f; + +std::vector ToFloatVector(OgaTensor& tensor) { + auto shape = tensor.Shape(); + int64_t count = std::accumulate(shape.begin(), shape.end(), int64_t{1}, std::multiplies()); + const float* data = static_cast(tensor.Data()); + return std::vector(data, data + count); +} + +int ArgMax(const std::vector& values) { + return static_cast(std::max_element(values.begin(), values.end()) - values.begin()); +} + +float Sum(const std::vector& values) { + return std::accumulate(values.begin(), values.end(), 0.0f); +} + +// Element-wise parity tolerance. The fixture runs on the CPU EP (its +// genai_config declares no provider), so the produced tensors come from the same +// ORT CPU MEA kernels that generated the golden in golden_io.npz; fp32 agreement +// is near-exact, well within 1e-3. This is ~4 orders tighter than the headline +// sum check yet still trivially passes same-kernel noise, while a genuinely wrong +// (but sum-preserving) output diverges by O(1) and fails. +constexpr float kParityAtol = 1e-3f; + +// Assert every element of *actual* matches *golden* within *atol*. A sum/argmax +// check alone can pass a wrong-but-sum-preserving output; this compares the full +// tensor and reports the single worst element for a readable failure. +void ExpectAllClose(const std::vector& actual, const float* golden, + size_t golden_count, float atol, const char* label) { + ASSERT_EQ(actual.size(), golden_count) << label << ": element-count mismatch"; + double max_abs_diff = 0.0; + size_t max_idx = 0; + for (size_t i = 0; i < golden_count; ++i) { + const double diff = std::abs(static_cast(actual[i]) - static_cast(golden[i])); + if (diff > max_abs_diff) { + max_abs_diff = diff; + max_idx = i; + } + } + EXPECT_LE(max_abs_diff, static_cast(atol)) + << label << ": max |actual-golden| = " << max_abs_diff << " at index " << max_idx + << " (actual=" << actual[max_idx] << ", golden=" << golden[max_idx] << ")"; +} + +// The slice-A fixture updates its KV cache with an opset-24 TensorScatter node. +// Some ORT packages used in CI (e.g. the DirectML and certain CUDA lanes) don't +// ship a TensorScatter(24) kernel yet, so the static-scatter graph cannot execute +// there and model run throws "Could not find an implementation for TensorScatter". +// Probe the runtime once so the cache tests skip cleanly where the op is absent +// (this path is intentionally not gated on a specific ORT release); lanes whose +// ORT has the kernel still run and assert full parity. +bool StaticScatterRuntimeAvailable() { + try { + auto model = OgaModel::Create(MODEL_PATH "static-scatter-bias-decoder"); + auto params = OgaGeneratorParams::Create(*model); + auto generator = OgaGenerator::Create(*model, *params); + generator->AppendTokens(std::vector{1}); + } catch (const std::exception& e) { + // Skip ONLY on the precise "kernel absent" signature. ORT phrases a missing + // kernel as: "Could not find an implementation for TensorScatter(24) node ...". + // Matching just "TensorScatter" would also swallow a GENUINE TensorScatter + // regression (a shape/type/dispatch bug, or a wrong-opset node) and silently + // skip parity, masking a real failure as "op not shipped". Require BOTH the + // missing-implementation phrase AND the opset-pinned op name; rethrow anything + // else as a real test error. + const std::string msg = e.what(); + if (msg.find("Could not find an implementation for") != std::string::npos && + msg.find("TensorScatter(24)") != std::string::npos) { + return false; + } + throw; // Unrelated failure: let it surface as a real test error. + } + return true; +} + +} // namespace + +// M1 fail-loud contract: StaticScatterKeyValueCache::RewindTo MUST throw rather +// than silently no-op. RewindTo cannot reset the write_indices/nonpad_kv_seqlen +// stream (it lives in InputIDs with no rewind hook), so a no-op would leave the +// index tracker stale -> wrong scatter slots + over-reported nonpad => silently +// wrong logits with no error. DecoderOnly_State::RewindTo reaches +// kv_cache_->RewindTo and OgaGenerator::RewindTo is a public API not blocked for +// these models, so the throw is the only guard against silent corruption. The +// happy-path e2e parity test never exercises rewind, so this pins it explicitly. +TEST(StaticScatterKeyValueCache, RewindToThrows) { + if (!StaticScatterRuntimeAvailable()) { + GTEST_SKIP() << "ORT runtime lacks the opset-24 TensorScatter kernel; " + "the static-scatter KV-cache path is unavailable on this build."; + } + auto model = OgaModel::Create(MODEL_PATH "static-scatter-bias-decoder"); + auto params = OgaGeneratorParams::Create(*model); + auto generator = OgaGenerator::Create(*model, *params); + + const std::vector prompt{1, 2, 3, 4}; + generator->AppendTokens(prompt); + + // Rewinding the static-scatter cache is unsupported and must fail loud. + EXPECT_THROW(generator->RewindTo(2), std::exception); +} + +// End-to-end parity test for the mobius slice-A fixture (#366, bias-aware +// external-KV static-cache decoder). It drives the real +// StaticScatterKeyValueCache C++ path through the public Oga generator API: +// AppendTokens runs the model, DefaultInputIDs feeds write_indices / +// nonpad_kv_seqlen via StaticScatterIndexTracker, and StaticScatterKeyValueCache +// binds the 3D in-place share-buffer KV cache. We force the exact prompt + decode +// token from golden_io.npz and compare the produced logits / updated caches to +// the frozen golden (ORT CPU MEA reference, opset 24). +TEST(StaticScatterKeyValueCache, EndToEndFixtureParity) { + if (!StaticScatterRuntimeAvailable()) { + GTEST_SKIP() << "ORT runtime lacks the opset-24 TensorScatter kernel; " + "the static-scatter KV-cache path is unavailable on this build."; + } + auto model = OgaModel::Create(MODEL_PATH "static-scatter-bias-decoder"); + auto params = OgaGeneratorParams::Create(*model); + auto generator = OgaGenerator::Create(*model, *params); + + // --- Prefill: force the 4-token golden prompt. StaticScatterIndexTracker + // must bind write_index=0, nonpad=4; TensorScatter writes rows 0..3. --- + const std::vector prompt{1, 2, 3, 4}; + generator->AppendTokens(prompt); + + auto prefill_logits = generator->GetLogits(); + auto prefill_logits_vec = ToFloatVector(*prefill_logits); // last token only: [1,1,256] + EXPECT_EQ(ArgMax(prefill_logits_vec), kPrefillArgmax); + EXPECT_NEAR(Sum(prefill_logits_vec), kPrefillLastTokLogitsSum, 1e-2f); + ExpectAllClose(prefill_logits_vec, kPrefillLastTokLogitsGolden, + std::size(kPrefillLastTokLogitsGolden), kParityAtol, "prefill logits"); + + auto prefill_cache = generator->GetOutput("updated_key_cache.0"); + auto prefill_cache_vec = ToFloatVector(*prefill_cache); + EXPECT_NEAR(Sum(prefill_cache_vec), kPrefillUpdatedKeyCache0Sum, 1e-2f); + ExpectAllClose(prefill_cache_vec, kPrefillUpdatedKeyCache0Golden, + std::size(kPrefillUpdatedKeyCache0Golden), kParityAtol, + "prefill updated_key_cache.0"); + + // --- Decode: force golden token 5. The tracker advances to write_index=4, + // nonpad=5; the decode step must read the slot prefill just wrote. --- + const std::vector decode_token{5}; + generator->AppendTokens(decode_token); + + auto decode_logits = generator->GetLogits(); + auto decode_logits_vec = ToFloatVector(*decode_logits); // [1,1,256] + EXPECT_EQ(ArgMax(decode_logits_vec), kDecodeArgmax); + EXPECT_NEAR(Sum(decode_logits_vec), kDecodeLogitsSum, 1e-2f); + ExpectAllClose(decode_logits_vec, kDecodeLastTokLogitsGolden, + std::size(kDecodeLastTokLogitsGolden), kParityAtol, "decode logits"); + + auto decode_cache = generator->GetOutput("updated_key_cache.0"); + auto decode_cache_vec = ToFloatVector(*decode_cache); + EXPECT_NEAR(Sum(decode_cache_vec), kDecodeUpdatedKeyCache0Sum, 1e-2f); + ExpectAllClose(decode_cache_vec, kDecodeUpdatedKeyCache0Golden, + std::size(kDecodeUpdatedKeyCache0Golden), kParityAtol, + "decode updated_key_cache.0"); +} +#endif // !USE_DML + +} // namespace Generators::test