Merge branch 'main' into dependabot/pip/requirements/lintrunner/ruff-0.15.9

Author: titaiwangms

.github/skills/multimodal-models/SKILL.md

@@ -29,9 +29,11 @@ input_ids ──────► Embedding ────────────
 | Component | File | Purpose |
 |-----------|------|---------|
 | `VisionModel` | `components/_vision.py` | SigLIP-style patch embedding + transformer encoder |
+| `PixtralVisionTower` | `components/_pixtral_vision.py` | Pixtral 2D RoPE vision encoder (bidirectional attention) |
 | `PatchEmbedding` | `components/_vision.py` | Conv2d → positional embedding |
 | `Gemma3MultiModalProjector` | `components/_multimodal.py` | AvgPool2d → RMSNorm → MatMul |
 | `MLPMultiModalProjector` | `components/_multimodal.py` | Linear → GELU → Linear |
+| `Mistral3MultiModalProjector` | `components/_pixtral_vision.py` | RMSNorm → spatial merge → Linear → GELU → Linear |
 | `LinearMultiModalProjector` | `components/_multimodal.py` | Single Linear |
 | `InputMixer` | `components/_multimodal.py` | Scatter vision embeddings at image-token positions |
 | `VisionLanguageTask` | `tasks/__init__.py` | ONNX I/O contract with `pixel_values` input |
@@ -44,7 +46,8 @@ embedding spaces.  Choose the one that matches the HuggingFace implementation:
 | Projector | Architecture | Models |
 |-----------|-------------|--------|
 | `Gemma3MultiModalProjector` | AvgPool2d → RMSNorm → MatMul | Gemma3 |
-| `MLPMultiModalProjector` | Linear → GELU → Linear | LLaVA, LLaVA-NeXT, VipLLaVA, Phi-4-MM, InternVL2, Pixtral, Molmo |
+| `MLPMultiModalProjector` | Linear → GELU → Linear | LLaVA, LLaVA-NeXT, VipLLaVA, Phi-4-MM, InternVL2, Molmo |
+| `Mistral3MultiModalProjector` | RMSNorm → spatial merge → Linear → GELU → Linear | Mistral-3, Pixtral |
 | `LinearMultiModalProjector` | Single Linear | PaliGemma, Qwen2-Audio, Idefics2, Florence2 |
 
 ### Gemma3MultiModalProjector

.github/workflows/benchmark.yml

@@ -133,12 +133,12 @@ jobs:
           python-version: "3.12"
 
       - name: Download baseline
-        uses: actions/download-artifact@v7
+        uses: actions/download-artifact@v8
         with:
           name: benchmark-baseline
 
       - name: Download current
-        uses: actions/download-artifact@v7
+        uses: actions/download-artifact@v8
         with:
           name: benchmark-current
 

CHANGELOG.md

@@ -7,6 +7,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Mistral-3 / Pixtral VLM Support
+
+#### Added
+
+- Support for Mistral-3 / Pixtral vision-language models (`mistral3` model type)
+  - Pixtral vision encoder with 2D RoPE, bidirectional attention, and spatial patch merging
+  - `Mistral3MultiModalProjector` for vision-to-text projection (RMSNorm → merge → MLP)
+  - `PixtralVisionTower` with precomputed 2D rotary caches
+  - Moved `mistral3` from CausalLM to VLM (LLaVA-style 3-model split: decoder, vision, embedding)
+  - FP8 quantization config handling (skip block quantization for fp8)
+  - Integration tests for `ministral3` (text-only) and `mistral3` (VLM)
+  - Config extraction for `PixtralVisionConfig.norm_eps` and `rope_parameters` fallback
+
 ### Static Cache Support
 
 #### Added

docs/design.md

@@ -81,7 +81,8 @@ Key components:
 | `DecoderLayer` | Pre-norm residual: LayerNorm → Attention → LayerNorm → MLP |
 | `MoELayer` | Mixture-of-Experts with pluggable gate |
 | `VisionModel` | SigLIP-style patch embedding + transformer encoder |
-| Projectors | `Gemma3MultiModalProjector`, `MLPMultiModalProjector`, `LinearMultiModalProjector` |
+| `PixtralVisionTower` | Pixtral 2D RoPE vision encoder with bidirectional attention |
+| Projectors | `Gemma3MultiModalProjector`, `MLPMultiModalProjector`, `Mistral3MultiModalProjector`, `LinearMultiModalProjector` |
 | `InputMixer` | Scatter vision embeddings into text positions |
 | RoPE variants | `DefaultRope`, `LinearRope`, `DynamicNTKRope`, `Llama3Rope`, `InterleavedMRope` |
 | `ALiBiAttention` | Attention with linear biases (Falcon, Bloom) |
