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

.squad/decisions/inbox/cohaagen-autoderive-io.md

@@ -0,0 +1,83 @@
+# Auto-derive `model.io` (kv + recurrent state_pairs) from graph ports
+
+Date: 2026-08-01
+By: Cohaagen
+Lane: #384 recurrent-state / native decode loader
+Status: PR opened (S increment)
+
+## What this unlocks
+
+Stock qwen3.x **linear-attention hybrid** exports (16 dense GQA layers + 48
+recurrent `conv_state`/`recurrent_state` layers with a `Scan` body), including
+Justin's named **qwen3.6-27b-int4**, now load + decode **natively** with **no
+hand-written `io:` overlay** — the loader derives the decoder I/O topology from
+the ONNX graph ports.
+
+Before this, the stock export (whose `inference_metadata.yaml` carries no `io`
+block) failed native load at token-input resolution:
+
+```
+cannot resolve model.io.token_input from tensor shape because 3 ports match:
+[input_ids, attention_mask, position_ids]
+```
+
+The recurrent hybrids also declare NO dense KV for the linear layers, so shape
+inference cannot classify their `conv_state`/`recurrent_state` running-state
+ports. During the #384 re-probe I proved a hand overlay (states in
+`io.state_pairs`) makes it load + decode byte-identically to the CPU fp32
+oracle. This PR removes the overlay requirement.
+
+## The change (blast radius: loader fallback + genai-config helper + tests + docs)
+
+1. **genai-config** (`compatibility.rs`): 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 a new public
+   `GenAiConfig::derive_decoder_io_from_graph(graph) -> Option<DerivedDecoderIo>`
+   that calls it with the conventional onnxruntime-genai patterns
+   (`past_key_values.%d.key`/`.value` → `present.%d.key`/`.value`). 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.
+
+2. **One relaxation, isolated to the fallback:** a `StateShapeMatch` param.
+   `strict_decoder_state` passes `Exact` (byte-for-byte unchanged). The new
+   fallback passes `AllowSymbolic`: a `present.*` state port whose exported shape
+   is fully symbolic (`[?,?,?]`) is accepted against its concrete `past_*` input
+   (`[?,10240,3]`). This was REQUIRED — the real 27b export leaves present
+   recurrent-state shapes fully symbolic, and `Exact` rejected them. A symbolic
+   axis means "unknown", not "different"; concrete-vs-concrete mismatch still
+   fails.
+
+3. **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 ≥1 recurrent state pair (the case
+   shape inference can't handle). Pure-dense decoders derive 0 state pairs → the
+   fallback declines → their existing `io=None` path is untouched. Declared `io`
+   always wins (fallback runs only when `io` is absent).
+
+## Correctness bar
+
+- **GPU e2e** (`native_autoderive_io_cuda_e2e.rs`, `#[ignore]`): STOCK 27b
+  metadata (no overlay) → native CUDA load + decode == native CPU fp32 oracle,
+  **byte-identical** token IDs. Proven parity token IDs:
+  `[11751, 13, 271, 248068, 271, 248069, 271, 4639, 369, 4252, 13, 11751, ...]`.
+- **Unit tests** (genai-config): 27b port layout → 32 kv entries + 96 state
+  pairs; qwen3-0.6b dense layout → dense kv + **0** state pairs (gate doesn't
+  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.91 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 model class (internal
+  `stl_vector` assertion), so there is no ORT reference to compare against; the
+  native CPU fp32 path (which already threads recurrent state) is the oracle.
+- Follow-up (separate): is the ORT-CUDA crash a mobius export bug or an ORT
+  limitation? Affects whether we ever get an ORT baseline for this class.

crates/onnx-genai-engine/src/native_decode/load.rs

@@ -277,11 +277,103 @@ impl NativeDecodeSession {
         Self::from_session_with_cuda_options_and_io(session, cuda_options, None)
     }
 
+    /// Best-effort auto-derived [`ModelIoSpec`] for a stock export whose sidecar
+    /// declares no `io` block, built purely from the session's graph ports.
+    ///
+    /// Reuses the guarded genai-config derivation
+    /// ([`GenAiConfig::derive_decoder_io_from_graph`]) so recurrent
+    /// `conv_state`/`recurrent_state` ports are classified as loop-carried
+    /// `state_pairs` and never confused with growable KV. Returns `None` (leaving
+    /// the caller's `io = None` shape-inference path untouched) unless the
+    /// derivation yields at least one recurrent state pair — the exact case the
+    /// shape-inference path cannot resolve. Non-KV ports (token/mask/position/
+    /// logits) are bound by conventional-name presence in the graph interface.
+    fn derive_fallback_io(session: &InferenceSession) -> Option<ModelIoSpec> {
+        let to_graph_tensor =
+            |meta: &onnx_runtime_session::IoMeta| onnx_genai_genai_config::GraphTensorInfo {
+                name: meta.name.clone(),
+                dtype: crate::engine::ir_dtype_name(meta.dtype).to_owned(),
+                dimensions: meta
+                    .shape
+                    .iter()
+                    .map(|dim| match dim {
+                        Dim::Static(value) => Some(*value),
+                        Dim::Symbolic(_) => None,
+                    })
+                    .collect(),
+            };
+        let graph = onnx_genai_genai_config::ModelGraphInfo {
+            inputs: session.inputs().iter().map(to_graph_tensor).collect(),
+            outputs: session.outputs().iter().map(to_graph_tensor).collect(),
+        };
+        let derived = onnx_genai_genai_config::GenAiConfig::derive_decoder_io_from_graph(&graph)?;
+        // Safety gate: only the recurrent-hybrid case the shape-inference path
+        // cannot handle. Pure-dense decoders (no state pairs) keep `io = None`.
+        if derived.state_pairs.is_empty() {
+            return None;
+        }
+        let input_names: HashSet<&str> = session
+            .inputs()
+            .iter()
+            .map(|meta| meta.name.as_str())
+            .collect();
+        let output_names: HashSet<&str> = session
+            .outputs()
+            .iter()
+            .map(|meta| meta.name.as_str())
+            .collect();
+        let present_input = |name: &str| input_names.contains(name).then(|| name.to_owned());
+        let present_output = |name: &str| output_names.contains(name).then(|| name.to_owned());
+        let state_pairs = derived
+            .state_pairs
+            .into_iter()
+            .map(|pair| LoopStatePair {
+                input: pair.input,
+                output: pair.output,
+                init: Some("zeros".to_owned()),
+                update: Some("replace".to_owned()),
+            })
+            .collect::<Vec<_>>();
+        Some(ModelIoSpec {
+            sequence_source: None,
+            kv_ownership: None,
+            token_input: present_input("input_ids"),
+            inputs_embeds_input: None,
+            attention_mask_input: present_input("attention_mask"),
+            position_ids_input: present_input("position_ids"),
+            logits_output: present_output("logits"),
+            hidden_output: None,
+            kv_inputs: (!derived.kv_inputs.is_empty()).then_some(derived.kv_inputs),
+            kv_outputs: (!derived.kv_outputs.is_empty()).then_some(derived.kv_outputs),
+            encoder_hidden_states_input: None,
+            audio_features_input: None,
+            cross_kv_inputs: None,
+            cross_kv_outputs: None,
+            kv_update: None,
+            state_pairs: Some(state_pairs),
+            optional_inputs: BTreeMap::new(),
+            static_cache: None,
+        })
+    }
+
     fn from_session_with_cuda_options_and_io(
         mut session: InferenceSession,
         cuda_options: NativeDecodeCudaOptions,
         io: Option<&ModelIoSpec>,
     ) -> anyhow::Result<Self> {
+        // Auto-derive a decoder I/O spec from the graph ports when the model
+        // package declares none. Declared `io` always wins; this fallback is
+        // additive and engages ONLY for hybrid linear-attention decoders (models
+        // exposing recurrent `conv_state`/`recurrent_state` state pairs the
+        // shape-inference path cannot classify). Pure-dense decoders derive no
+        // state pairs and keep their existing `io = None` load path unchanged, so
+        // no currently-loadable model changes behavior. See #384.
+        let derived_io: Option<ModelIoSpec> = if io.is_none() {
+            Self::derive_fallback_io(&session)
+        } else {
+            None
+        };
+        let io = io.or(derived_io.as_ref());
         let mut io_span = onnx_genai_ort::prof_span!("native.inspect_decode_io");
         let input_names = session
             .inputs()

crates/onnx-genai-engine/src/native_decode/mod.rs

@@ -6,15 +6,17 @@ use crate::decode_loop::{DecodeLoopBackend, DecodeLoopState, run_decode_loop};
 use crate::logits::{ProcessorChain, TokenId};
 use crate::sampling::sample_greedy;
 use anyhow::{Context, bail};
-use onnx_genai_metadata::{KvOwnership, ModelIoSpec, SequenceInputKind, SharedKvGroup};
+use onnx_genai_metadata::{
+    KvOwnership, LoopStatePair, ModelIoSpec, SequenceInputKind, SharedKvGroup,
+};
 use onnx_genai_ort::Tokenizer;
 use onnx_runtime_ir::{DataType, DeviceType, Dim, SymbolId};
 use onnx_runtime_session::{
     CaptureDeclineReport, DecodePrecision, DeviceAllocationCounts, DeviceBindingTransferStats,
     DeviceGraphCaptureResult, DeviceIoBinding, DevicePreference, InferenceSession, Tensor,
 };
 use onnx_runtime_tracer::{Args, TraceContext, capture_rejected};
-use std::collections::{HashMap, HashSet};
+use std::collections::{BTreeMap, HashMap, HashSet};
 use std::path::Path;
 use std::sync::Arc;
 

crates/onnx-genai-engine/src/native_decode/tests.rs

@@ -467,6 +467,185 @@ fn tiny_decoder(last_token_logits: bool) -> InferenceSession {
     InferenceSession::from_graph(graph).expect("build tiny decoder")
 }
 
+/// A tiny hybrid decoder mirroring `tiny_decoder` but adding a recurrent
+/// linear-attention layer: layer 0 carries `conv_state`/`recurrent_state`
+/// (fixed loop-carried state, threaded past→present by `Identity`), layer 1 is
+/// dense GQA (`key`/`value`). All ports use the conventional
+/// `past_key_values.%d.*` / `present.%d.*` names, so the loader's graph-derived
+/// I/O fallback can classify them without a declared `io:` block.
+fn tiny_hybrid_decoder() -> InferenceSession {
+    let mut graph = Graph::new();
+    graph.opset_imports.insert(String::new(), 11);
+    let batch = graph.intern_symbol("batch");
+    let sequence = graph.intern_symbol("sequence");
+    let total = graph.intern_symbol("total");
+    let past = graph.intern_symbol("past");
+    let shape = |dims: &[Dim]| -> Shape { dims.to_vec() };
+
+    let input_ids = graph.create_named_value(
+        "input_ids",
+        DataType::Int64,
+        shape(&[batch.into(), sequence.into()]),
+    );
+    let attention_mask = graph.create_named_value(
+        "attention_mask",
+        DataType::Int64,
+        shape(&[batch.into(), total.into()]),
+    );
+    let position_ids = graph.create_named_value(
+        "position_ids",
+        DataType::Int64,
+        shape(&[batch.into(), sequence.into()]),
+    );
+    // Layer 0: recurrent state ports (fixed-extent, replaced each step).
+    let conv_state = graph.create_named_value(
+        "past_key_values.0.conv_state",
+        DataType::Float32,
+        shape(&[batch.into(), 4.into(), 3.into()]),
+    );
+    let recurrent_state = graph.create_named_value(
+        "past_key_values.0.recurrent_state",
+        DataType::Float32,
+        shape(&[batch.into(), 2.into(), 4.into(), 4.into()]),
+    );
+    // Layer 1: dense GQA KV ports (growable along seq).
+    let past_key = graph.create_named_value(
+        "past_key_values.1.key",
+        DataType::Float32,
+        shape(&[batch.into(), 1.into(), past.into(), 1.into()]),
+    );
+    let past_value = graph.create_named_value(
+        "past_key_values.1.value",
+        DataType::Float32,
+        shape(&[batch.into(), 1.into(), past.into(), 1.into()]),
+    );
+    for input in [
+        input_ids,
+        attention_mask,
+        position_ids,
+        conv_state,
+        recurrent_state,
+        past_key,
+        past_value,
+    ] {
+        graph.add_input(input);
+    }
+
+    let cast = graph.create_value(DataType::Float32, shape(&[batch.into(), sequence.into()]));
+    insert_op(
+        &mut graph,
+        "Cast",
+        vec![input_ids],
+        cast,
+        &[("to", Attribute::Int(1))],
+    );
+    let current_kv = graph.create_value(
+        DataType::Float32,
+        shape(&[batch.into(), 1.into(), sequence.into(), 1.into()]),
+    );
+    insert_op(
+        &mut graph,
+        "Unsqueeze",
+        vec![cast],
+        current_kv,
+        &[("axes", Attribute::Ints(vec![1, 3]))],
+    );
+
+    let logits = graph.create_named_value(
+        "logits",
+        DataType::Float32,
+        shape(&[batch.into(), sequence.into(), 1.into()]),
+    );
+    insert_op(
+        &mut graph,
+        "Unsqueeze",
+        vec![cast],
+        logits,
+        &[("axes", Attribute::Ints(vec![2]))],
+    );
+
+    // Recurrent state: replaced wholesale (Identity pass-through here).
+    let present_conv = graph.create_named_value(
+        "present.0.conv_state",
+        DataType::Float32,
+        shape(&[batch.into(), 4.into(), 3.into()]),
+    );
+    insert_op(&mut graph, "Identity", vec![conv_state], present_conv, &[]);
+    let present_recurrent = graph.create_named_value(
+        "present.0.recurrent_state",
+        DataType::Float32,
+        shape(&[batch.into(), 2.into(), 4.into(), 4.into()]),
+    );
+    insert_op(
+        &mut graph,
+        "Identity",
+        vec![recurrent_state],
+        present_recurrent,
+        &[],
+    );
+    // Dense KV: appended along the sequence axis.
+    let present_key = graph.create_named_value(
+        "present.1.key",
+        DataType::Float32,
+        shape(&[batch.into(), 1.into(), total.into(), 1.into()]),
+    );
+    insert_op(
+        &mut graph,
+        "Concat",
+        vec![past_key, current_kv],
+        present_key,
+        &[("axis", Attribute::Int(2))],
+    );
+    let present_value = graph.create_named_value(
+        "present.1.value",
+        DataType::Float32,
+        shape(&[batch.into(), 1.into(), total.into(), 1.into()]),
+    );
+    insert_op(
+        &mut graph,
+        "Concat",
+        vec![past_value, current_kv],
+        present_value,
+        &[("axis", Attribute::Int(2))],
+    );
+    for output in [
+        logits,
+        present_conv,
+        present_recurrent,
+        present_key,
+        present_value,
+    ] {
+        graph.add_output(output);
+    }
+    InferenceSession::from_graph(graph).expect("build tiny hybrid decoder")
+}
+
+#[test]
+fn native_decoder_auto_derives_hybrid_state_io() {
+    // A hybrid decoder with NO declared `io:` must load via the graph-derived
+    // fallback: token/mask/position/logits bound by conventional name, dense KV
+    // for layer 1, and the layer-0 recurrent ports classified as loop-carried
+    // state pairs. This is the exact gap that blocked qwen3.x hybrids (#384).
+    NativeDecodeSession::from_session(tiny_hybrid_decoder())
+        .expect("hybrid decoder auto-derives its decoder-state io from the graph");
+}
+
+#[test]
+fn native_decoder_auto_derive_skips_dense_ambiguous_decoder() {
+    // The fallback's safety gate: a pure-dense decoder derives ZERO state pairs,
+    // so auto-derive declines and the model keeps its existing shape-inference
+    // path — which still fails on the ambiguous token input exactly as before.
+    // This proves the fallback never over-fires on models that load today.
+    let error = match NativeDecodeSession::from_session(tiny_decoder(false)) {
+        Ok(_) => panic!("dense ambiguous decoder must still require explicit token_input"),
+        Err(error) => error,
+    };
+    assert!(
+        error.to_string().contains("model.io.token_input"),
+        "{error:#}"
+    );
+}
+
 fn target_io(sequence_source: SequenceInputKind) -> ModelIoSpec {
     ModelIoSpec {
         sequence_source: Some(sequence_source),