Cast vision/audio inputs f32 to model dtype for GenAI compatibility (#265)

## Summary

Keep vision/audio encoder graph inputs as float32 (matching GenAI
image/audio processor output) and add a Cast(f32→model_dtype) at the
graph entry. Encoder weights use the requested f16/bf16 dtype for memory
efficiency.

## Problem

GenAI multimodal processor outputs float32 pixel values and audio
features. When building with `--dtype f16` or `--dtype bf16`, the
vision/audio encoder graph inputs were also cast to f16/bf16, creating a
type mismatch: GenAI sends f32, encoder expects f16, ORT produces
all-zero output (silent failure).

## Solution: Cast at graph entry

Instead of keeping entire vision/audio encoders at f32 (wasteful), add a
`Cast` op at the start of each encoder graph:

```
pixel_values (f32, from GenAI) → Cast(f32→f16) → vision encoder (f16 weights) → image_features
```

This gives both **GenAI compatibility** (f32 input) and **memory
efficiency** (f16/bf16 weights).

## Changes

- `_gemma4.py`: Vision/audio encoder inputs always `ir.DataType.FLOAT`,
with `op.Cast` to `config.dtype` when dtype ≠ f32
- `_builder.py`: No changes needed — all parameters cast to requested
dtype as before

## Testing

13 gemma4 tests pass. Verified f16 build: vision input=f32, first
node=Cast, weights=f16.

---------

Signed-off-by: Justin Chu <justinchu@microsoft.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: justinchuby

src/mobius/_builder.py

@@ -192,6 +192,9 @@ def forward(self, op, input_ids, attention_mask,
     if hasattr(config, "validate"):
         config.validate()
     dtype = getattr(config, "dtype", ir.DataType.FLOAT)
+    # Cast all parameters to the target dtype. Vision/audio encoder weights
+    # are included — their graph inputs are kept at f32 (matching GenAI's
+    # image processor output) with a Cast at the graph entry.
     _cast_module_dtype(module, dtype)
     resolved_task = get_task(task)
     capabilities = ep_registry.require(execution_provider)

src/mobius/tasks/_fun_asr_speech_language.py

@@ -91,14 +91,19 @@ def _build_audio_encoder(
         input_dim = (config.audio.input_size if config.audio else None) or 560
 
         graph, builder = _make_graph(name="audio_encoder")
+        op = builder.op
 
+        # Audio encoder input is always f32 (matching audio processor output).
+        # Cast at graph entry for f16/bf16 builds.
         input_features = builder.input(
             "input_features",
-            dtype=config.dtype,
+            dtype=ir.DataType.FLOAT,
             shape=[batch, seq_len, input_dim],
         )
+        if config.dtype and config.dtype != ir.DataType.FLOAT:
+            input_features = op.Cast(input_features, to=config.dtype)
 
-        audio_features = audio_encoder(builder.op, input_features)
+        audio_features = audio_encoder(op, input_features)
 
         builder.add_output(audio_features, "audio_features")
         return _make_model(graph)

src/mobius/tasks/_gemma4.py

@@ -298,11 +298,16 @@ def _build_vision(
         graph, builder = _make_graph(name="vision_encoder")
         op = builder.op
 
+        # Vision encoder input is always f32 (matching GenAI's image processor
+        # output). When the model uses f16/bf16, add a Cast at the graph entry
+        # so weights can stay at the requested dtype for memory efficiency.
         pixel_values = builder.input(
             "pixel_values",
-            dtype=config.dtype,
+            dtype=ir.DataType.FLOAT,
             shape=[batch, num_patches, pixel_dim],
         )
+        if config.dtype and config.dtype != ir.DataType.FLOAT:
+            pixel_values = op.Cast(pixel_values, to=config.dtype)
         pixel_position_ids = builder.input(
             "pixel_position_ids",
             dtype=ir.DataType.INT64,
@@ -352,9 +357,11 @@ def _build_audio(
 
         input_features = builder.input(
             "input_features",
-            dtype=config.dtype,
+            dtype=ir.DataType.FLOAT,  # Always f32 (matching audio processor output)
             shape=[batch, time, input_size],
         )
+        if config.dtype and config.dtype != ir.DataType.FLOAT:
+            input_features = op.Cast(input_features, to=config.dtype)
         input_features_mask = builder.input(
             "input_features_mask",
             dtype=ir.DataType.BOOL,

src/mobius/tasks/_phi4mm_multimodal.py

@@ -85,19 +85,24 @@ def _build_vision(
         image_size = (config.vision.image_size if config.vision else None) or 448
 
         graph, builder = _make_graph(name="vision_encoder")
+        op = builder.op
 
+        # Vision encoder input is always f32 (matching GenAI's image processor
+        # output). Cast at graph entry for f16/bf16 builds.
         pixel_values = builder.input(
             "pixel_values",
-            dtype=config.dtype,
+            dtype=ir.DataType.FLOAT,
             shape=[batch, 3, image_size, image_size],
         )
+        if config.dtype and config.dtype != ir.DataType.FLOAT:
+            pixel_values = op.Cast(pixel_values, to=config.dtype)
         image_sizes = builder.input(
             "image_sizes",
             dtype=ir.DataType.INT64,
             shape=[num_images, 2],
         )
 
-        image_features = vision(builder.op, pixel_values, image_sizes=image_sizes)
+        image_features = vision(op, pixel_values, image_sizes=image_sizes)
 
         builder.add_output(image_features, "image_features")
         return _make_model(graph)
@@ -118,12 +123,17 @@ def _build_speech(
         input_size = (config.audio.input_size if config.audio else None) or 80
 
         graph, builder = _make_graph(name="audio_encoder")
+        op = builder.op
 
+        # Audio encoder input is always f32 (matching audio processor output).
+        # Cast at graph entry for f16/bf16 builds.
         audio_embeds = builder.input(
             "audio_embeds",
-            dtype=config.dtype,
+            dtype=ir.DataType.FLOAT,
             shape=[batch, audio_seq_len, input_size],
         )
+        if config.dtype and config.dtype != ir.DataType.FLOAT:
+            audio_embeds = op.Cast(audio_embeds, to=config.dtype)
         audio_sizes = builder.input(
             "audio_sizes",
             dtype=ir.DataType.INT64,
@@ -136,7 +146,7 @@ def _build_speech(
         )
 
         speech_out = speech(
-            builder.op,
+            op,
             audio_embeds,
             audio_sizes=audio_sizes,
             audio_projection_mode=audio_projection_mode,

src/mobius/tasks/_speech_language.py

@@ -83,14 +83,19 @@ def _build_audio_encoder(
         n_mels = (config.audio.num_mel_bins if config.audio else None) or 128
 
         graph, builder = _make_graph(name="audio_encoder")
+        op = builder.op
 
+        # Audio encoder input is always f32 (matching audio processor output).
+        # Cast at graph entry for f16/bf16 builds.
         input_features = builder.input(
             "input_features",
-            dtype=config.dtype,
+            dtype=ir.DataType.FLOAT,
             shape=[batch, n_mels, mel_seq],
         )
+        if config.dtype and config.dtype != ir.DataType.FLOAT:
+            input_features = op.Cast(input_features, to=config.dtype)
 
-        audio_features = audio_encoder(builder.op, input_features)
+        audio_features = audio_encoder(op, input_features)
 
         builder.add_output(audio_features, "audio_features")
         return _make_model(graph)

src/mobius/tasks/_vision_language_3model.py

@@ -83,12 +83,17 @@ def _build_vision(
         image_size = (config.vision.image_size if config.vision else None) or 224
 
         graph, builder = _make_graph(name="vision_encoder")
+        op = builder.op
+        # Vision encoder input is always f32 (matching GenAI's image processor
+        # output). Cast at graph entry for f16/bf16 builds.
         pixel_values = builder.input(
             "pixel_values",
-            dtype=config.dtype,
+            dtype=ir.DataType.FLOAT,
             shape=[batch, 3, image_size, image_size],
         )
-        image_features = vision(builder.op, pixel_values=pixel_values)
+        if config.dtype and config.dtype != ir.DataType.FLOAT:
+            pixel_values = op.Cast(pixel_values, to=config.dtype)
+        image_features = vision(op, pixel_values=pixel_values)
 
         builder.add_output(image_features, "image_features")
         return _make_model(graph)
@@ -134,19 +139,24 @@ def _build_vision(
         pixel_dim = in_channels * temporal_patch_size * patch_size * patch_size
 
         graph, builder = _make_graph(name="vision_encoder")
+        op = builder.op
+        # Vision encoder input is always f32 (matching GenAI's image processor
+        # output). Cast at graph entry for f16/bf16 builds.
         pixel_values = builder.input(
             "pixel_values",
-            dtype=config.dtype,
+            dtype=ir.DataType.FLOAT,
             shape=[total_patches, pixel_dim],
         )
+        if config.dtype and config.dtype != ir.DataType.FLOAT:
+            pixel_values = op.Cast(pixel_values, to=config.dtype)
         image_grid_thw = builder.input(
             "image_grid_thw",
             dtype=ir.DataType.INT64,
             shape=[num_images, 3],
         )
 
         image_features = vision(
-            builder.op,
+            op,
             pixel_values=pixel_values,
             image_grid_thw=image_grid_thw,
         )
@@ -209,12 +219,16 @@ def _build_vision(
         width = ir.SymbolicDim("width")
 
         graph, builder = _make_graph(name="vision_encoder")
+        op = builder.op
+        # Vision encoder input is always f32 (matching GenAI's image processor
+        # output). Cast at graph entry for f16/bf16 builds.
         pixel_values = builder.input(
             "pixel_values",
-            dtype=config.dtype,
+            dtype=ir.DataType.FLOAT,
             shape=[batch, 3, height, width],
         )
-        op = builder.op
+        if config.dtype and config.dtype != ir.DataType.FLOAT:
+            pixel_values = op.Cast(pixel_values, to=config.dtype)
 
         image_features = vision(
             op,

tests/build_graph_test.py

@@ -1857,6 +1857,117 @@ def test_dtype_casts_float_initializers(self, dtype_str, expected):
                 f"Initializer '{name}' dtype is {init.dtype}, expected {expected_dtype}"
             )
 
+    @pytest.mark.parametrize(
+        "dtype_str",
+        ["f16", "bf16"],
+    )
+    def test_multimodal_encoder_inputs_are_float32(self, dtype_str):
+        """Vision/audio encoder graph inputs stay f32 with Cast at entry.
+
+        When building multimodal models with f16/bf16, encoder graph inputs
+        (pixel_values, input_features) must remain FLOAT because ORT GenAI's
+        image/audio processors output f32. A Cast node at graph entry converts
+        to the target dtype for the encoder's internal computation.
+        """
+        # Use a VL model with 3-model split (vision_encoder is separate)
+        config = _base_config(
+            vision=VisionConfig(
+                hidden_size=32,
+                intermediate_size=64,
+                num_hidden_layers=1,
+                num_attention_heads=2,
+                image_size=28,
+                patch_size=14,
+                norm_eps=1e-6,
+            ),
+            image_token_id=32000,
+        )
+        config.dtype = DTYPE_MAP[dtype_str]
+        model_cls = registry.get("llava")
+        module = model_cls(config)
+        task = get_task("vision-language")
+        pkg = task.build(module, config)
+
+        # Vision encoder pixel_values input must be FLOAT
+        vision_model = pkg["vision_encoder"]
+        pixel_values_input = vision_model.graph.inputs[0]
+        assert pixel_values_input.name == "pixel_values"
+        assert pixel_values_input.dtype == ir.DataType.FLOAT, (
+            f"Vision encoder input dtype is {pixel_values_input.dtype}, "
+            f"expected FLOAT (Cast should handle conversion to {dtype_str})"
+        )
+
+        # First non-input node should be Cast to target dtype
+        first_node = next(iter(vision_model.graph))
+        assert first_node.op_type == "Cast", (
+            f"Expected Cast as first node, got {first_node.op_type}"
+        )
+
+    @pytest.mark.parametrize(
+        "dtype_str",
+        ["f16", "bf16"],
+    )
+    def test_gemma4_encoder_inputs_are_float32(self, dtype_str):
+        """Gemma4 vision and audio encoder inputs stay f32 in bf16/f16 builds."""
+        from mobius._configs import Gemma4AudioConfig, Gemma4Config
+
+        config = Gemma4Config(
+            num_hidden_layers=2,
+            hidden_size=64,
+            intermediate_size=128,
+            num_attention_heads=4,
+            num_key_value_heads=1,
+            head_dim=16,
+            vocab_size=256,
+            rms_norm_eps=1e-6,
+            hidden_act="silu",
+            attn_qk_norm=True,
+            layer_types=["sliding_attention", "sliding_attention"],
+            sliding_window=8,
+            global_head_dim=16,
+            global_rope_theta=10_000.0,
+            global_partial_rotary_factor=0.25,
+            final_logit_softcapping=0.0,
+            hidden_size_per_layer_input=0,
+            image_token_id=255999,
+            pad_token_id=0,
+            tie_word_embeddings=True,
+            num_kv_shared_layers=1,
+            vision=VisionConfig(
+                hidden_size=32,
+                intermediate_size=64,
+                num_hidden_layers=1,
+                num_attention_heads=2,
+                patch_size=16,
+                norm_eps=1e-6,
+            ),
+            audio=Gemma4AudioConfig(
+                input_size=16,
+                hidden_size=32,
+                num_layers=1,
+                output_dim=64,
+                output_proj_dims=64,
+                audio_token_id=255998,
+            ),
+            dtype=DTYPE_MAP[dtype_str],
+        )
+        model_cls = registry.get("gemma4")
+        module = model_cls(config)
+        task = get_task("gemma4")
+        pkg = task.build(module, config)
+
+        # Vision encoder pixel_values must be FLOAT
+        vision_model = pkg["vision_encoder"]
+        pv_input = vision_model.graph.inputs[0]
+        assert pv_input.name == "pixel_values"
+        assert pv_input.dtype == ir.DataType.FLOAT
+
+        # Audio encoder input_features must be FLOAT
+        audio_model = pkg["audio_encoder"]
+        af_input = audio_model.graph.inputs[0]
+        assert af_input.name == "input_features"
+        assert af_input.dtype == ir.DataType.FLOAT
+
 
 class TestBuildGraphMultiModal:
     """Verify Phi4MM builds with Phi4MMMultiModalTask (4-model split)."""