Skip to content

Design: Execution Provider aware model building #101

Closed
Closed
Design: Execution Provider aware model building#101

Description

@justinchuby
justinchuby
opened on Apr 2, 2026

Design: Execution Provider aware model building

See #99 for the raw EP analysis of ORT GenAI's model builders, the adversarial design debate, and critical review findings.

Summary

Mobius currently builds EP-agnostic ONNX graphs. To replace ORT GenAI's Python model builders, we need EP-aware optimization that produces graphs tuned for specific execution providers (CPU, CUDA, DML, WebGPU, TRT-RTX). This document proposes a hybrid design: of the 28 EP branch points identified in ORT GenAI, 13 are orthogonal concerns (dtype, quantization, config metadata) handled by existing infrastructure, 12 are addressed by rewrite rules (fuse or lower), and 3 require build-time structural flags. Of the 15 that need new EP-specific code, rewrite rules handle 80% and build-time flags handle 20%.

The design preserves mobius's core invariant: components and models are EP-agnostic. EP logic lives exclusively in the builder's optimization pipeline and, for a small number of structural cases, in the task layer.

1. Custom Logic Survey: All 28 EP Branch Points in ORT GenAI

Analysis of all 20 Python files in onnxruntime-genai/src/python/py/models/. 14 files contain zero EP logic. The 28 branch points concentrate in 6 files: base.py (17), builder.py (5), gptoss.py (3), phi.py (1), qwen.py (1), gemma.py (1).

Category 1: Dtype and Precision (3 branch points)

ID Description EPs Affected Proposed Mechanism
BP-1 set_io_dtype() — selects model I/O dtype based on EP+precision All Orthogonal — already a build() parameter (dtype)
BP-2 I/O binding dtype for KV cache entries cuda, dml Orthogonal — derived from model dtype
BP-3 Mixed-precision attention (FP32 accumulation) cuda, trt-rtx Orthogonal — attribute on op.Attention

Category 2: Attention Op Selection (6 branch points)

ID Description EPs Affected Proposed Mechanism
BP-4 GQA vs MHA selection based on (EP, dtype) support matrix All Rewrite rule — GQAFusionPass gated by (ep, dtype) ∈ GQA_SUPPORT
BP-5 PackedAttention vs standard attention cuda, dml Rewrite rule — PackedAttentionPass gated by support matrix
BP-6 Fused RotaryEmbedding vs separate Q/K rotation All except dml Rewrite rule — SeparateRoPEPass for DML lowering
BP-7 DML requires separate Q/K/V projections (no packed QKV) dml Rewrite rule — UnpackQKVPass for DML lowering
BP-8 Multi-head vs grouped attention cache layout cuda, trt-rtx Rewrite rule — cache layout derived from attention op type
BP-9 Attention mask format (2D vs 4D, causal vs custom) All Rewrite rule — mask reformatting pass

Category 3: Graph Structure (6 branch points)

ID Description EPs Affected Proposed Mechanism
BP-10 If operator → Where substitution dml, webgpu Rewrite rule — DecomposeIfPass
BP-11 Shape operator elimination (constant reshape targets [0,0,N,H]) webgpu Build-time — task emits concrete dims when use_concrete_dims=True
BP-12 int64 → int32 input casting webgpu Build-time — task emits int32 input types
BP-13 Shape→ReduceMax(ReduceSum) for attention mask dims webgpu Rewrite rule — EliminateShapePass
BP-14 TRT-RTX split-If (separate then/else subgraphs) trt-rtx Rewrite rule — SplitIfPass
BP-15 Graph capture constraints (no dynamic control flow) webgpu, trt-rtx Build-time — task avoids Scan/Loop when unsupported

Category 4: Quantization (5 branch points)

ID Description EPs Affected Proposed Mechanism
BP-16 MatMulNBits vs MatMulFpQ4 selection cuda, dml, webgpu Orthogonal — separate quantization pass
BP-17 Block size constraints (TRT-RTX requires 128) trt-rtx Orthogonal — quantization config parameter
BP-18 Zero-point handling (TRT-RTX: no zero_points) trt-rtx Orthogonal — quantization pass attribute
BP-19 QMoE weight layout (CUDA-specific block ordering) cuda Orthogonal — quantization pass + weight preprocessing
BP-20 Quantization axis and group size defaults All Orthogonal — quantization config

Category 5: Operator Selection (3 branch points)

ID Description EPs Affected Proposed Mechanism
BP-21 SkipLayerNorm → Add+LayerNorm decomposition trt-rtx Rewrite rule — DecomposeSkipLayerNormPass
BP-22 Fused MoE → decomposed expert routing dml Rewrite rule — DecomposeMoEPass
BP-23 LayerNorm → primitive ops (Div+Sqrt+ReduceMean) trt-rtx Rewrite rule — DecomposeLayerNormPass

Category 6: Config and Metadata (3 branch points)

ID Description EPs Affected Proposed Mechanism
BP-24 genai_config.json EP-specific fields All Orthogonal — config generation step
BP-25 Cache type naming (sliding_window vs default) trt-rtx Orthogonal — config generation
BP-26 Session options and provider-specific settings All Orthogonal — config generation

