diff --git a/.agents/skills/ort-genai-config/SKILL.md b/.agents/skills/ort-genai-config/SKILL.md index 6ab46aa10..9a1cf368d 100644 --- a/.agents/skills/ort-genai-config/SKILL.md +++ b/.agents/skills/ort-genai-config/SKILL.md @@ -97,11 +97,22 @@ Key decoder fields: `filename`, `hidden_size`, `head_size`, ### LLM (decoder-only → `DecoderOnly_Model`) ``` -chatglm, decoder, ernie4_5, gemma, gemma2, gemma3_text, gpt2, -gptoss, granite, internlm2, llama, mistral, nemotron, olmo, -phi, phimoe, phi3, phi3small, qwen2, qwen3, smollm3 +decoder ``` +Mobius emits `model.type = "decoder"` for every graph-representable, +single-model decoder-only package, regardless of the HuggingFace architecture +name. The optimized graph and config mappings define the runtime contract. + +Specialized decoder-only exceptions are explicit: + +- `gpt2`: selects `Gpt_Model`; Mobius rejects its separate key/value cache ABI + because the runtime requires combined rank-5 cache tensors. +- `lfm2`: selects `LFM2_Model` and its convolution-cache implementation. +- `phi3`, `phi3small`, `phimoe`: used only for LongRoPE configs whose runtime + must recompute caches after crossing the short-context threshold. Non-LongRoPE + exports use `decoder`. + ### VLM (vision-language → `MultiModalLanguageModel`) ``` @@ -141,7 +152,7 @@ phi3small_pipeline, qwen2_5_vl_pipeline ```json { "model": { - "type": "llama", + "type": "decoder", "vocab_size": 32000, "context_length": 4096, "eos_token_id": 2, diff --git a/.agents/skills/ort-genai-config/references/genai-config-fields.md b/.agents/skills/ort-genai-config/references/genai-config-fields.md index 5555d7efe..fcbb4b465 100644 --- a/.agents/skills/ort-genai-config/references/genai-config-fields.md +++ b/.agents/skills/ort-genai-config/references/genai-config-fields.md @@ -10,7 +10,7 @@ parent [SKILL.md](../SKILL.md). | Field | Type | Required | Description | |---|---|---|---| -| `type` | string | **yes** | Model type identifier (see registry in SKILL.md) | +| `type` | string | **yes** | `decoder` for graph-representable decoder-only packages; a specialized type only for the explicit runtime/topology exceptions in SKILL.md | | `vocab_size` | int | yes | Vocabulary size | | `context_length` | int | **yes** | Maximum context length; must be > 0 | | `bos_token_id` | int | no | Beginning-of-sequence token | diff --git a/docs/cli_reference.md b/docs/cli_reference.md index a1e5bd589..93c691778 100644 --- a/docs/cli_reference.md +++ b/docs/cli_reference.md @@ -161,8 +161,10 @@ tokenizer files to the output directory: - With `--config` (local directory): tokenizer files are copied from that directory. -For a single-model decoder-only text graph, Mobius emits the architecture-neutral -`model.type: "decoder"` contract supported by onnxruntime-genai 0.14.0 and newer. +For a graph-representable, single-model decoder-only text graph, Mobius emits the +architecture-neutral `model.type: "decoder"` contract supported by +onnxruntime-genai 0.14.0 and newer. Validation targets only the latest stable +release, currently 0.15.2. The graph determines the exact semantic input names, output names, cache templates, and global cache indices, so dense, MoE, tied-weight, quantized, and unknown architecture names do not need a runtime registry entry. @@ -171,9 +173,10 @@ Architecture-specific types remain only where the runtime selects different behavior. `lfm2` uses its legacy convolution-cache implementation. `gpt2` selects `Gpt_Model`, but Mobius's separate rank-4 key/value cache ABI does not match that runtime's rank-5 combined-cache contract, so config generation currently fails -closed. `phi3`, `phimoe`, and `phi3small` retain their names because the released -generator uses them to recompute LongRoPE caches when generation crosses the -short-context threshold. Multimodal, audio, encoder-decoder, special-position-ID, +closed. `phi3`, `phimoe`, and `phi3small` retain their names only when their config selects +LongRoPE, because the released generator uses those names to recompute caches when +generation crosses the short-context threshold. Ordinary Phi-3-family graphs use +`decoder`. Multimodal, audio, encoder-decoder, special-position-ID, and split pipeline packages remain outside the generic path and require their dedicated types and schemas. These exceptions follow the [v0.15.2 runtime model factory](https://github.com/microsoft/onnxruntime-genai/blob/v0.15.2/src/models/model.cpp#L874-L907), diff --git a/examples/gemma4/ort_genai/README.md b/examples/gemma4/ort_genai/README.md index 7007875fa..cda45a668 100644 --- a/examples/gemma4/ort_genai/README.md +++ b/examples/gemma4/ort_genai/README.md @@ -127,16 +127,16 @@ Output: `audio_features [batch, time/4, 1536]` — projected to text hidden_size ## ORT GenAI support required -These configs use the `gemma4` model type which is not yet in a released -ORT GenAI build: +The decoder-only config uses the generic `decoder` type. Multimodal configs +retain `gemma4` because that value selects the vision/audio pipeline: | Config | `model.type` | Pipeline variant | |---|---|---| -| `text/genai_config.json` | `gemma4` | Decoder-only (text) | +| `text/genai_config.json` | `decoder` | Decoder-only (text) | | `vlm/genai_config.json` | `gemma4` | Multimodal (vision + audio + text) | -The runtime auto-detects the pipeline variant from which ONNX files are -present. +The specialized multimodal runtime selects its pipeline variant from which +ONNX files are present. --- diff --git a/examples/gemma4/ort_genai/text/genai_config.json b/examples/gemma4/ort_genai/text/genai_config.json index 553c6e771..c49766431 100644 --- a/examples/gemma4/ort_genai/text/genai_config.json +++ b/examples/gemma4/ort_genai/text/genai_config.json @@ -1,6 +1,6 @@ { "model": { - "type": "gemma4", + "type": "decoder", "vocab_size": 262144, "context_length": 131072, "bos_token_id": 2, diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index da8d0fc4b..55171953a 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -144,10 +144,8 @@ def _revision_kwargs(revision: str | None) -> dict[str, str]: "gpt2": "gpt2", "lfm2": "lfm2", "lfm2_vl": "lfm2", - "phi3": "phi3", - "phi3small": "phi3small", - "phimoe": "phimoe", } +_LONGROPE_TEXT_TYPES = frozenset({"phi3", "phi3small", "phimoe"}) _GENERIC_DECODER_MIN_VERSION = (0, 14, 0) _GENERIC_DECODER_TESTED_VERSIONS = ("0.15.2",) _DECODER_SEMANTIC_INPUTS = frozenset( @@ -269,14 +267,15 @@ def _select_ort_model_type( hf_model_type: str | None, *, is_decoder_only: bool, + rope_type: str | None = None, ) -> str: """Choose the ORT-GenAI model type for an exported package. - Released ORT GenAI 0.14+ dispatches ``decoder`` to its generic + Released ORT GenAI dispatches ``decoder`` to its generic ``DecoderOnly_Model``. Decoder-only packages therefore use that type unless the runtime has genuinely different behavior: ``gpt2`` selects ``Gpt_Model``, - ``lfm2`` selects ``LFM2_Model``/``LFM2Cache``, and Phi-3 family names enable - LongRoPE cache recomputation after the short-context threshold. + ``lfm2`` selects ``LFM2_Model``/``LFM2Cache``, and Phi-3 family names are + retained only for LongRoPE cache recomputation after the short-context threshold. Multimodal and encoder-decoder packages retain their architecture-specific type because those values select distinct runtime pipelines and position-ID @@ -287,6 +286,8 @@ def _select_ort_model_type( resolved = _resolve_ort_genai_model_type(source_type or "unknown") if resolved in _ARCHITECTURE_SPECIFIC_TEXT_TYPES: return _ARCHITECTURE_SPECIFIC_TEXT_TYPES[resolved] + if resolved in _LONGROPE_TEXT_TYPES and rope_type == "longrope": + return resolved return "decoder" return _resolve_ort_genai_model_type(hf_model_type or "unknown") @@ -1546,6 +1547,7 @@ def _write_genai_config( else _count_cache_layer_slots(decoder_model) ), sliding_window=sliding_window, + has_specialized_topology=not _is_single_model_decoder_package(pkg), ) generator.with_special_tokens( **_special_token_ids_from_tokenizer_config(output_dir, config.vocab_size) @@ -1865,7 +1867,10 @@ def write_ort_genai_config( # See _select_ort_model_type: decoder-only packages prefer the package's # own config.model_type; multimodal packages keep the HF parent type. ort_model_type = _select_ort_model_type( - cfg_model_type, model_type, is_decoder_only=is_decoder_only + cfg_model_type, + model_type, + is_decoder_only=is_decoder_only, + rope_type=getattr(config, "rope_type", None), ) # Token IDs may live on the parent config or the text sub-config # (e.g. Gemma4Config has text_config with bos_token_id=2). @@ -1901,7 +1906,10 @@ def write_ort_genai_config( ort_model_type = "gemma3n" else: ort_model_type = _select_ort_model_type( - raw_type, raw_type, is_decoder_only=is_decoder_only + raw_type, + raw_type, + is_decoder_only=is_decoder_only, + rope_type=getattr(config, "rope_type", None), ) if ort_model_type == "unknown": logger.warning( @@ -1948,7 +1956,6 @@ def write_ort_genai_config( "Generic ORT GenAI decoder packages require onnxruntime-genai >= 0.14.0; " f"requested {runtime_version}" ) - result: dict[str, str] = {} if "mtp" in pkg: @@ -2022,9 +2029,11 @@ def write_ort_genai_config( has_speech=has_speech, ) result["genai_config"] = genai_path + with open(genai_path, encoding="utf-8") as handle: + emitted_model_type = json.load(handle)["model"]["type"] compatibility_path = _write_runtime_compatibility( directory, - model_type=ort_model_type, + model_type=emitted_model_type, runtime_version=runtime_version, ) result["runtime_compatibility"] = compatibility_path diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 1fc301542..dcf80958b 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -202,14 +202,29 @@ def test_decoder_only_preserves_specialized_hf_fallback(self): ("gpt2", "gpt2"), ("lfm2", "lfm2"), ("lfm2_vl", "lfm2"), - ("phi3", "phi3"), - ("phi3small", "phi3small"), - ("phimoe", "phimoe"), ], ) def test_specialized_decoder_type_is_preserved(self, model_type, expected): assert _select_ort_model_type(model_type, model_type, is_decoder_only=True) == expected + @pytest.mark.parametrize("model_type", ["phi3", "phi3small", "phimoe"]) + def test_phi3_family_uses_generic_decoder_without_longrope(self, model_type): + assert ( + _select_ort_model_type(model_type, model_type, is_decoder_only=True) == "decoder" + ) + + @pytest.mark.parametrize("model_type", ["phi3", "phi3small", "phimoe"]) + def test_phi3_family_preserves_specialized_type_for_longrope(self, model_type): + assert ( + _select_ort_model_type( + model_type, + model_type, + is_decoder_only=True, + rope_type="longrope", + ) + == model_type + ) + class TestWriteProcessorConfig: def test_no_vision_returns_none(self, tmp_path): @@ -3409,6 +3424,20 @@ def test_generic_decoder_runtime_compatibility_metadata(tmp_path): } +def test_decoder_sidecar_preserves_type_and_matching_compatibility_metadata(tmp_path): + pkg = _make_fake_llm_pkg("qwen2") + pkg["mtp"] = _mock_model(inputs=["hidden_states"], outputs=["draft_logits"]) + + result = write_ort_genai_config(pkg, str(tmp_path)) + with open(result["genai_config"], encoding="utf-8") as handle: + config = json.load(handle) + with open(result["runtime_compatibility"], encoding="utf-8") as handle: + compatibility = json.load(handle) + + assert config["model"]["type"] == "qwen2" + assert compatibility["model_type"] == config["model"]["type"] + + @pytest.mark.parametrize( ("model_type", "config_overrides"), [ @@ -3470,12 +3499,22 @@ def test_gpt2_specialized_runtime_rejects_separate_cache_graph(tmp_path): @pytest.mark.parametrize("model_type", ["phi3", "phi3small", "phimoe"]) def test_phi3_longrope_runtime_type_is_preserved(tmp_path, model_type): - result = write_ort_genai_config(_make_fake_llm_pkg(model_type), str(tmp_path)) + pkg = _make_fake_llm_pkg(model_type) + pkg.config.rope_type = "longrope" + result = write_ort_genai_config(pkg, str(tmp_path)) with open(result["genai_config"], encoding="utf-8") as handle: config = json.load(handle) assert config["model"]["type"] == model_type +@pytest.mark.parametrize("model_type", ["phi3", "phi3small", "phimoe"]) +def test_phi3_without_longrope_emits_generic_decoder(tmp_path, model_type): + result = write_ort_genai_config(_make_fake_llm_pkg(model_type), str(tmp_path)) + with open(result["genai_config"], encoding="utf-8") as handle: + config = json.load(handle) + assert config["model"]["type"] == "decoder" + + class TestGemma4RealModel: """Build a real tiny Gemma4 model and verify genai config inputs.""" @@ -3819,7 +3858,7 @@ def test_auto_export_produces_genai_config(self, tmp_path): genai_config = gen.generate() assert "model" in genai_config - assert genai_config["model"]["type"] == "qwen2" + assert genai_config["model"]["type"] == "decoder" assert genai_config["model"]["vocab_size"] == 256 assert genai_config["model"]["decoder"]["num_hidden_layers"] == 2 @@ -3834,7 +3873,7 @@ def test_auto_export_produces_genai_config(self, tmp_path): with open(os.path.join(output_dir, "genai_config.json")) as f: saved = json.load(f) - assert saved["model"]["type"] == "qwen2" + assert saved["model"]["type"] == "decoder" def test_phi4mm_detection_and_config(self, tmp_path): """Simulate phi4mm auto-export: verify detection and config.""" diff --git a/src/mobius/integrations/ort_genai/genai_config.py b/src/mobius/integrations/ort_genai/genai_config.py index b231d13b5..de3cb241c 100644 --- a/src/mobius/integrations/ort_genai/genai_config.py +++ b/src/mobius/integrations/ort_genai/genai_config.py @@ -15,6 +15,13 @@ import os from typing import Any +_SPECIALIZED_DECODER_MODEL_TYPES = { + "gpt2": "gpt2", + "lfm2": "lfm2", + "lfm2_vl": "lfm2", +} +_LONGROPE_DECODER_MODEL_TYPES = frozenset({"phi3", "phi3small", "phimoe"}) + def _default_decoder_inputs( *, @@ -138,8 +145,10 @@ class GenaiConfigGenerator: and assembles the nested dict structure that ORT-GenAI expects. Args: - model_type: The ORT-GenAI model type string (e.g. ``"qwen2"``, - ``"llama"``, ``"qwen2_5_vl"``). + model_type: The source architecture or specialized ORT-GenAI model type. + Decoder-only configs emit ``"decoder"`` unless this value identifies + a runtime-specific state ABI, or LongRoPE is explicitly requested. + Multimodal configs retain the supplied pipeline type. vocab_size: Model vocabulary size. hidden_size: Decoder hidden dimension. num_hidden_layers: Number of decoder transformer layers. @@ -163,6 +172,10 @@ class GenaiConfigGenerator: ``past_value_names``). decoder_outputs: Explicit decoder output mapping derived from the graph, including logits and present-cache templates. + uses_longrope: Preserve a Phi-3-family specialized type because the + runtime must recompute LongRoPE caches across the context threshold. + has_specialized_topology: Preserve the supplied type for packages with + auxiliary graphs or runtime-managed pipelines. """ def __init__( @@ -188,6 +201,8 @@ def __init__( layer_types: list[str] | None = None, conv_cache_size: int | None = None, sliding_window: dict[str, Any] | None = None, + uses_longrope: bool = False, + has_specialized_topology: bool = False, ): self.model_type = model_type self.vocab_size = vocab_size @@ -217,6 +232,8 @@ def __init__( self._layer_types = layer_types self._conv_cache_size = conv_cache_size self._sliding_window = sliding_window + self._uses_longrope = uses_longrope + self._has_specialized_topology = has_specialized_topology # Optional VLM fields (set via with_vision()) self._vision: dict[str, Any] | None = None @@ -247,6 +264,7 @@ def from_config( supports_in_place_kv_cache: bool | None = None, num_cache_layer_slots: int | None = None, sliding_window: dict[str, Any] | None = None, + has_specialized_topology: bool = False, ) -> GenaiConfigGenerator: """Create a generator from a BaseModelConfig-like dataclass. @@ -302,6 +320,11 @@ def from_config( else None ), sliding_window=sliding_window, + uses_longrope=( + model_type in _LONGROPE_DECODER_MODEL_TYPES + and getattr(config, "rope_type", None) == "longrope" + ), + has_specialized_topology=has_specialized_topology, ) def with_vision( @@ -524,6 +547,14 @@ def with_special_tokens(self, **token_ids: int) -> GenaiConfigGenerator: def generate(self) -> dict[str, Any]: """Generate the full genai_config.json dict.""" is_multimodal = self._vision is not None or self._audio is not None + if is_multimodal or self._has_specialized_topology: + emitted_model_type = self.model_type + elif self.model_type in _SPECIALIZED_DECODER_MODEL_TYPES: + emitted_model_type = _SPECIALIZED_DECODER_MODEL_TYPES[self.model_type] + elif self.model_type in _LONGROPE_DECODER_MODEL_TYPES and self._uses_longrope: + emitted_model_type = self.model_type + else: + emitted_model_type = "decoder" # Decoder section — use explicit inputs when available (from # graph introspection), otherwise fall back to defaults. @@ -562,7 +593,7 @@ def generate(self) -> dict[str, Any]: # Model section model: dict[str, Any] = { - "type": self.model_type, + "type": emitted_model_type, "vocab_size": self.vocab_size, "context_length": self.context_length, "decoder": decoder, diff --git a/src/mobius/integrations/ort_genai/genai_config_test.py b/src/mobius/integrations/ort_genai/genai_config_test.py index d85e2ba78..d9a70e8e8 100644 --- a/src/mobius/integrations/ort_genai/genai_config_test.py +++ b/src/mobius/integrations/ort_genai/genai_config_test.py @@ -32,7 +32,7 @@ def test_minimal_llm_config(self): ) config = gen.generate() - assert config["model"]["type"] == "llama" + assert config["model"]["type"] == "decoder" assert config["model"]["vocab_size"] == 32000 assert config["model"]["context_length"] == 4096 @@ -44,6 +44,65 @@ def test_minimal_llm_config(self): assert decoder["head_size"] == 128 assert decoder["filename"] == "model.onnx" + @pytest.mark.parametrize("model_type", ["llama", "qwen2", "gemma4_text", "custom"]) + def test_decoder_only_types_are_normalized(self, model_type): + gen = GenaiConfigGenerator( + model_type, + vocab_size=256, + hidden_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + ) + assert gen.generate()["model"]["type"] == "decoder" + + @pytest.mark.parametrize( + ("model_type", "expected"), + [("gpt2", "gpt2"), ("lfm2", "lfm2"), ("lfm2_vl", "lfm2")], + ) + def test_specialized_decoder_types_are_preserved(self, model_type, expected): + gen = GenaiConfigGenerator( + model_type, + vocab_size=256, + hidden_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + ) + assert gen.generate()["model"]["type"] == expected + + def test_auxiliary_graph_topology_preserves_runtime_type(self): + gen = GenaiConfigGenerator( + "qwen2", + vocab_size=256, + hidden_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + has_specialized_topology=True, + ) + assert gen.generate()["model"]["type"] == "qwen2" + + def test_phi3_type_is_preserved_only_for_longrope(self): + common = { + "vocab_size": 256, + "hidden_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 16, + } + assert GenaiConfigGenerator("phi3", **common).generate()["model"]["type"] == "decoder" + assert ( + GenaiConfigGenerator("phi3", uses_longrope=True, **common).generate()["model"][ + "type" + ] + == "phi3" + ) + def test_llm_decoder_inputs_have_input_ids(self): """LLM decoders receive input_ids, not inputs_embeds.""" gen = GenaiConfigGenerator( @@ -739,7 +798,7 @@ def test_write_creates_valid_json(self, tmp_path): with open(path) as f: loaded = json.load(f) - assert loaded["model"]["type"] == "llama" + assert loaded["model"]["type"] == "decoder" assert "search" in loaded def test_write_roundtrips_vlm(self, tmp_path):