feat(models): add Cosmos3-Edge vision-language model (3-model split)

Extend the cosmos3_edge support from a text-only reasoner to the full
vision-language model (nvidia/Cosmos3-Edge,
Cosmos3EdgeForConditionalGeneration) as a LLaVA-style 3-model
onnxruntime-genai split (decoder + vision_encoder + embedding):

- decoder: squared-ReLU GQA text reasoner taking inputs_embeds, built
  with 3D multimodal RoPE (mrope_section=[24,20,20]).
- vision_encoder: SigLIP vision tower + new Cosmos3EdgeMultiModalProjector
  (pre-shuffle LayerNorm -> 2x2 pixel-shuffle -> linear_fc1 -> GELU ->
  linear_fc2).
- embedding: token embedding + image-feature fusion at image_token_id=19.

preprocess_weights routes the single HF checkpoint to the three
sub-models: model.visual.* / model.projector.* -> vision (SigLIP
mlp.fc1/fc2 -> up_proj/down_proj), embed_tokens -> embedding, the
top-level text tower -> decoder (self_attn.to_{q,k,v,out} ->
{q,k,v,o}_proj), and drops the generator-tower k_norm_und_for_gen
key-norm. Built via a new Cosmos3EdgeVLTask ("cosmos3-edge-vl").

A vision config hook reconstructs image_size from num_patches and pulls
the projector's merger_intermediate_size from projector_config. The
decoder-only text reasoner remains available as cosmos3_edge_text.

L1 graph-build tested only: NVIDIA publishes no modeling code for
cosmos3_edge (not in transformers, no remote-code module), so exact
pixel-shuffle ordering and numerical parity are unverifiable.

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

Author: justinchuby

CHANGELOG.md

@@ -7,22 +7,35 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
-### NVIDIA Cosmos 3 Edge text reasoner (`cosmos3_edge`)
+### NVIDIA Cosmos 3 Edge vision-language model (`cosmos3_edge`)
 
 #### Added
 
-- Support for the **text reasoner backbone** of the `cosmos3_edge`
-  vision-language checkpoint (`nvidia/Cosmos3-Edge`). The language tower is a
-  grouped-query-attention decoder with a **non-gated squared-ReLU FFN**
-  (`hidden_act="relu2"`, `up_proj → relu2 → down_proj`) and 3D multimodal RoPE
-  (`mrope_section=[24, 20, 20]`, equivalent to 1D RoPE for text-only). Mapped
-  onto the standard `CausalLMModel` backbone + `FCMLP`; `preprocess_weights`
-  renames the `self_attn.to_{q,k,v,out}` projections, nests the top-level text
-  tower under `model.`, and drops the vision encoder, projector, and the
-  generator-tower `k_norm_und_for_gen` key-norm. Registered as `cosmos3_edge`
-  / `cosmos3_edge_text`; L1 graph-build tested. The `cosmos3_omni` variants
-  (`Cosmos3-Nano`/`-Super`) are two-tower diffusion world models and remain out
-  of scope for this decoder-only path.
+- Support for the **full `cosmos3_edge` vision-language model**
+  (`nvidia/Cosmos3-Edge`, `Cosmos3EdgeForConditionalGeneration`) as a 3-model
+  onnxruntime-genai split (`decoder` + `vision_encoder` + `embedding`):
+  - **decoder**: grouped-query-attention text reasoner with a **non-gated
+    squared-ReLU FFN** (`hidden_act="relu2"`, `up_proj → relu2 → down_proj`)
+    and 3D multimodal RoPE (`mrope_section=[24, 20, 20]`); takes
+    `inputs_embeds`.
+  - **vision_encoder**: SigLIP vision tower + a new
+    `Cosmos3EdgeMultiModalProjector` (pre-shuffle `LayerNorm` → 2×2
+    pixel-shuffle → `linear_fc1` → GELU → `linear_fc2`).
+  - **embedding**: token embedding + image-feature fusion at
+    `image_token_id=19`.
+  `preprocess_weights` routes the single HF checkpoint to the three
+  sub-models: `model.visual.*` / `model.projector.*` → vision (with SigLIP
+  `mlp.fc1/fc2` → `up_proj/down_proj`), `embed_tokens` → embedding, the
+  top-level text tower (`layers.*` / `norm` / `lm_head`) → decoder (renaming
+  `self_attn.to_{q,k,v,out}` → `{q,k,v,o}_proj`), and drops the
+  generator-tower `k_norm_und_for_gen` key-norm. Built via a new
+  `Cosmos3EdgeVLTask` (`cosmos3-edge-vl`). The decoder-only text reasoner
+  remains available as `cosmos3_edge_text`.
+- **L1 graph-build tested only.** NVIDIA does not publish modeling code for
+  `cosmos3_edge` (not in `transformers`, no remote-code module), so the exact
+  pixel-shuffle ordering and numerical parity are unverifiable; L4/L5 parity
+  is deferred. The `cosmos3_omni` variants (`Cosmos3-Nano`/`-Super`) are
+  two-tower diffusion world models tracked separately.
 
 ### Text-only export for multimodal Gemma 4 (`--text-only`)
 

