From 8736a06caf0c3bd40c7f9a8c3e12ba0cb5e9adb2 Mon Sep 17 00:00:00 2001
From: Justin Chu <justinchuby@users.noreply.github.com>
Date: Fri, 31 Jul 2026 05:58:27 +0000
Subject: [PATCH] profile_native: resolve genai_config decoder io so capture
 counters print

The simple `profile_native` generate path (the only path that prints
`cuda_graph: captures/replays/fallbacks` and supports `--trace`
capture-reject reasons) loaded decoders via `NativeDecodeSession::load`
with `io = None`. For onnxruntime-genai `genai_config.json` decoders whose
token/attention/position inputs are all `Int64` rank-2, shape-based I/O
inference is ambiguous, so the load failed:

    cannot resolve model.io.token_input from tensor shape because 2 ports
    match: [("input_ids", ...), ("attention_mask", ...)]

and capture observability was unreachable for those models.

Add a public `NativeDecodeSession::load_with_resolved_io` that resolves the
model directory's `ModelIoSpec` from an adjacent
`inference_metadata.{yaml,yml,json}` sidecar, else onnxruntime-genai
`genai_config.json` compatibility synthesis (same precedence as the engine
directory loader), and threads it into the existing io-aware load path.
`profile_native`'s simple path now calls it.

Observability/plumbing only: capture semantics are unchanged
(`graph_capture` stays auto-decided, no defaults change, no kernel touched).

After: genai_config decoders load in the simple path and print real
counters, e.g. `cuda_graph: enabled=true captures=3 replays=87 fallbacks=0`,
and `--trace` surfaces per-op `capture_status` / ARG_CAPTURE_REJECTED reasons.

Refs #87 (capture observability).

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
---
 .../src/bin/profile_native.rs                 |  2 +-
 .../src/native_decode/load.rs                 | 57 +++++++++++++++++++
 2 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/crates/onnx-genai-bench/src/bin/profile_native.rs b/crates/onnx-genai-bench/src/bin/profile_native.rs
index 5e1004920a..69d23b9a9f 100644
--- a/crates/onnx-genai-bench/src/bin/profile_native.rs
+++ b/crates/onnx-genai-bench/src/bin/profile_native.rs
@@ -629,7 +629,7 @@ fn main() -> Result<()> {
             .context("build synthetic native session")?;
         NativeDecodeSession::from_session(native).context("wrap synthetic native decoder")?
     } else {
-        NativeDecodeSession::load(&model, device)
+        NativeDecodeSession::load_with_resolved_io(&model, device)
             .with_context(|| format!("load native decoder {}", model.display()))?
     };
 
diff --git a/crates/onnx-genai-engine/src/native_decode/load.rs b/crates/onnx-genai-engine/src/native_decode/load.rs
index 57ea0ad822..64da80358e 100644
--- a/crates/onnx-genai-engine/src/native_decode/load.rs
+++ b/crates/onnx-genai-engine/src/native_decode/load.rs
@@ -14,7 +14,64 @@ fn native_metadata_max_len_from_model_path(path: &Path) -> Option<usize> {
     .and_then(|metadata| metadata.model.and_then(|model| model.max_sequence_length))
 }
 
+/// Resolve a model directory's [`InferenceMetadata`] using the same precedence
+/// as the engine's directory loader: a native `inference_metadata.{yaml,yml,json}`
+/// sidecar first, then onnxruntime-genai `genai_config.json` compatibility
+/// synthesis. Returns `None` when neither is present so callers fall back to
+/// shape-based I/O inference exactly as before.
+fn resolve_io_metadata_from_model_path(
+    path: &Path,
+) -> Option<onnx_genai_metadata::InferenceMetadata> {
+    let root = if path.is_dir() { path } else { path.parent()? };
+    if let Some(metadata) = [
+        "inference_metadata.yaml",
+        "inference_metadata.yml",
+        "inference_metadata.json",
+    ]
+    .iter()
+    .map(|name| root.join(name))
+    .find(|path| path.is_file())
+    .and_then(|path| onnx_genai_metadata::load_metadata(&path).ok())
+    {
+        return Some(metadata);
+    }
+    let genai_config = root.join("genai_config.json");
+    if genai_config.is_file() {
+        return crate::engine::genai_config_compat_metadata_from_model_path(
+            Some(genai_config.as_path()),
+            path,
+        )
+        .ok()
+        .flatten();
+    }
+    None
+}
+
 impl NativeDecodeSession {
+    /// Load a decoder-with-past model, resolving its [`ModelIoSpec`] from an
+    /// adjacent `inference_metadata.{yaml,yml,json}` sidecar or (for
+    /// onnxruntime-genai packages) `genai_config.json`, so genai_config decoders
+    /// bind their token input from metadata instead of guessing from ambiguous
+    /// tensor shapes. This is the observability entry point used by measurement
+    /// tools so CUDA-graph capture counters and `--trace` reject reasons surface
+    /// for genai_config decoders. Capture semantics are identical to [`load`]:
+    /// `graph_capture` stays auto-decided and no defaults change; only I/O
+    /// resolution is threaded.
+    ///
+    /// [`load`]: NativeDecodeSession::load
+    pub fn load_with_resolved_io(
+        path: impl AsRef<Path>,
+        device: NativeDecodeDevice,
+    ) -> anyhow::Result<Self> {
+        let path = path.as_ref();
+        let metadata = resolve_io_metadata_from_model_path(path);
+        let io = metadata
+            .as_ref()
+            .and_then(|metadata| metadata.model.as_ref())
+            .and_then(|model| model.io.as_ref());
+        Self::load_with_cuda_options_and_io(path, device, NativeDecodeCudaOptions::default(), io)
+    }
+
     pub(crate) fn load_with_weight_offload_host_cache(
         path: impl AsRef<Path>,
         device: NativeDecodeDevice,