@@ -229,8 +230,10 @@ the correct gate variant per layer.
 3. `InputMixer` — scatter vision embeddings at image-token positions
 4. `Gemma3CausalLMModel` — standard text decoder
 
-`LLaVAModel` follows a similar pattern with MLPMultiModalProjector.
-Many VL models (InternVL2, Pixtral, Idefics, Molmo, etc.) reuse the LLaVA
+`LLaVAModel` follows a similar pattern with `MLPMultiModalProjector`.
+Mistral-3 / Pixtral uses `Mistral3MultiModalProjector` with a dedicated
+`PixtralVisionTower` (2D RoPE, bidirectional attention) instead of CLIP/SigLIP.
+Many other VL models (InternVL2, Idefics, Molmo, etc.) reuse the LLaVA
 pattern with CLIP/SigLIP vision encoder + projector + LLM.
 
 Uses `VisionLanguageTask` which adds `pixel_values` to the ONNX graph inputs.
@@ -241,6 +244,7 @@ Three projector variants are available for different model families:
 |-----------|-------------|---------|
 | `Gemma3MultiModalProjector` | AvgPool2d → RMSNorm → MatMul | Gemma3 |
 | `MLPMultiModalProjector` | Linear → GELU → Linear | LLaVA, Phi4MM |
+| `Mistral3MultiModalProjector` | RMSNorm → spatial merge → Linear → GELU → Linear | Mistral-3, Pixtral |
 | `LinearMultiModalProjector` | Single Linear | PaliGemma |
 
 ### Encoder-only models (BERT, RoBERTa)

src/mobius/_configs.py

@@ -95,6 +95,10 @@ class VisionConfig:
     norm_eps: float = 1e-6
     mm_tokens_per_image: int | None = None
     image_token_id: int | None = None
+    # Pixtral / Mistral-3 vision fields
+    model_type: str | None = None
+    head_dim: int | None = None
+    rope_theta: float | None = None
     # Qwen VL-specific
     out_hidden_size: int | None = None
     in_channels: int = 3
@@ -273,7 +277,13 @@ def _extract_rope_config(config) -> RoPEConfig:
             _nested_rope_theta(rope_scaling, "full_attention"),
             default=10_000.0,
         ),
-        rope_scaling=(_normalize_rope_scaling(rope_scaling) or None),
+        # Some models (e.g. Ministral-3) store YaRN config under
+        # rope_parameters instead of rope_scaling; fall back accordingly.
+        rope_scaling=(
+            _normalize_rope_scaling(rope_scaling)
+            or _normalize_rope_scaling(rope_parameters)
+            or None
+        ),
         partial_rotary_factor=_first_not_none(
             getattr(config, "partial_rotary_factor", None),
             rope_scaling.get("partial_rotary_factor", None),
@@ -284,12 +294,11 @@ def _extract_rope_config(config) -> RoPEConfig:
             getattr(config, "rope_local_base_freq", None),
             _nested_rope_theta(rope_scaling, "sliding_attention"),
         ),
-        original_max_position_embeddings=(
-            getattr(
-                config,
-                "original_max_position_embeddings",
-                rope_scaling.get("original_max_position_embeddings", None),
-            )
+        original_max_position_embeddings=_first_not_none(
+            getattr(config, "original_max_position_embeddings", None),
+            rope_scaling.get("original_max_position_embeddings", None),
+            # Also check rope_parameters (see rope_scaling comment above).
+            rope_parameters.get("original_max_position_embeddings", None),
         ),
     )
 
@@ -351,10 +360,22 @@ def _extract_vision_config(config, parent_config, model_type: str) -> dict:
             ),
             image_size=getattr(vc, "image_size", None),
             patch_size=getattr(vc, "patch_size", None),
-            norm_eps=getattr(vc, "layer_norm_eps", 1e-6),
+            norm_eps=_first_not_none(
+                getattr(vc, "layer_norm_eps", None),
+                getattr(vc, "norm_eps", None),
+                default=1e-6,
+            ),
+            # Pixtral / Mistral-3 vision fields
+            model_type=getattr(vc, "model_type", None),
+            head_dim=getattr(vc, "head_dim", None),
+            rope_theta=getattr(vc, "rope_theta", None),
             # Qwen VL-specific vision fields
             out_hidden_size=getattr(vc, "out_hidden_size", None),
