[Mistral3] Add VLM support with multi-image inference by titaiwangms · Pull Request #2077 · microsoft/onnxruntime-genai

titaiwangms · 2026-04-08T23:15:30Z

Summary

Adds Mistral3/Pixtral VLM support to onnxruntime-genai with multi-image inference. Includes C++ image processor, PixtralVisionState for per-image vision processing, Python export support, and comprehensive tests.

Changes

C++ Runtime

Mistral3 image processor — [IMG]/[IMG_BREAK]/[IMG_END] token expansion based on image resolution and patch geometry, multi-image support
PixtralVisionState — per-image vision processing loop with bounds checks and overflow guard; slices from padded batch tensor using image_sizes metadata from ort-extensions PixtralImageSizes op
Virtual SetExtraInputs — proper polymorphic dispatch for vision state subclasses
IsPixtralFamily() model type detection — enables Pixtral-specific codepath
processor_config.json — with PixtralImageSizes preprocessing step
context_length / max_length separation — context_length controls KV cache allocation while max_length controls generation stopping, preventing premature EOS with large image token counts
INT32 input_ids — token IDs above 32767 (Pixtral [IMG]=128011) require int32

Python Export Support

Mistral3 model classes (Mistral3Config, Mistral3ForConditionalGeneration)
FP8 dtype promotion for checkpoint loading
get_user_content() handler for Mistral3 prompt formatting

Multi-Image Architecture

Pixtral uses dynamic image sizes (28×28 to 1540×1540) so images can't be batched in the vision encoder. PixtralVisionState processes each image individually by:

Reading image_sizes tensor from ort-extensions PixtralImageSizes op (provides per-image H×W)
Slicing the padded [N, C, max_H, max_W] batch tensor to extract each image's actual pixels
Running vision encoder on each image separately
Concatenating vision embeddings for the decoder

Dependencies

onnxruntime-extensions PR Fix C# phi-2 test #1050 — PixtralImageSizes custom op for image size metadata
PR Fix YaRN RoPE bugs in model builder and add parity tests #2076 — YaRN RoPE parity fixes (merged) ✅

Testing

5 multi-image token expansion tests (various image sizes and counts)
Virtual dispatch verification tests for SetExtraInputs
YaRN RoPE parity tests (from merged Fix YaRN RoPE bugs in model builder and add parity tests #2076)
E2E verified: multi-image (fish.jpg + challenge.jpg) correctly describes both images

Copilot

Pull request overview

Adds end-to-end support for Pixtral/Ministral-3 (Mistral3ForConditionalGeneration) vision-language models across the C++ runtime (VLM registration + image processor) and the Python ModelBuilder (architecture dispatch + FP8 weight handling) so the text decoder and vision inputs can be exported/consumed by ORT GenAI.

Changes:

Register mistral3 as a VLM type and wire up a Mistral3ImageProcessor in the C++ processor factory.
Implement Pixtral-specific prompt/image preprocessing (special token expansion, NHWC→NCHW transpose, spatial-merge-aware token counting).
Add Python ModelBuilder support for Mistral3ForConditionalGeneration via Mistral3TextModel and builder dispatch/config flattening.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
src/models/model_type.h	Adds `mistral3` to the VLM type set so it routes through multimodal model codepaths.
src/models/model.cpp	Registers `Mistral3ImageProcessor` in `MultiModalProcessor` factory.
src/models/mistral3_image_processor.h	Declares the Pixtral image processor used by the runtime.
src/models/mistral3_image_processor.cpp	Implements Pixtral prompt expansion, pixel preprocessing, and NHWC→NCHW transpose.
src/python/py/models/builder.py	Adds architecture dispatch for `Mistral3ForConditionalGeneration` and config adjustments for text-only export.
src/python/py/models/builders/mistral.py	Introduces `Mistral3TextModel` with FP8 linear weight dequantization when loading HF weights.
src/python/py/models/builders/init.py	Exports `Mistral3TextModel` for builder imports.

Copilot

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 7 comments.

titaiwangms · 2026-04-10T22:36:43Z

⚠️ Merge Dependency: This PR is stacked on top of PR #2076 (pr/fix-yarn-config-resolution). PR #2076 must be merged first, then this PR can be rebased onto main and merged.

Add full Mistral3/Pixtral VLM support to onnxruntime-genai: Model builder (src/python/py/models/builders/mistral.py): - Mistral3TextModel builder with FP8 dequantization for quantized checkpoints - image_token_id injection into genai_config.json Image preprocessor (src/models/mistral3_image_processor.cpp): - ProcessPixtralPrompt: per-image token sequence with [IMG], [IMG_BREAK], [IMG_END] - Smart resize dimensions from image_sizes tensor with shape validation - NCHW transpose for ort-extensions NHWC output Multi-image vision loop (src/models/multi_modal.cpp): - PixtralVisionState: per-image vision.onnx execution with zero-padded slicing - SetExtraInputs override reads image_sizes for per-image dimensions - Uses actual tensor memory info for GPU-safe sub-tensor views - Explicit rejection of multi-image without image_sizes metadata - Overflow guard on feature buffer offsets Factory and config: - IsPixtralFamily() in model_type.h for vision state dispatch - YaRN RoPE parity test with GPT-OSS-20B config - Simplified type hierarchy test (compile-time static_asserts) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot

Pull request overview

Copilot reviewed 15 out of 16 changed files in this pull request and generated 5 comments.

…ment cleanup - multi_modal.cpp: Add image_widths_ size validation matching heights check - multi_modal.cpp: Add h_i/w_i bounds checks before memcpy (0 < h_i <= h_max) - common.py: Add mistral3 model_type elif for multi-image [IMG] tags - test_mistral3_tokens.py: Clarify token IDs are test-only, add del generator comment - test_yarn_rope_parity.py: Remove duplicate comment header Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

…ral config Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

…mpty stub - mistral3_image_processor.cpp: Keep soft if/null fallback for image_sizes tensor. PixtralImageSizes step is not yet in processor_config.json (olive-recipes needs updating), so CheckResult would break single-image inference. Added clarifying comment explaining the dependency. - builders/base.py: Add pass to empty update_genai_config() method stub. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

The runtime should crash on misconfigured models (missing PixtralImageSizes) rather than silently producing wrong multi-image results. Models must be exported with PixtralImageSizes in processor_config.json. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Add comment noting opportunity to explore batched image processing for improved throughput in PixtralVisionState. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

…161) ## Summary Add an end-to-end olive-recipe demo for Ministral-3-3B VLM and enhance the mobius genai integration with Pixtral processor config support and tokenizer class remapping. ## Changes ### Core: genai integration enhancements (`src/mobius/integrations/ort_genai/`) - **`_write_processor_config`**: Enhanced to generate full ORT-extensions image transform pipeline for VL models. Dispatches by model type (Pixtral: longest_edge resize; generic: smart_resize). Derives normalization params from HF processor config. - **`_fix_tokenizer_config`**: New function with `_TOKENIZER_CLASS_REMAP` dict to remap unsupported tokenizer classes (e.g. `TokenizersBackend` → `LlamaTokenizer`). Called automatically from `write_ort_genai_config()`. - 4 new tests (31 total passing) ### Example: olive-recipe demo (`examples/olive/ministral-3-3b-vlm/`) - **optimize.py**: Pure mobius export (`build()` → `save()` → `write_ort_genai_config()`) with optional Olive quantization via `--olive-config` - **inference.py**: ORT GenAI multimodal inference (text-only, image+text, interactive) - **eval.py**: AI2D benchmark evaluation (ONNX vs PyTorch comparison) - Olive configs for CPU (INT4) and CUDA (FP16) quantization - README with setup, export, inference, and evaluation instructions ### Skills - New `olive-recipe` skill documenting the mobius + Olive hybrid pipeline pattern - Updated `ort-genai-config` skill with Pixtral processor config and tokenizer remap info ## References - [Issue #158](#158) - [PR #130](#130) (Pixtral/Ministral3 VLM support) - [olive-recipes PR #352](microsoft/olive-recipes#352) - [onnxruntime-genai PR #2077](microsoft/onnxruntime-genai#2077) Closes #158 --------- Signed-off-by: Ti-Tai Wang <titaiwang@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

titaiwangms marked this pull request as draft April 8, 2026 23:17

titaiwangms force-pushed the pr/mistral3-vlm-support branch from 09b475c to accc0ee Compare April 8, 2026 23:41

titaiwangms requested a review from Copilot April 9, 2026 01:32

Copilot started reviewing on behalf of titaiwangms April 9, 2026 01:33 View session

Copilot AI reviewed Apr 9, 2026

View reviewed changes

Comment thread src/models/mistral3_image_processor.cpp Outdated

titaiwangms force-pushed the pr/fix-yarn-config-resolution branch from dbb6147 to d786f35 Compare April 9, 2026 20:28

titaiwangms force-pushed the pr/mistral3-vlm-support branch 10 times, most recently from 3600ddf to ffae0a0 Compare April 10, 2026 00:03

titaiwangms requested a review from Copilot April 10, 2026 20:48

Copilot started reviewing on behalf of titaiwangms April 10, 2026 20:49 View session

Copilot AI reviewed Apr 10, 2026

View reviewed changes

titaiwangms force-pushed the pr/fix-yarn-config-resolution branch from d786f35 to 216066d Compare April 10, 2026 21:10

titaiwangms marked this pull request as ready for review April 10, 2026 21:15

titaiwangms changed the title ~~Add Mistral3/Pixtral VLM support~~ Add Mistral3/Pixtral VLM support: image processor, model builder, parity tests Apr 10, 2026

titaiwangms force-pushed the pr/mistral3-vlm-support branch from ffae0a0 to 078d398 Compare April 10, 2026 21:46

titaiwangms requested review from baijumeswani, kunal-vaishnavi, tianleiwu and xadupre April 10, 2026 21:51

This was referenced Apr 10, 2026

Add Ministral-3-3B VLM olive-recipe E2E demo to examples/olive/ onnxruntime/mobius#158

Closed

Add Ministral-3-3B VLM olive-recipe demo + enhance genai integration onnxruntime/mobius#161

Merged

github-advanced-security AI found potential problems Apr 20, 2026

View reviewed changes

Comment thread test/python/test_mistral3_tokens.py Fixed

titaiwangms force-pushed the pr/mistral3-vlm-support branch from b9bd58f to f84c0f6 Compare April 20, 2026 20:49

titaiwangms changed the title ~~feat: Add Mistral3 (Pixtral) VLM support with context_length/max_length separation~~ [Mistral3] Add VLM support with multi-image inference Apr 20, 2026

kunal-vaishnavi reviewed Apr 20, 2026

View reviewed changes

Comment thread src/models/multi_modal.cpp

kunal-vaishnavi reviewed Apr 20, 2026

View reviewed changes

Comment thread src/python/py/models/builders/mistral.py Outdated

titaiwangms requested a review from Copilot April 21, 2026 19:59

Copilot started reviewing on behalf of titaiwangms April 21, 2026 19:59 View session

Copilot AI reviewed Apr 21, 2026

View reviewed changes

Comment thread test/python/test_yarn_rope_parity.py Outdated

Comment thread src/models/mistral3_image_processor.cpp

Comment thread src/models/multi_modal.cpp

Comment thread src/models/multi_modal.cpp

Comment thread test/python/test_mistral3_tokens.py Outdated

titaiwangms and others added 2 commits April 21, 2026 20:29

Refactor: add update_genai_config hook to base builder, simplify mist…

d89140c

…ral config Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

kunal-vaishnavi reviewed Apr 21, 2026

View reviewed changes

Comment thread src/models/mistral3_image_processor.cpp Outdated

kunal-vaishnavi reviewed Apr 21, 2026

View reviewed changes

Comment thread src/python/py/models/builders/base.py

kunal-vaishnavi previously approved these changes Apr 21, 2026

View reviewed changes

titaiwangms dismissed kunal-vaishnavi’s stale review via a6f120e April 21, 2026 22:32

titaiwangms force-pushed the pr/mistral3-vlm-support branch from a6f120e to 84978b2 Compare April 21, 2026 22:50

titaiwangms requested a review from kunal-vaishnavi April 23, 2026 17:50

titaiwangms enabled auto-merge (squash) April 23, 2026 18:06

Add TODO for vision encoder batching optimization

9cbe1a0

Add comment noting opportunity to explore batched image processing for improved throughput in PixtralVisionState. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

kunal-vaishnavi approved these changes Apr 24, 2026

View reviewed changes

titaiwangms merged commit 19c2f7b into main Apr 26, 2026
16 of 18 checks passed

titaiwangms deleted the pr/mistral3-vlm-support branch April 26, 2026 18:47

Copilot AI mentioned this pull request Apr 28, 2026

extend modelbuilder to build Olmo3, SmolLM3 and other models #2078

Open

dependabot Bot mentioned this pull request May 30, 2026

Bump Microsoft.ML.OnnxRuntimeGenAI from 0.13.1 to 0.14.0 dvai-global/dvai-bridge#33

Closed

dependabot Bot mentioned this pull request Jun 7, 2026

Bump Microsoft.ML.OnnxRuntimeGenAI.DirectML and Microsoft.ML.OnnxRuntimeGenAI.Managed EricCogen/GauntletCI#226

Merged

This was referenced Aug 1, 2026

Bump Microsoft.ML.OnnxRuntimeGenAI and Microsoft.ML.OnnxRuntimeGenAI.Managed trackdubllc/Trackdub#53

Merged

Bump Microsoft.ML.OnnxRuntime.Gpu and Microsoft.ML.OnnxRuntimeGenAI.Cuda trackdubllc/Trackdub#54

Closed

Uh oh!

Conversation

titaiwangms commented Apr 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

C++ Runtime

Python Export Support

Multi-Image Architecture

Dependencies

Testing

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

titaiwangms commented Apr 10, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

titaiwangms commented Apr 8, 2026 •

edited

Loading