mobius

ONNX model definitions for GenAI using the onnxscript.nn API.

Overview

This package provides model definitions for generative AI architectures — LLMs, MoE, multimodal, encoder-only, encoder-decoder, vision, audio, and diffusion models — built directly as ONNX graphs using onnxscript.nn.Module. Rather than tracing or exporting PyTorch models, it constructs the ONNX graph declaratively, then applies pretrained HuggingFace weights.

Supports building ONNX models from HuggingFace model IDs with automatic weight downloading, dtype casting (including bfloat16 via ir.LazyTensor), and multi-component export for pipelines.

📖 Documentation · 📦 Supported Models

Highlighted Models

Mage-VL supports direct three-model ONNX export. ORT GenAI export is currently rejected because the runtime cannot supply its required patch_positions input or Mage-VL's 1D decoder positions.

Supports 290+ Transformers model types and 10 Diffusers component types across 40+ task types and 100+ reusable components.

See the model documentation for the complete list.

Installation

pip install -e .

For running tests:

pip install -e ".[testing]"

Quick Start

Python API

from mobius import build

# Build a model package with weights
pkg = build("meta-llama/Llama-3.2-1B")
pkg.save("output/llama-3.2-1b/")

Static cache (opt-in) pre-allocates fixed-size KV cache buffers, which is useful when you know the maximum sequence length up front:

from mobius import build, CausalLMTask

task = CausalLMTask(static_cache=True, max_seq_len=2048)
pkg = build("meta-llama/Llama-3.2-1B", task=task)
pkg.save("output/llama-3.2-1b-static/")

EP-aware optimization generates graphs tuned for a specific runtime execution provider. Pass execution_provider to target CUDA, DirectML, WebGPU, and more — each with the right set of fused kernels and lowering passes applied automatically:

from mobius import build

# CUDA: GQA fusion, SkipLayerNorm, PackQKV
pkg = build("meta-llama/Llama-3.2-1B",
            execution_provider="cuda", dtype="f16")

# WebGPU: GQA fusion, Shape ops replaced with portable alternatives
pkg = build("meta-llama/Llama-3.2-1B",
            execution_provider="webgpu", dtype="f16")

See the EP quickstart and full EP reference for all supported EPs and options.

CLI

mobius build --model Qwen/Qwen2.5-0.5B --output output_dir/

# Build for CUDA with f16
mobius build --model meta-llama/Llama-3.2-1B --output output_dir/ --ep cuda --dtype f16

# Build a diffusers pipeline (all components)
mobius build --model Qwen/Qwen-Image-2512 --output output_dir/

# Build encoder-decoder model (produces encoder/model.onnx + decoder/model.onnx)
mobius build --model openai/whisper-tiny --output output_dir/

Build-mode toggles use the cargo-style --features option. Available features are static-cache, fp8-kv-cache, prune-prefill-prefix, and text-only. Pass them as a comma-separated list or repeat the option:

mobius build --model meta-llama/Llama-3.2-1B --output output_dir/ \
      --features static-cache,prune-prefill-prefix --max-seq-len 2048

Use --release with either build or build-gguf to potentially reduce saved model size by stripping build-time debug and provenance metadata. Functional metadata with keys prefixed by mobius. is preserved:

mobius build --model meta-llama/Llama-3.2-1B --output output_dir/ --release
mobius build-gguf model.gguf --output output_dir/ --release

See the CLI Reference for all subcommands and flags.

Examples

• examples/build_and_save.py — Build and save ONNX models (simplest workflow)
• examples/text_generation.py — Greedy text generation with a causal LM
• examples/static_cache_generation.py — Text generation with static KV cache
• examples/multimodal_generation.py — Image captioning with a multimodal model

Architecture

HuggingFace Hub
       │
       ▼
 ArchitectureConfig ◄── from_transformers() / from_diffusers()
       │
       ▼
 Model Module ◄── Reusable Components (Attention, MLP, RMSNorm, RoPE, …)
       │
       ▼
 Task ◄── CausalLMTask, VisionLanguageTask, VAETask, DenoisingTask, …
       │
       ▼
 ONNX Model ◄── preprocess_weights() + apply_weights()

The package is organised into four layers:

• Components — onnxscript.nn.Module building blocks (Attention, MLP, DecoderLayer, RoPE, VisionEncoder, MoELayer, …)
• Models — Full architectures composed from components
• Tasks — Define the ONNX graph I/O contract (inputs, outputs, KV cache)
• Registry — Maps HuggingFace model_type strings to model classes

See the design document for details.

Development

# Unit tests (fast, no network needed)
pytest tests/build_graph -v

# Integration tests (downloads models)
pytest tests/integration -m integration -v

# All unit tests (components, configs, tasks, models)
pytest src tests -m "not integration" -v

# Linting
lintrunner f --all-files

Adding a new model

To request a new model, or ask an AI coding agent to implement a new model, use the resources below:

See the AI-assisted model support strategy and the developer skills in .agents/skills/: