Add Pixtral/Ministral3 VLM support: dynamic vision, FP8 dequant, PatchMerger fix, MHA

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,56 @@ 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`` and ``activation_scale`` tensors.
+
+    Returns:
+        A new dict with FP8 weights dequantized to bfloat16 and auxiliary
+        scale tensors removed.
+    """
+    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 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 +221,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

@@ -320,6 +320,12 @@ def forward(
                 rotary_embedding_dim=self.rotary_embedding_dim,
                 interleaved=self._rope_interleave,
             )
+            # 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.
+            if len(position_embeddings) > 2:
+                attn_scale = position_embeddings[2]
+                query_states = op.Mul(query_states, attn_scale)
 
         attn_output, present_key, present_value = _apply_attention(
             op,

src/mobius/components/_rotary_embedding.py

@@ -67,8 +67,11 @@ def apply_rotary_pos_emb(
     Args:
         op: The OpBuilder.
         x: Input tensor of shape ``(batch_size, seq_length, num_heads * head_dim)``.
-        position_embeddings: Tuple of ``(cos, sin)`` embeddings, each
-            ``(batch_size, seq_length, rotary_dim)``.
+        position_embeddings: Tuple of ``(cos, sin)`` or ``(cos, sin, attn_scale)``
+            embeddings. The cos/sin tensors have shape
+            ``(batch_size, seq_length, rotary_dim)``. The optional attn_scale
+            (used by Ministral3/Mistral4) is not consumed here — it is applied
+            separately in the attention module after RoPE.
         num_heads: Number of attention heads.
         rotary_embedding_dim: Dimension for partial RoPE (0 = full embedding).
         interleaved: If True, use interleaved RoPE layout where real/imag
@@ -78,7 +81,7 @@ def apply_rotary_pos_emb(
     Returns:
         Tensor with RoPE applied, same shape as input.
     """
-    cos, sin = position_embeddings
+    cos, sin = position_embeddings[0], position_embeddings[1]
     return op.RotaryEmbedding(
         x,
         cos,
@@ -223,8 +226,13 @@ def forward(self, op: builder.OpBuilder, position_ids: ir.Value):
 class YarnRope(BaseRope):
     """YaRN (Yet another RoPE extensioN) rotary embeddings.
 
-    Used by DeepSeek-V2/V3. Blends interpolated and extrapolated
-    frequencies with a linear ramp, and applies mscale attention factor.
+    Used by DeepSeek-V2/V3 and Ministral3 (Pixtral). Blends interpolated
+    and extrapolated frequencies with a linear ramp, and applies mscale
+    attention factor.
+
+    For Ministral3/Mistral4 models with ``llama_4_scaling_beta`` in
+    rope_scaling, ``forward()`` returns a 3-tuple ``(cos, sin, attn_scale)``
+    where ``attn_scale`` is a position-dependent query scaling factor.
 
     Reference: https://huggingface.co/papers/2309.00071
     """
@@ -291,6 +299,32 @@ def find_correction_dim(num_rotations):
         )
         super().__init__(cos_cache, sin_cache)
 
