Fix DecoderState input_ids check regression introduced in #2103

[Ti-Tai Wang (titaiwangms)]

Problem

PR #2103 (Add Gemma4 multimodal support) introduced an optional decoder_input_ids_ in DecoderState for models like Gemma4 that require input_ids alongside inputs_embeds. The guard condition was:

if (model_.session_info_.HasInput(model_.config_->model.decoder.inputs.input_ids)) {

However, model_.session_info_ is a combined SessionInfo aggregating inputs from ALL sessions — decoder, embedding, vision, and speech. Since the embedding session always takes input_ids as its primary input, HasInput("input_ids") returns true for every VLM that has an embedding model, even when the decoder ONNX itself has no input_ids input.

Impact

Any VLM using the 3-model pipeline (embedding + decoder) where the decoder does not have input_ids will fail at inference with:

RuntimeError: Invalid input name: input_ids

Confirmed broken: mistral3 (Ministral-3B / Pixtral family), whose decoder only accepts inputs_embeds.

Fix

Scope the check to a decoder-only SessionInfo so it only fires when the decoder ONNX itself declares input_ids as an input:

// Use a decoder-only SessionInfo to avoid false positives from the embedding session
SessionInfo decoder_only_info;
decoder_only_info.Add(*model_.decoder_session_);
if (decoder_only_info.HasInput(model_.config_->model.decoder.inputs.input_ids)) {
    decoder_input_ids_ = std::make_unique<DefaultInputIDs>(*this);
    decoder_input_ids_->Add();
}

Gemma4 is unaffected — its decoder ONNX genuinely has input_ids as an input, so the narrowed check still fires correctly for that model.

Verification

Tested with the exported Ministral-3B-Instruct-2512 (mistral3 type, 3-model VLM):

• generator.set_inputs(inputs) no longer throws
• Token generation succeeds end-to-end

Regressed by: #2103

[Ti-Tai Wang (titaiwangms)]

Akshay Sonawane (@apsonawane) Can you check if this makes sense?

[Copilot]

Pull request overview

Fixes a regression in the multimodal decoding pipeline where DecoderState could incorrectly assume the decoder ONNX model accepts input_ids, due to using an aggregated SessionInfo that also includes embedding-session inputs.

Changes:

• Narrow the input_ids-presence check in DecoderState to a decoder-only SessionInfo built from decoder_session_.
• Prevent passing an invalid input_ids input to decoders that only accept inputs_embeds (e.g., mistral3 / Pixtral-family decoders).

[Justin Chu (justinchuby)]

Thanks- also a good point raised by kunal-vaishnavi where decoder potentially should not take input ids. Akshay Sonawane (@apsonawane) do you have a suggestion on what the proper contract for the components should be? Is the current contract with input ids in the decoder ok?

[Akshay Sonawane (apsonawane)]

Justin Chu (@justinchuby) is input_ids in gemma4 used for any other reason other than embedding lookup? If not, we should remove it from decoder and keep input_embeds as input. Embedding model already resolves tokens to embeddings so ideally decoder would not need input_ids. Looks like mistral3/pixtral, phi3-vision, qwen3.5 follows the same contract of keeping input_ids out of decoder.

[Akshay Sonawane (apsonawane)]

Ti-Tai Wang (@titaiwangms) Thank you for this fix. Shouldn't this issue be caught by mistral tests you added?

[Justin Chu (justinchuby)]

The key is its per layer embedding architecture. If we are not careful there would be multiple slicing and duplication. Or a more complicated decoder signature. I can share more details here in a bit

[Ti-Tai Wang (titaiwangms)]

Justin Chu (@justinchuby) Akshay Sonawane (@apsonawane) The fix is correct for both models, and the design debate is a separate concern from this bug fix.

By narrowing the check to decoder-only SessionInfo, the logic now correctly reflects each model's actual decoder contract:

• Gemma4: The decoder ONNX genuinely declares input_ids as an input (required for per-layer image embedding integration). Decoder-only HasInput("input_ids") → true. Behavior is unchanged from before PR Add Gemma4 multimodal support (vision + audio) #2103.
• Mistral3/Pixtral: The decoder only accepts inputs_embeds. Decoder-only HasInput("input_ids") → false. input_ids is no longer incorrectly injected.

The bug in PR #2103 was that the combined session_info_ includes all sessions (decoder + embedding + vision). The embedding session always has input_ids. So the old check was effectively true for every VLM with a separate embedding model — even those whose decoder doesn't declare it.

Whether Gemma4's decoder should eventually be refactored to accept only inputs_embeds (removing input_ids) is a valid design question, but that's a separate PR. Justin Chu (@justinchuby)'s point about per-layer embedding architecture is the key complexity there — it's not a simple swap. This fix restores correct behavior for Mistral3 without touching the Gemma4 contract.

[Ti-Tai Wang (titaiwangms)]

Akshay Sonawane (@apsonawane) Good catch on the test gap. The existing tests likely cover each session type in isolation and don't exercise DecoderState construction with all three sessions (decoder + embedding + vision) loaded simultaneously — which is the exact scenario where the combined session_info_ returns a false positive for input_ids.

I'll add a unit test that:

1. Constructs a mock multi-modal pipeline with separate decoder, embedding, and vision sessions
2. Verifies input_ids is injected only when the decoder session genuinely declares it
3. Verifies input_ids is NOT injected for a decoder that only declares inputs_embeds, even when an embedding session (which does have input_ids) is present

This would have caught the PR #2103 regression at the time it was introduced.

[Justin Chu (justinchuby)]

I created onnxruntime/mobius#296 and a discussion on #2149. Please take a look.

[Akshay Sonawane (apsonawane)]

If we can add a negative test that would be good