Add onnx-genai InferenceMetadata export (--runtime onnx-genai) (#398)

justinchuby · Copilot · web-flow · commit 40dbf1fb47e2 · 2026-07-16T08:19:41.000-07:00
Adds a Mobius integration that emits onnx-genai's own
`inference_metadata.yaml` (the InferenceMetadata schema) instead of
ORT-GenAI's `genai_config.json`, so onnx-genai models carry the config
that runtime needs.

- New `--runtime onnx-genai` build target; emitter
`src/mobius/integrations/onnx_genai/inference_metadata.py`.
- Emits `model.attention` (type/heads/head_dim),
`model.max_sequence_length`, `kv_cache.native_dtype`, and
`required_capabilities` using the runtime's supported strings
(`grouped_query_attention` / `multi_head_attention`).
- KV pre-size capacity defaults to `min(4096, model limit)` (avoids
exceeding WebGPU's 256 MiB buffer limit for large context windows);
optional `--max-length` override.
- **Round-trip verified:** onnx-genai loads a GQA WebGPU model from the
emitted `inference_metadata.yaml` alone (no genai_config.json) and
produces coherent output.

lintrunner clean; emitter pytest 8 passed.

Co-authored-by: Copilot &lt;223556219+Copilot@users.noreply.github.com&gt;
diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py
@@ -127,6 +127,11 @@ def _cmd_build(args: argparse.Namespace) -> None:
     if args.max_seq_len is not None and args.max_seq_len <= 0:
         raise SystemExit("Error: --max-seq-len must be a positive integer.")
 
+    if args.max_length is not None and args.runtime != "onnx-genai":
+        raise SystemExit("Error: --max-length can only be used with --runtime onnx-genai.")
+    if args.max_length is not None and args.max_length <= 0:
+        raise SystemExit("Error: --max-length must be a positive integer.")
+
     # Validate --static-cache + --task compatibility
     if args.static_cache and args.task is not None:
         raise SystemExit(
@@ -293,6 +298,13 @@ def _save_package(
         )
         for name, path in artifacts.items():
             print(f"  {name}: {path}")
+    elif runtime == "onnx-genai":
+        from mobius.integrations.onnx_genai import write_inference_metadata
+
+        path = write_inference_metadata(
+            pkg, output_dir, max_sequence_length=getattr(args, "max_length", None)
+        )
+        print(f"  inference_metadata: {path}")
 
 
 def _cmd_list(args: argparse.Namespace) -> None:
@@ -555,14 +567,25 @@ def main(argv: list[str] | None = None) -> None:
     build_parser.add_argument(
         "--runtime",
         default=None,
-        choices=["ort-genai"],
+        choices=["onnx-genai", "ort-genai"],
         metavar="RUNTIME",
         help=(
             "Generate runtime-specific config files after building. "
-            "Currently supports: 'ort-genai' (writes genai_config.json and "
-            "copies tokenizer files). When used with --model, tokenizer files "
-            "are downloaded from HuggingFace. When used with --config (local "
-            "directory), tokenizer files are copied from that directory."
+            "Supports: 'onnx-genai' (writes inference_metadata.yaml) and "
+            "'ort-genai' (writes genai_config.json and copies tokenizer files). "
+            "For ort-genai, --model downloads tokenizer files from HuggingFace "
+            "and --config copies them from the local directory."
+        ),
+    )
+    build_parser.add_argument(
+        "--max-length",
+        type=int,
+        default=None,
+        metavar="N",
+        help=(
+            "Serving KV capacity written to onnx-genai inference metadata. "
+            "Only used with --runtime onnx-genai. Defaults to the smaller of "
+            "4096 and the model's max_position_embeddings."
         ),
     )
     build_parser.add_argument(
diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py
@@ -77,8 +77,10 @@ def save(
             This method writes ONNX files only.  If you need a directory that
             ``onnxruntime-genai`` can load (i.e. with ``genai_config.json`` and
             tokenizer files), use
-            :func:`mobius.integrations.ort_genai.export_package` instead — it
-            wraps :meth:`save` with the ORT-GenAI config-generation step.
+            :func:`mobius.integrations.ort_genai.export_package` instead. For
+            onnx-genai, call
+            :func:`mobius.integrations.onnx_genai.write_inference_metadata`
+            after saving, or use ``mobius build --runtime onnx-genai``.
 
         Args:
             directory: Path to the output directory (created if needed).
diff --git a/src/mobius/integrations/onnx_genai/__init__.py b/src/mobius/integrations/onnx_genai/__init__.py
@@ -0,0 +1,11 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""onnx-genai integration for inference metadata generation."""
+
+from mobius.integrations.onnx_genai.inference_metadata import (
+    generate_inference_metadata,
+    write_inference_metadata,
+)
+
+__all__ = ["generate_inference_metadata", "write_inference_metadata"]
diff --git a/src/mobius/integrations/onnx_genai/inference_metadata.py b/src/mobius/integrations/onnx_genai/inference_metadata.py
@@ -0,0 +1,151 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Generate onnx-genai ``inference_metadata.yaml`` sidecars."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+import onnx_ir as ir
+
+from mobius._model_package import ModelPackage
+
+_DTYPE_NAMES = {
+    ir.DataType.FLOAT: "float32",
+    ir.DataType.FLOAT16: "float16",
+    ir.DataType.BFLOAT16: "bfloat16",
+}
+_DEFAULT_MAX_SEQUENCE_LENGTH = 4096
+
+
+def _positive_int(config: object, name: str) -> int:
+    value = getattr(config, name, None)
+    if not isinstance(value, int) or value <= 0:
+        raise ValueError(
+            f"onnx-genai inference metadata requires a positive {name}, got {value!r}."
+        )
+    return value
+
+
+def _kv_dtype(config: object) -> str:
+    dtype = getattr(config, "dtype", None)
+    try:
+        return _DTYPE_NAMES[dtype]
+    except KeyError:
+        raise ValueError(
+            "onnx-genai inference metadata supports float32, float16, or bfloat16 "
+            f"KV caches, got {dtype!r}."
+        ) from None
+
+
+def _max_sequence_length(config: object, requested: int | None) -> int:
+    model_max = _positive_int(config, "max_position_embeddings")
+    if requested is None:
+        return min(model_max, _DEFAULT_MAX_SEQUENCE_LENGTH)
+    if not isinstance(requested, int) or requested <= 0:
+        raise ValueError(
+            f"onnx-genai max_sequence_length must be a positive integer, got {requested!r}."
+        )
+    if requested > model_max:
+        raise ValueError(
+            f"onnx-genai max_sequence_length {requested} exceeds the model limit {model_max}."
+        )
+    return requested
+
+
+def generate_inference_metadata(
+    config: object, *, max_sequence_length: int | None = None
+) -> dict[str, Any]:
+    """Map a decoder config to metadata with a conservative serving KV capacity."""
+    num_attention_heads = _positive_int(config, "num_attention_heads")
+    num_kv_heads = _positive_int(config, "num_key_value_heads")
+    head_dim = _positive_int(config, "head_dim")
+    max_sequence_length = _max_sequence_length(config, max_sequence_length)
+    kv_dtype = _kv_dtype(config)
+
+    is_gqa = num_kv_heads != num_attention_heads
+    capabilities = ["grouped_query_attention" if is_gqa else "multi_head_attention"]
+
+    attention: dict[str, Any] = {
+        "type": "group_query_attention" if is_gqa else "multi_head_attention",
+        "num_kv_heads": num_kv_heads,
+        "num_attention_heads": num_attention_heads,
+        "head_dim": head_dim,
+    }
+    sliding_window = getattr(config, "sliding_window", None)
+    if isinstance(sliding_window, int) and sliding_window > 0:
+        attention["sliding_window"] = sliding_window
+
+    return {
+        "required_capabilities": capabilities,
+        "model": {
+            "attention": attention,
+            "max_sequence_length": max_sequence_length,
+            "runtime_configurable": {"kv_cache": {"dtype": [kv_dtype]}},
+        },
+        "kv_cache": {"native_dtype": kv_dtype},
+    }
+
+
+def _to_yaml(metadata: dict[str, Any]) -> str:
+    capabilities = metadata["required_capabilities"]
+    attention = metadata["model"]["attention"]
+    kv_dtypes = metadata["model"]["runtime_configurable"]["kv_cache"]["dtype"]
+
+    lines = ["required_capabilities:"]
+    if capabilities:
+        lines.extend(f"  - {capability}" for capability in capabilities)
+    else:
+        lines[-1] += " []"
+
+    lines.extend(
+        [
+            "model:",
+            "  attention:",
+            f"    type: {attention['type']}",
+            f"    num_kv_heads: {attention['num_kv_heads']}",
+            f"    num_attention_heads: {attention['num_attention_heads']}",
+            f"    head_dim: {attention['head_dim']}",
+        ]
+    )
+    if "sliding_window" in attention:
+        lines.append(f"    sliding_window: {attention['sliding_window']}")
+    lines.extend(
+        [
+            f"  max_sequence_length: {metadata['model']['max_sequence_length']}",
+            "  runtime_configurable:",
+            "    kv_cache:",
+            "      dtype:",
+            *(f"        - {dtype}" for dtype in kv_dtypes),
+            "kv_cache:",
+            f"  native_dtype: {metadata['kv_cache']['native_dtype']}",
+        ]
+    )
+    return "\n".join(lines) + "\n"
+
+
+def write_inference_metadata(
+    pkg: ModelPackage,
+    directory: str,
+    *,
+    max_sequence_length: int | None = None,
+) -> str:
+    """Write ``inference_metadata.yaml`` for an already-built model package."""
+    config = getattr(pkg, "config", None)
+    if config is None:
+        raise ValueError(
+            "write_inference_metadata requires ModelPackage.config to be set. "
+            "This is set automatically when building with mobius.build()."
+        )
+
+    os.makedirs(directory, exist_ok=True)
+    path = os.path.join(directory, "inference_metadata.yaml")
+    with open(path, "w", encoding="utf-8") as file:
+        file.write(
+            _to_yaml(
+                generate_inference_metadata(config, max_sequence_length=max_sequence_length)
+            )
+        )
+    return path
diff --git a/src/mobius/integrations/onnx_genai/inference_metadata_test.py b/src/mobius/integrations/onnx_genai/inference_metadata_test.py
@@ -0,0 +1,106 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import onnx_ir as ir
+import pytest
+
+from mobius._configs import ArchitectureConfig
+from mobius._model_package import ModelPackage
+from mobius.integrations.onnx_genai import (
+    generate_inference_metadata,
+    write_inference_metadata,
+)
+
+
+def _config(**kwargs) -> ArchitectureConfig:
+    values = {
+        "num_attention_heads": 14,
+        "num_key_value_heads": 2,
+        "head_dim": 64,
+        "max_position_embeddings": 32768,
+        "dtype": ir.DataType.FLOAT16,
+    }
+    values.update(kwargs)
+    return ArchitectureConfig(**values)
+
+
+def test_generate_gqa_fp16_metadata() -> None:
+    metadata = generate_inference_metadata(_config(sliding_window=4096))
+
+    assert metadata == {
+        "required_capabilities": ["grouped_query_attention"],
+        "model": {
+            "attention": {
+                "type": "group_query_attention",
+                "num_kv_heads": 2,
+                "num_attention_heads": 14,
+                "head_dim": 64,
+                "sliding_window": 4096,
+            },
+            "max_sequence_length": 4096,
+            "runtime_configurable": {"kv_cache": {"dtype": ["float16"]}},
+        },
+        "kv_cache": {"native_dtype": "float16"},
+    }
+
+
+def test_generate_mha_float32_metadata() -> None:
+    metadata = generate_inference_metadata(
+        _config(
+            num_attention_heads=8,
+            num_key_value_heads=8,
+            dtype=ir.DataType.FLOAT,
+            sliding_window=None,
+        )
+    )
+
+    assert metadata["required_capabilities"] == ["multi_head_attention"]
+    assert metadata["model"]["attention"] == {
+        "type": "multi_head_attention",
+        "num_kv_heads": 8,
+        "num_attention_heads": 8,
+        "head_dim": 64,
+    }
+    assert metadata["kv_cache"]["native_dtype"] == "float32"
+
+
+def test_write_inference_metadata_yaml(tmp_path) -> None:
+    pkg = ModelPackage(config=_config())
+
+    path = write_inference_metadata(pkg, str(tmp_path))
+
+    assert path == str(tmp_path / "inference_metadata.yaml")
+    assert (tmp_path / "inference_metadata.yaml").read_text() == (
+        "required_capabilities:\n"
+        "  - grouped_query_attention\n"
+        "model:\n"
+        "  attention:\n"
+        "    type: group_query_attention\n"
+        "    num_kv_heads: 2\n"
+        "    num_attention_heads: 14\n"
+        "    head_dim: 64\n"
+        "  max_sequence_length: 4096\n"
+        "  runtime_configurable:\n"
+        "    kv_cache:\n"
+        "      dtype:\n"
+        "        - float16\n"
+        "kv_cache:\n"
+        "  native_dtype: float16\n"
+    )
+
+
+def test_max_sequence_length_override() -> None:
+    metadata = generate_inference_metadata(_config(), max_sequence_length=2048)
+
+    assert metadata["model"]["max_sequence_length"] == 2048
+
+
+@pytest.mark.parametrize("max_sequence_length", [0, -1, 32769])
+def test_invalid_max_sequence_length(max_sequence_length: int) -> None:
+    with pytest.raises(ValueError, match="max_sequence_length"):
+        generate_inference_metadata(_config(), max_sequence_length=max_sequence_length)
+
+
+def test_write_requires_package_config(tmp_path) -> None:
+    with pytest.raises(ValueError, match=r"ModelPackage\.config"):
+        write_inference_metadata(ModelPackage(), str(tmp_path))
