EP Support: Hybrid rewrite-rule + build-time architecture (28 branch points, Phase 1-3 plan)

[justinchuby]

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.

---

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 local function? Instead of rewrite rules fusing standard ops into GQA, could mobius emit GroupQueryAttention directly as a com.microsoft ONNX local function (like LinearAttention)? The function body would be standard ops (RotaryEmbedding + Attention) serving as automatic fallback. This eliminates pattern fragility entirely. Key open question: does the seqlens_k computation (derived from attention_mask) belong inside the function body or remain a graph-level responsibility? If solvable, GQA-as-function could replace the GQA rewrite rule in Phase 2.

---

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

[justinchuby]

ONNX Function Approach for Attention Fusion (Addendum)

An alternative to rewrite rules for attention fusion: emit com.microsoft::GroupQueryAttention as an ir.Function whose body contains the fallback implementation (RotaryEmbedding + standard Attention). This follows the same pattern as LinearAttention.

How it works:

• Component emits op.GroupQueryAttention(q, k, v, cos, sin, past_key, past_value) directly
• The ir.Function body contains RotaryEmbedding + Attention as standard ops
• Runtime with native GQA kernel: executes fused kernel, ignores body
• Runtime without native kernel: expands function body, executes standard ops

Impact on the design:

• Eliminates pattern fragility for 4 of 12 rewrite BPs (BP-4, BP-5, BP-6, BP-8)
• Silent failure becomes impossible for attention fusion — function body is guaranteed fallback
• Remaining 8 lowering rules (If→Where, Shape elimination, SkipLayerNorm decomposition, etc.) still need rewrite rules

Blocker: The function must correctly derive seqlens_k from attention_mask and match ORT native kernel input contract. Requires validation.

Timeline: Phase 2 deliverable. Phase 1 uses rewrite rules as the lower-risk path.

[justinchuby]

Debuggability: Optimization Tracing

Concern raised: rewrite rules are harder to debug than inline EP logic.

Response

Rewrite rules ARE harder to debug when they go wrong — with inline logic you set a breakpoint, with rules you need to trace pattern matching. This is a real cost.

However, it is a one-time cost per rule (written once, debugged once), not an ongoing cost per model (80+ models × 5 EPs = 400 code paths to debug with inline logic).

Mitigation: trace_optimization mode

Added trace_optimization: bool = False to build() and build_from_module(). When enabled, produces a step-by-step diagnostic log:

[EP Trace] Target: cuda, dtype: float16, role: decoder
[EP Trace] Stage 1: Cleanup (9 passes)
[EP Trace]   ConstantFolding: 12 nodes folded
[EP Trace] Stage 2: Fusion (3 rules)
[EP Trace]   GQAFusion: 24/24 Attention nodes → GroupQueryAttention
[EP Trace]   SkipLayerNorm: 24/24 Add+LayerNorm → SkipLayerNormalization
[EP Trace]   BiasGelu: 0 matches (model uses SiLU, not GELU)
[EP Trace] Stage 3: Lowering (0 rules for cuda)
[EP Trace] Stage 4: Constant folding
[EP Trace] Summary:
[EP Trace]   Rule           | Matched | +Nodes | -Nodes
[EP Trace]   GQAFusion      |      24 |     24 |     24
[EP Trace]   SkipLayerNorm  |      24 |     24 |     48
[EP Trace]   BiasGelu       |       0 |      0 |      0

What the trace shows

• Which rules ran per stage (cleanup → fuse → lower → fold)
• How many nodes each rule matched vs total candidates
• Nodes added/removed per rule (net graph change)
• Rules that matched 0 nodes (potential pattern fragility)

Usage

pkg = build("Qwen/Qwen2.5-7B", execution_provider="cuda", trace_optimization=True)

This addresses the debuggability concern: when a rewrite rule misbehaves, the trace shows exactly which rule, which stage, and how many nodes were affected. Default off (no performance impact), opt-in for diagnosis.

[justinchuby]

Updated Design Direction: ONNX Functions + EP Capabilities

Based on human team feedback, the design direction is refined:

Primary mechanism: ONNX functions with decomposition fallback

For fusions (GQA, PackedAttention, fused RoPE, etc.), use ONNX ir.Function as the primary mechanism:

