diff --git a/examples/gemma4/ort_genai/vlm/genai_config.json b/examples/gemma4/ort_genai/vlm/genai_config.json index eaf76be1..c47657bf 100644 --- a/examples/gemma4/ort_genai/vlm/genai_config.json +++ b/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", diff --git a/examples/gemma4/ort_genai/vlm/processor_config.json b/examples/gemma4/ort_genai/vlm/processor_config.json index 63a50fa2..96198946 100644 --- a/examples/gemma4/ort_genai/vlm/processor_config.json +++ b/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 } } diff --git a/src/mobius/_optimizations.py b/src/mobius/_optimizations.py index 2612c10c..fd9f972e 100644 --- a/src/mobius/_optimizations.py +++ b/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) diff --git a/src/mobius/_passes/__init__.py b/src/mobius/_passes/__init__.py index 589e0bc4..f9258e03 100644 --- a/src/mobius/_passes/__init__.py +++ b/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 diff --git a/src/mobius/_passes/_remove_dead_inputs.py b/src/mobius/_passes/_remove_dead_inputs.py new file mode 100644 index 00000000..8056735e --- /dev/null +++ b/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) diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 71f79510..307908a0 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/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, @@ -150,10 +168,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) @@ -164,6 +204,7 @@ def _write_genai_config( config: Any, output_dir: str, *, + pkg: ModelPackage, ort_model_type: str, ep: str, context_length: int, @@ -175,10 +216,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, @@ -187,20 +243,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: @@ -331,6 +411,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, diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index c7abce3c..f6d0a950 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -15,7 +15,9 @@ from mobius.integrations.ort_genai.auto_export import ( _copy_tokenizer_files, _copy_tokenizer_files_from_local, + _graph_input_names, _resolve_ort_genai_model_type, + _write_genai_config, _write_processor_config, write_ort_genai_config, ) @@ -482,9 +484,295 @@ class FakeConfig: assert data["model"]["eos_token_id"] == [1, 106] -@pytest.mark.integration -class TestAutoExportEndToEnd: - """Integration test: auto_export with a tiny model (no real download).""" +class TestGemma4GenaiConfig: + """Tests for Gemma4-specific genai_config generation via graph introspection.""" + + @staticmethod + def _make_gemma4_pkg(): + """Build a mock Gemma4 VLM package with graph inputs.""" + import dataclasses + + from mobius._model_package import ModelPackage + + @dataclasses.dataclass + class FakeVision: + image_size: int = 448 + patch_size: int = 16 + mm_tokens_per_image: int = 256 + + @dataclasses.dataclass + class FakeConfig: + model_type: str = "gemma4" + vocab_size: int = 262144 + hidden_size: int = 2048 + num_hidden_layers: int = 26 + num_attention_heads: int = 8 + num_key_value_heads: int = 4 + head_dim: int = 256 + max_position_embeddings: int = 8192 + image_token_id: int = 255999 + vision: FakeVision = dataclasses.field(default_factory=FakeVision) + + # Mock graph inputs for each sub-model + def _mock_model_with_inputs(names): + inputs = [] + for n in names: + inp = mock.MagicMock() + inp.name = n + inputs.append(inp) + m = mock.MagicMock() + m.graph.inputs = inputs + return m + + decoder = _mock_model_with_inputs( + [ + "inputs_embeds", + "input_ids", + "attention_mask", + "position_ids", + "past_key_values.0.key", + "past_key_values.0.value", + ] + ) + vision = _mock_model_with_inputs( + [ + "pixel_values", + "pixel_position_ids", + ] + ) + embedding = _mock_model_with_inputs( + [ + "input_ids", + "image_features", + ] + ) + + return ModelPackage( + { + "decoder": decoder, + "vision": vision, + "embedding": embedding, + }, + config=FakeConfig(), + ) + + def test_gemma4_vision_inputs(self, tmp_path): + """Gemma4 vision uses pixel_values + pixel_position_ids.""" + pkg = self._make_gemma4_pkg() + path = _write_genai_config( + pkg.config, + str(tmp_path), + pkg=pkg, + ort_model_type="gemma4", + ep="cpu", + context_length=4096, + bos_token_id=2, + eos_token_id=1, + pad_token_id=0, + is_vlm=True, + has_speech=False, + ) + with open(path) as f: + data = json.load(f) + vision_inputs = data["model"]["vision"]["inputs"] + assert "pixel_values" in vision_inputs + assert "pixel_position_ids" in vision_inputs + assert "image_grid_thw" not in vision_inputs + assert "spatial_merge_size" not in data["model"]["vision"] + + def test_gemma4_decoder_has_input_ids_and_inputs_embeds(self, tmp_path): + """Gemma4 decoder has both inputs_embeds and input_ids.""" + pkg = self._make_gemma4_pkg() + path = _write_genai_config( + pkg.config, + str(tmp_path), + pkg=pkg, + ort_model_type="gemma4", + ep="cpu", + context_length=4096, + bos_token_id=2, + eos_token_id=1, + pad_token_id=0, + is_vlm=True, + has_speech=False, + ) + with open(path) as f: + data = json.load(f) + decoder_inputs = data["model"]["decoder"]["inputs"] + assert "inputs_embeds" in decoder_inputs + assert "input_ids" in decoder_inputs + # KV cache templates are present + assert decoder_inputs["past_key_names"] == "past_key_values.%d.key" + + def test_gemma4_embedding_inputs(self, tmp_path): + """Gemma4 embedding inputs discovered from graph.""" + pkg = self._make_gemma4_pkg() + path = _write_genai_config( + pkg.config, + str(tmp_path), + pkg=pkg, + ort_model_type="gemma4", + ep="cpu", + context_length=4096, + bos_token_id=2, + eos_token_id=1, + pad_token_id=0, + is_vlm=True, + has_speech=False, + ) + with open(path) as f: + data = json.load(f) + emb_inputs = data["model"]["embedding"]["inputs"] + assert "input_ids" in emb_inputs + assert "image_features" in emb_inputs + + +class TestGraphInputNames: + """Tests for _graph_input_names() helper.""" + + @staticmethod + def _mock_model(names): + inputs = [] + for n in names: + inp = mock.MagicMock() + inp.name = n + inputs.append(inp) + m = mock.MagicMock() + m.graph.inputs = inputs + return m + + def test_filters_kv_cache_inputs(self): + """KV cache inputs (past_key_values.*) are filtered out.""" + model = self._mock_model( + [ + "input_ids", + "attention_mask", + "past_key_values.0.key", + "past_key_values.0.value", + "past_key_values.1.key", + "past_key_values.1.value", + ] + ) + result = _graph_input_names(model) + assert result == ["input_ids", "attention_mask"] + + def test_filters_past_prefix(self): + """Inputs starting with 'past_' are also filtered out.""" + model = self._mock_model( + [ + "input_ids", + "past_something", + ] + ) + result = _graph_input_names(model) + assert result == ["input_ids"] + + def test_skips_none_names(self): + """Inputs with name=None are skipped.""" + inp_good = mock.MagicMock() + inp_good.name = "input_ids" + inp_none = mock.MagicMock() + inp_none.name = None + m = mock.MagicMock() + m.graph.inputs = [inp_good, inp_none] + result = _graph_input_names(m) + assert result == ["input_ids"] + + def test_returns_all_semantic_inputs(self): + """All non-KV-cache inputs are returned in order.""" + model = self._mock_model( + [ + "inputs_embeds", + "input_ids", + "attention_mask", + "position_ids", + ] + ) + result = _graph_input_names(model) + assert result == [ + "inputs_embeds", + "input_ids", + "attention_mask", + "position_ids", + ] + + +class TestGemma4RealModel: + """Build a real tiny Gemma4 model and verify genai config inputs.""" + + def test_gemma4_genai_config_from_real_model(self, tmp_path): + """Build tiny Gemma4 VLM, generate genai config, verify inputs.""" + from mobius._builder import build_from_module + from mobius._config_resolver import _default_task_for_model + from mobius._configs import Gemma4Config, VisionConfig + from mobius._registry import registry + from mobius.tasks import get_task + + config = Gemma4Config( + model_type="gemma4", + 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", "full_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, + vision=VisionConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=2, + patch_size=16, + norm_eps=1e-6, + ), + ) + model_cls = registry.get("gemma4") + module = model_cls(config) + task_name = _default_task_for_model("gemma4") + task = get_task(task_name) + pkg = build_from_module(module, config, task=task) + pkg.config = config + + result = write_ort_genai_config(pkg, str(tmp_path)) + with open(result["genai_config"]) as f: + data = json.load(f) + + # Decoder inputs introspected from graph + decoder_inputs = data["model"]["decoder"]["inputs"] + assert "inputs_embeds" in decoder_inputs + assert "input_ids" in decoder_inputs + assert "attention_mask" in decoder_inputs + assert "position_ids" in decoder_inputs + assert decoder_inputs["past_key_names"] == ("past_key_values.%d.key") + + # Vision inputs introspected from graph + vision_inputs = data["model"]["vision"]["inputs"] + assert "pixel_values" in vision_inputs + assert "pixel_position_ids" in vision_inputs + assert "image_grid_thw" not in vision_inputs + + # Embedding inputs introspected from graph + emb_inputs = data["model"]["embedding"]["inputs"] + assert "input_ids" in emb_inputs + assert "image_features" in emb_inputs + + # Config-level properties are still present + assert data["model"]["image_token_id"] == 255999 + assert "spatial_merge_size" not in data["model"]["vision"] + assert data["model"]["vision"]["config_filename"] == "processor_config.json" def test_auto_export_produces_genai_config(self, tmp_path): """Mock build() to return a tiny package, verify genai_config.""" diff --git a/src/mobius/integrations/ort_genai/genai_config.py b/src/mobius/integrations/ort_genai/genai_config.py index b70c5d13..e35b15db 100644 --- a/src/mobius/integrations/ort_genai/genai_config.py +++ b/src/mobius/integrations/ort_genai/genai_config.py @@ -123,6 +123,12 @@ class GenaiConfigGenerator: bos_token_id: Beginning-of-sequence token ID. eos_token_id: End-of-sequence token ID(s). pad_token_id: Padding token ID. + decoder_inputs: Explicit decoder input name mapping. When + provided (e.g. from ONNX graph introspection), used + directly instead of the default mapping from + :func:`_default_decoder_inputs`. Must already include KV + cache template entries (``past_key_names``, + ``past_value_names``). """ def __init__( @@ -140,6 +146,7 @@ def __init__( bos_token_id: int | None = None, eos_token_id: int | list[int] | None = None, pad_token_id: int | None = None, + decoder_inputs: dict[str, str] | None = None, ): self.model_type = model_type self.vocab_size = vocab_size @@ -154,6 +161,9 @@ def __init__( self.eos_token_id = eos_token_id self.pad_token_id = pad_token_id + # Explicit decoder inputs (from graph introspection); None → use defaults + self._decoder_inputs = decoder_inputs + # Optional VLM fields (set via with_vision()) self._vision: dict[str, Any] | None = None self._embedding: dict[str, Any] | None = None @@ -173,6 +183,7 @@ def from_config( bos_token_id: int | None = None, eos_token_id: int | list[int] | None = None, pad_token_id: int | None = None, + decoder_inputs: dict[str, str] | None = None, ) -> GenaiConfigGenerator: """Create a generator from a BaseModelConfig-like dataclass. @@ -204,6 +215,7 @@ def from_config( bos_token_id=bos_token_id, eos_token_id=eos_token_id, pad_token_id=pad, + decoder_inputs=decoder_inputs, ) def with_vision( @@ -216,6 +228,7 @@ def with_vision( config_filename: str = "processor_config.json", input_names: dict[str, str] | None = None, output_names: dict[str, str] | None = None, + embedding_input_names: dict[str, str] | None = None, vision_start_token_id: int | None = None, video_token_id: int | None = None, ) -> GenaiConfigGenerator: @@ -234,6 +247,10 @@ def with_vision( Defaults to pixel_values + image_grid_thw. output_names: Override vision model output name mapping. Defaults to image_features. + embedding_input_names: Override embedding model input name + mapping. When provided (e.g. from ONNX graph + introspection), used directly. Defaults to + input_ids + image_features. vision_start_token_id: Token ID for ``<|vision_start|>``. video_token_id: Token ID for video placeholders. @@ -248,6 +265,11 @@ def with_vision( output_names = { "image_features": "image_features", } + if embedding_input_names is None: + embedding_input_names = { + "input_ids": "input_ids", + "image_features": "image_features", + } self._vision = { "filename": filename, @@ -261,10 +283,7 @@ def with_vision( self._embedding = { "filename": embedding_filename, - "inputs": { - "input_ids": "input_ids", - "image_features": "image_features", - }, + "inputs": embedding_input_names, "outputs": { "inputs_embeds": "inputs_embeds", }, @@ -328,13 +347,18 @@ def generate(self) -> dict[str, Any]: """Generate the full genai_config.json dict.""" is_multimodal = self._vision is not None or self._speech is not None - # Decoder section + # Decoder section — use explicit inputs when available (from + # graph introspection), otherwise fall back to defaults. + if self._decoder_inputs is not None: + decoder_inputs = dict(self._decoder_inputs) + else: + decoder_inputs = _default_decoder_inputs(is_vlm=is_multimodal) decoder: dict[str, Any] = { "session_options": _make_session_options(self.ep), "filename": "model.onnx", "head_size": self.head_dim, "hidden_size": self.hidden_size, - "inputs": _default_decoder_inputs(is_vlm=is_multimodal), + "inputs": decoder_inputs, "outputs": _default_decoder_outputs(), "num_attention_heads": self.num_attention_heads, "num_hidden_layers": self.num_hidden_layers, @@ -360,8 +384,10 @@ def generate(self) -> dict[str, Any]: if self._vision is not None: model["vision"] = self._vision if self._embedding is not None: - # Add audio_features to embedding inputs when speech is enabled - if self._speech is not None: + # Add audio_features to embedding inputs when speech is + # enabled and not already present (graph-introspected + # inputs already include it). + if self._speech is not None and "audio_features" not in self._embedding["inputs"]: self._embedding["inputs"]["audio_features"] = "audio_features" model["embedding"] = self._embedding if self._speech is not None: diff --git a/src/mobius/integrations/ort_genai/genai_config_test.py b/src/mobius/integrations/ort_genai/genai_config_test.py index d96c6b59..c59e4538 100644 --- a/src/mobius/integrations/ort_genai/genai_config_test.py +++ b/src/mobius/integrations/ort_genai/genai_config_test.py @@ -336,6 +336,29 @@ def test_image_token_id_required(self): with pytest.raises(TypeError): gen.with_vision() # missing image_token_id + def test_custom_embedding_input_names(self): + """embedding_input_names overrides the default embedding inputs.""" + gen = GenaiConfigGenerator( + "gemma4", + vocab_size=262144, + hidden_size=2048, + num_hidden_layers=26, + num_attention_heads=8, + num_key_value_heads=4, + head_dim=256, + ).with_vision( + image_token_id=255999, + embedding_input_names={ + "input_ids": "input_ids", + "image_features": "image_features", + "custom_input": "custom_input", + }, + ) + config = gen.generate() + emb = config["model"]["embedding"] + assert emb["inputs"]["custom_input"] == "custom_input" + assert emb["inputs"]["input_ids"] == "input_ids" + class TestGenaiConfigFromConfig: """Test from_config() factory method.""" @@ -443,6 +466,56 @@ def test_write_roundtrips_vlm(self, tmp_path): assert loaded["model"]["image_token_id"] == 151655 +class TestExplicitDecoderInputs: + """Test decoder_inputs parameter overrides defaults.""" + + def test_explicit_decoder_inputs_used(self): + """When decoder_inputs is provided, it replaces the defaults.""" + decoder_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", + "past_value_names": "past_key_values.%d.value", + } + gen = GenaiConfigGenerator( + "gemma4", + vocab_size=262144, + hidden_size=2048, + num_hidden_layers=26, + num_attention_heads=8, + num_key_value_heads=4, + head_dim=256, + decoder_inputs=decoder_inputs, + ).with_vision( + image_token_id=255999, + spatial_merge_size=None, + ) + config = gen.generate() + result = config["model"]["decoder"]["inputs"] + assert "input_ids" in result + assert "inputs_embeds" in result + assert result["past_key_names"] == "past_key_values.%d.key" + + def test_default_used_when_decoder_inputs_none(self): + """When decoder_inputs is None, default mapping is used.""" + gen = GenaiConfigGenerator( + "llama", + vocab_size=32000, + hidden_size=4096, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=8, + head_dim=128, + ) + config = gen.generate() + inputs = config["model"]["decoder"]["inputs"] + # LLM default: input_ids, not inputs_embeds + assert "input_ids" in inputs + assert "inputs_embeds" not in inputs + + class TestGenaiConfigGeneratorMultimodal: """Test genai_config generation for multimodal (vision + speech)."""