fix: Gemma4 genai config - vision inputs, decoder input_ids, processor (#199)

## Summary

Fix Gemma4 ORT GenAI config generation to produce correct
genai_config.json and processor_config.json.

## Changes

### 1. Vision inputs (auto_export.py)
Gemma4 uses `pixel_values + pixel_position_ids` (not `image_grid_thw`).
Added Gemma4-specific branch in `_write_genai_config()` that sets the
correct input names and disables `spatial_merge_size`.

### 2. Decoder input_ids (genai_config.py)
Gemma4 decoders need `input_ids` alongside `inputs_embeds` for per-layer
token embeddings (E2B architecture). Added `with_extra_decoder_inputs()`
method to `GenaiConfigGenerator` and wired it for Gemma4 in
`_write_genai_config()`.

### 3. Processor config (auto_export.py)
Updated `_write_processor_config()` to detect Gemma4 and write the
correct format with `name`, `tokens_per_image`, `mean`, and `std` fields
wrapped under a `processor` key, matching the ort-extensions expected
format.

### 4. Reference config updated
Added `input_ids` to decoder inputs in
`examples/gemma4/ort_genai/vlm/genai_config.json`.

## Testing
- All 85 ORT GenAI tests pass
- All 12 Gemma4 build graph tests pass
- Full test suite: 2557 passed, 41 skipped
- Linter clean

---------

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

Author: justinchuby

examples/gemma4/ort_genai/vlm/genai_config.json

@@ -20,6 +20,7 @@
             "num_hidden_layers": 15,
             "inputs": {
                 "inputs_embeds": "inputs_embeds",
+                "input_ids": "input_ids",
                 "attention_mask": "attention_mask",
                 "position_ids": "position_ids",
                 "past_key_names": "past_key_values.%d.key",

examples/gemma4/ort_genai/vlm/processor_config.json

@@ -3,8 +3,6 @@
         "name": "gemma4_image_processor",
         "image_size": 448,
         "patch_size": 16,
-        "tokens_per_image": 280,
-        "mean": [0.5, 0.5, 0.5],
-        "std": [0.5, 0.5, 0.5]
+        "tokens_per_image": 280
     }
 }

src/mobius/_optimizations.py

@@ -54,7 +54,11 @@
 
 from mobius._execution_providers import EpCapabilities, ep_registry
 from mobius._flags import flags
-from mobius._passes import FoldConcatInitializersPass, FoldTransposedInitializerPass
+from mobius._passes import (
+    FoldConcatInitializersPass,
+    FoldTransposedInitializerPass,
+    RemoveDeadGraphInputsPass,
+)
 from mobius.functions import register_function_bodies
 from mobius.rewrite_rules import (
     gelu_fusion_rules,
@@ -445,7 +449,8 @@ def _should_inline(func: ir.Function) -> bool:
         for _, ir_pass in lower_ir_passes:
             ir_pass(model)
 
-    # Stage 4: Final dead-node removal and constant folding after rewrites.
+    # Stage 4: Final dead-node removal, constant folding, and dead input
+    # cleanup after rewrites.
     if trace:
         before_fold = sum(_count_all_ops(model).values())
         logger.info("[EP Trace] Stage 4: Constant folding")
@@ -461,6 +466,9 @@ def _should_inline(func: ir.Function) -> bool:
                 input_size_limit=8192,
                 output_size_limit=_FOLD_OUTPUT_SIZE_LIMIT,
             ),
+            # Remove graph inputs whose consumers were all eliminated by
+            # fusion (e.g. position_ids when GQA absorbs RoPE).
+            RemoveDeadGraphInputsPass(),
         ]
     )
     fold_pass(model)

src/mobius/_passes/__init__.py

@@ -42,7 +42,9 @@
 __all__ = [
     "FoldTransposedInitializerPass",
     "FoldConcatInitializersPass",
+    "RemoveDeadGraphInputsPass",
 ]
 
 from mobius._passes._fold_concat import FoldConcatInitializersPass
 from mobius._passes._fold_transpose import FoldTransposedInitializerPass