-            in_channels=getattr(vc, "in_channels", 3),
+            in_channels=_first_not_none(
+                getattr(vc, "in_channels", None),
+                getattr(vc, "num_channels", None),
+                default=3,
+            ),
             spatial_merge_size=getattr(vc, "spatial_merge_size", 2),
             temporal_patch_size=getattr(vc, "temporal_patch_size", 2),
             num_position_embeddings=getattr(vc, "num_position_embeddings", None),
@@ -573,6 +594,11 @@ def from_transformers(cls, hf_config) -> QuantizationConfig | None:
         method = qc.get("quant_method", "none")
         if method == "none":
             return None
+        # FP8 per-tensor quantization (float8_e4m3fn + scalar scale)
+        # is handled by dtype casting in _assign_weight(), not by
+        # QuantizedLinear block quantization.
+        if method == "fp8":
+            return None
         return cls(
             bits=qc.get("bits", 4),
             group_size=qc.get("group_size", 128),

src/mobius/_configs_test.py

@@ -628,6 +628,15 @@ def test_from_transformers_method_none_returns_none(self):
         hf = type("HFConfig", (), {"quantization_config": {"quant_method": "none"}})()
         assert QuantizationConfig.from_transformers(hf) is None
 
+    def test_from_transformers_fp8_returns_none(self):
+        """FP8 per-tensor quantization is not block quantization; returns None."""
+        hf = type(
+            "HFConfig",
+            (),
+            {"quantization_config": {"quant_method": "fp8", "bits": 8}},
+        )()
+        assert QuantizationConfig.from_transformers(hf) is None
+
     def test_from_transformers_to_dict_object(self):
         """HF QuantizationConfig objects have a to_dict() method."""
         inner = type(

src/mobius/_registry.py

@@ -350,7 +350,7 @@ def _create_default_registry() -> ModelRegistry:
         "ministral",
         "ministral3",
         "mistral",
-        "mistral3",
+        "nanochat",
         "open-llama",
         "openelm",
         "qwen2",
@@ -485,6 +485,7 @@ def _create_default_registry() -> ModelRegistry:
         "llava_next",
         "llava_next_video",
         "llava_onevision",
+        "mistral3",
         "molmo",
         "paligemma",
         "pixtral",
@@ -734,7 +735,6 @@ def _create_default_registry() -> ModelRegistry:
     "helium": "kyutai/helium-1-preview-2b",
     "minicpm": "optimum-intel-internal-testing/tiny-random-minicpm",
     "minicpm3": "openbmb/MiniCPM3-4B",
-    "mistral3": "mistralai/Mistral-Small-3.2-24B-Instruct-2506",
     "ministral3": "Aratako/Ministral-3-3B-Instruct-2512-BF16-TextOnly",
     "nanochat": "nanochat-students/nanochat-d20",
     "olmo3": "allenai/Olmo-3-7B-Instruct",
@@ -843,6 +843,7 @@ def _create_default_registry() -> ModelRegistry:
     "instructblip": "Salesforce/instructblip-flan-t5-xl",
     "llava_onevision": "llava-hf/llava-onevision-qwen2-0.5b-ov-hf",
     "molmo": "allenai/MolmoE-1B-0924",
+    "mistral3": "mistralai/Ministral-3-3B-Instruct-2512",
 
     # --- Speech ---
     "whisper": "openai/whisper-tiny",

src/mobius/components/__init__.py

@@ -46,6 +46,7 @@
     "LinearMultiModalProjector",
     "LoRALinear",
     "MLP",
+    "Mistral3MultiModalProjector",
     "MLPMultiModalProjector",
     "Mamba2Block",
     "Mamba2Scan",
@@ -54,6 +55,7 @@
     "OffsetRMSNorm",
     "PatchEmbed",
     "PatchEmbedding",
+    "PixtralVisionTower",
     "PostGatedRMSNorm",
     "PostNormDecoderLayer",
     "QFormer",
@@ -186,6 +188,10 @@
     LinearMultiModalProjector,
     MLPMultiModalProjector,
 )
+from mobius.components._pixtral_vision import (
+    Mistral3MultiModalProjector,
+    PixtralVisionTower,
+)
 from mobius.components._qformer import (
     QFormer,
     QFormerAttention,