From b2e4a43d829642b419d4650c476da1484822cac3 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Sat, 1 Aug 2026 08:18:30 +0000 Subject: [PATCH] feat(native-decode): auto-derive io kv/state_pairs from graph ports (#384) Stock qwen3.x linear-attention hybrid exports (dense GQA + recurrent conv_state/recurrent_state layers under a Scan body), incl. qwen3.6-27b-int4, declare no model.io block, so native load failed at token-input resolution and could not classify the recurrent running-state ports. Add an additive loader fallback: when io is None, derive the decoder I/O topology from the graph ports. - genai-config: extract decoder_state_from_patterns (pure core of strict_decoder_state) and add pub GenAiConfig::derive_decoder_io_from_graph, reusing the guarded suffix-based recurrent-state derivation. A StateShapeMatch param keeps strict_decoder_state Exact (unchanged) while the fallback allows a fully-symbolic present.* state shape to pair with its concrete past_* input. - engine loader: build ModelGraphInfo from session ports and engage the derived spec ONLY when it yields >=1 recurrent state pair. Declared io always wins; pure-dense decoders (0 state pairs) keep their existing path unchanged. Correctness: auto-derived native CUDA decode of the STOCK 27b (no overlay) is byte-identical to the native CPU fp32 oracle (ORT-CUDA crashes on this class), reproducing the exact #384 re-probe sequence [11751,13,271,248068,...]. Runs eager (~5.9 tok/s): Scan declines CUDA-graph capture (separate perf lane). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../decisions/inbox/cohaagen-autoderive-io.md | 83 ++++++++ .../src/native_decode/load.rs | 92 +++++++++ .../src/native_decode/mod.rs | 6 +- .../src/native_decode/tests.rs | 179 ++++++++++++++++++ .../tests/native_autoderive_io_cuda_e2e.rs | 124 ++++++++++++ .../src/compatibility.rs | 136 ++++++++++++- crates/onnx-genai-genai-config/src/lib.rs | 1 + crates/onnx-genai-genai-config/src/tests.rs | 171 +++++++++++++++++ docs/MODEL_METADATA.md | 27 +++ 9 files changed, 816 insertions(+), 3 deletions(-) create mode 100644 .squad/decisions/inbox/cohaagen-autoderive-io.md create mode 100644 crates/onnx-genai-engine/tests/native_autoderive_io_cuda_e2e.rs diff --git a/.squad/decisions/inbox/cohaagen-autoderive-io.md b/.squad/decisions/inbox/cohaagen-autoderive-io.md new file mode 100644 index 0000000000..ccf044c68c --- /dev/null +++ b/.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` + 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. diff --git a/crates/onnx-genai-engine/src/native_decode/load.rs b/crates/onnx-genai-engine/src/native_decode/load.rs index 64da80358e..cc4034a6b4 100644 --- a/crates/onnx-genai-engine/src/native_decode/load.rs +++ b/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 { + 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::>(); + 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 { + // 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 = 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() diff --git a/crates/onnx-genai-engine/src/native_decode/mod.rs b/crates/onnx-genai-engine/src/native_decode/mod.rs index c2d29b4305..408bbeeeb6 100644 --- a/crates/onnx-genai-engine/src/native_decode/mod.rs +++ b/crates/onnx-genai-engine/src/native_decode/mod.rs @@ -6,7 +6,9 @@ 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::{ @@ -14,7 +16,7 @@ use onnx_runtime_session::{ 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; diff --git a/crates/onnx-genai-engine/src/native_decode/tests.rs b/crates/onnx-genai-engine/src/native_decode/tests.rs index 3000fb5540..8554ede1db 100644 --- a/crates/onnx-genai-engine/src/native_decode/tests.rs +++ b/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), diff --git a/crates/onnx-genai-engine/tests/native_autoderive_io_cuda_e2e.rs b/crates/onnx-genai-engine/tests/native_autoderive_io_cuda_e2e.rs new file mode 100644 index 0000000000..88f170715c --- /dev/null +++ b/crates/onnx-genai-engine/tests/native_autoderive_io_cuda_e2e.rs @@ -0,0 +1,124 @@ +//! Live GPU correctness lock for auto-derived decoder I/O (issue #384). +//! +//! Proves that a STOCK qwen3.6-27b int4 export — a hybrid linear-attention +//! decoder (16 dense GQA layers + 48 recurrent `conv_state`/`recurrent_state` +//! layers) whose `inference_metadata.yaml` declares NO `io` block — loads and +//! decodes natively purely from the graph-derived I/O fallback, with a token +//! stream byte-identical to the native CPU fp32 oracle. +//! +//! ORT-CUDA crashes on this model (internal `stl_vector` assertion), so there is +//! no ORT reference; the trusted oracle is our own native CPU backend, which +//! already threads recurrent state correctly. Any divergence between the CUDA +//! and CPU native runs is therefore a real correctness bug, not an export gap. +//! +//! This is the zero-overlay counterpart of the hand-overlay parity that was +//! proven manually during the #384 re-probe: with auto-derive, the stock export +//! "just works" with no sidecar `io:` block. +//! +//! Note: the CPU oracle run of a 27B model is SLOW (multi-minute prefill + +//! decode); this test is `#[ignore]` and meant to be run deliberately. +//! +//! ```bash +//! QWEN3_6_27B_CUDA_E2E_DIR=/home/justinchu/mary-models/qwen3.6-27b-int4-cuda \ +//! CUDA_VISIBLE_DEVICES=2 cargo test -p onnx-genai-engine \ +//! --features cuda,native-backend --test native_autoderive_io_cuda_e2e \ +//! -- --ignored --nocapture +//! ``` +#![cfg(all(feature = "cuda", feature = "native-backend"))] + +use std::path::{Path, PathBuf}; +use std::time::Instant; + +use onnx_genai_engine::{ + Engine, EngineConfig, EngineDecodeBackend, GenerateRequest, GenerateResult, NativeDecodeDevice, +}; + +const DEFAULT_MODEL_DIR: &str = "/home/justinchu/mary-models/qwen3.6-27b-int4-cuda"; +const PROMPT: &str = "The capital of France is"; +const MAX_NEW_TOKENS: usize = 16; + +fn model_dir() -> Option { + let dir = std::env::var_os("QWEN3_6_27B_CUDA_E2E_DIR") + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from(DEFAULT_MODEL_DIR)); + let required = [ + "model.onnx", + "model.onnx.data", + "inference_metadata.yaml", + "tokenizer.json", + ]; + let missing: Vec<_> = required + .iter() + .filter(|name| !dir.join(name).is_file()) + .collect(); + if missing.is_empty() { + Some(dir) + } else { + eprintln!( + "skipping native auto-derive io e2e: model directory {} is missing {:?}", + dir.display(), + missing + ); + None + } +} + +fn generate(dir: &Path, device: NativeDecodeDevice) -> anyhow::Result { + let mut engine = Engine::from_dir( + dir, + EngineConfig { + decode_backend: EngineDecodeBackend::Native, + native_device: Some(device), + ..EngineConfig::default() + }, + )?; + let mut request = GenerateRequest::new(PROMPT.to_string()); + request.options.max_new_tokens = MAX_NEW_TOKENS; + request.options.temperature = 0.0; + request.options.greedy = true; + request.options.stop_on_eos = false; + engine.generate(request) +} + +#[test] +#[ignore = "requires the real Qwen3.6-27B int4 hybrid export and a CUDA device (slow CPU oracle)"] +fn stock_export_auto_derives_io_and_matches_cpu_oracle() -> anyhow::Result<()> { + let Some(dir) = model_dir() else { + return Ok(()); + }; + if let Err(error) = onnx_runtime_ep_cuda::CudaExecutionProvider::new(0) { + eprintln!("skipping native auto-derive io e2e: CUDA unavailable: {error}"); + return Ok(()); + } + + // Native CUDA decode. The stock export has no declared `io:` block, so this + // only loads if the loader's graph-derived fallback classifies the recurrent + // conv/recurrent-state ports as loop-carried state pairs (not growable KV). + let cuda_start = Instant::now(); + let cuda = generate(&dir, NativeDecodeDevice::Cuda { index: Some(0) })?; + let cuda_elapsed = cuda_start.elapsed(); + + // Native CPU fp32 oracle over the same stock export. + let cpu_start = Instant::now(); + let cpu = generate(&dir, NativeDecodeDevice::Cpu)?; + let cpu_elapsed = cpu_start.elapsed(); + + eprintln!( + "native auto-derive io e2e:\n cuda tokens={:?} ({:?})\n cpu tokens={:?} ({:?})", + cuda.token_ids, cuda_elapsed, cpu.token_ids, cpu_elapsed + ); + + assert_eq!( + cuda.token_ids, cpu.token_ids, + "auto-derived native CUDA decode diverged from the CPU fp32 oracle — \ + graph-derived io must be byte-exact.\ncuda={:?}\ncpu={:?}", + cuda.token_ids, cpu.token_ids + ); + assert_eq!( + cuda.token_ids.len(), + MAX_NEW_TOKENS, + "expected {MAX_NEW_TOKENS} generated tokens, got {}", + cuda.token_ids.len() + ); + Ok(()) +} diff --git a/crates/onnx-genai-genai-config/src/compatibility.rs b/crates/onnx-genai-genai-config/src/compatibility.rs index 3a4e1c1ebb..4a855c4f80 100644 --- a/crates/onnx-genai-genai-config/src/compatibility.rs +++ b/crates/onnx-genai-genai-config/src/compatibility.rs @@ -33,6 +33,65 @@ struct DecoderStateMetadata { kv_dtype: String, } +/// How strictly a fixed-state input/output port pair's ONNX shapes must agree +/// during decoder-state derivation. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum StateShapeMatch { + /// Every axis must match exactly, including symbolic vs concrete. Used by the + /// config-driven `strict_decoder_state` path, whose contract is unchanged. + Exact, + /// Same rank, with symbolic (unknown) axes treated as wildcards: a concrete + /// axis matches an unknown axis, only two differing concrete extents fail. + /// Stock exports frequently leave a `present.*` recurrent-state shape fully + /// symbolic (`[?,?,?]`) even though the paired `past_*` input carries the + /// concrete running-state extent, so the config-free graph fallback must not + /// reject that legitimate pairing. + AllowSymbolic, +} + +impl StateShapeMatch { + fn shapes_pair(self, input: &[Option], output: &[Option]) -> bool { + match self { + StateShapeMatch::Exact => input == output, + StateShapeMatch::AllowSymbolic => { + input.len() == output.len() + && input.iter().zip(output).all(|(a, b)| match (a, b) { + (Some(a), Some(b)) => a == b, + _ => true, + }) + } + } + } +} + +/// A single loop-carried recurrent state port pair (`conv_state`/ +/// `recurrent_state`), threaded in→out and replaced wholesale each decode step +/// rather than appended along the sequence axis like growable KV. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct DerivedStatePair { + /// Graph input port carrying the state into the step. + pub input: String, + /// Graph output port carrying the updated state out of the step. + pub output: String, +} + +/// Decoder KV/state topology derived purely from an ONNX graph's port +/// inventory, without a `genai_config.json`. Returned by +/// [`GenAiConfig::derive_decoder_io_from_graph`]; consumed by the native +/// loader's auto-derive fallback when the sidecar declares no `io` block. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct DerivedDecoderIo { + /// Growable dense KV input ports (`past_key_values.%d.key`/`.value`), paired + /// positionally with `kv_outputs`. + pub kv_inputs: Vec, + /// Growable dense KV output ports (`present.%d.key`/`.value`). + pub kv_outputs: Vec, + /// Fixed loop-carried recurrent state pairs (empty for pure-dense decoders). + pub state_pairs: Vec, + /// Canonical dtype spelling shared by every KV cache tensor. + pub kv_dtype: String, +} + /// Coarse structural family a `genai_config.json` describes. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) enum ModelShape { @@ -1485,6 +1544,36 @@ impl GenAiConfig { "model.decoder.outputs.present_value_names", )?; + Self::decoder_state_from_patterns( + graph, + past_key, + past_value, + present_key, + present_value, + StateShapeMatch::Exact, + ) + } + + /// Derive a decoder's KV/state topology from the ONNX graph interface using + /// the four indexed `%d` key/value name patterns. + /// + /// This is the pure, config-independent core of + /// [`strict_decoder_state`](Self::strict_decoder_state): it depends only on + /// the graph inventory and the four pattern strings, so callers that lack a + /// `genai_config.json` (e.g. stock native exports carrying only an + /// `inference_metadata.yaml`) can reuse the exact same guarded derivation via + /// [`derive_decoder_io_from_graph`](Self::derive_decoder_io_from_graph). The + /// guard is inherited unchanged: only `suffix_tensor_map` (not + /// `strict_indexed_kv`) finds the non-KV recurrent `state_pairs`, so + /// cross-attention/Whisper KV is never misclassified as running state. + fn decoder_state_from_patterns( + graph: &ModelGraphInfo, + past_key: &str, + past_value: &str, + present_key: &str, + present_value: &str, + state_shape_match: StateShapeMatch, + ) -> Result { let past_key_names = match_indexed_tensors(&graph.inputs, past_key)?; let past_value_names = match_indexed_tensors(&graph.inputs, past_value)?; let present_key_names = match_indexed_tensors(&graph.outputs, present_key)?; @@ -1562,7 +1651,7 @@ impl GenAiConfig { for (suffix, input) in state_inputs { let output = state_outputs[&suffix]; require_same_dtype(input, output, "fixed-state input/output")?; - if input.dimensions != output.dimensions { + if !state_shape_match.shapes_pair(&input.dimensions, &output.dimensions) { return Err(incomplete(format!( "fixed-state pair '{}'/'{}' has different ONNX shapes", input.name, output.name @@ -1584,6 +1673,51 @@ impl GenAiConfig { }) } + /// Derive a decoder's KV/state I/O topology directly from an ONNX graph's + /// port inventory, using the conventional onnxruntime-genai key/value name + /// patterns (`past_key_values.%d.key`/`.value` → `present.%d.key`/`.value`). + /// + /// This is the config-free entry point for stock native exports that ship an + /// `inference_metadata.yaml` without an explicit `io` block (and no + /// `genai_config.json`): the native loader calls it as an additive fallback + /// so hybrid linear-attention decoders (which carry recurrent + /// `conv_state`/`recurrent_state` ports the shape-inference path cannot + /// classify) still load. It reuses the exact guarded derivation + /// ([`decoder_state_from_patterns`](Self::decoder_state_from_patterns)), so + /// dense cross-attention/Whisper KV is never misclassified as running state. + /// + /// Structural derivation failures are swallowed into `None` (mirroring + /// [`graph_decoder_state`](Self::graph_decoder_state)) so a caller can fall + /// back to its existing shape-inference path without regressing any model + /// that loads today. + pub fn derive_decoder_io_from_graph(graph: &ModelGraphInfo) -> Option { + let metadata = Self::decoder_state_from_patterns( + graph, + "past_key_values.%d.key", + "past_key_values.%d.value", + "present.%d.key", + "present.%d.value", + StateShapeMatch::AllowSymbolic, + ) + .ok()?; + let state_pairs = metadata + .state_pairs + .iter() + .map(|pair| { + Some(DerivedStatePair { + input: pair.get("input")?.as_str()?.to_owned(), + output: pair.get("output")?.as_str()?.to_owned(), + }) + }) + .collect::>>()?; + Some(DerivedDecoderIo { + kv_inputs: metadata.kv_inputs, + kv_outputs: metadata.kv_outputs, + state_pairs, + kv_dtype: metadata.kv_dtype, + }) + } + /// Best-effort graph-derived decoder KV/state topology for the single-decoder /// compatibility path. /// diff --git a/crates/onnx-genai-genai-config/src/lib.rs b/crates/onnx-genai-genai-config/src/lib.rs index 6ec3aa9a8b..bfe1599ded 100644 --- a/crates/onnx-genai-genai-config/src/lib.rs +++ b/crates/onnx-genai-genai-config/src/lib.rs @@ -33,6 +33,7 @@ mod json_builders; mod loading; mod wire_types; +pub use compatibility::{DerivedDecoderIo, DerivedStatePair}; pub(crate) use compatibility::{ ModelShape, incomplete, required_str, transducer_unsupported, unrepresentable_preprocessing, }; diff --git a/crates/onnx-genai-genai-config/src/tests.rs b/crates/onnx-genai-genai-config/src/tests.rs index 26cae80924..c12c16c110 100644 --- a/crates/onnx-genai-genai-config/src/tests.rs +++ b/crates/onnx-genai-genai-config/src/tests.rs @@ -326,6 +326,177 @@ fn hybrid_decoder_derives_sparse_kv_and_state_pairs() { } } +/// A qwen3.6-27b-shaped hybrid graph: 64 decoder layers where every 4th layer +/// (indices 3, 7, ... 63) is dense GQA full-attention (`key`/`value`) and the +/// remaining 48 layers are linear-attention recurrent +/// (`conv_state`/`recurrent_state`), using the conventional onnxruntime-genai +/// `past_key_values.%d.*` / `present.%d.*` port names the config-free +/// derivation keys off. +fn qwen27b_hybrid_graph() -> ModelGraphInfo { + let dense = [Some(1), Some(8), None, Some(128)]; + let conv = [Some(1), Some(10240), Some(3)]; + let recur = [Some(1), Some(16), Some(128), Some(128)]; + let mut inputs = vec![ + hybrid_graph_tensor("input_ids", "int64", &[Some(1), None]), + hybrid_graph_tensor("attention_mask", "int64", &[Some(1), None]), + hybrid_graph_tensor("position_ids", "int64", &[Some(1), None]), + ]; + let mut outputs = vec![hybrid_graph_tensor( + "logits", + "float32", + &[Some(1), None, Some(151936)], + )]; + for layer in 0..64 { + if layer % 4 == 3 { + inputs.push(hybrid_graph_tensor( + &format!("past_key_values.{layer}.key"), + "float16", + &dense, + )); + inputs.push(hybrid_graph_tensor( + &format!("past_key_values.{layer}.value"), + "float16", + &dense, + )); + outputs.push(hybrid_graph_tensor( + &format!("present.{layer}.key"), + "float16", + &dense, + )); + outputs.push(hybrid_graph_tensor( + &format!("present.{layer}.value"), + "float16", + &dense, + )); + } else { + inputs.push(hybrid_graph_tensor( + &format!("past_key_values.{layer}.conv_state"), + "float16", + &conv, + )); + inputs.push(hybrid_graph_tensor( + &format!("past_key_values.{layer}.recurrent_state"), + "float16", + &recur, + )); + // Stock exports leave the present recurrent-state shapes fully + // symbolic even though the paired past input carries the concrete + // running-state extent; the config-free fallback must accept this. + outputs.push(hybrid_graph_tensor( + &format!("present.{layer}.conv_state"), + "float16", + &[None, None, None], + )); + outputs.push(hybrid_graph_tensor( + &format!("present.{layer}.recurrent_state"), + "float16", + &[None, None, None, None], + )); + } + } + ModelGraphInfo { inputs, outputs } +} + +/// A dense qwen3-0.6b-shaped graph: 28 decoder layers, all dense GQA +/// (`key`/`value`), no recurrent state ports. +fn qwen06b_dense_graph() -> ModelGraphInfo { + let dense = [Some(1), Some(8), None, Some(128)]; + let mut inputs = vec![ + hybrid_graph_tensor("input_ids", "int64", &[Some(1), None]), + hybrid_graph_tensor("attention_mask", "int64", &[Some(1), None]), + hybrid_graph_tensor("position_ids", "int64", &[Some(1), None]), + ]; + let mut outputs = vec![hybrid_graph_tensor( + "logits", + "float32", + &[Some(1), None, Some(151936)], + )]; + for layer in 0..28 { + inputs.push(hybrid_graph_tensor( + &format!("past_key_values.{layer}.key"), + "float16", + &dense, + )); + inputs.push(hybrid_graph_tensor( + &format!("past_key_values.{layer}.value"), + "float16", + &dense, + )); + outputs.push(hybrid_graph_tensor( + &format!("present.{layer}.key"), + "float16", + &dense, + )); + outputs.push(hybrid_graph_tensor( + &format!("present.{layer}.value"), + "float16", + &dense, + )); + } + ModelGraphInfo { inputs, outputs } +} + +#[test] +fn derive_decoder_io_from_graph_splits_hybrid_kv_and_state() { + // The 27b hybrid layout: 16 dense GQA layers (2 ports each = 32 kv entries) + // and 48 recurrent layers (conv_state + recurrent_state = 96 state pairs). + let derived = GenAiConfig::derive_decoder_io_from_graph(&qwen27b_hybrid_graph()) + .expect("hybrid graph derives decoder io"); + assert_eq!(derived.kv_inputs.len(), 32); + assert_eq!(derived.kv_outputs.len(), 32); + assert_eq!(derived.state_pairs.len(), 96); + assert_eq!(derived.kv_dtype, "float16"); + + // KV lists must contain ONLY the dense full-attention ports. + assert!( + derived + .kv_inputs + .contains(&"past_key_values.3.key".to_string()) + ); + assert!( + derived + .kv_inputs + .contains(&"past_key_values.63.value".to_string()) + ); + assert!( + !derived + .kv_inputs + .iter() + .any(|name| name.contains("conv_state") || name.contains("recurrent_state")) + ); + + // Every recurrent port is a loop-carried state pair, past→present paired. + assert!(derived.state_pairs.iter().any(|pair| { + pair.input == "past_key_values.0.conv_state" && pair.output == "present.0.conv_state" + })); + assert!(derived.state_pairs.iter().any(|pair| { + pair.input == "past_key_values.62.recurrent_state" + && pair.output == "present.62.recurrent_state" + })); + assert!( + derived + .state_pairs + .iter() + .all(|pair| pair.input.contains("_state") && pair.output.contains("_state")) + ); +} + +#[test] +fn derive_decoder_io_from_graph_dense_has_no_state_pairs() { + // A pure-dense GQA decoder must derive KV ports but NEVER over-derive state + // pairs — the loader's safety gate leaves such models on the existing + // shape-inference path unchanged. + let derived = GenAiConfig::derive_decoder_io_from_graph(&qwen06b_dense_graph()) + .expect("dense graph derives decoder io"); + assert_eq!(derived.kv_inputs.len(), 56); + assert_eq!(derived.kv_outputs.len(), 56); + assert!( + derived.state_pairs.is_empty(), + "dense decoder must not gain recurrent state pairs, got {:?}", + derived.state_pairs + ); +} + #[test] fn uniform_decoder_graph_matches_pattern_expansion() { // A dense-KV model must produce the SAME kv_inputs whether or not the diff --git a/docs/MODEL_METADATA.md b/docs/MODEL_METADATA.md index 66014ca89f..95d80a4f82 100644 --- a/docs/MODEL_METADATA.md +++ b/docs/MODEL_METADATA.md @@ -225,6 +225,33 @@ attention masks, and position IDs all having rank two, or multiple same-shaped KV tensors—the corresponding metadata fields are required and loading fails with an error naming the missing field. +### Graph-derived `model.io` fallback for hybrid decoders (#384) + +When a package declares **no** `model.io` block at all, the native loader +attempts one additive fallback before shape inference: it derives the decoder +I/O topology from the graph's port inventory using the conventional +onnxruntime-genai key/value names (`past_key_values.%d.key`/`.value` → +`present.%d.key`/`.value`). This engages **only** when the derivation finds at +least one recurrent state pair — i.e. a hybrid linear-attention decoder that +also exposes non-KV `conv_state`/`recurrent_state` ports (qwen3.x, incl. the +27B). Those ports carry fixed loop-carried state that shape inference cannot +classify, so without this fallback such stock exports fail to load natively. + +A declared `model.io` always wins; the fallback never overrides one. Pure-dense +decoders derive zero state pairs, so the fallback declines and they keep the +existing shape-inference path unchanged — no model that loads today changes its +behavior. The derivation reuses the same guarded logic as the +`genai_config.json` compatibility path (recurrent ports are matched by suffix, +never misclassified as KV), with one relaxation: a `present.*` state port whose +exported shape is fully symbolic is accepted against its concrete `past_*` input +(a symbolic axis is unknown, not proof of a different shape). + +These hybrids run **eager**: their `Scan` control-flow body declines CUDA-graph +capture, so this fallback unblocks correctness (native decode of the model +class) independent of the capture perf lane. Correctness is validated +byte-for-byte against the native CPU fp32 oracle (ORT-CUDA crashes on this model +class, so there is no ORT reference). + Attention representation metadata is interpreted independently of `model.attention.type`. In particular, `key_sequence_lengths.scalar_broadcast` controls the representation contract rather than selecting an implementation,