+from mobius._passes._remove_dead_inputs import RemoveDeadGraphInputsPass

src/mobius/_passes/_remove_dead_inputs.py

@@ -0,0 +1,46 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Pass that removes unused graph inputs.
+
+After EP-aware optimization (e.g. GQA fusion absorbs RoPE), some graph
+inputs may have zero consumers.  For example, ``position_ids`` becomes
+dead when all attention layers use ``GroupQueryAttention`` with
+``do_rotary=1``.  Removing dead inputs produces cleaner models and
+avoids requiring the runtime to provide dummy feed values.
+
+KV cache inputs (``past_key_values.*``) are always retained even if
+they appear unused in the graph, because ORT GenAI manages them
+externally via the KV cache protocol.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import onnx_ir as ir
+
+logger = logging.getLogger(__name__)
+
+
+class RemoveDeadGraphInputsPass(ir.passes.InPlacePass):
+    """Remove graph inputs that have no consumers.
+
+    Skips inputs whose name starts with ``past_key_values.`` (KV cache
+    entries managed by the runtime) and inputs with ``None`` names.
+    """
+
+    def call(self, model: ir.Model) -> ir.passes.PassResult:
+        dead = [
+            inp
+            for inp in model.graph.inputs
+            if inp.name is not None
+            and not inp.name.startswith("past_key_values.")
+            and len(inp.uses()) == 0
+        ]
+        for inp in dead:
+            model.graph.inputs.remove(inp)
+            logger.debug("Removed dead graph input: %s", inp.name)
+
+        modified = len(dead) > 0
+        return ir.passes.PassResult(model, modified=modified)

src/mobius/integrations/ort_genai/auto_export.py

@@ -42,6 +42,8 @@
 from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
+    import onnx_ir as ir
+
     from mobius._model_package import ModelPackage
 
 logger = logging.getLogger(__name__)
@@ -69,6 +71,22 @@ def _resolve_ort_genai_model_type(model_type: str) -> str:
     return _ORT_GENAI_MODEL_TYPE.get(model_type, model_type)
 
 
