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

Category	Examples
Text Generation	Llama 2/3/4, Mistral, Qwen 2/2.5/3/3.5/3.6, Phi-3/3.5, Gemma 1/2/3/4, Granite, GPT-2, OPT, OLMo, SmolLM3, and many more
Mixture of Experts	PhiMoE, GPTOSS, Mixtral, OLMoE, DeepSeek-V2/V3, Qwen2-MoE, Qwen3-MoE, Qwen3-Next, GLM-4-MoE, Arctic, DBRX, Jamba
Multimodal	Gemma 3/4, Phi-4MM (vision + audio + LoRA), Nemotron Parse, LLaVA, InternVL2, Mage-VL (image + streaming video), MiniCPM-V 4.6, Qwen2.5-VL, Qwen3-VL, Qwen3.5/3.6-VL, Pixtral
Encoder-only	BERT, RoBERTa, ALBERT, DeBERTa, DistilBERT, ELECTRA, XLNet
Encoder-Decoder	BART, T5/mT5, Marian, M2M-100, Pegasus, BigBird-Pegasus
Speech-to-Text	Whisper, Moonshine, Moonshine Streaming, FastConformer-RNNT, FunASR, GLM-ASR, Qwen3-ASR, SenseVoice
Audio	Wav2Vec2, HuBERT, WavLM, SpeechT5
Vision	ViT, BEiT, DeiT, DINOv2, Swin, CLIP, SigLIP
Diffusion	Stable Diffusion (UNet + VAE + ControlNet), Flux, SD3, DiT, QwenImage / Qwen-Image-Edit-2509, HunyuanDiT, CogVideoX
Adapters	T2I-Adapter, IP-Adapter

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.

Projection-specific QMoE export

Olive's integer-affine, fused, K-last expert checkpoints can independently use INT2, INT3, INT4, or INT8 for FC1 and FC2, with a supported model-wide INT4 fallback for non-expert projections. All 16 bit-width pairs are accepted when their layouts satisfy the checks below; no experimental flag is required. INT3 uses the proposed raw QMoE contract from microsoft/onnxruntime#32657.

from mobius import build

package = build("./local-olive-qwen3-moe", dtype="f16")
package.save("./qmoe-export")

Use an Olive model-wide INT4 configuration with exact source-name overrides for model.layers.<i>.mlp.experts.gate_up_proj and model.layers.<i>.mlp.experts.down_proj. The generic MoE adapter, including Qwen3-MoE, validates fused checkpoint dimensions and binds packed bytes without requantization. INT3 is symmetric unsigned offset-4, row-local LSB-first uint8; FC1 gate/up rows and scales are interleaved together, while FC2 retains its K-last down-projection layout. Explicit all-4 INT3 zero points are validated and omitted in the graph.

Both projections require a shared power-of-two block size of at least 16, dividing hidden and intermediate sizes. Pairs containing INT3 require block size 32/64/128, symmetric INT3, and FP16/BF16 activations. For these pairs both FC scale dtypes must match the activations. Every projection's scales must be finite and nonnegative. When the adapter supplies the activation dtype, scales must remain finite after conversion. Non-INT3 layouts retain the standard weight loader's dtype conversion. The legacy uniform INT4 path receives the same numeric scale checks. INT2, INT4, and INT8 projections retain their existing zero-point semantics; asymmetric zero points must be packed uint8 sidecars. Dense INT3 MatMulNBits, embeddings, other source layouts, and schema-unapproved tails are not enabled. Keep non-expert projections at supported precisions.

build() defers INT3-layout FC1 interleaving to initializer serialization, using uncached lazy tensors instead of retaining a second complete FC1 payload. The original checkpoint is still loaded eagerly; this bounds additional interleaving storage, not total checkpoint memory. Direct preprocess_weights() calls retain eager behavior unless an output weight_transforms mapping is provided. Pass that mapping to ModelPackage.apply_weights() with the returned weights to preserve the deferred transformations.

Export support is not a runtime qualification matrix. INT3 QMoE nodes carry a draft-format/runtime-warning marker; they do not invent an approved schema version. Execution of every pair depends on the target ORT version and execution provider. Tests cover bytes, interleaving, binding, and external data save/reload, not INT3 kernel execution, model quality, or speed. Final schema/version review, a complete quantization provenance manifest, and checkpoint-to-ORT parity remain necessary before production use. This is not completion of the runtime plan's P5 acceptance gate.

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/:

Skill	Use when
adding-a-new-model	Adding any new HuggingFace model architecture
reusable-components	Creating or extending components
moe-models	Adding a Mixture-of-Experts model
multimodal-models	Adding a vision-language model
writing-tests	Writing unit or integration tests
writing-rewrite-rules	Adding ONNX graph rewrite rules