src/mobius/_configs/_sub_configs.py

@@ -74,6 +74,10 @@ class VisionConfig:
     pooling_kernel_size: int | None = None
     # MLP activation for vision encoder layers (e.g. "gelu_pytorch_tanh" for Gemma4 SigLIP)
     hidden_act: str | None = None
+    # Cosmos3-Edge pixel-shuffle projector intermediate size
+    # (HF: projector_config.merger_intermediate_size). ``None`` means the model
+    # does not use a Cosmos-style merger projector.
+    projector_intermediate_size: int | None = None
     # CLIP-style feature extraction: which ``hidden_states`` index to output
     # (HuggingFace convention, e.g. -2 for Phi-3.5-Vision). ``None`` means use
     # the final hidden state (all layers + post_layernorm).

src/mobius/_configs/per_model/__init__.py

@@ -23,6 +23,7 @@
 # defaults run separately as an explicit first pass, so `ruff` / `isort`
 # may freely re-sort this block.
 from mobius._configs.per_model import (  # noqa: F401
+    _cosmos3_edge_vision,
     _gemma4_audio,
     _gemma4_unified_audio,
     _gemma4_unified_vision,

src/mobius/_configs/per_model/_cosmos3_edge_vision.py

@@ -0,0 +1,62 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Vision extractor hook for NVIDIA Cosmos3-Edge.
+
+Cosmos3-Edge (``cosmos3_edge`` / ``Cosmos3EdgeForConditionalGeneration``) pairs
+a SigLIP-style vision encoder (``cosmos3_edge_vision``) with a pixel-shuffle
+merger projector (``cosmos3_edge_projector``). Two config quirks need bridging
+into :class:`~mobius._configs.VisionConfig`:
+
+- The vision config declares ``num_patches`` (256) instead of ``image_size``.
+  The standard :class:`~mobius.components.PatchEmbedding` derives the patch
+  count from ``image_size // patch_size``, so we reconstruct
+  ``image_size = sqrt(num_patches) * patch_size`` (16 * 16 = 256).
+- The projector's intermediate width lives in a sibling ``projector_config``
+  (``merger_intermediate_size``), not in ``vision_config``.
+"""
+
+from __future__ import annotations
+
+import math
+
+from mobius._configs._extractors import register_vision_hook
+
+
+@register_vision_hook("cosmos3_edge", "cosmos3_edge_vision")
+def _cosmos3_edge_vision(config, parent_config, model_type: str, fields: dict):
+    vision_source = parent_config or config
+    hf_vision = getattr(vision_source, "vision_config", None) or getattr(
+        config, "vision_config", None
+    )
+
+    # Reconstruct image_size from num_patches (256 -> 16x16 grid -> 256 px).
+    if fields.get("image_size") is None and hf_vision is not None:
+        num_patches = getattr(hf_vision, "num_patches", None)
+        patch_size = getattr(hf_vision, "patch_size", None) or fields.get("patch_size")
+        if num_patches is not None and patch_size is not None:
+            grid = round(math.sqrt(num_patches))
+            fields["image_size"] = grid * patch_size
+
+    # Pixel-shuffle projector intermediate size from projector_config.
+    projector_cfg = getattr(vision_source, "projector_config", None) or getattr(
+        config, "projector_config", None
+    )
+    if projector_cfg is not None:
+        get = (
+            projector_cfg.get
+            if isinstance(projector_cfg, dict)
+            else lambda k, d=None: getattr(projector_cfg, k, d)
+        )
+        fields["projector_intermediate_size"] = get("merger_intermediate_size")
+        # out_hidden_size drives the projector output (text hidden size).
+        if fields.get("out_hidden_size") is None:
+            fields["out_hidden_size"] = get("out_hidden_size")
+        merge = get("spatial_merge_size")
+        if merge is not None:
+            fields["spatial_merge_size"] = merge
+
+    # Cosmos3-Edge places image_token_id at the top level (default 19).
+    if fields.get("image_token_id") is None:
+        fields["image_token_id"] = getattr(vision_source, "image_token_id", 19)
+    return None

src/mobius/_registry.py

@@ -36,6 +36,7 @@
     CausalLMModel,
     ChatGLMCausalLMModel,
     Cosmos3EdgeTextModel,
+    Cosmos3EdgeVLModel,
     DeepSeekOCR2CausalLMModel,
     DeepSeekV3CausalLMModel,
     DeepSeekV4CausalLMModel,
@@ -398,7 +399,7 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None:
     "codegen": ModelRegistration(CodeGenCausalLMModel),
     "cohere": ModelRegistration(CohereCausalLMModel),
     "cohere2": ModelRegistration(CohereCausalLMModel),
-    "cosmos3_edge": ModelRegistration(Cosmos3EdgeTextModel),
+    "cosmos3_edge": ModelRegistration(Cosmos3EdgeVLModel),
     "cosmos3_edge_text": ModelRegistration(Cosmos3EdgeTextModel),
     "diffllama": ModelRegistration(DiffLlamaCausalLMModel),
     "doge": ModelRegistration(DogeCausalLMModel),

src/mobius/components/__init__.py

@@ -47,6 +47,7 @@
     "LoRALinear",
     "MLP",
     "MLPMultiModalProjector",
+    "Cosmos3EdgeMultiModalProjector",
     "MoELayer",
     "OffsetRMSNorm",
     "PatchEmbed",
@@ -169,6 +170,9 @@
     SparseMixerGate,
     TopKGate,
 )
+from mobius.components._multimodal import (
+    Cosmos3EdgeMultiModalProjector as Cosmos3EdgeMultiModalProjector,
+)
 from mobius.components._multimodal import (
     Gemma3MultiModalProjector as Gemma3MultiModalProjector,
 )

src/mobius/components/_multimodal.py

@@ -23,6 +23,7 @@
 
 from mobius.components._common import Linear
 from mobius.components._rms_norm import RMSNorm
+from mobius.components._vision import VisionLayerNorm
 
 if TYPE_CHECKING:
     import onnx_ir as ir
@@ -117,6 +118,81 @@ def forward(self, op: OpBuilder, vision_features: ir.Value):
         return hidden
 
 
+class Cosmos3EdgeMultiModalProjector(nn.Module):
+    """Cosmos3-Edge pixel-shuffle merger projector.
+
+    ``LayerNorm → spatial 2x2 pixel-shuffle → Linear(fc1) → GELU → Linear(fc2)``
+
+    The SigLIP vision encoder emits a fixed ``grid x grid`` patch grid
+    (``num_patches`` patches, e.g. 16x16 = 256). ``use_postshuffle_norm=false``
+    means the ``LayerNorm`` is applied on the raw ``vision_hidden_size`` (1152)
+    features **before** the spatial merge. The merge concatenates each
+    ``spatial_merge_size x spatial_merge_size`` block of adjacent patches into a
+    single ``spatial_merge_size**2 * vision_hidden_size`` (4608) vector, which
+    ``linear_fc1`` maps to ``intermediate_size`` (11520) and ``linear_fc2`` maps
+    to ``text_hidden_size`` (2048).
+
+    HF weights (``model.projector.*``):
+    - ``norm.{weight,bias}`` (pre-shuffle LayerNorm)
+    - ``linear_fc1.{weight,bias}``
+    - ``linear_fc2.{weight,bias}``
+    """
+
+    def __init__(
+        self,
+        vision_hidden_size: int,
+        text_hidden_size: int,
+        intermediate_size: int,
+        grid_size: int,
+        spatial_merge_size: int = 2,
+        norm_eps: float = 1e-6,
+    ):
+        super().__init__()
+        self._grid = grid_size
+        self._ms = spatial_merge_size
+        self._vision_hidden = vision_hidden_size
+        merged_dim = vision_hidden_size * spatial_merge_size * spatial_merge_size
+        # Pre-shuffle LayerNorm over the raw vision hidden size.
+        self.norm = VisionLayerNorm(vision_hidden_size, eps=norm_eps)
+        self.linear_fc1 = Linear(merged_dim, intermediate_size, bias=True)
+        self.linear_fc2 = Linear(intermediate_size, text_hidden_size, bias=True)
+
+    def forward(self, op: OpBuilder, vision_features: ir.Value):
+        # vision_features: [batch, grid*grid, vision_hidden]
+        ms = self._ms
+        g = self._grid
+        gm = g // ms
+        d = self._vision_hidden
+
+        # Pre-shuffle LayerNorm (use_postshuffle_norm=false).
+        x = self.norm(op, vision_features)
+
+        batch = op.Shape(vision_features, start=0, end=1)  # dynamic [1]
+
+        # [B, g*g, D] -> [B, g/ms, ms, g/ms, ms, D]
+        shape_6d = op.Concat(
+            batch,
+            op.Constant(value_ints=[gm, ms, gm, ms, d]),
+            axis=0,
+        )
+        x = op.Reshape(x, shape_6d)
+        # Group hidden dim outermost per merged block (HF F.unfold ordering):
+        # [B, g/ms, ms, g/ms, ms, D] -> [B, g/ms, g/ms, D, ms, ms]
+        x = op.Transpose(x, perm=[0, 1, 3, 5, 2, 4])
+        # Flatten to [B, (g/ms)^2, D*ms*ms]
+        shape_3d = op.Concat(
+            batch,
+            op.Constant(value_ints=[gm * gm, d * ms * ms]),
+            axis=0,
+        )
+        x = op.Reshape(x, shape_3d)
+
+        x = self.linear_fc1(op, x)
+        x = op.Gelu(x)
+        x = self.linear_fc2(op, x)
+        return x
+
+
 class LinearMultiModalProjector(nn.Module):
     """Single linear projection (PaliGemma, Qwen2-Audio).
 

src/mobius/models/__init__.py

@@ -24,6 +24,7 @@
     "CohereCausalLMModel",
     "ControlNetModel",
     "Cosmos3EdgeTextModel",
+    "Cosmos3EdgeVLModel",
     "DeepSeekOCR2CausalLMModel",
     "DeepSeekV3CausalLMModel",
     "DeepSeekV4CausalLMModel",
@@ -165,7 +166,7 @@
 from mobius.models.cogvideox import CogVideoXTransformer3DModel
 from mobius.models.cohere import CohereCausalLMModel
 from mobius.models.controlnet import ControlNetModel
-from mobius.models.cosmos import Cosmos3EdgeTextModel
+from mobius.models.cosmos import Cosmos3EdgeTextModel, Cosmos3EdgeVLModel
 from mobius.models.ctrl import CTRLCausalLMModel
 from mobius.models.deepseek import DeepSeekV3CausalLMModel
 from mobius.models.deepseek_ocr2 import DeepSeekOCR2CausalLMModel