Add Pixtral/Ministral3 VLM support: dynamic vision, FP8 dequant, PatchMerger fix, MHA (#130)

## Summary

Adds Pixtral/Ministral3 VLM support with 3-model split
(vision/embedding/decoder) for onnxruntime-genai.

## Key Changes

### Vision encoder (`_pixtral_vision.py`)
- `PixtralVisionTower`: Conv2d patch embedding → RMSNorm → 24-layer
transformer with 2D RoPE
- `PixtralRoPE2D`: Precomputed 2D rotary position embeddings over
spatial grid
- `PixtralAttention`: Bidirectional MHA with 2D RoPE
(`com.microsoft.MultiHeadAttention`; `op.Attention` fallback for
onnx-standard EP)
- `Mistral3PatchMerger`: Spatial 2×2 patch merging matching HF
`F.unfold` dim-major ordering
- `Mistral3MultiModalProjector`: norm → merge → GELU MLP projection

### PatchMerger fix (cosine sim 0.007 → 0.999973)
The original transpose permutation `[0,1,3,2,4,5]` produced patch-major
ordering, but HF `F.unfold` produces dim-major `[D, ms_h, ms_w]`. The
learned `merging_layer` weights expect HF ordering. Fixed to
`[0,1,3,5,2,4]`.

### FP8 weight dequantization (`_weight_loading.py`)
- Detects FP8 weights, dequantizes via `fp8.to(bf16) * scale`
- Always returns new dict (non-mutating, addresses review comment)
- Docstring updated to list all removed suffixes

### Attention scale optimization (`_attention.py`)
- Moved `attn_scale` Mul before RoPE (scaling commutes with rotation) to
preserve `RotaryEmbedding → Attention` pattern for GQA rewrite rule
matching

### Task and infrastructure
- `PixtralVLTask`: Dynamic H×W vision input with Squeeze for rank-2
output
- Registered as `pixtral-vl` task in `_vision_language_3model.py`
- Olive-recipe `optimize.py` integration verified

## Testing
- Cosine sim vs HuggingFace PyTorch: **0.999973**
- E2E `model-mm.py` with fish.jpg and challenge.jpg: correct
descriptions
- Olive-recipe pipeline export + E2E verified
- `test_patch_merger_matches_hf_unfold_ordering`: regression test
against HF F.unfold
- All 1261 unit tests pass, lintrunner clean

## Files Changed
- `src/mobius/components/_pixtral_vision.py` — PatchMerger fix + MHA
- `src/mobius/components/_pixtral_vision_test.py` — F.unfold ordering
test
- `src/mobius/tasks/_vision_language_3model.py` — Squeeze vision output
- `src/mobius/_weight_loading.py` — Docstring + mutation fix
- `src/mobius/components/_attention.py` — attn_scale before RoPE

---------

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

Author: titaiwangms

src/mobius/_registry.py

@@ -476,14 +476,14 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None:
     "llava_next": ModelRegistration(LLaVAModel, task="vision-language"),
     "llava_next_video": ModelRegistration(LLaVAModel, task="vision-language"),
     "llava_onevision": ModelRegistration(LLaVAModel, task="vision-language"),
-    "mistral3": ModelRegistration(LLaVAModel, task="vision-language"),
+    "mistral3": ModelRegistration(LLaVAModel, task="pixtral-vl"),
     "mllama": ModelRegistration(MllamaCausalLMModel, task="mllama-vision-language"),
     "molmo": ModelRegistration(LLaVAModel, task="vision-language"),
     "ovis2": ModelRegistration(LLaVAModel, task="vision-language"),
     "paligemma": ModelRegistration(LLaVAModel, task="vision-language"),
     "phi4_multimodal": ModelRegistration(Phi4MMMultiModalModel, task="phi4mm-multimodal"),
     "phi4mm": ModelRegistration(Phi4MMMultiModalModel, task="phi4mm-multimodal"),
-    "pixtral": ModelRegistration(LLaVAModel, task="vision-language"),
+    "pixtral": ModelRegistration(LLaVAModel, task="pixtral-vl"),
     "qwen2_5_vl": ModelRegistration(Qwen25VLCausalLMModel, task="qwen-vl"),
     "qwen2_5_vl_text": ModelRegistration(Qwen25VLTextModel),
     "qwen2_vl": ModelRegistration(Qwen25VLCausalLMModel, task="qwen-vl"),

src/mobius/_weight_loading.py

@@ -147,6 +147,58 @@ def _parallel_download(
     return [path_map[f] for f in filenames]
 
 
+def _dequantize_fp8_weights(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    """Dequantize FP8 weights and return a new dict with float tensors.
+
+    Some HuggingFace checkpoints (e.g. Ministral-3-3B) store linear layer
+    weights as float8_e4m3fn with a scalar ``weight_scale_inv`` tensor.
+    The real weight value is ``fp8_weight.to(bfloat16) * weight_scale_inv``.
+
+    Dequantization targets bfloat16 because that is the native training
+    dtype for FP8-quantized checkpoints — the FP8 values represent
+    bfloat16 values scaled into the FP8 range.
+
+    This function detects FP8 tensors, applies the scale, and removes the
+    auxiliary ``weight_scale_inv``, ``activation_scale``, and
+    ``input_scale`` tensors.
+
+    Returns:
+        A new dict with FP8 weights dequantized to bfloat16 and the
+        auxiliary scale tensors removed.  Always returns a new dict,
+        even when no FP8 weights are found.
+    """
+    fp8_dtypes = {torch.float8_e4m3fn, torch.float8_e5m2}
+    fp8_keys = [k for k, v in state_dict.items() if v.dtype in fp8_dtypes]
+    if not fp8_keys:
+        return dict(state_dict)
+
+    # Work on a copy to avoid mutating the caller's dict
+    result = dict(state_dict)
+
+    logger.info("Dequantizing %d FP8 weights", len(fp8_keys))
+    for key in fp8_keys:
+        # Derive scale key from weight key using suffix replacement to avoid
+        # replacing '.weight' substrings that appear in the middle of the key
+        # (e.g. 'model.weight_proj.weight' → 'model.weight_proj.weight_scale_inv').
+        if key.endswith(".weight"):
+            scale_key = key[: -len(".weight")] + ".weight_scale_inv"
+        else:
+            scale_key = key + "_scale_inv"
+
+        if scale_key in result:
+            # Cast scale to bfloat16 to guarantee the output dtype is bfloat16,
+            # even when weight_scale_inv is stored as FP32 in the checkpoint.
+            scale = result[scale_key].to(torch.bfloat16)
+            result[key] = result[key].to(torch.bfloat16) * scale
+        else:
+            logger.warning("FP8 weight '%s' has no scale_inv — casting without scaling", key)
+            result[key] = result[key].to(torch.bfloat16)
+
+    # Remove auxiliary FP8 tensors (not needed in the ONNX graph)
+    aux_suffixes = (".weight_scale_inv", ".activation_scale", ".input_scale")
+    return {k: v for k, v in result.items() if not any(k.endswith(s) for s in aux_suffixes)}
+
+
 def _download_weights(model_id: str) -> dict[str, torch.Tensor]:
     """Download weights from HuggingFace and return as a state dict.
 
@@ -171,4 +223,6 @@ def _download_weights(model_id: str) -> dict[str, torch.Tensor]:
     state_dict: dict[str, torch.Tensor] = {}
     for path in tqdm.tqdm(paths, desc="Loading weights"):
         state_dict.update(safetensors.torch.load_file(path))
+
+    state_dict = _dequantize_fp8_weights(state_dict)
     return state_dict

src/mobius/_weight_loading_test.py

@@ -518,3 +518,108 @@ def build(self, module, config):
         module = CausalLMModel(config)
         pkg = build_from_module(module, config, task=StubTask())
         assert pkg["model"].graph.name == "stub"
+
+
+class TestDequantizeFP8Weights:
+    """Tests for _dequantize_fp8_weights."""
+
+    def test_no_fp8_returns_unchanged(self):
+        """Non-FP8 state dicts pass through unchanged."""
+        from mobius._weight_loading import _dequantize_fp8_weights
+
+        state_dict = {
+            "layer.weight": torch.randn(4, 4),
+            "layer.bias": torch.randn(4),
+        }
+        result = _dequantize_fp8_weights(state_dict)
+        assert set(result.keys()) == set(state_dict.keys())
+        assert torch.equal(result["layer.weight"], state_dict["layer.weight"])
+
+    def test_fp8_e4m3fn_dequantized(self):
+        """FP8 weights are multiplied by weight_scale_inv."""
+        from mobius._weight_loading import _dequantize_fp8_weights
+
+        fp8_weight = torch.tensor([1.0, 2.0, -1.0, 0.5], dtype=torch.float32).to(
+            torch.float8_e4m3fn
+        )
+        scale_inv = torch.tensor(0.5, dtype=torch.bfloat16)
+        state_dict = {
+            "proj.weight": fp8_weight,
+            "proj.weight_scale_inv": scale_inv,
+        }
+        result = _dequantize_fp8_weights(state_dict)
+        assert "proj.weight" in result
+        assert "proj.weight_scale_inv" not in result  # aux tensor removed
+        assert result["proj.weight"].dtype == torch.bfloat16
+        # Verify dequant: fp8→bf16 * scale_inv
+        expected = fp8_weight.to(torch.bfloat16) * scale_inv
+        assert torch.allclose(result["proj.weight"], expected)
+
+    def test_activation_scale_removed(self):
+        """Auxiliary activation_scale tensors are removed."""
+        from mobius._weight_loading import _dequantize_fp8_weights
+
+        state_dict = {
+            "proj.weight": torch.tensor([1.0], dtype=torch.float32).to(torch.float8_e4m3fn),
+            "proj.weight_scale_inv": torch.tensor(1.0, dtype=torch.bfloat16),
+            "proj.activation_scale": torch.tensor(1.0, dtype=torch.bfloat16),
+        }
+        result = _dequantize_fp8_weights(state_dict)
+        assert "proj.activation_scale" not in result
+
+    def test_suffix_replace_not_greedy(self):
+        """The scale key derivation uses suffix replacement, not global replace.
+
+        For a key like 'model.weight_proj.weight', the scale key should be
+        'model.weight_proj.weight_scale_inv' (not 'model.weight_scale_inv_proj.weight_scale_inv').
+        """
+        from mobius._weight_loading import _dequantize_fp8_weights
+
+        fp8_weight = torch.tensor([1.0], dtype=torch.float32).to(torch.float8_e4m3fn)
+        scale = torch.tensor(2.0, dtype=torch.bfloat16)
+        state_dict = {
+            "model.weight_proj.weight": fp8_weight,
+            "model.weight_proj.weight_scale_inv": scale,
+        }
+        result = _dequantize_fp8_weights(state_dict)
+        assert "model.weight_proj.weight" in result
+        assert result["model.weight_proj.weight"].dtype == torch.bfloat16
+
+    def test_missing_scale_casts_without_scaling(self):
+        """FP8 weight without scale_inv is cast to bfloat16 without scaling."""
+        from mobius._weight_loading import _dequantize_fp8_weights
+
+        fp8_weight = torch.tensor([1.0, 2.0], dtype=torch.float32).to(torch.float8_e4m3fn)
+        state_dict = {"orphan.weight": fp8_weight}
+        result = _dequantize_fp8_weights(state_dict)
+        assert result["orphan.weight"].dtype == torch.bfloat16
+
+    def test_fp32_scale_produces_bf16_output(self):
+        """FP32 weight_scale_inv should still produce bfloat16 output."""
+        from mobius._weight_loading import _dequantize_fp8_weights
+
+        fp8_weight = torch.tensor([1.0, 2.0], dtype=torch.float32).to(torch.float8_e4m3fn)
+        # Scale stored as FP32 (common for scalar scales in real checkpoints)
+        scale_inv = torch.tensor(0.5, dtype=torch.float32)
+        state_dict = {
+            "proj.weight": fp8_weight,
+            "proj.weight_scale_inv": scale_inv,
+        }
+        result = _dequantize_fp8_weights(state_dict)
+        assert result["proj.weight"].dtype == torch.bfloat16, (
+            f"Expected bfloat16, got {result['proj.weight'].dtype}"
+        )
+
+    def test_does_not_mutate_input(self):
+        """_dequantize_fp8_weights should not mutate the input dict."""
+        from mobius._weight_loading import _dequantize_fp8_weights
+
+        fp8_weight = torch.tensor([1.0], dtype=torch.float32).to(torch.float8_e4m3fn)
+        scale = torch.tensor(1.0, dtype=torch.bfloat16)
+        original = {
+            "proj.weight": fp8_weight,
+            "proj.weight_scale_inv": scale,
+        }
+        original_keys = set(original.keys())
+        _dequantize_fp8_weights(original)
+        assert set(original.keys()) == original_keys, "Input dict was mutated"

src/mobius/components/_attention.py

@@ -304,6 +304,17 @@ def forward(
 
         # Apply rotary position embeddings (skip when not provided)
         if position_embeddings is not None:
+            # Apply llama_4_attn_scale if present (Ministral3/Mistral4).
+            # The scale is computed from position_ids by the RoPE module
+            # and passed as the 3rd element of position_embeddings.
+            # Applied BEFORE RoPE so the graph keeps the
+            # RotaryEmbedding → Attention pattern that the
+            # RotaryAttentionToGQA rewrite rule matches. Scaling
+            # commutes with rotation: scale(RoPE(q)) == RoPE(scale(q)).
+            if len(position_embeddings) > 2:
+                attn_scale = position_embeddings[2]
+                query_states = op.Mul(query_states, attn_scale)
+
             query_states = apply_rotary_pos_emb(
                 op,
                 x=query_states,

src/mobius/components/_pixtral_vision.py

@@ -24,6 +24,7 @@
 from onnxscript import nn
 from onnxscript._internal import builder
 
+from mobius._build_context import ep_capabilities
 from mobius._configs import ArchitectureConfig
 from mobius.components._common import Linear
 from mobius.components._conv import Conv2dNoBias
@@ -174,17 +175,31 @@ def forward(
             num_heads=self._num_heads,
         )
 
-        # Bidirectional attention (is_causal=0, no KV cache)
+        # Bidirectional attention (no causal mask, no KV cache).
+        # Use com.microsoft.MultiHeadAttention for all EPs that support
+        # custom-domain ops (it has fused kernels on CUDA/DML and runs
+        # correctly on CPU). Fall back to standard opset-23 Attention
+        # for onnx-standard EP which prohibits custom-domain ops.
         scale = float(1.0 / (self._head_dim**0.5))
-        attn_output = op.Attention(
-            q,
-            k,
-            v,
-            q_num_heads=self._num_heads,
-            kv_num_heads=self._num_heads,
-            scale=scale,
-            is_causal=0,
-        )
+        if ep_capabilities().name == "onnx-standard":
+            attn_output = op.Attention(
+                q,
+                k,
+                v,
+                q_num_heads=self._num_heads,
+                kv_num_heads=self._num_heads,
+                scale=scale,
+                is_causal=0,
+            )
+        else:
+            attn_output = op.MultiHeadAttention(
+                q,
+                k,
+                v,
+                num_heads=self._num_heads,
+                scale=scale,
+                _domain="com.microsoft",
+            )
         return self.o_proj(op, attn_output)
 
 
@@ -349,10 +364,15 @@ def forward(
         )
         x = op.Reshape(hidden_states, shape_6d)
 
-        # Transpose: [batch, H/ms, W/ms, ms, ms, D]
-        x = op.Transpose(x, perm=[0, 1, 3, 2, 4, 5])
+        # Transpose to match HuggingFace F.unfold ordering (dim-major).
+        # F.unfold groups elements as [D, ms_h, ms_w] per spatial position,
+        # meaning the hidden dim D is the outermost loop. To reproduce
+        # this with reshape + transpose:
+        #   from: [batch, H/ms, ms, W/ms, ms, D]
+        #     to: [batch, H/ms, W/ms, D, ms, ms]
+        x = op.Transpose(x, perm=[0, 1, 3, 5, 2, 4])
 
-        # Flatten: [batch, (H/ms)*(W/ms), ms*ms*D]
+        # Flatten: [batch, (H/ms)*(W/ms), D*ms*ms]
         merged_count = op.Mul(h_m_1d, w_m_1d)
         shape_3d = op.Concat(
             batch,

src/mobius/components/_pixtral_vision_test.py

@@ -6,6 +6,8 @@
 from __future__ import annotations
 
 import numpy as np
+import onnx_ir as ir
+import onnxruntime as ort
 
 from mobius._configs import ArchitectureConfig, VisionConfig
 from mobius.components._pixtral_vision import (
@@ -103,3 +105,98 @@ def test_patch_merger_builds():
     merger = Mistral3PatchMerger(hidden_size=32, spatial_merge_size=2)
     # input_dim = 32 * 2 * 2 = 128, output_dim = 32
     assert list(merger.merging_layer.weight.shape) == [32, 128]
+
+
+def test_patch_merger_matches_hf_unfold_ordering():
+    """PatchMerger element ordering matches HuggingFace F.unfold (dim-major).
+
+    HF uses ``F.unfold(image_grid, kernel_size=ms, stride=ms)`` which
+    groups elements as ``[D, ms_h, ms_w]`` per spatial position (dim is
+    the outermost loop). The ONNX implementation must reproduce this
+    ordering so the learned ``merging_layer`` projection is correct.
+    """
+    import tempfile
+
+    import torch
+    from onnxscript._internal.builder import GraphBuilder
+
+    hidden_size = 8
+    ms = 2
+    grid_h, grid_w = 4, 4
+    seq_len = grid_h * grid_w
+
+    rng = np.random.default_rng(42)
+    x = rng.standard_normal((1, seq_len, hidden_size)).astype(np.float32)
+
+    # HF reference: F.unfold ordering
+    x_torch = torch.from_numpy(x.squeeze(0))  # (seq_len, D)
+    image_grid = x_torch.view(grid_h, grid_w, hidden_size).permute(2, 0, 1).unsqueeze(0)
+    grid = torch.nn.functional.unfold(image_grid, kernel_size=ms, stride=ms)
+    hf_merged = grid.view(hidden_size * ms * ms, -1).t().numpy()  # (num_merged, D*ms*ms)
+
+    # Build ONNX model that performs only the reshape+transpose+flatten
+    # (no linear projection) so we can compare the raw merge ordering.
+    x_input = ir.Value(
+        name="x",
+        shape=ir.Shape([1, seq_len, hidden_size]),
+        type=ir.TensorType(ir.DataType.FLOAT),
+    )
+    gh_input = ir.Value(
+        name="grid_h",
+        shape=ir.Shape([]),
+        type=ir.TensorType(ir.DataType.INT64),
+    )
+    gw_input = ir.Value(
+        name="grid_w",
+        shape=ir.Shape([]),
+        type=ir.TensorType(ir.DataType.INT64),
+    )
+    graph = ir.Graph(
+        inputs=[x_input, gh_input, gw_input],
+        outputs=[],
+        nodes=[],
+        name="test_merge_ordering",
+        opset_imports={"": 23},
+    )
+    gb = GraphBuilder(graph)
+    op = gb.op
+
+    # Reproduce the PatchMerger reshape+transpose+flatten logic
+    batch = op.Shape(x_input, start=0, end=1)
+    d = op.Shape(x_input, start=2, end=3)
+    ms_scalar = op.Constant(value_int=ms)
+    h_m = op.Div(gh_input, ms_scalar)
+    w_m = op.Div(gw_input, ms_scalar)
+    ms_1d = op.Constant(value_ints=[ms])
+    h_m_1d = op.Reshape(h_m, op.Constant(value_ints=[1]))
+    w_m_1d = op.Reshape(w_m, op.Constant(value_ints=[1]))
+    shape_6d = op.Concat(batch, h_m_1d, ms_1d, w_m_1d, ms_1d, d, axis=0)
+    merged = op.Reshape(x_input, shape_6d)
+    merged = op.Transpose(merged, perm=[0, 1, 3, 5, 2, 4])
+    merged_count = op.Mul(h_m_1d, w_m_1d)
+    shape_3d = op.Concat(batch, merged_count, op.Constant(value_ints=[-1]), axis=0)
+    result = op.Reshape(merged, shape_3d)
+    result.name = "output"
+    graph.outputs.append(result)
+
+    model = ir.Model(graph, ir_version=11)
+
+    with tempfile.NamedTemporaryFile(suffix=".onnx", delete=True) as f:
+        ir.save(model, f.name)
+        sess = ort.InferenceSession(f.name, providers=["CPUExecutionProvider"])
+        onnx_out = sess.run(
+            None,
+            {
+                "x": x,
+                "grid_h": np.array(grid_h, dtype=np.int64),
+                "grid_w": np.array(grid_w, dtype=np.int64),
+            },
+        )[0]
+
+    np.testing.assert_allclose(
+        onnx_out.squeeze(0),
+        hf_merged,
+        atol=1e-5,
+        rtol=1e-5,
+        err_msg="PatchMerger ordering does not match HF F.unfold",
+    )