+        # Store llama_4_scaling_beta for Ministral3 position-dependent query scaling.
+        # When set, forward() returns (cos, sin, attn_scale) instead of (cos, sin).
+        self._llama4_beta = rope_scaling.get("llama_4_scaling_beta")
+        self._llama4_original_max_pos = float(original_max_pos)
+
+    def forward(self, op: builder.OpBuilder, position_ids: ir.Value):
+        cos_sin = get_rotary_pos_emb(op, position_ids, self.cos_cache, self.sin_cache)
+        if self._llama4_beta is None:
+            return cos_sin
+
+        # Compute position-dependent attention scale for Ministral3/Mistral4:
+        #   scale = 1 + beta * log(1 + floor(position_ids / original_max_pos))
+        # For pos < original_max_pos: floor(pos / max) = 0 → scale = 1.0
+        # Computed in FP32 for precision, then cast to match model dtype.
+        pos_float = op.Cast(position_ids, to=ir.DataType.FLOAT)
+        floored = op.Floor(op.Div(pos_float, float(self._llama4_original_max_pos)))
+        log_term = op.Log(op.Add(floored, 1.0))
+        attn_scale = op.Add(op.Mul(log_term, float(self._llama4_beta)), 1.0)
+        # Cast to match model dtype (e.g. FP16) — cos_cache has the target dtype
+        cos_dtype = self.cos_cache.dtype
+        if cos_dtype is not None and cos_dtype != ir.DataType.FLOAT:
+            attn_scale = op.Cast(attn_scale, to=cos_dtype)
+        # Unsqueeze to [batch, seq_len, 1] for broadcasting with 3D query states
+        attn_scale = op.Unsqueeze(attn_scale, [-1])
+        return (cos_sin[0], cos_sin[1], attn_scale)
+
 
 class _MRopeBase(BaseRope):
     """Base class for multi-dimensional RoPE variants.

src/mobius/components/_rotary_embedding_test.py

@@ -254,3 +254,70 @@ def test_get_rotary_pos_emb(self):
         cos, sin = get_rotary_pos_emb(op, pos_ids, cos_cache, sin_cache)
         assert cos is not None
         assert sin is not None
+
+
+class TestYarnRopeAttnScale:
+    """Tests for YarnRope with llama_4_attn_scale (Ministral3)."""
+
+    def test_yarn_without_attn_scale_returns_2tuple(self):
+        """Standard YaRN (no llama_4_scaling_beta) returns (cos, sin)."""
+        from mobius.components._rotary_embedding import YarnRope
+
+        config = make_config(
+            head_dim=128,
+            max_position_embeddings=16384,
+            rope_theta=1000000.0,
+            rope_scaling={
+                "rope_type": "yarn",
+                "factor": 16.0,
+                "beta_fast": 32.0,
+                "beta_slow": 1.0,
+                "mscale": 1.0,
+                "mscale_all_dim": 1.0,
+                "original_max_position_embeddings": 16384,
+            },
+        )
+        rope = YarnRope(config)
+        builder, op, _graph = create_test_builder()
+        pos_ids = create_test_input(builder, "pos_ids", [1, 4])
+        result = rope.forward(op, pos_ids)
+        assert len(result) == 2, "Without llama_4_scaling_beta, should return (cos, sin)"
+
+    def test_yarn_with_attn_scale_returns_3tuple(self):
+        """YaRN with llama_4_scaling_beta returns (cos, sin, attn_scale)."""
+        from mobius.components._rotary_embedding import YarnRope
+
+        config = make_config(
+            head_dim=128,
+            max_position_embeddings=262144,
+            rope_theta=1000000.0,
+            rope_scaling={
+                "rope_type": "yarn",
+                "factor": 16.0,
+                "beta_fast": 32.0,
+                "beta_slow": 1.0,
+                "mscale": 1.0,
+                "mscale_all_dim": 1.0,
+                "original_max_position_embeddings": 16384,
+                "llama_4_scaling_beta": 0.1,
+            },
+        )
+        rope = YarnRope(config)
+        builder, op, _graph = create_test_builder()
+        pos_ids = create_test_input(builder, "pos_ids", [1, 4])
+        result = rope.forward(op, pos_ids)
+        assert len(result) == 3, (
+            "With llama_4_scaling_beta, should return (cos, sin, attn_scale)"
+        )
+
+    def test_apply_rotary_pos_emb_ignores_3rd_element(self):
+        """apply_rotary_pos_emb should work with both 2-tuple and 3-tuple."""
+        builder, op, _graph = create_test_builder()
+        x = create_test_input(builder, "x", [1, 4, 64])
+        cos = create_test_input(builder, "cos", [1, 4, 8])
+        sin = create_test_input(builder, "sin", [1, 4, 8])
+        scale = create_test_input(builder, "scale", [1, 4, 1])
+
+        # 3-tuple should work — apply_rotary_pos_emb only uses [0] and [1]
+        result = apply_rotary_pos_emb(op, x, (cos, sin, scale), num_heads=4)
+        assert result is not None

src/mobius/models/llava.py

@@ -129,7 +129,20 @@ def forward(self, op: builder.OpBuilder, input_ids: ir.Value, image_features: ir
         indices = op.Sub(cumsum, op.Constant(value_int=1))
         indices = op.Clip(indices, op.Constant(value_int=0))
 
-        gathered = op.Gather(image_features, indices, axis=0)
+        # Pad image_features with one zero row so Gather is valid even when
+        # image_features is empty (text-only input: num_image_tokens == 0).
+        # The Where mask ensures the padding row is never used in the output.
+        pad_row = op.Expand(
+            op.CastLike(op.Constant(value_float=0.0), image_features),
+            op.Concat(
+                op.Constant(value_ints=[1]),
+                op.Shape(image_features, start=1, end=2),
+                axis=0,
+            ),
+        )
+        padded_features = op.Concat(image_features, pad_row, axis=0)
+
+        gathered = op.Gather(padded_features, indices, axis=0)
         return op.Where(image_mask_3d, gathered, text_embeds)
 
     def preprocess_weights(

src/mobius/tasks/__init__.py

@@ -35,6 +35,7 @@
     "OPSET_VERSION",
     "ObjectDetectionTask",
     "Phi4MMMultiModalTask",
+    "PixtralVLTask",
     "Qwen3VLVisionLanguageTask",
     "QwenImageVAETask",
     "QwenVLTask",
@@ -79,6 +80,7 @@
 from mobius.tasks._vision_language_3model import (
     HybridQwenVLTask,
     MllamaVisionLanguageTask,
+    PixtralVLTask,
     QwenVLTask,
     VisionLanguageTask,
 )
@@ -102,6 +104,7 @@
     "vae": VAETask,
     "qwen-image-vae": QwenImageVAETask,
     "vision-language": VisionLanguageTask,
+    "pixtral-vl": PixtralVLTask,
     "mllama-vision-language": MllamaVisionLanguageTask,
     "qwen-vl": QwenVLTask,
     "hybrid-qwen-vl": HybridQwenVLTask,