Category 7: Forced Overrides and Validation (2 branch points)

ID Description EPs Affected Proposed Mechanism
BP-27 Phi3MoE forced CUDA override (MoE only works on CUDA) cuda Orthogonal — registry validation
BP-28 EP compatibility validation (reject unsupported model+EP combos) All Orthogonal — registry validation

Summary by mechanism

Mechanism Count Branch Points
Rewrite rules (fuse or lower) 12 BP-4 through BP-10, BP-13, BP-14, BP-21 through BP-23
Build-time structural 3 BP-11, BP-12, BP-15
Orthogonal (dtype/quant/config/validation) 13 BP-1 through BP-3, BP-16 through BP-20, BP-24 through BP-28

2. EP Capability Matrix

GQA (GroupQueryAttention) Support

EP FLOAT FLOAT16 BFLOAT16
cpu ✅ ❌ ❌
cuda ❌ ✅ ✅
dml ❌ ✅ ❌
webgpu ✅ ✅ ❌
trt-rtx ❌ ✅ ✅

PackedAttention Support

EP FLOAT FLOAT16 BFLOAT16
cpu ✅ ❌ ❌
cuda ✅ ✅ ✅
dml ✅ ✅ ❌
webgpu ✅ ✅ ❌
trt-rtx ✅ ✅ ✅

Op Support Matrix

Capability cpu cuda dml webgpu trt-rtx
Fused RotaryEmbedding ✅ ✅ ❌ ✅ ✅
Packed QKV ✅ ✅ ❌ ✅ ✅
If operator ✅ ✅ ❌ ❌ ⚠️ split
Shape operator ✅ ✅ ✅ ❌ ✅
int64 types ✅ ✅ ✅ ❌ ✅
SkipLayerNorm ✅ ✅ ✅ ✅ ❌
Fused MoE ✅ ✅ ❌ ❌ ✅
SkipNorm (Add+RMSNorm) ✅ ✅ ✅ ✅ ❌

3. Design Approach Analysis

School A: Generate EP differences in modeling code

EP-aware graph construction — components and models emit different ops at build time based on the target EP.

Pros:

  • Correct by construction — no post-processing step can fail to match
  • Simpler debugging: wrong graph → look at forward(), not rewrite-rule interaction chains
  • Handles the WebGPU 'no Shape ops' case cleanly (concrete dims at build time — the one case that genuinely cannot be post-processed)
  • No pattern fragility — rewrite rules can't 'silently not match'

Cons:

  • Maintenance explosion: O(models × EPs) — with 81 models × 5 EPs this becomes unmanageable
  • Violates component-model separation: _attention.py starts knowing about deployment targets
  • Testing burden multiplies: every model needs an EP-axis in its test matrix
  • Adding EP Bump ruff from 0.15.4 to 0.15.6 in /requirements/lintrunner #6 requires auditing every component that might be affected
  • ORT GenAI tried this approach: their most complex model files (phi.py, qwen.py) have 57-200 lines of EP-specific subgraph builders

Best for: BP-11, BP-12, BP-15 — structural constraints where the required information doesn't exist in the completed graph (~3 of 28 branch points).

School B: Rewrite after generation

Build a single generic graph, then apply EP-specific rewrite rules to transform it.

