Add StaticScatterKeyValueCache: drive mobius bias-aware external-KV (TensorScatter + nonpad_kv_seqlen) static cache

[titaiwangms]

What

Adds StaticScatterKeyValueCache, a new KV-cache variant that lets onnxruntime-genai
drive a mobius-exported bias-aware external-KV static-cache decoder directly. It:

• Reuses the existing past_present_share_buffer lifecycle: each key_cache.{i} /
  value_cache.{i} is a pre-allocated 3D [batch, max_seq_len, kv_hidden] buffer that
  the graph updates in place via opset-24 TensorScatter, so past and present alias
  one OrtValue and never need rebinding between steps.
• Adds a small [B] int64 write_indices / nonpad_kv_seqlen producer in
  input_ids (driven by a standalone StaticScatterIndexTracker): write_indices is the
  scatter row offset (valid tokens before this step), nonpad_kv_seqlen is the valid
  cached length after this step (Attention's per-batch seqlens_k).
• Factory auto-detect: CreateKeyValueCache selects this variant when the model
  declares both write_indices and nonpad_kv_seqlen inputs (kept in lockstep with
  the producer gate so a half-declared model can't create the cache with unbound indices).
• Consumes mobius's 3D external-KV emission directly (vs the 4D
  [batch, num_kv_heads, seq, head_dim] default layout), and supports per-layer-varying
  kv_hidden (e.g. Gemma-4 sliding GQA vs global MQA), read from each layer's own
  declared input shape. The batch dim is taken from BatchBeamSize() (symbolic in the
  graph); only max_seq_len / kv_hidden must be static.

Why

genai's existing GQA past_present_share_buffer path cannot carry an arbitrary float
attention bias. mobius now emits float-bias decoders (causal + sliding-window built
in-graph) whose KV is an external TensorScatter cache — that path can carry the bias.
This PR is the host-side consumer for that graph, recovering near-GQA decode behaviour
on the memory-efficient-attention (MEA) path for bias-aware models like Gemma-4.

Relationship to onnxruntime-genai#2204

#2204 proposed Path A = a prefill/decode GQA graph-split. This PR implements Path B
= a single TensorScatter external-KV graph, which sidesteps the cross-stage KV handoff
that #2204 flagged as the crux. Referenced as related; this PR does not close #2204.

Mobius side

The bias-aware external-KV graph this cache consumes is emitted by mobius
(onnxruntime/mobius#367, issue onnxruntime/mobius#366). This PR is the genai-side consumer.

Testing

• StaticScatterKeyValueCache.EndToEndFixtureParity — bit-exact vs an ORT-CPU-MEA golden
  on a committed slice-A fixture (prefill seq=4 @ write 0/nonpad 4, then decode @ write
  4/nonpad 5): logits argmax/sum + updated_key_cache.0 sums.
• StaticScatterKeyValueCache.RewindToThrows — deny-by-default guard: RewindTo() throws
  (the write_indices/nonpad_kv_seqlen stream has no rewind hook, so a silent no-op would
  desync the tracker; matches the LFM2Cache / WindowedKeyValueCache siblings).
• 5 StaticScatterIndexTracker producer tests pinning the off-by-one / init contract.
• 7/7 green. CPU / MEA path — not gated on onnxruntime#28958.

Scope / limits

• batch == 1 (beam search unsupported, throws).
• MEA ≈ GQA decode ceiling, not Flash: the arbitrary float bias precludes the Flash path.
• Per-layer kv_hidden supported (Gemma-4 sliding/global mix).
• The slice-A fixture is committed under test/models/ per the existing test-model
  convention, tracked via a .gitignore allowlist entry matching its siblings.

[Copilot]

Pull request overview

Adds a new KV-cache implementation to let onnxruntime-genai drive mobius-exported “static-scatter” (opset-24 TensorScatter) external-KV decoders, including per-step production of write_indices and nonpad_kv_seqlen and an end-to-end fixture-backed parity test.

Changes:

• Introduces StaticScatterKeyValueCache (3D [B, max_seq_len, kv_hidden] per-layer KV buffers updated in-place) and auto-detection via declared model inputs.
• Adds StaticScatterIndexTracker plus DefaultInputIDs plumbing to produce/bind write_indices and nonpad_kv_seqlen.
• Adds a committed slice-A test fixture + new unit tests (producer contract tests, rewind guard, and end-to-end parity).

Reviewed changes

Copilot reviewed 9 out of 11 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
test/static_scatter_kv_cache_test.cpp	New unit tests for index-tracker semantics and end-to-end parity against a committed fixture.
test/models/static-scatter-bias-decoder/genai_config.json	Adds a minimal test model config declaring write_indices/nonpad_kv_seqlen and cache I/O names.
src/models/static_scatter_indices.h	Defines the StaticScatterIndices pair and StaticScatterIndexTracker helper.
src/models/kv_cache.h	Declares StaticScatterKeyValueCache interface and members.
src/models/kv_cache.cpp	Implements StaticScatterKeyValueCache, shape discovery/allocation, and factory auto-detect.
src/models/input_ids.h	Extends DefaultInputIDs to hold static-scatter driver tensors + tracker.
src/models/input_ids.cpp	Creates/binds/updates write_indices and nonpad_kv_seqlen when declared by the model.
src/config.h	Adds default config names for write_indices and nonpad_kv_seqlen.
src/config.cpp	Parses the new decoder input names from genai_config.json.
.gitignore	Allowlists the new committed test fixture directory under test/models/.

[titaiwangms]

Review summary — StaticScatterKeyValueCache (Path B consumer)

Reviewed at 2fd84e8 by a multi-model review team plus independent call-trace verification. Mergeable as scoped infrastructure; no Critical correctness defect found.

Verified correct ✅

• Prefill index ordering is sound (I traced decoder_only.cpp): input_ids_.Update() → StaticScatterIndexTracker::Advance(N) runs before State::Run, so prefill binds write=0, nonpad=N — the ctor's 0/0 default is dead, and the feared one-step lag is not present.
• Index semantics agree with mobius: write=valid_before, nonpad=valid_after, satisfying mobius's nonpad == write + S_q for unpadded chunks. The tracker unit tests pin the off-by-one contract well.
• RewindTo fail-loud throw, factory predicate in lockstep with the producer gate (both inputs required), per-layer-varying kv_hidden, and batch==1 enforced on both sides — all consistent.

Major — factory hard-fails instead of falling back

IsStaticScatterCache() selects on presence of both driver inputs before the other cache types in CreateKeyValueCache. A model that declares those inputs but uses 4D KV would hit the constructor's rank != 3 throw rather than falling back to Default/Windowed. Name collision is unlikely (low real risk), but consider tightening the predicate (require the rank-3 static layout) or making the factory attempt-and-fall-back. (kv_cache.cpp:711-719, 1069-1076)

Minor

• Layer-index parse: std::stoi(idx_str) accepts malformed names ("0.bad" → 0) and doesn't reject duplicates. Use std::from_chars / stoi with pos == size and validate non-negative + unique. (kv_cache.cpp:740-746)
• Test runtime probe: StaticScatterRuntimeAvailable() skips on any exception message containing "TensorScatter". A real TensorScatter regression could be misclassified as "kernel absent" and silently skip parity. Match the precise missing-kernel signature ("Could not find an implementation" + "TensorScatter(24)") and rethrow others. (test/static_scatter_kv_cache_test.cpp)
• E2E parity strength: the fixture asserts argmax + sum (1e-2 tol). sum can pass wrong-but-sum-preserving outputs (permutation/cancellation). Add full-tensor or slice-wise parity against the golden.

Open question

• Decode token equal to pad_token_id → Advance(0) leaves write_index/nonpad unchanged; two such steps alias a cache slot. Mirrors the existing current_sequence_length logic and is benign if pad==eos. Confirm pad_token_id can't be a legal mid-stream generated token for the targeted models.

Scope clarity (not a defect)

• batch==1 is an enforced, documented MVP constraint (driver tensors allocated [1], throws on BatchBeamSize()!=1), even though the mobius graph + tests support [batch]. Intentional on both sides — worth tracking as a known limitation, not a break.
• KV name mapping: this consumer reads names from genai_config.json, while mobius emits key_cache.{i} / updated_key_cache.{i}. The fixture maps them manually; production needs the genai_config generator to emit matching names.

Nits

• input_index_ / output_index_ / layer_shapes_ are set in Add() but never read back — annotate clearly or remove.
• kv_cache.h comment "differs only in the allocator" erases the throw-vs-silent RewindTo split vs DefaultKeyValueCache; the struct field write_index (singular) vs the tensor write_indices (plural) is a small cross-boundary translation.

Praise: static_scatter_indices.h is a textbook extraction of a tricky off-by-one contract into a minimal, unit-testable type — the "crux mobius and genai must agree on" framing is exactly right. The two-tier test layout (pure-C++ tracker tests on every EP + fixture-backed e2e guarded by #if !USE_DML + runtime probe) is thoughtfully done.

[titaiwangms]

Re-review at 72e7199 — prior findings addressed; one deferred with rationale ✅

Re-checked the new commit "Address review feedback on static-scatter KV cache". I independently verified the new golden by running the committed fixture model.onnx fresh on ORT 1.27 CPU: it matches golden_io.npz (≤1e-4), the scalar sum/argmax constants match, and the C++ header arrays (static_scatter_golden.h) match the fresh run exactly (maxdiff = 0.0) for all four tensors (256-elem logits ×2, 512-elem cache ×2).

Prior finding	Status
Minor — std::stoi accepted malformed layer indices ("0.bad" → 0), no dedup	✅ Factored into a model-free DiscoverKvLayerIndices() using std::from_chars with a complete-consume check (parse_end == end), non-negative guard, and duplicate rejection. New unit tests cover well-formed/sort, throw-on-malformed, throw-on-duplicate.
Minor — runtime probe skipped on any "TensorScatter" substring (could mask a real regression)	✅ Tightened to require both "Could not find an implementation for" and "TensorScatter(24)"; rethrows anything else.
Minor — e2e parity used only argmax + sum (passes wrong-but-sum-preserving outputs)	✅ Added full-tensor ExpectAllClose @ 1e-3 against a generated, provenance-documented golden header (kept the sum/argmax as headline checks). Verified the golden is authentic and reproducible.
Nit — dead input_index_ / output_index_ / layer_shapes_	✅ Removed.
Nit — "differs only in the allocator" erased the RewindTo throw-vs-reshape split	✅ Comment rewritten to document both the 3D-vs-4D layout and the RewindTo difference.
Open question — decode token == pad_token_id → Advance(0) slot aliasing	✅ Documented the assumption in StaticScatterIndexTracker::Advance (safe because pad_token_id == eos_token_id ends the sequence before the stalled slot is reread; would only break if a model made pad a legal mid-stream token).
Major — factory hard-fails (rank != 3 throw) instead of falling back to Default/Windowed	⏳ Deferred as a tracked follow-up with an in-code rationale (the rank-3 tightening touches the sensitive factory-selection path and risks layer-discovery drift for sparse/hybrid models lacking a layer-0 KV input; collision considered unlikely). Reasonable disposition — acceptable to land now and follow up.

DiscoverKvLayerIndices logic reviewed by reading: the size() > prefix+suffix guard keeps idx_str non-empty, and from_chars/negative/duplicate checks are correct. I did not rebuild the genai C++ (expensive) but verified the golden provenance directly, which is the strongest available check. LGTM with the factory fallback tracked separately.