• Components emit high-level function ops (e.g., com.microsoft::GroupQueryAttention)
• The function body contains the decomposed/fallback implementation (RotaryEmbedding + standard Attention)
• Runtime with native kernel: executes fused kernel, ignores body
• Runtime without native kernel: expands function body, executes standard ops
• Components remain completely EP-agnostic — they always emit the same function op

This replaces rewrite rules for fusion-type EP optimizations. No GQA rewrite rule needed — the function IS both the fused and unfused form.

EP capability registry drives decomposition

Register capabilities per EP:

_EP_CAPABILITIES = {
    "cuda": {"GroupQueryAttention": {FLOAT16, BFLOAT16}, "PackedAttention": {FLOAT16, BFLOAT16, FLOAT}, ...},
    "dml": {"GroupQueryAttention": {FLOAT16}, ...},  # no fused RoPE variant
    "cpu": {"GroupQueryAttention": {FLOAT}, ...},
}

The builder queries capabilities to determine if function bodies need expansion. If an EP does NOT support a function natively, the body is expanded to standard ops during optimization. Components never see this logic.

Rewrite rules remain for graph-structural changes

For non-fusion EP differences that are graph transforms (not op-level):

• If → Where decomposition (DML, WebGPU)
• Shape elimination (WebGPU graph capture)
• int64 → int32 casting (WebGPU)
• SkipLayerNorm decomposition (TRT-RTX)

These dont fit the function pattern because they operate on graph structure, not individual ops.

Summary of mechanism split

EP difference type	Mechanism	Why
Attention fusion (GQA, packed QKV, fused RoPE)	ONNX ir.Function	Function body is the fallback; runtime picks native kernel
Op-level fusion (BiasGelu, SkipLayerNorm)	ONNX ir.Function	Same pattern — function body is decomposed form
Graph-structural (If elimination, Shape removal)	Rewrite rules	Not an op substitution — graph topology change
Type constraints (int64→int32)	Rewrite rules	Global graph transform, not per-op
Quantization	Separate pass	Orthogonal concern
Config generation	Metadata	Not a graph change

Key principle

Components never know about EPs. They emit ONNX functions. The EP capability registry and the runtime together determine the execution path. Rewrite rules handle the remaining graph-structural changes that cannot be expressed as functions.

[justinchuby]

DML Insight: ir.Function Eliminates Two Rules

Migrating GQA to an ir.Function eliminates both the fusion rule AND the DML separate-RoPE decomposition rule:

• GQA function body includes RoPE (RotaryEmbedding + standard Attention)
• CUDA/WebGPU/TRT-RTX: Runtime has native GQA+RoPE kernel → uses it directly, ignores body
• DML: No native GQA+RoPE kernel → expands function body → gets RotaryEmbedding + Attention (which IS the separate-RoPE pattern DML needs)
• CPU: No native kernel → expands body → runs standard ops

Net result: the function body expansion IS the DML-compatible pattern. No need for a separate _separate_rope.py rewrite rule. Two fewer rules for the price of one function body.

Remaining rewrite rules after full ir.Function migration:

• If → Where (DML, WebGPU) — graph-structural
• Shape elimination (WebGPU) — graph-structural
• int64 → int32 (WebGPU) — type constraint
• SkipLayerNorm decomposition (TRT-RTX) — function body expansion
• SimplifiedLayerNorm decomposition (TRT-RTX) — function body expansion

The last two (SkipLayerNorm, SimplifiedLayerNorm for TRT-RTX) are effectively "function body expansions" and could also become ir.Functions long-term.

[justinchuby]

Design Invariant: When to Use Functions vs Rewrite Rules

Converged principle for future contributors:

• Use ir.Function when an op has a local, self-contained fallback (function body). The component emits the function op; the runtime uses a native kernel or expands the body. Examples: GQA, PackedAttention, BiasGelu, SkipLayerNorm fusion.

• Use rewrite rules when the transform is global — spanning multiple ops, changing control flow topology, or applying a cross-graph constraint. Examples: If→Where (DML/WebGPU), Shape elimination (WebGPU), int64→int32 casting (WebGPU), SkipLayerNorm decomposition (TRT-RTX).

• Use neither (config/task parameters) when it is a build-time configuration decision. Examples: concrete dims for WebGPU, static cache, model-specific overrides.

Pipeline Evolution

Phase 1 (current scaffolding):
  Stage 1: Cleanup → Stage 2: Fuse (rewrite rules) → Stage 3: Lower (rewrite rules) → Stage 4: Fold

