feat(native-decode): auto-derive io kv/state_pairs from graph ports (#384)

[justinchuby]

What this unlocks

Stock qwen3.x linear-attention hybrid exports (dense GQA layers interleaved with recurrent conv_state/recurrent_state layers under a Scan body) — including Justin's named qwen3.6-27b-int4 — now load and decode natively with NO hand-written io: overlay. The native loader auto-derives the decoder I/O topology from the ONNX graph ports when the package declares no model.io block.

Before this, the stock 27b export (whose inference_metadata.yaml carries no io block) failed native load at token-input resolution ("3 ports match: input_ids, attention_mask, position_ids"), and its recurrent layers expose no dense KV, so shape inference cannot classify the running-state ports. Lineage: #384 recurrent-state lane.

The change (blast radius: loader fallback + genai-config helper + tests + docs)

• genai-config: extracted the pure, config-free core of strict_decoder_state into decoder_state_from_patterns(graph, past_key, past_value, present_key, present_value, shape_match). Added pub GenAiConfig::derive_decoder_io_from_graph(graph) -> Option that calls it with the conventional onnxruntime-genai patterns. This is true reuse of the guarded derivation — recurrent ports are found by suffix_tensor_map (never strict_indexed_kv), so cross-attention / Whisper KV is never misclassified as running state.
• One isolated relaxation via a StateShapeMatch param: strict_decoder_state passes Exact (its contract is byte-for-byte unchanged); the new fallback passes AllowSymbolic. Real stock exports leave the present.* recurrent-state shapes fully symbolic while the paired past_* input carries the concrete extent — a symbolic axis is "unknown", not "different", so the fallback accepts it; concrete-vs-concrete mismatch still fails.
• Engine loader (native_decode/load.rs): in from_session_with_cuda_options_and_io, when io is None, build a ModelGraphInfo from the session ports and call the helper. Safety gate: engage the derived spec ONLY when it yields at least one recurrent state pair. Pure-dense decoders derive 0 state pairs, the fallback declines, and their existing io=None shape-inference path is untouched. A declared io always wins (fallback runs only when io is absent).

Correctness bar (byte-identical vs CPU fp32 oracle)

ORT-CUDA crashes on this model class (internal stl_vector assertion), so there is no ORT reference; the oracle is our own native CPU fp32 backend, which already threads recurrent state.

• GPU e2e (native_autoderive_io_cuda_e2e.rs, #[ignore], skips cleanly when the model dir is absent): STOCK 27b metadata (no overlay) native CUDA load + decode == native CPU fp32 oracle, byte-identical token IDs. Manually confirmed the auto-derived CUDA decode reproduces the exact oracle-validated sequence proven during the Large-model (27B) native E2E offload blocked: Unsqueeze rank bug + missing recurrent_state + mobius#432 #384 re-probe: [11751, 13, 271, 248068, 271, 248069, 271, 4639, 369, 4252, 13, 11751, 369, 279, 6511, 321] ("Paris ... That is correct. Paris is the capital").
• Unit tests (genai-config): 27b port layout derives 32 kv entries + 96 state pairs; qwen3-0.6b dense layout derives dense kv + 0 state pairs (gate does not over-derive). Non-vacuous both directions; the 27b fixture uses fully-symbolic present-state shapes to lock the AllowSymbolic path.
• No-regression (engine): a hybrid decoder with io=None now auto-derives and loads; a dense ambiguous decoder with io=None still fails with the same model.io.token_input error (fallback declines). Existing declared-io tests unchanged and green.

Honest caveats

• Eager only, ~5.9 tok/s. The Scan control-flow body declines CUDA-graph capture, so these hybrids run eager. Capturing the Scan body is a separate, parked perf lane; this PR is a correctness unblock, not a perf win.
• CPU is the oracle, not ORT (ORT-CUDA crashes on this class). Whether that crash is a mobius export bug or an ORT limitation is a separate follow-up.

Decision note: .squad/decisions/inbox/cohaagen-autoderive-io.md

Working as Cohaagen (EP/runtime engineer). Please review before merge (auto-merge intentionally NOT enabled).

[codecov]

Codecov Report

❌ Patch coverage is 90.00000% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.74%. Comparing base (25d41f4) to head (b2e4a43).

Files with missing lines	Patch %	Lines
...rates/onnx-genai-genai-config/src/compatibility.rs	90.00%	1 Missing and 4 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #573      +/-   ##
==========================================
- Coverage   81.33%   80.74%   -0.59%     
==========================================
  Files         315      315              
  Lines      123688   123734      +46     
  Branches   123688   123734      +46     
==========================================
- Hits       100597    99913     -684     
- Misses      19018    19749     +731     
+ Partials     4073     4072       -1     

Flag	Coverage Δ
cli-ort-linux	86.69% <ø> (ø)
cli-ort-windows	?
mlas	77.91% <ø> (-0.82%)	⬇️
offline	80.54% <90.00%> (-0.61%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
...rates/onnx-genai-genai-config/src/compatibility.rs	80.68% <90.00%> (+0.36%)	⬆️

... and 8 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

✅ Benchmarks — No Regression

Comparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).

ℹ️ Absolute times are informational only — they vary with runner load. The % change column is the reliable signal because both sides ran under identical conditions.

Status	Scenario	Base	PR	Change
✅	add/small_f32_threads=1-internal/1024	226.8 ns	245.8 ns	+8.4%
✅	matmul/medium_generic_f32_threads=8/32x512x512	1.62 ms	1.74 ms	+6.9%
✅	matmul/medium_generic_f32_threads=1/32x512x512	2.73 ms	2.92 ms	+6.8%
✅	matmul/medium_generic_f16_threads=8/32x512x512	47.93 µs	50.37 µs	+5.1%
✅	add/small_f16_threads=1-internal/1024	15.33 µs	16.05 µs	+4.7%
✅	matmul/small_generic_bf16_threads=8/1x256x256	41.78 µs	42.67 µs	+2.1%
✅	matmul/small_generic_bf16_threads=1/1x256x256	41.31 µs	41.02 µs	-0.7%
✅	add/small_bf16_threads=1-internal/1024	14.93 µs	14.83 µs	-0.7%
✅	matmul/medium_generic_bf16_threads=8/32x512x512	663.25 µs	654.58 µs	-1.3%
✅	matmul/large_generic_f32_threads=1/32x1024x1024	11.61 ms	11.14 ms	-4.0%
✅	matmul/medium_generic_f16_threads=1/32x512x512	41.84 µs	39.73 µs	-5.0%
✅	matmul/small_generic_f16_threads=8/1x256x256	41.24 µs	39.07 µs	-5.3%
✅	matmul/large_generic_f16_threads=8/32x1024x1024	112.00 µs	105.12 µs	-6.1%
✅	logit_processing/seven_processor_chain_per_step	1.48 ms	1.39 ms	-6.4%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	651.01 µs	606.95 µs	-6.8%
✅	add/medium_f32_threads=1-internal/262144	3.24 ms	3.01 ms	-6.9%
✅	sampling_latency/top_p_per_token	1.11 ms	1.04 ms	-7.0%
✅	matmul/small_generic_f32_threads=1/1x256x256	48.48 µs	44.70 µs	-7.8%
✅	grammar_masking/llguidance_compute_mask/32	98.13 µs	90.04 µs	-8.2%
✅	add/large_f16_threads=1-internal/4194304	53.81 ms	49.10 ms	-8.7%
✅	sampling_latency/top_k_per_token	539.50 µs	486.29 µs	-9.9%
✅	sampling_latency/greedy_per_token	3.72 µs	3.33 µs	-10.4%
✅	gather/small_bf16_threads=1-internal/4096	665.2 ns	589.5 ns	-11.4%
✅	sampling_latency/min_p_per_token	398.37 µs	352.31 µs	-11.6%
✅	gather/small_f32_threads=1-internal/4096	890.9 ns	785.6 ns	-11.8%
✅	reduce_mean/small_f32_threads=1-internal/4096	19.49 µs	17.09 µs	-12.3%
✅	reduce_mean/medium_f32_threads=1-internal/65536	327.69 µs	286.94 µs	-12.4%
✅	add/large_f32_threads=1-internal/4194304	55.96 ms	48.58 ms	-13.2%
✅	gather/small_f16_threads=1-internal/4096	640.3 ns	553.8 ns	-13.5%
✅	add/medium_f16_threads=1-internal/262144	3.46 ms	2.99 ms	-13.6%
✅	tokenization/encode_tokens_per_second	449.16 µs	386.42 µs	-14.0%
✅	tokenization/decode_tokens_per_second	7.78 ms	6.68 ms	-14.1%
✅	reduce_mean/large_f32_threads=1-internal/262144	1.30 ms	1.11 ms	-14.7%
🟢	add/large_bf16_threads=1-internal/4194304	58.43 ms	49.52 ms	-15.3%
🟢	kv_cache/alloc_dealloc_pages	43.85 µs	37.04 µs	-15.5%
🟢	matmul/small_generic_f16_threads=1/1x256x256	40.09 µs	33.76 µs	-15.8%
🟢	gather/medium_f16_threads=1-internal/32768	3.82 µs	3.21 µs	-15.9%
🟢	gather/medium_f32_threads=1-internal/32768	6.27 µs	5.25 µs	-16.2%
🟢	matmul/large_generic_f16_threads=1/32x1024x1024	98.75 µs	82.03 µs	-16.9%
🟢	gather/large_bf16_threads=1-internal/131072	18.22 µs	14.56 µs	-20.1%
🟢	matmul/large_generic_bf16_threads=1/32x1024x1024	2.54 ms	1.99 ms	-21.6%
🟢	add/medium_bf16_threads=1-internal/262144	3.92 ms	3.07 ms	-21.9%
🟢	matmul/large_generic_bf16_threads=8/32x1024x1024	2.08 ms	1.55 ms	-25.4%
🟢	gather/large_f32_threads=1-internal/131072	58.44 µs	39.21 µs	-32.9%
🟢	matmul/small_generic_f32_threads=8/1x256x256	63.26 µs	42.30 µs	-33.1%
🟢	gather/medium_bf16_threads=1-internal/32768	3.61 µs	2.34 µs	-35.1%
🟢	gather/large_f16_threads=1-internal/131072	21.98 µs	13.61 µs	-38.1%
🟢	matmul/large_generic_f32_threads=8/32x1024x1024	7.01 ms	3.92 ms	-44.1%

Visual flags: ⚠️ ≥ 15% slower, 🔴 ≥ 30% slower — calibrated against measured runner noise (~27% worst-case on multi-threaded matmul)

Host info

CPU: Apple M1 (Virtual)
Cores: 3
OS: Darwin 25.4.0 arm64
Rust: rustc 1.97.1 (8bab26f4f 2026-07-14)
Load avg: { 4.06 3.85 4.94 }

What this cannot catch

• Regressions in code paths not covered by these benchmarks (e.g., end-to-end decode with a real model)
• Sub-threshold regressions that compound over multiple PRs
• Performance changes that only manifest under GPU execution
• Latency changes in the ORT integration path (these benchmarks exercise the native Rust kernels)

[justinchuby]

VERDICT: APPROVE

Independent review of #573 (auto-derive decoder io kv/state_pairs from graph ports). Reviewed detached at head b2e4a43 in an isolated worktree; read-only, nothing committed.

1. Diff scope + additive-fallback proof — PASS

git diff origin/main...HEAD: 9 files, +816/-3. Changes confined to native_decode/load.rs (fallback), native_decode/mod.rs (imports only), genai-config/compatibility.rs + lib.rs (derivation), tests, docs, decision note.

• The loader engages the derived spec ONLY when io.is_none() AND the derivation yields >=1 recurrent state pair: let io = io.or(derived_io.as_ref()); with derive_fallback_io returning None when derived.state_pairs.is_empty(). Declared io always wins (short-circuits, derived_io never computed).
• No edits to the decode loop, KV/state binder internals, capture, weight-offload, or GAP-3. Confirmed by the diff stat — none of those files are touched.

2. AllowSymbolic relaxation — PASS (cannot over-match, cannot leak into strict path)

StateShapeMatch::AllowSymbolic is passed ONLY by derive_decoder_io_from_graph; strict_decoder_state passes StateShapeMatch::Exact explicitly (compatibility.rs:1553). The strict path is byte-unchanged (only the input == output shape check was factored into shapes_pair, and Exact returns exactly input == output).

Crucially, AllowSymbolic can NOT cause over-derivation: which ports become state pairs is decided by suffix_tensor_map NAME matching (share the past_key_values./present. prefix, and are excluded if they are the indexed .key/.value KV ports). AllowSymbolic only relaxes the shape-equality check applied to pairs already selected by name — it broadens accept/reject, never selection. So a symbolic present.* state shape can pair with a concrete past_* input (the real stock-export case), but dense/cross-attn KV can never be pulled into state by it.

3. Over-derivation guard is non-vacuous — PASS (mutation has teeth)

Unit tests: 27b layout -> 32 kv-pairs + 96 state_pairs; qwen3-0.6b dense -> kv only, 0 state_pairs. All 29 genai-config tests pass (incl. whisper_strict_encoder_decoder_synth_routes_cross_kv).

Independent mutation: I disabled the KV exclusion in suffix_tensor_map (if false && excluded.contains(...)). Result: derive_decoder_io_from_graph_dense_has_no_state_pairs FAILED — the dense 0.6b graph produced 56 bogus state pairs (all .key/.value misclassified as recurrent state). Guard confirmed non-vacuous. Reverted; worktree clean (git status empty).

4. Byte-identical 27B native vs CPU oracle — RAN, PASSED

Ran native_autoderive_io_cuda_e2e on GPU 6 (CUDA_VISIBLE_DEVICES=6, ONNX_GENAI_REQUIRE_CUDA=1) against the stock /home/justinchu/mary-models/qwen3.6-27b-int4-cuda export. Verified its inference_metadata.yaml has NO io: block, so loading depends entirely on the fallback.

Result (test passed, 5552s; CPU oracle ~92min, CUDA 19s):
cuda tokens = [11751, 13, 271, 248068, 271, 248069, 271, 4639, 369, 4252, 13, 11751, 369, 279, 6511, 321]
cpu tokens = [11751, 13, 271, 248068, 271, 248069, 271, 4639, 369, 4252, 13, 11751, 369, 279, 6511, 321]

Byte-identical, and matches the expected sequence exactly. Note: the test asserts cuda==cpu and len==16; it does not itself hard-code the expected IDs, but the observed IDs match the brief's expectation.

5. No-regression + fmt — PASS

native_decoder_auto_derive_skips_dense_ambiguous_decoder passes: a pure-dense decoder with io=None derives 0 state pairs, the fallback declines, and the model still errors on model.io.token_input exactly as before — no silent wrong-spec auto-derive. native_decoder_auto_derives_hybrid_state_io passes. cargo fmt --all --check clean.

Summary

Strictly-additive fallback, gated to the recurrent-hybrid case declared-io cannot already cover; declared-io path untouched. AllowSymbolic is scoped to the graph fallback, relaxes only shape-equality on name-selected pairs, and cannot leak into the Exact strict path or over-derive dense/cross-attn KV into state. Guard proven to have teeth by mutation. 27B stock native CUDA decode is byte-identical to the CPU fp32 oracle. Ships the token-parity Justin cares about with zero regression risk to currently-loadable models.