+def _graph_input_names(model: ir.Model) -> list[str]:
+    """Return non-KV-cache input names from an ONNX model graph.
+
+    Filters out KV cache inputs (``past_key_values.*`` and ``past_*``)
+    since those are represented as template patterns in genai_config.json,
+    not as literal graph input names.
+    """
+    return [
+        inp.name
+        for inp in model.graph.inputs
+        if inp.name is not None
+        and not inp.name.startswith("past_key_values.")
+        and not inp.name.startswith("past_")
+    ]
+
+
 def _copy_tokenizer_files(
     model_id: str,
     output_dir: str,
@@ -152,10 +170,32 @@ def _write_processor_config(
     if vision is None:
         return None
 
-    processor: dict[str, Any] = {
-        "image_size": getattr(vision, "image_size", 448),
-        "patch_size": getattr(vision, "patch_size", 14),
-    }
+    model_type = getattr(config, "model_type", "")
+
+    if model_type in ("gemma4", "gemma4_text"):
+        # Gemma4 needs a processor wrapper with model-specific fields
+        tokens_per_image = (
+            getattr(vision, "mm_tokens_per_image", None)
+            or getattr(config, "mm_tokens_per_image", None)
+            or getattr(vision, "max_soft_tokens", None)
+            or 280
+        )
+        image_size = getattr(vision, "image_size", None) or 448
+        patch_size = getattr(vision, "patch_size", None) or 16
+        processor: dict[str, Any] = {
+            "processor": {
+                "name": "gemma4_image_processor",
+                "image_size": image_size,
+                "patch_size": patch_size,
+                "tokens_per_image": tokens_per_image,
+            }
+        }
+    else:
+        processor = {
+            "image_size": getattr(vision, "image_size", None) or 448,
+            "patch_size": getattr(vision, "patch_size", None) or 14,
+        }
+
     path = os.path.join(output_dir, "processor_config.json")
     with open(path, "w", encoding="utf-8") as f:
         json.dump(processor, f, indent=4)
@@ -166,6 +206,7 @@ def _write_genai_config(
     config: Any,
     output_dir: str,
     *,
+    pkg: ModelPackage,
     ort_model_type: str,
     ep: str,
     context_length: int,
@@ -177,10 +218,25 @@ def _write_genai_config(
 ) -> str:
     """Generate and write genai_config.json.
 
+    Input names for each sub-model (decoder, vision, embedding) are
+    introspected from the ONNX graphs in *pkg* rather than hard-coded
+    per model type.
+
     Returns the path to the written file.
     """
     from mobius.integrations.ort_genai.genai_config import GenaiConfigGenerator
 
+    # --- Discover decoder inputs from the ONNX graph ---
+    decoder_model = pkg.get("decoder") or pkg.get("model")
+    if decoder_model is not None:
+        decoder_input_names = _graph_input_names(decoder_model)
+        decoder_inputs: dict[str, str] | None = {name: name for name in decoder_input_names}
+        # KV cache entries are template-based, not per-input
+        decoder_inputs["past_key_names"] = "past_key_values.%d.key"
+        decoder_inputs["past_value_names"] = "past_key_values.%d.value"
+    else:
+        decoder_inputs = None  # fall back to defaults
+
     generator = GenaiConfigGenerator.from_config(
         config,
         ort_model_type,
@@ -189,20 +245,44 @@ def _write_genai_config(
         bos_token_id=bos_token_id,
         eos_token_id=eos_token_id,
         pad_token_id=pad_token_id,
+        decoder_inputs=decoder_inputs,
     )
 
     if is_vlm:
         image_token_id = getattr(config, "image_token_id", None)
         if image_token_id is not None:
+            # Discover vision inputs from the graph
+            vision_model = pkg.get("vision")
+            if vision_model is not None:
+                names = _graph_input_names(vision_model)
+                vision_input_mapping: dict[str, str] | None = {n: n for n in names}
+            else:
+                vision_input_mapping = None
+
+            # Discover embedding inputs from the graph
+            embedding_model = pkg.get("embedding")
+            if embedding_model is not None:
+                names = _graph_input_names(embedding_model)
+                embedding_input_mapping: dict[str, str] | None = {n: n for n in names}
+            else:
+                embedding_input_mapping = None
+
+            # spatial_merge_size and config_filename are config-level
+            # properties that cannot be inferred from the graph.
             vision_kwargs: dict[str, Any] = {}
+            model_type = getattr(config, "model_type", "")
             if has_speech:
-                # Phi4MM uses different vision inputs than Qwen2.5-VL
                 vision_kwargs["spatial_merge_size"] = None
                 vision_kwargs["config_filename"] = "vision_processor.json"
-                vision_kwargs["input_names"] = {
-                    "pixel_values": "pixel_values",
-                    "image_sizes": "image_sizes",
-                }
+            elif model_type in ("gemma4", "gemma4_text"):
+                vision_kwargs["spatial_merge_size"] = None
+                vision_kwargs["config_filename"] = "processor_config.json"
+
+            if vision_input_mapping is not None:
+                vision_kwargs["input_names"] = vision_input_mapping
+            if embedding_input_mapping is not None:
+                vision_kwargs["embedding_input_names"] = embedding_input_mapping
+
             generator.with_vision(image_token_id=image_token_id, **vision_kwargs)
 
     if has_speech:
@@ -333,6 +413,7 @@ def write_ort_genai_config(
     genai_path = _write_genai_config(
         config,
         directory,
+        pkg=pkg,
         ort_model_type=ort_model_type,
         ep=ep,
         context_length=context_length,