Phase 2+ (ir.Function migration):
  Stage 1: Cleanup → Stage 2: Expand unsupported functions → Stage 3: Lower (rewrite rules) → Stage 4: Fold

Stage 2 evolves from "apply fusion rewrite rules" to "expand function bodies for EPs without native kernels." Stage 3 operates on the expanded graph — any If/Shape/int64 nodes introduced by function expansion are handled automatically.

Rule Count Reduction

Phase	Fusion rules	Graph-structural rules	Total
Phase 1 (scaffolding)	4	4	8
Phase 2+ (ir.Function)	0	4	4

56% reduction. The 4 remaining rules share a common property: they are global transforms that cannot be localized to a single op or function boundary.

[justinchuby]

Final Design: EpCapabilities + Function-Based Portability

Support Strategy

Mobius uses a three-tier support strategy for execution providers:

Tier 1: ONNX Functions (portability by default)

• Components emit custom ops (e.g., com.microsoft::GroupQueryAttention) with model-local function bodies
• The function body contains the decomposed standard-ONNX implementation
• Any ONNX-conformant runtime can execute the model by expanding function bodies
• Custom ops with function bodies ARE portable ONNX — the function body IS the specification
• No decomposition happens by default — the graph keeps function ops as-is

Tier 2: EP-Specific Optimizations (optional, declared via EpCapabilities)

• When targeting a specific EP (e.g., execution_provider="cuda"), the builder queries EpCapabilities to determine:
  • Which function ops have native kernel support (keep the function call, runtime uses fast kernel)
  • Which standard ops to fuse into vendor-specific function ops (e.g., RotaryEmbedding + Attention → GQA)
• Optimizations are additive — they improve performance but the model works without them

Tier 3: EP-Specific Constraints (required, graph-structural)

• Some EPs have hard constraints that require graph transformation:
  • DML: no If operator → rewrite to Where
  • WebGPU: no Shape ops in graph capture mode → rewrite to ReduceMax
  • WebGPU: no int64 → cast to int32
  • TRT-RTX: no SkipLayerNorm → decompose to primitives
• These are NOT optional — the model fails without them on the target EP
• Implemented as rewrite rules that run after function expansion

EpCapabilities Class

@dataclass(frozen=True)
class EpCapabilities:
    """Declares what an execution provider supports.
    Adding a new EP = adding one entry to _EP_REGISTRY.
    """
    name: str
    gqa_dtypes: frozenset[DataType]      # GQA fusion for these dtypes
    packed_attn_dtypes: frozenset[DataType]  # PackedAttention for these dtypes
    supports_fused_rope: bool = True     # Can handle RoPE fused in attention
    supports_if: bool = True             # If operator supported
    supports_shape: bool = True          # Shape operator supported
    supports_int64: bool = True          # INT64 types supported
    supports_skip_layer_norm: bool = True  # com.microsoft.SkipLayerNorm
    supports_simplified_layer_norm: bool = True
    supports_fused_moe: bool = True
    ...

Default EP: Pure ONNX Conformant

The default execution provider ("default") represents any ONNX-conformant runtime:

• Custom ops with function bodies are allowed — the function body is the portable fallback
• No decomposition — keep all ops as-is
• No vendor-specific fusions — no GQA optimization, no packed attention
• All standard ONNX features supported (If, Shape, int64)

build(model_id) → portable ONNX that runs on any conformant runtime
build(model_id, execution_provider="cuda") → CUDA-optimized with fused ops

Adding a New EP

To add EP #6, a contributor:

1. Adds ONE entry to _EP_REGISTRY declaring capabilities
2. Writes any new rewrite rules for hard constraints
3. No changes to components, models, or tasks

Pipeline

build(model_id, execution_provider="cuda")
  → Components emit ONNX ops + function ops (EP-agnostic)
  → Task builds graph (EP-agnostic)
  → _optimize(model, caps, dtype, model_role):
      Stage 1: Cleanup (dead nodes, constant fold)
      Stage 2: Expand unsupported functions (per EpCapabilities)
      Stage 3: Graph-structural lowering (If→Where, Shape→ReduceMax, int64→int32)
      Stage 4: Final constant folding
  → ModelPackage

Design Invariant

• ir.Function: when an op has a local, self-contained fallback (function body)
• Rewrite rules: when the transform is global (control flow, type casting, cross-graph)
• Neither: when it is a build-time config decision (task parameters)
• Components never know about EPs