feat(minimax_m3): serve image input on MiniMax-M3

docs/models.md

@@ -17,6 +17,7 @@ for them; other checkpoints of the same architectures work too.
 | gpt-oss | [openai/gpt-oss-120b](https://huggingface.co/openai/gpt-oss-120b), [openai/gpt-oss-20b](https://huggingface.co/openai/gpt-oss-20b) |
 | Gemma-4 | [google/gemma-4-26B-A4B-it](https://huggingface.co/google/gemma-4-26B-A4B-it), [nvidia/Gemma-4-26B-A4B-NVFP4](https://huggingface.co/nvidia/Gemma-4-26B-A4B-NVFP4), [google/gemma-4-12B-it](https://huggingface.co/google/gemma-4-12B-it), [nvidia/Gemma-4-31B-IT-NVFP4](https://huggingface.co/nvidia/Gemma-4-31B-IT-NVFP4) .. |
 | MiniMax-M2.5 | [nvidia/MiniMax-M2.5-NVFP4](https://huggingface.co/nvidia/MiniMax-M2.5-NVFP4) |
+| MiniMax-M3 | [nvidia/MiniMax-M3-NVFP4](https://huggingface.co/nvidia/MiniMax-M3-NVFP4) |
 | Muse-Glimmer | [meta-models/Muse-Glimmer-30B](https://huggingface.co/meta-models/Muse-Glimmer-30B), [RedHatAI/Muse-Glimmer-30B-NVFP4](https://huggingface.co/RedHatAI/Muse-Glimmer-30B-NVFP4) |
 
 ### Image input
@@ -31,6 +32,7 @@ These families accept image input by default; pass `--text-model-only` to skip t
 | Gemma-4 12B (`gemma4_unified`: linear patch embedder, resident under either placement) | same budgets, one 48x48 super-patch per soft token | same as the tower releases | same |
 | GLM-5.3-Flash (`glm5_next`: ViT tower, streamed under `--mm-encoder-weights host`) | one token per 28x28 pixels of the resized image, dynamic resolution on a canvas zero-padded to a 28-multiple | token counts, passed through as the processor's `min_image_tokens` / `max_image_tokens`; checkpoint defaults 16 to 8000 tokens | `{"max_image_tokens": 2048}` |
 | Muse-Glimmer-30B (windowed ViT tower, streamed under `--mm-encoder-weights host`) | one token per 28x28 pixels of the resized image, aspect ratio kept under a token cap; checkpoint default 4096 tokens | the maximum is the cap (`max_image_tokens`); the minimum has no effect | `{"max_image_tokens": 1024}` |
+| MiniMax-M3 (`minimax_m3`: CLIP-style ViT tower, streamed under `--mm-encoder-weights host`) | one token per 28x28 pixels of the resized image, dynamic resolution | pixel areas in `size.shortest_edge` / `longest_edge`; checkpoint defaults 4 to 576 tokens | `{"size": {"longest_edge": 1048576}}` |
 
 ## MoE strategies
 

python/freetoken/attention/m3_sparse.py

@@ -100,7 +100,7 @@ def _page_ok(name: str) -> bool:
 
     from freetoken.engine.engine import _resolve_auto_attention_backend
 
-    name = _resolve_auto_attention_backend(frozenset({AttnType.FULL}), False)
+    name = _resolve_auto_attention_backend(frozenset({AttnType.FULL}))
     if not _page_ok(name):
         # trtllm (the sm_100 first pick) pins 16/32/64-token pages; walk the rest
         # of the SAME tree (same arch gates, same requirement probes) with the

python/freetoken/mm/processors/minimax_m3.py

@@ -0,0 +1,92 @@
+"""MiniMax-M3 image processor: patch-grid ViT with a 2x2 merger, a pixel-area budget and 1-D rope."""
+
+from __future__ import annotations
+
+import struct
+from typing import Any
+
+import torch
+
+from freetoken.message import MMItem
+from freetoken.mm import mm_pad_value
+from freetoken.mm.config import MultimodalConfig
+from freetoken.mm.processor import MMProcessor, PromptReplacement, content_hash
+
+# the config names only the image token; the reference processor reads these two from the tokenizer's fixed added tokens
+IMAGE_START_ID = 200029  # ]<]start of image[>[
+IMAGE_END_ID = 200030  # ]<]end of image[>[
+
+
+def _compression(vc: Any, key: str) -> int:
+    """The native vision config carries the merge sizes flat; the checkpoint's own config nests them under img_token_compression_config."""
+    value = getattr(vc, key, None)
+    return int(vc.img_token_compression_config[key] if value is None else value)
+
+
+class MiniMaxM3MMProcessor(MMProcessor):
+    def __init__(self, hf_config: Any, model_path: str, mm: MultimodalConfig) -> None:
+        super().__init__(model_path, mm)
+        vc = hf_config.vision_config
+        # the native config maps image_token_id onto the checkpoint's image_token_index
+        self.image_token_id = getattr(hf_config, "image_token_id", getattr(hf_config, "image_token_index", None))
+        # the template renders one image token per image; the wrappers come from the replacement
+        self.placeholder = [self.image_token_id]
+        self.merge = _compression(vc, "spatial_merge_size")
+        self.pixels_per_token = (vc.patch_size * self.merge) ** 2
+        self.patch_dim = vc.num_channels * _compression(vc, "temporal_patch_size") * vc.patch_size**2
+
+    def get_mm_processor_kwargs(self, mm: MultimodalConfig) -> dict[str, Any]:
+        kwargs: dict[str, Any] = {"return_tensors": "pt"}
+        if mm.image_min_tokens is not None or mm.image_max_tokens is not None:
+            # the budget travels as pixel areas in size.shortest_edge / longest_edge
+            size = dict(self._image_processor().size)
+            if mm.image_min_tokens is not None:
+                size["shortest_edge"] = mm.image_min_tokens * self.pixels_per_token
+            if mm.image_max_tokens is not None:
+                size["longest_edge"] = mm.image_max_tokens * self.pixels_per_token
+            kwargs["size"] = size
+        return {**kwargs, **mm.processor_kwargs}
+
+    def process(self, images: list[Any]) -> list[MMItem]:
+        processor = self._image_processor()
+        kwargs = self.get_mm_processor_kwargs(self.mm)
+        items: list[MMItem] = []
+        for pil in images:
+            out = processor(images=pil, **kwargs)
+            pixel_values = out["pixel_values"].to(torch.bfloat16)
+            grid = out["image_grid_thw"][0].tolist()
+            t, h, w = int(grid[0]), int(grid[1]), int(grid[2])
+            if t != 1:
+                raise ValueError("video/temporal input is not supported")
+            # hash the bf16 wire buffer: half the bytes, and the exact tensor the cache and radix key agree on
+            item_hash = content_hash(pixel_values, struct.pack("<3i", t, h, w))
+            items.append(
+                MMItem(
+                    modality="image",
+                    hash=item_hash,
+                    pad_value=mm_pad_value(item_hash),
+                    offsets=[],
+                    feature=pixel_values,
+                    model_specific_data={"grid_thw": [t, h, w]},
+                )
+            )
+        return items
+
+    def prompt_replacement(self, item: MMItem) -> PromptReplacement:
+        t, h, w = item.grid_thw
+        full = [IMAGE_START_ID] + [self.image_token_id] * ((t * h * w) // (self.merge * self.merge)) + [IMAGE_END_ID]
+        return PromptReplacement.select_token_id(full, self.image_token_id)
+
+    def dummy_items(self, dtype: torch.dtype, device: torch.device) -> list[MMItem]:
+        merge = self.merge
+        return [MMItem(
+            modality="image",
+            hash=0,
+            pad_value=0,
+            offsets=[[0, 1]],
+            feature=torch.zeros(merge * merge, self.patch_dim, dtype=dtype, device=device),
+            model_specific_data={"grid_thw": [1, merge, merge]},
+        )]
+
+
+__all__ = ["IMAGE_END_ID", "IMAGE_START_ID", "MiniMaxM3MMProcessor"]

python/freetoken/models/minimax_m3/__init__.py

@@ -1,5 +1,6 @@
-from .config import parse_config
-from .model import MiniMaxM3ForCausalLM
+from .config import VisionConfig, parse_config, parse_vision_config
+from .model import MiniMaxM3ForCausalLM, MiniMaxM3ForConditionalGeneration
+from .vision import MiniMaxM3VisionModel
 from .weight import (
     nvfp4_expert_spec,
     iter_weights,
@@ -8,6 +9,10 @@
 __all__ = [
     "nvfp4_expert_spec",
     "MiniMaxM3ForCausalLM",
+    "MiniMaxM3ForConditionalGeneration",
+    "MiniMaxM3VisionModel",
+    "VisionConfig",
     "parse_config",
+    "parse_vision_config",
     "iter_weights",
 ]

python/freetoken/models/minimax_m3/config.py

@@ -1,10 +1,9 @@
 """Engine-facing config for MiniMax-M3 (``minimax_m3``).
 
 The checkpoint is a multimodal wrapper (``model_type=minimax_m3_vl``): the text tower
-lives in ``text_config`` and the weights carry a ``language_model.`` prefix. FreeToken
-serves the text tower; the ViT vision stack (``vision_tower.`` / projector) is skipped
-at load like the other VL checkpoints' (``VISION_KEY_PREFIXES``) -- multimodal input
-is future work, so ``ModelConfig.vision_config`` stays None here.
+lives in ``text_config`` and the weights carry a ``language_model.`` prefix. Its vision
+section becomes ``ModelConfig.vision_config`` for the tower in vision.py; an engine that
+builds no vision encoder hands parse_config a config without that section.
 
 Attention is GQA with block-sparse selection on the trailing layers: parse_config
 declares ONE full-attention group over all layers carrying ``mla=False`` plus the
@@ -27,6 +26,7 @@
 from __future__ import annotations
 
 import os
+from dataclasses import dataclass
 from typing import Any
 
 from freetoken.models.config import (
@@ -41,6 +41,54 @@ def _text_config(hf_config: Any) -> Any:
     return getattr(hf_config, "text_config", None) or hf_config
 
 
+@dataclass(frozen=True)
+class VisionConfig:
+    hidden_size: int
+    num_layers: int
+    num_heads: int
+    intermediate_size: int
+    num_channels: int
+    patch_size: int
+    temporal_patch_size: int
+    spatial_merge_size: int
+    layer_norm_eps: float
+    rope_theta: float
+    projector_hidden_size: int
+    text_hidden_size: int
+
+
+def _compression(vc: Any, key: str) -> int:
+    """The native vision config carries the merge sizes flat; the checkpoint's own config nests them under img_token_compression_config."""
+    value = getattr(vc, key, None)
+    return int(vc.img_token_compression_config[key] if value is None else value)
+
+
+def parse_vision_config(hf_config: Any) -> VisionConfig | None:
+    """None when the config carries no vision section, which is how a text-only engine asks for no tower."""
+    vc = getattr(hf_config, "vision_config", None)
+    if vc is None:
+        return None
+    if vc.hidden_act != "gelu":
+        raise NotImplementedError(f"minimax_m3 vision tower activation {vc.hidden_act!r}; only gelu is implemented")
+    # the native config moves rope_theta into rope_parameters; the checkpoint's own config keeps it flat
+    rope_params = getattr(vc, "rope_parameters", None) or {}
+    text_hidden_size = _text_config(hf_config).hidden_size
+    return VisionConfig(
+        hidden_size=vc.hidden_size,
+        num_layers=vc.num_hidden_layers,
+        num_heads=vc.num_attention_heads,
+        intermediate_size=vc.intermediate_size,
+        num_channels=vc.num_channels,
+        patch_size=vc.patch_size,
+        temporal_patch_size=_compression(vc, "temporal_patch_size"),
+        spatial_merge_size=_compression(vc, "spatial_merge_size"),
+        layer_norm_eps=vc.layer_norm_eps,
+        rope_theta=float(rope_params["rope_theta"] if "rope_theta" in rope_params else vc.rope_theta),
+        projector_hidden_size=getattr(hf_config, "projector_hidden_size", None) or text_hidden_size,
+        text_hidden_size=text_hidden_size,
+    )
+
+
 # House rule: an env switch that changes WHAT IS SERVED must leave a server-log
 # trace. parse_config runs a few times per process (engine + weight loader), so
 # each resolved-mode line is logged once per distinct value.
@@ -179,7 +227,10 @@ def parse_config(hf_config: Any) -> ModelConfig:
         dense_quant="mxfp8" if dense_mxfp8 else "none",
         lm_head_quant="none",
         m3_args=args,
+        vision_config=parse_vision_config(hf_config),
+        # the native config maps image_token_id onto the checkpoint's image_token_index
+        image_token_id=getattr(hf_config, "image_token_id", getattr(hf_config, "image_token_index", None)),
     )
 
 
-__all__ = ["parse_config"]
+__all__ = ["VisionConfig", "parse_config", "parse_vision_config"]

python/freetoken/models/minimax_m3/model.py

@@ -11,14 +11,16 @@
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from freetoken.models.blocks import BaseLLMModel
+from freetoken.models.blocks import BaseLLMModel, embed_input_ids
 from freetoken.utils import nvtx_annotate
 
 from .attention import MiniMaxM3Attention
 from .mlp import MiniMaxM3MLP
 from .moe import MiniMaxM3SparseMoeBlock
+from .vision import MiniMaxM3VisionModel
 
 if TYPE_CHECKING:
+    from freetoken.message import MMItem
     from freetoken.models.config import ModelConfig
 
 
@@ -79,7 +81,7 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"):
         self.norm = GemmaPlusOneRMSNormFused(size=config.hidden_size, eps=config.rms_norm_eps)
 
     def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
-        x = self.embed_tokens.forward(input_ids)
+        x = embed_input_ids(self.embed_tokens, input_ids, get_global_ctx().batch)
         residual: torch.Tensor | None = None
         for layer in self.layers.op_list:
             x, residual = layer.forward(x, residual)
@@ -105,4 +107,17 @@ def forward(self) -> torch.Tensor:
         return self.lm_head.forward(output)
 
 
-__all__ = ["MiniMaxM3ForCausalLM"]
+class MiniMaxM3ForConditionalGeneration(MiniMaxM3ForCausalLM):
+    def __init__(self, config: ModelConfig):
+        super().__init__(config)
+        if config.is_multimodal:
+            self.vision_tower = MiniMaxM3VisionModel(config.vision_config)
+
+    def place_encoder_weights(self, mode: str) -> None:
+        self.vision_tower.place_weights(mode)
+
+    def encode(self, item: MMItem) -> torch.Tensor:
+        return self.vision_tower.forward(item.feature, item.grid_thw)
+
+
+__all__ = ["MiniMaxM3ForCausalLM", "MiniMaxM3ForConditionalGeneration"]