Pros:

  • O(models + EPs) scaling instead of O(models × EPs) — new EP = one rule file, zero model changes
  • Components and models stay EP-agnostic (mobius's core design principle)
  • 8 existing rewrite rules already written and tested — Phase 1 partially activates existing code
  • Each rule independently testable: input pattern → output pattern, no shared state
  • Generic graph is inspectable as the 'canonical' form, separate from any deployment target

Cons:

  • Pattern fragility: rules can silently fail to match if graph structure varies
  • Cannot handle cases where required information doesn't exist in the graph (WebGPU no-Shape)
  • Rule ordering sensitivity: fuse-before-lower ordering must be enforced
  • Information loss: generic op.Attention doesn't encode whether QKV was packed or separate

Best for: BP-4 through BP-10, BP-13, BP-14, BP-21 through BP-23 — op substitutions and fusions expressible as graph transformations (~12 of 28 branch points).

Recommended: Hybrid approach with con mitigations

School B Con Mitigation in Hybrid Design Residual Risk
Silent rule failure _optimize() returns PassResult with match counts. Fusion assertions: count_ops('GQA') > 0 or count_ops('Attention') == 0. Audit logging opt-in via --verbose-optimization. Partially mitigated — new model variants may need assertion updates
Cannot handle WebGPU no-Shape Explicit School A for this case: task-level use_concrete_dims flag, ~1 implementation in base task class inherited by all tasks. Fully mitigated
Rule ordering sensitivity Architecturally enforced: cleanup → fuse → lower → fold. Fusions never see lowered ops; lowering never sees unfused primitives. Mostly mitigated — ~2-3 intra-stage ordering constraints to document
Intermediate invalid state Generic graph is valid ONNX (runs on CPU). Only 'suboptimal for EP X,' never broken. Pipeline is atomic. Fully mitigated
Information loss (packed QKV) Graph topology encodes this — single MatMul vs three separate MatMuls. Rule checks q.producer.op_type. Partially mitigated — topology inference is more brittle than explicit flag
School A Con Mitigation in Hybrid Design Residual Risk
O(models × EPs) School A used for ~3 branch points in base task class, not 81 models × 5 EPs. Fully mitigated
Components lose EP-agnostic guarantee Hard rule: EP logic banned from components and models. School A touch points live only in task layer, accessed via named boolean flags (use_concrete_dims), never EP strings. Fully mitigated
Testing multiplies 1 task variant needs EP-axis testing (base class), not 81 models. Fully mitigated
Ignores existing rewrite rules School B (primary path) activates all existing rules. Fully mitigated
DML restrictions scatter across code All DML constraints in rule registry. Zero DML logic in components or models. Fully mitigated

The boundary principle

If the information needed to make the decision exists in the completed graph, post-process it (School B). If the information must be injected during construction, do it at build time (School A).

This gives School A exactly 3 branch points (BP-11, BP-12, BP-15) and School B the remaining 12 rule-based points. The 13 orthogonal concerns are handled by existing parameters (dtype, quantization config, genai_config generation, registry validation).

4. Recommended Design

API design

# Public API — single new parameter
pkg = build("Qwen/Qwen2.5-7B-Instruct", execution_provider="cuda")
pkg = build("meta-llama/Llama-3.1-8B", execution_provider="dml", dtype="float16")

# Valid values: "cpu" (default), "cuda", "dml", "webgpu", "trt-rtx"

Architecture

┌─────────────────────────────────────────────────────────────────┐
│  build(model_id, execution_provider="cuda", dtype="float16")    │
│                                                                 │
│  1. Registry lookup → module_class, task, config                │
│  2. Module instantiation (EP-agnostic)                          │
│  3. Task.build(module, config, use_concrete_dims=False)         │
│     └─ Produces generic ONNX graph (valid for any EP)           │
│  4. _optimize(model, ep="cuda", dtype=FLOAT16, role="decoder")  │
│     ├─ Stage 1: Base cleanup passes (current _DEFAULT_PASSES)   │
│     ├─ Stage 2: Fusion passes (GQA, SkipNorm, BiasGelu...)     │
│     │   └─ Only fusions supported by target EP                  │
│     ├─ Stage 3: Lowering passes (decompose unsupported ops)     │
│     │   └─ SeparateRoPE, DecomposeIf, UnpackQKV...              │
│     └─ Stage 4: Constant folding                                │
│  5. Return ModelPackage with optimized models                   │
└─────────────────────────────────────────────────────────────────┘

Per-layer responsibilities

Layer EP Awareness Responsibility
Components (components/) ❌ None Emit standard ONNX ops. Never see EP string.
Models (models/) ❌ None Compose components. Never see EP string.
Tasks (tasks/) ⚠️ Minimal Accept use_concrete_dims: bool for WebGPU structural case. No EP string — builder translates.
Builder (_builder.py) ✅ Full Translates EP string → structural flags for tasks + optimization passes. Owns the _optimize() pipeline.

Rule registry design

# Per-EP fuse/lower lists — never fuse what you'll decompose
_EP_RULE_REGISTRY: dict[str, dict[str, list]] = {
    "cpu": {
        "fuse": [],  # CPU runs generic graph
        "lower": [],
    },
    "cuda": {
        "fuse": [
            ("gqa", GQA_DTYPES_CUDA),      # {FLOAT16, BFLOAT16}
            ("packed_attn", PACKED_CUDA),     # {FLOAT16}
            ("skip_norm", ALL_DTYPES),
            ("skip_layer_norm", ALL_DTYPES),
            ("gelu_fusion", ALL_DTYPES),
        ],
        "lower": [],  # CUDA supports all fused ops
    },
    "dml": {
        "fuse": [
            ("gqa", GQA_DTYPES_DML),        # {FLOAT16}
            ("skip_layer_norm", ALL_DTYPES),
            ("gelu_fusion", ALL_DTYPES),
        ],
        "lower": [
            SeparateRoPEPass,
            UnpackQKVPass,
            DecomposeIfPass,
        ],
    },
    "webgpu": {
        "fuse": [
            ("gqa", GQA_DTYPES_WEBGPU),     # {FLOAT, FLOAT16}
            ("skip_norm", ALL_DTYPES),
            ("gelu_fusion", ALL_DTYPES),
        ],
        "lower": [
            DecomposeIfPass,
            EliminateShapePass,
            CastInt64ToInt32Pass,
        ],
    },
    "trt-rtx": {
        "fuse": [
            ("gqa", GQA_DTYPES_TRT),        # {FLOAT16, BFLOAT16}
            ("gelu_fusion", ALL_DTYPES),
            # NOT skip_layer_norm — TRT-RTX decomposes it
        ],
        "lower": [
            DecomposeSkipLayerNormPass,
            DecomposeLayerNormPass,
            SplitIfPass,
        ],
    },
}

model_role for multi-model tasks

# VisionLanguageTask produces 3 models with different roles
pkg = task.build(module, config)
for name, model in pkg.items():
    role = _MODEL_ROLE_MAP.get(name, "decoder")
    # role: "decoder" for pkg["model"], "vision" for pkg["vision"],
    #        "embedding" for pkg["embedding"]
    _optimize(model, ep=ep, dtype=dtype, model_role=role)

GQA fusion only fires for role="decoder". ViT-specific fusions only fire for role="vision". BiasGelu and normalization fusions fire for all roles.

Base-class concrete dims for WebGPU

# In ModelTask base class
class ModelTask:
    def _create_dims(self, config, use_concrete_dims: bool = False):
        if use_concrete_dims:
            batch_size = ir.Dim(1)  # Concrete
            seq_len = ir.Dim(1)
        else:
            batch_size = ir.SymbolicDim("batch_size")
            seq_len = ir.SymbolicDim("seq_len")
        return batch_size, seq_len

# Builder translates EP → structural flag
use_concrete_dims = (ep == "webgpu")
pkg = task.build(module, config, use_concrete_dims=use_concrete_dims)

All 29 task classes inherit from ModelTask. The concrete-dims logic is implemented once.

ONNX function approach for attention fusion (Phase 2)

Mobius already uses ir.Function for LinearAttention: the component emits a com.microsoft::LinearAttention node whose function body contains standard ONNX ops (Reshape, Scan, MatMul). Runtimes with a native kernel execute it directly; others expand the body as a fallback. This pattern could replace rewrite rules for attention fusion.

A com.microsoft::GroupQueryAttention function would contain RotaryEmbedding + Attention as its body. The Attention component would emit this function node directly — no rewrite rule needed. The runtime decides whether to use its fused GQA kernel or fall back to the standard ops in the body.

Impact on the 12 rewrite-rule branch points:

  • 4 BPs eliminated (BP-4 GQA, BP-5 PackedAttention, BP-6 fused RoPE, BP-8 cache layout) — replaced by ir.Function, zero pattern fragility
  • 8 BPs remain as rewrite rules (BP-7, BP-9, BP-10, BP-13, BP-14, BP-21, BP-22, BP-23) — these are lowering/decomposition rules (If→Where, Shape elimination, SkipLayerNorm decomposition) with no fused-op fallback pattern

Blocker: ORT's GQA kernel requires seqlens_k derived from attention_mask. The function must correctly compute this internally to trigger the native kernel. This needs validation before replacing the rewrite rule.

Timeline: Phase 1 uses rewrite rules (lower risk, activates existing code). Phase 2 prototypes the ir.Function approach for GQA. If validated, it replaces the GQA rewrite rule and eliminates the "silent failure" concern for the highest-impact optimization.

5. Example Implementation Snippet

# src/mobius/_builder.py — proposed changes

from __future__ import annotations

import dataclasses
from collections.abc import Sequence

import onnxscript.ir as ir
from onnxscript.ir import passes as common_passes

from mobius.rewrite_rules import (
    gelu_fusion_rules,
    group_query_attention_rules,
    packed_attention_rules,
    skip_layer_norm_rules,
    skip_norm_rules,
)


# ---------------------------------------------------------------------------
# EP support matrices
# ---------------------------------------------------------------------------

_GQA_SUPPORT: frozenset[tuple[str, ir.DataType]] = frozenset([
    ("cpu", ir.DataType.FLOAT),
    ("cuda", ir.DataType.FLOAT16),
    ("cuda", ir.DataType.BFLOAT16),
    ("dml", ir.DataType.FLOAT16),
    ("webgpu", ir.DataType.FLOAT),
    ("webgpu", ir.DataType.FLOAT16),
    ("trt-rtx", ir.DataType.FLOAT16),
    ("trt-rtx", ir.DataType.BFLOAT16),
])

_PACKED_ATTN_SUPPORT: frozenset[tuple[str, ir.DataType]] = frozenset([
    ("cpu", ir.DataType.FLOAT),
    ("cuda", ir.DataType.FLOAT),
    ("cuda", ir.DataType.FLOAT16),
    ("cuda", ir.DataType.BFLOAT16),
    ("dml", ir.DataType.FLOAT),
    ("dml", ir.DataType.FLOAT16),
    ("webgpu", ir.DataType.FLOAT),
    ("webgpu", ir.DataType.FLOAT16),
    ("trt-rtx", ir.DataType.FLOAT),
    ("trt-rtx", ir.DataType.FLOAT16),
    ("trt-rtx", ir.DataType.BFLOAT16),
])


# ---------------------------------------------------------------------------
# Pass result tracking
# ---------------------------------------------------------------------------

@dataclasses.dataclass
class PassResult:
    """Tracks which optimization passes fired and how many nodes matched."""
    pass_name: str
    nodes_matched: int
    nodes_total: int


# ---------------------------------------------------------------------------
# Optimization pass factory
# ---------------------------------------------------------------------------

def _get_optimization_passes(
    ep: str,
    dtype: ir.DataType,
    model_role: str = "decoder",
) -> tuple[list, list]:
    """Return (fusion_rules, lowering_rules) for the given EP and dtype.

    Returns lists of pattern rewrite rules (for onnxscript.rewriter.rewrite),
    not ir.passes.Pass objects. Fusion rules are applied first, then lowering.
    Principle: never fuse what you'll decompose.
    """
    fuse: list[RewriteRuleSet] = []
    lower: list[RewriteRuleSet] = []

    # --- Fusion passes (only for supported EP+dtype) ---
    if model_role == "decoder":
        if (ep, dtype) in _GQA_SUPPORT:
            fuse.append(group_query_attention_rules())
        if (ep, dtype) in _PACKED_ATTN_SUPPORT:
            fuse.append(packed_attention_rules())

    # Normalization fusions — all roles, all dtypes
    if ep != "trt-rtx":  # TRT-RTX decomposes these
        fuse.append(skip_layer_norm_rules())
        fuse.append(skip_norm_rules())

    # Activation fusions — all roles, all dtypes
    fuse.append(gelu_fusion_rules())

    # --- Lowering passes (decompose unsupported ops) ---
    if ep == "dml":
        # lower.append(separate_rope_rules())       # Phase 2
        # lower.append(unpack_qkv_rules())           # Phase 2
        # lower.append(decompose_if_rules())         # Phase 2
        pass
    elif ep == "webgpu":
        # lower.append(decompose_if_rules())         # Phase 2
        # lower.append(eliminate_shape_rules())       # Phase 2
        # lower.append(cast_int64_to_int32_rules())   # Phase 2
        pass
    elif ep == "trt-rtx":
        # lower.append(decompose_skip_layer_norm_rules())  # Phase 2
        # lower.append(split_if_rules())                    # Phase 2
        pass

    return fuse, lower


# ---------------------------------------------------------------------------
# Optimization pipeline
# ---------------------------------------------------------------------------

def _optimize(
    model: ir.Model,
    ep: str = "cpu",
    dtype: ir.DataType = ir.DataType.FLOAT,
    model_role: str = "decoder",
) -> list[PassResult]:
    """Apply EP-aware optimization passes to a model in-place.

    Pipeline stages:
      1. Base cleanup (identity elimination, CSE, dead code, etc.)
      2. Fusion (promote generic ops to fused EP-supported ops)
      3. Lowering (decompose fused ops the EP doesn't support)
      4. Constant folding

    Returns a list of PassResult for audit logging.
    """
    results: list[PassResult] = []

    # Stage 1: Base cleanup (current _DEFAULT_PASSES)
    base_pass = ir.passes.PassManager(_DEFAULT_PASSES, steps=2)
    base_pass(model)

    # Stage 2+3: Fusion and lowering (each rule set applied separately)
    fuse_rule_sets, lower_rule_sets = _get_optimization_passes(ep, dtype, model_role)

    from onnxscript.rewriter import rewrite
    for rule_set in fuse_rule_sets:
        rewrite(model, pattern_rewrite_rules=rule_set)

    for rule_set in lower_rule_sets:
        rewrite(model, pattern_rewrite_rules=rule_set)

    # Stage 4: Final constant folding
    fold_pass = ir.passes.PassManager([
        common_passes.RemoveUnusedNodesPass(),
        onnxscript.optimizer._constant_folding.FoldConstantsPass(
            shape_inference=False,
            input_size_limit=8192,
            output_size_limit=512 * 512,
        ),
    ])
    fold_pass(model)

    # --- Fusion assertions (hard errors) ---
    if model_role == "decoder" and (ep, dtype) in _GQA_SUPPORT:
        gqa_count = _count_ops(model, "GroupQueryAttention")
        attn_count = _count_ops(model, "Attention")
        if gqa_count == 0 and attn_count > 0:
            raise RuntimeError(
                f"GQA fusion expected for ep={ep}/dtype={dtype} but "
                f"found 0 GroupQueryAttention and {attn_count} Attention "
                f"nodes. This indicates a pattern matching failure in "
                f"the GQA rewrite rule."
            )
        results.append(PassResult("GQAFusion", gqa_count, gqa_count + attn_count))

    return results


def _count_ops(model: ir.Model, op_type: str) -> int:
    """Count nodes of a given op_type in the model graph."""
    return sum(1 for node in model.graph if node.op_type == op_type)


# ---------------------------------------------------------------------------
# Updated build_from_module
# ---------------------------------------------------------------------------

_MODEL_ROLE_MAP = {
    "model": "decoder",
    "decoder": "decoder",
    "vision": "vision",
    "embedding": "embedding",
    "encoder": "encoder",
}


def build_from_module(
    module,
    config,
    task="text-generation",
    execution_provider: str = "cpu",
) -> "ModelPackage":
    """Build ONNX models from a module, with EP-aware optimization."""
    if hasattr(config, "validate"):
        config.validate()

    dtype = getattr(config, "dtype", ir.DataType.FLOAT)
    _cast_module_dtype(module, dtype)

    resolved_task = get_task(task)

    # Translate EP → structural flags for the task
    use_concrete_dims = (execution_provider == "webgpu")
    pkg = resolved_task.build(
        module, config, use_concrete_dims=use_concrete_dims
    )

    # EP-aware optimization per model with role awareness
    for name, model in pkg.items():
        role = _MODEL_ROLE_MAP.get(name, "decoder")
        _optimize(model, ep=execution_provider, dtype=dtype, model_role=role)

    return pkg


# ---------------------------------------------------------------------------
# Updated build (public API)
# ---------------------------------------------------------------------------

def build(
    model_id: str,
    task=None,
    *,
    module_class=None,
    dtype=None,
    load_weights: bool = True,
    trust_remote_code: bool = False,
    execution_provider: str = "cpu",
) -> "ModelPackage":
    """Build ONNX model(s) from a HuggingFace model ID.

    Args:
        model_id: HuggingFace model repository ID.
        task: Task name or ModelTask instance.
        module_class: Custom module class override.
        dtype: Model dtype (e.g., "float16", ir.DataType.FLOAT16).
        load_weights: Whether to load and apply HuggingFace weights.
        trust_remote_code: Trust remote code from HuggingFace.
        execution_provider: Target execution provider. One of:
            "cpu" (default), "cuda", "dml", "webgpu", "trt-rtx".

    Returns:
        ModelPackage containing optimized ONNX model(s).
    """
    # ... (existing registry lookup, config creation, module instantiation)
    # Then:
    return build_from_module(
        module, config, task,
        execution_provider=execution_provider,
    )

Pipeline test pattern

# tests/ep_optimization_test.py

import pytest
from mobius import build_from_module
from mobius._builder import _count_ops

_EP_FUSION_EXPECTATIONS = [
    # (model_type, ep, dtype, expected_ops)
    ("llama", "cuda", FLOAT16, {"GroupQueryAttention": ">0"}),
    ("llama", "dml", FLOAT16, {"GroupQueryAttention": ">0"}),
    ("llama", "cpu", FLOAT, {}),  # No fusions expected
    ("qwen2", "cuda", FLOAT16, {"GroupQueryAttention": ">0"}),
    ("bert", "cuda", FLOAT, {"GroupQueryAttention": "==0"}),  # MHA, not GQA
    ("qwen2", "webgpu", FLOAT16, {"GroupQueryAttention": ">0", "Shape": "==0"}),
]


def _check_op_constraint(actual: int, constraint: str) -> bool:
    """Check an op count against a constraint like '>0' or '==0'."""
    op_str, val_str = constraint.split(None, 1)
    val = int(val_str)
    if op_str == ">":
        return actual > val
    if op_str == "==":
        return actual == val
    if op_str == ">=":
        return actual >= val
    raise ValueError(f"Unknown constraint operator: {op_str}")


@pytest.mark.parametrize("model_type,ep,dtype,expectations", _EP_FUSION_EXPECTATIONS)
def test_ep_produces_expected_ops(model_type, ep, dtype, expectations):
    """Verify EP-specific optimization produces expected fused ops."""
    config = _get_tiny_config(model_type)
    module = _create_module(model_type, config)
    pkg = build_from_module(module, config, execution_provider=ep)

    model = pkg["model"]
    for op_type, constraint in expectations.items():
        actual = _count_ops(model, op_type)
        assert _check_op_constraint(actual, constraint), (
            f"{model_type}/{ep}/{dtype}: {op_type} expected {constraint}, "
            f"got {actual}"
        )


def test_ep_does_not_fuse_vision_encoder_with_gqa():
    """GQA should NOT be applied to vision encoder models."""
    config = _get_tiny_vl_config("llava")
    module = _create_module("llava", config)
    pkg = build_from_module(module, config, execution_provider="cuda")

    vision_model = pkg["vision"]
    assert _count_ops(vision_model, "GroupQueryAttention") == 0, (
        "Vision encoder should not have GQA fusion"
    )

6. Implementation Phases

Phase 1: Wire EP parameter + activate existing rules

Deliverables:

  • Add execution_provider: str parameter to build() and build_from_module()
  • Implement _get_optimization_passes(ep, dtype, model_role) factory
  • Implement _EP_RULE_REGISTRY with per-EP fuse/lower lists (fuse only in Phase 1; lower stubs)
  • Wire _optimize(model, ep, dtype, model_role) with 4-stage pipeline
  • Add model_role parameter derived from ModelPackage keys
  • Fusion assertion hard errors (GQA, SkipNorm) in _optimize()
  • Pipeline tests: tests/ep_optimization_test.py with expectations table
  • Base task class: _create_dims(use_concrete_dims) helper

Phase 1 must-haves (from critical review):

  1. ✅ Audit logging + fusion assertions — not deferred
  2. ✅ Pipeline tests ship with execution_provider parameter
  3. ✅ Per-EP fuse/lower rule registry (not just inclusions)
  4. ✅ model_role parameter on _optimize()
  5. ✅ Base-class concrete-dims helper for WebGPU

Files changed: _builder.py (primary), tasks/_base.py (concrete dims helper)
Files added: tests/ep_optimization_test.py
Existing rules activated: group_query_attention_rules, skip_norm_rules, skip_layer_norm_rules, gelu_fusion_rules

Phase 2: EP-conditional rule registry + missing lowering rules

Deliverables:

  • SeparateRoPEPass — decompose fused RoPE for DML
  • UnpackQKVPass — split packed QKV projections for DML
  • DecomposeIfPass — If → Where for DML/WebGPU
  • EliminateShapePass — Shape → constant for WebGPU
  • CastInt64ToInt32Pass — int64 → int32 for WebGPU
  • DecomposeSkipLayerNormPass — SkipLayerNorm → Add+LayerNorm for TRT-RTX
  • Packed attention activation for CUDA
  • Rule interaction test suite (verify GQA + SkipNorm + BiasGelu compose correctly)
  • Integration tests: parity check (generic graph logits ≈ EP graph logits)
  • Promote EpCapabilities to internal implementation if conditional constraints emerge

Phase 3: Advanced EP features

Deliverables:

  • TRT-RTX: SplitIfPass, DecomposeLayerNormPass, sliding window cache naming
  • Quantization integration: EP-specific quantization configs (block size, zero points)
  • genai_config.json generation with EP-specific fields
  • Registry validation: reject unsupported (model_type, EP) combinations
  • Golden graph snapshots for regression detection (nightly CI)
  • Formal EP dialect lowering (fuse-up/lower-down as declarative op-set membership) — if EP set expands significantly

7. Coverage and Limitations

Honest coverage numbers

Optimization Coverage Notes
GQA fusion ~60% of attention models Models with kv_num_heads < q_num_heads using standard op.Attention. Llama, Qwen, Mistral, Phi, Gemma2, etc.
SkipNorm (Add+RMSNorm) ~52% of transformers Models using RMSNorm with residual connections
SkipLayerNorm (Add+LayerNorm) ~49% of transformers Models using LayerNorm with residual connections
BiasGelu fusion ~15% of transformers Only models with approximate='tanh' GeLU activation (GPT-2-like). Most modern LLMs use SiLU with gated MLP.
Gelu fusion (decomposed→native) ~25% of transformers Models that emit decomposed GeLU (Div→Erf→Add→Mul)
At least one EP optimization ~85% of transformers Combined coverage of normalization + activation + attention fusions

Known gaps

Gap Models Affected Mitigation Timeline
Custom attention (MLA, Lightning, etc.) DeepSeek, Qwen3.5 linear attention, ~23 model files Phase 3 — per-architecture fusion rules
ALiBi attention (no RoPE) Falcon, MPT, BLOOM Phase 2 — ALiBi-aware GQA variant
Encoder-only attention (no KV cache) BERT, RoBERTa, DeBERTa Low priority — encoder attention already runs as standard MHA; GQA not applicable
Cross-attention in encoder-decoder T5, BART, Whisper Phase 2 — cross-attention GQA rule
QK-norm with PackedQKV Qwen3, models with pre-attention Q/K normalization Phase 2 — QK-norm-aware packed attention variant
SSM models (Mamba, RWKV) ~5 models N/A — no attention ops to fuse; EP optimization limited to normalization and activation fusions
Diffusion models ~15 models Separate concern — diffusion uses different EP strategies (UNet/DiT-specific)

What "no EP optimization" means in practice

A model that gets zero EP-specific fusions still produces a valid, correct ONNX graph. It runs on any EP — just without hardware-specific operator fusions. The model is functionally correct; it's a performance optimization gap, not a correctness gap. The generic graph benefits from the base cleanup passes (identity elimination, CSE, dead code removal, constant folding) which are EP-agnostic and apply to all models.

8. New Model Onboarding: Model #81

Step-by-step for a standard LLM

Step What Where Lines
1 Create model file src/mobius/models/newmodel.py ~200-300
2 Implement __init__, forward, preprocess_weights Same file (compose existing components) Included above
3 Export from models/__init__.py 1 line from .newmodel import NewModelCausalLM
4 Register in _registry.py 1 line reg.register("newmodel", NewModelCausalLM)
5 Add tiny config to tests/_test_configs.py ~5 lines ("newmodel", {"hidden_size": 64, ...}, True)
6 Run tests CLI pytest tests/build_graph_test.py -k "newmodel"
7 EP support Nothing 0 lines

The model uses standard Attention component → emits op.Attention with q_num_heads/kv_num_heads → GQA fusion rule matches automatically for CUDA/DML/WebGPU. The contributor never writes if ep == ....

When EP cost is non-zero

If model #81 uses a genuinely new attention pattern not covered by existing rules (e.g., a novel sparse attention variant), one new rewrite rule is needed. That rule is written once in rewrite_rules/, tested independently, and benefits all future models with the same pattern. Cost: ~50-100 lines for the rule + test, written by the EP team, not the model contributor.

Comparison with ORT GenAI

Dimension ORT GenAI Mobius (proposed)
Model file ~300-500 lines (uses base.py builder) ~200-300 lines (composes components)
EP-specific code per model 20-50 lines (EP dispatch in model file + base.py) 0 lines
Adding EP #6 Audit 6+ model files with EP logic Add 1 rule file, 0 model changes
Testing per model Each model × EP tested via builder.py orchestration Model tested once; pipeline test validates EP optimization for all models
Lines-per-model (EP included) ~340 ~25 (registration + test config only)
Total EP code for 81 models ~4800 lines across base.py + model files ~500 lines in rule registry + rules (shared across all models)

9. Regression Testing Strategy

Test layers

Layer Name What It Catches When It Runs Cost
L1 Graph build × EP Graph construction crashes for EP-specific builds Every PR +~30s (90 models × 2 EPs, <0.5s each)
L1 Rewrite rule unit tests Individual rule breakage Every PR Existing (~5s)
L2 Fusion coverage Missed fusions (performance regression) — the silent-failure mitigation Every PR +~20s
L3 EP parity Numerical accuracy: generic_logits ≈ ep_logits (atol=1e-5) PR (affected models) +~5min
L4 Golden × EP End-to-end regression with real weights on GPU Nightly + release +~10min
L5 Golden graph snapshots Unexpected op histogram changes (diff alert, not blocking) Nightly +~2min

L2 fusion coverage test design

# Concrete assertions — the key defense against silent rule failures

def test_cuda_fp16_has_gqa(model_type, tiny_config):
    """Models with GQA-compatible attention MUST produce GroupQueryAttention on CUDA+FP16."""
    pkg = build_from_module(module, config, execution_provider="cuda")
    gqa = count_ops(pkg["model"], "GroupQueryAttention")
    attn = count_ops(pkg["model"], "Attention")
    # If any Attention nodes remain unfused, that's a regression
    assert gqa > 0 or attn == 0, f"GQA expected but got {gqa} GQA, {attn} Attention"


def test_webgpu_has_no_shape(model_type, tiny_config):
    """WebGPU graphs must not contain Shape operators."""
    pkg = build_from_module(module, config, execution_provider="webgpu")
    shapes = count_ops(pkg["model"], "Shape")
    assert shapes == 0, f"WebGPU graph has {shapes} Shape ops"


def test_dml_has_no_if(model_type, tiny_config):
    """DML graphs must not contain If operators."""
    pkg = build_from_module(module, config, execution_provider="dml")
    ifs = count_ops(pkg["model"], "If")
    assert ifs == 0, f"DML graph has {ifs} If ops"

CI integration

  • L1 + L2: Run in existing main.yml L1 Smoke job (10min budget, <1min additional)
  • L3: Run in main.yml L3 Synthetic Parity job with affected-model detection
  • L4: Run in gpu_tests.yml on A10 self-hosted runner
  • L5: Run in nightly_l2.yml

10. Open Questions

Confirmed decisions

Decision Resolution
Primary mechanism Rewrite rules (School B) for 12/28 branch points
Secondary mechanism Build-time structural flags for 3/28 branch points
Orthogonal concerns Existing parameters for 13/28 branch points
Component/model EP awareness Banned — EP logic in builder + task layers only
API surface execution_provider: str parameter on build()
Rule ordering Cleanup → Fuse → Lower → Fold (architecturally enforced)

Open questions for human team

  1. Replacing or coexisting? Is mobius intended to fully replace ORT GenAI's Python model builders, or coexist alongside them? The design assumes replacement — if coexisting, porting ORT GenAI's EP logic directly would have lower risk.

  2. Phase 1 scope. Should Phase 1 include all 5 EPs or start with CPU + CUDA only? Starting with 2 EPs simplifies the initial rule registry and testing matrix.

  3. Quantization integration timing. Quantization is listed as orthogonal (Phase 3), but some quantization decisions interact with EP (e.g., TRT-RTX block_size=128). Should quantization be a Phase 2 deliverable?

  4. EpCapabilities promotion timing. The flat (ep, dtype) matrix is sufficient for Phase 1. If a conditional constraint surfaces (e.g., 'GQA works only if RoPE is also fused'), we promote to EpCapabilities dataclass. Should we preemptively build the dataclass in Phase 2 regardless?

  5. Diffusion model EP support. The 28 branch points are from text-generation models. Diffusion models (UNet, DiT, Flux) have different EP concerns (e.g., FP16 accumulation, attention memory optimization). Should diffusion EP support be a separate design document?

  6. Rule coverage target. Is 85% coverage (at least one optimization) acceptable for Phase 1-2, with custom attention architectures deferred to Phase 3? Or is higher coverage required before shipping EP support?

  7. GQA as ONNX function: Could GroupQueryAttention be implemented as an ir.Function (like LinearAttention) instead of via rewrite rules? The function body would contain the fallback implementation (RoPE + standard attention), while runtimes with native GQA kernels skip the body. This would eliminate pattern fragility entirely. Blocker: the function must correctly derive seqlens_k from attention_mask to trigger ORT's native kernel. See also: op.Attention already uses this pattern for standard attention.

See #99 for the raw EP analysis, adversarial debate, and critical review.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions