Add Mage-VL image and video support - #472
Open
justinchuby wants to merge 4 commits into
Open
Conversation
Implement the Mage-ViT vision tower, Qwen3 decoder package, explicit sampled-frame RoPE, four-frame packed attention, processor/config extraction, and exact checkpoint weight routing for microsoft/Mage-VL. Add synthetic multi-dtype parity, batch-safe media embedding coverage, real mixed image/video L4 and 24-token L5 goldens, CUDA and runtime export integration, and multimodal golden infrastructure. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
State the ORT GenAI position and media-input limitations precisely, and resolve lint findings in the synthetic parity reference. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Apply the repository formatter to the newly tracked model module. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds first-class support for the microsoft/Mage-VL (model_type: mage_vl) vision-language model as a standardized 3-model package (decoder / vision_encoder / embedding), including streaming-video-specific vision inputs (patch_positions) and updates to the test + export infrastructure to exercise real image+video pipelines.
Changes:
- Introduces the Mage-VL model implementation, task split, and HF-config extraction for the custom Mage-ViT vision tower.
- Extends golden/e2e infrastructure and fixtures to support ordered image+video inputs (including deterministic frame sampling and pixel budgets).
- Updates ORT GenAI auto-export to propagate
trust_remote_code, emit Mage-VL’s processor config filename, and keep configs schema-loadable when the runtime can’t acceptpatch_positions.
Reviewed changes
Copilot reviewed 28 out of 30 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/weight_alignment_test.py | Adds Mage-VL HF checkpoint → package-parameter alignment test. |
| tests/mage_vl_parity_test.py | Adds synthetic parity coverage for Mage-VL vision + embedding behaviors (incl. packed-attn CUDA build). |
| tests/e2e_golden_test.py | Extends VLM golden runner to preprocess videos and cast vision feeds to session input dtypes. |
| tests/cli_test.py | Verifies --trust-remote-code propagates into ORT GenAI config generation. |
| tests/_test_configs.py | Adds tiny mage_vl config to VL test configs. |
| testdata/golden/vision-language/mage-vl.json | Adds Mage-VL golden logits/token IDs fixture. |
| testdata/golden/vision-language/mage-vl_generation.json | Adds Mage-VL golden deterministic generation fixture. |
| testdata/cases/vision-language/mage-vl.yaml | Adds real image+video Mage-VL golden case (CI-skipped) with frame/pixel controls. |
| testdata/cases/schema.json | Adds schema support for videos, video_num_frames, media_max_pixels. |
| src/mobius/tasks/_vision_language_3model.py | Adds MageVLTask vision builder with patch_positions input. |
| src/mobius/tasks/init.py | Exports/registers MageVLTask under task name mage-vl. |
| src/mobius/models/mage_vl.py | Implements Mage-VL decoder, vision encoder, embedding mixer, and weight routing. |
| src/mobius/models/init.py | Publicly exports MageVLForConditionalGeneration. |
| src/mobius/integrations/ort_genai/auto_export.py | Adds mage_vl mapping, trust_remote_code plumb-through, Mage-VL processor filename, and runtime-compat config handling. |
| src/mobius/integrations/ort_genai/auto_export_test.py | Adds tests for Mage-VL processor config emission + trust_remote_code propagation + patch_positions omission in runtime config. |
| src/mobius/components/init.py | Exposes build_packed_token_offset in the public components API. |
| src/mobius/_testing/torch_reference.py | Adds Mage-VL-specific optional import shim for HF remote-code loading in test reference. |
| src/mobius/_testing/golden.py | Extends GoldenTestCase to include video inputs and media pixel/frame controls. |
| src/mobius/_testing/golden_test.py | Updates golden test expectations for new GoldenTestCase fields. |
| src/mobius/_registry.py | Registers mage_vl model_type → Mage-VL module/task association. |
| src/mobius/_configs/per_model/_mage_vl_vision.py | Adds vision-config extraction hook for Mage-VL’s custom Mage-ViT config. |
| src/mobius/_configs/per_model/init.py | Imports the Mage-VL vision hook for registration side effects. |
| src/mobius/_configs/_sub_configs.py | Adds video/vision token IDs + windowing fields to VisionConfig. |
| src/mobius/_configs/_extractors.py | Propagates new shared vision fields into extracted VisionConfig. |
| src/mobius/_configs/_base.py | Adds new top-level multimodal fields and ensures Mage-VL model_type retention in some hierarchical cases. |
| src/mobius/main.py | Propagates CLI --trust-remote-code into ORT GenAI artifact generation. |
| scripts/generate_golden.py | Extends golden generation to include videos + media pixel/frame controls and float-casts logits before extraction. |
| README.md | Lists Mage-VL as a supported multimodal model. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+895
to
+916
| image_processor = getattr(processor, "image_processor", None) | ||
| video_processor = getattr(processor, "video_processor", None) | ||
| saved_image_max = getattr(image_processor, "max_pixels", None) | ||
| image_size = getattr(image_processor, "size", None) | ||
| saved_image_longest = getattr(image_size, "longest_edge", None) | ||
| saved_video_max = getattr(video_processor, "max_pixels", None) | ||
| try: | ||
| if case.media_max_pixels is not None: | ||
| if image_processor is not None: | ||
| image_processor.max_pixels = case.media_max_pixels | ||
| if image_size is not None and saved_image_longest is not None: | ||
| image_size.longest_edge = case.media_max_pixels | ||
| if video_processor is not None: | ||
| video_processor.max_pixels = case.media_max_pixels | ||
| processed_pt = processor(**kwargs) | ||
| finally: | ||
| if image_processor is not None and saved_image_max is not None: | ||
| image_processor.max_pixels = saved_image_max | ||
| if image_size is not None and saved_image_longest is not None: | ||
| image_size.longest_edge = saved_image_longest | ||
| if video_processor is not None and saved_video_max is not None: | ||
| video_processor.max_pixels = saved_video_max |
Comment on lines
+482
to
+487
| if image_processor is not None and saved_image_max is not None: | ||
| image_processor.max_pixels = saved_image_max | ||
| if image_size is not None and saved_image_longest is not None: | ||
| image_size.longest_edge = saved_image_longest | ||
| if video_processor is not None and saved_video_max is not None: | ||
| video_processor.max_pixels = saved_video_max |
Register the public Mage-VL checkpoint as the architecture-validation model so registry coverage and remote-config graph validation include the new architecture. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Member
Author
CI triage updateThe first CI run on
No Bamba source, configuration, tolerance, or test code is changed by this PR, and I will not modify unrelated Bamba behavior to make this PR green. Baseline reference: the PR base is The follow-up CI run for |
Member
Author
Follow-up CI result (
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds first-class support for
microsoft/Mage-VL(model_type: mage_vl) as a standardized three-model package:decoder: Qwen3 text decoder with ordinary 1D RoPEvision_encoder: custom 24-layer Mage-ViT for images and streaming videoembedding: batch-safe text/visual feature mixerThe vision graph matches Mage-VL's fused biased QKV layout, exact 4:6:6 temporal/height/width RoPE, original sampled-frame temporal positions, 2x2 patch merger, and independent four-frame attention windows. CUDA/DML builds use packed attention with dynamically derived window boundaries; portable builds retain a standard-ONNX block-diagonal fallback.
This also adds remote config extraction, registry/public exports, exact checkpoint weight routing, ordered mixed image/video processing, ORT config hooks, real media fixtures, and L1-L5 coverage.
Validation
mage_vlregistry/config/model support was present.d88b153285f1633a61b2f693c59c8576693af185.[0, 180, 360, 539, 719], and 1,138 prompt tokens.python -m pytest tests/build_graph_test.py tests/weight_alignment_test.py tests/mage_vl_parity_test.py -k "mage_vl" -q: 10 passed.atol/rtol=1e-4), FP16 parity (1e-2), four-frame-boundary video coverage, empty-media decode, multi-row media ordering, and packed-CUDA graph coverage.python -m pytest tests/yaml_schema_test.py src/mobius/_testing/golden_test.py src/mobius/integrations/ort_genai/auto_export_test.py -q: 345 passed.python -m pytest tests/build_graph_test.py tests/cli_test.py src/ -q -k "not phi4mm and not apply_weights_unknown" --tb=short -n auto: 3,639 passed, 57 skipped.python -m pytest tests/quantization_integration_test.py -v: 6 passed.lintrunner -a: clean.Real checkpoint / GPU
Hardware: NVIDIA RTX A1000, 8,188 MiB; ORT 1.28.0 CUDA EP.
CUDAExecutionProvider: 1 passed in 324.12s.The image depicts a group of four men standing around a table in a stadium....CLI exports
Both completed successfully with real weights and emitted
decoder/,vision_encoder/,embedding/, tokenizer files,genai_config.json, andimage_processor.json:Olive
Olive 0.10.2 NF4 quantization completed on the real FP16 decoder in 74.46s, reducing external weights from 8,112,046,080 bytes to 2,204,314,624 bytes. The quantized decoder loaded in ORT CPU, produced finite
(1, 3, 151936)logits, and generated:Waivers / known runtime limits
qwen2_5_vlpath supplies 3D MRoPEposition_ids; Mage-VL's Qwen3 decoder correctly requires 2D 1D-RoPE positions, so generation fails withInvalid rank for input: position_ids Got: 3 Expected: 2. Mapping toqwen3loads but expects decoderinput_ids, incompatible with the standardized multimodalinputs_embedspipeline.patch_positionscontaining original sampled frame indices. ORT GenAI 0.15.2 rejects that vision input mapping as unknown, and its Qwen processor does not produce it. The exporter emits a precise warning and a schema-loadable config; direct Mobius ONNX image/video execution is covered by the passing real CUDA L4/L5 tests.