From PyTorch and ONNX to an Accelerator Model

Topics We'll Cover

• Why a Simulator Needs a Front End
• One Trace Format, Four Front Ends
• Real Models, No Weights: the Meta Device
• Front End 1: the Dispatch Trace
• Meta or Fake Tensors: What the Trace Means
• Front Ends 2 and 3: torch.export and a torch.compile Backend
• Front End 4: ONNX Without the Weights
• Four Front Ends, One Answer
• Checking the Trace Against the Closed Form
• Costing the Trace
• Interactive: Operator Coverage, Counted Three Ways
• Operator Coverage as an Engineering Metric
• What to Take Away

Why a Simulator Needs a Front End

An accelerator simulator is only as useful as the workloads it can run. Hand-written workloads ("a 4096 × 4096 matmul, then a softmax") answer the questions the architect already thought of. Real users arrive with a PyTorch model or an ONNX file. A front end turns either into the operator stream the cost model consumes, so the simulator sees what the framework would actually execute.

LLM Inference Simulators 09 surveys the integration routes. This deck builds four of them in Torch_Sim_Frontend, runs them on published model configurations (Llama-3-8B and -70B, Mistral-7B, Qwen2.5-0.5B, GPT-2) on a CPU-only machine without a single weight, and checks them against each other, against closed forms and against PyTorch's own FLOP counter.

One Trace Format, Four Front Ends

Every front end produces the same record per operator. Tensor identities let a cost model follow data from producer to consumer, which is what an offload model needs to charge a transfer. The weight flag separates parameter traffic from activations, and is propagated through views (W.t() before mm is still a weight).

class Op:
    """One operator call. ``name`` is canonical: ``aten.mm`` or ``onnx.MatMul``."""

    name: str
    inputs: list[TensorMeta]
    outputs: list[TensorMeta]
    attrs: dict = field(default_factory=dict)
    category: str = "unknown"
    flops: int = 0
    bytes_read: int = 0
    bytes_written: int = 0
    weight_bytes: int = 0      # the part of bytes_read that is parameters

Real Models, No Weights: the Meta Device

A tensor on the meta device has a shape, a dtype and a stride, but no storage. Building a model inside torch.device("meta") gives every parameter that form, so a 70-billion-parameter model is instantiated in under a second and costs no memory. The configurations are the architectural fields of the published Hugging Face config.json files, bundled with their source URLs.

def build(cfg, device: str = "meta", dtype: torch.dtype = torch.bfloat16, seed: int = 0) -> torch.nn.Module:
    """``AutoModelForCausalLM.from_config`` on ``device``; ``cfg`` is a config or a registry name."""
    from transformers import AutoModelForCausalLM

    if isinstance(cfg, str):
        cfg = config(cfg)
    torch.manual_seed(seed)
    with torch.device(device):
        model = AutoModelForCausalLM.from_config(cfg, dtype=dtype)
    return model.eval()

Source: examples/results.md in Torch_Sim_Frontend

Every parameter was checked to be a meta tensor. The whole results script, which traces all of these and more, peaks at 647 MB of resident memory.

Front End 1: the Dispatch Trace

Every PyTorch operator passes through the dispatcher. A TorchDispatchMode sees each call after autograd and Python-level decompositions, with the real argument tensors, so it records exactly what would execute. The mode below also gives each tensor an identity and marks views of parameters as weights:

def __torch_dispatch__(self, func, types, args=(), kwargs=None):
    kwargs = kwargs or {}
    out = func(*args, **kwargs)
    name = str(func.overloadpacket)                # e.g. "aten.mm"
    flat_in, _ = tree_flatten((args, kwargs))
    flat_out, _ = tree_flatten(out)
    ins = [self._meta(t) for t in flat_in if isinstance(t, torch.Tensor)]
    # A view of a parameter (W.t() before mm) is still weight traffic.
    is_view = name.split(".", 1)[1] in VIEW_PACKETS
    param_view = is_view and any(t.param for t in ins)
    outs = [self._meta(t, param=param_view) for t in flat_out if isinstance(t, torch.Tensor)]
    attrs = {k: v for k, v in kwargs.items() if isinstance(v, (int, float, bool, str))}
    self.ops.append(apply_aten_rule(Op(name, ins, outs, attrs)))
    return out

• A gotcha with real models. With the KV cache off, transformers 5 checks for packed sequences by calling .item() on a computed tensor, and the meta device cannot produce a value: RuntimeError: Tensor.item() cannot be called on meta tensors. With the cache on, that check is skipped, so meta prefill traces keep the cache on.
• Decode steps are traced by running the prompt first (outside the mode, at no cost on the meta device) and then tracing one step with the cache it produced.

Meta or Fake Tensors: What the Trace Means

The dispatch trace shows what PyTorch would run on the device the tensors claim to be on. On the meta device, scaled_dot_product_attention has no fused kernel and decomposes to bmm, _safe_softmax and a mask over the whole 2,048 × 2,048 score matrix per head. Under a FakeTensorMode the tensors claim to be on the CPU, and attention is one fused flash-attention operator per layer.

Source: examples/results.md in Torch_Sim_Frontend

• The arithmetic is identical (QK^T and AV FLOPs match exactly). The bytes are not: in decode the meta trace moves 25.02 GB per step against 15.84 GB. 4.30 GB of that is transformers' repeat_kv copying 8 KV heads into 32 for the math path.
• transformers changes path when it detects tracing. In eager mode on real data it inspects the mask with .item() and uses SDPA's native grouped-query attention. Under fake tensors, torch.export or torch.compile, it builds an explicit mask and repeats the KV heads. A test pins this down.
• On plain PyTorch (no data-dependent Python), fake tensors record exactly the operators and shapes of a real run. That is also tested.

The lesson: a trace is evidence about one code path on one device. Say which.

Front Ends 2 and 3: torch.export and a torch.compile Backend

torch.export traces the whole model into one FX graph with fake-tensor metadata on every node, and run_decompositions() lowers it to Core ATen (InfSim 09, slide 03). It works on the meta device, with one restriction: transformers' DynamicCache is not a registered pytree, so prefill is exported with use_cache=False. That is safe because transformers recognises the fake tensors and takes its tracing path.

A torch.compile backend is any callable that takes a GraphModule and returns a callable (InfSim 09, slide 04). Dynamo's graphs are at the torch level (linear, scaled_dot_product_attention), so this backend hands them to AOTAutograd for an ATen graph, walks it with the same rules as the other front ends, and returns the graph's own forward. The user's program runs unchanged and gives identical outputs (tested).

class SimBackend:
    """``torch.compile(model, backend=SimBackend())``: records an operator trace per captured graph."""

    def __init__(self):
        self.graphs: list[list[Op]] = []

    def __call__(self, gm: torch.fx.GraphModule, example_inputs):
        # Dynamo names lifted parameters after their module path ("..._parameters_weight_");
        # AOTAutograd keeps the input order, so the i-th placeholder below is the same tensor.
        is_param = [n.op == "placeholder" and "_parameters_" in n.name for n in gm.graph.nodes
                    if n.op == "placeholder"]

        def fw_compiler(fgm: torch.fx.GraphModule, _inputs):
            ph = [n for n in fgm.graph.nodes if n.op == "placeholder"]
            params = {n.name for n, p in zip(ph, is_param, strict=True) if p}
            self.graphs.append(walk_graph(fgm.graph, params))
            return fgm.forward

        return aot_autograd(fw_compiler=fw_compiler)(gm, example_inputs)

Lifted parameters keep Dynamo's module-path names (..._parameters_weight_), and AOTAutograd keeps the input order, which is how the walk knows which graph inputs are weights.

Front End 4: ONNX Without the Weights

ONNX is the framework-neutral route: whatever produced the model, the cost model sees standard operators with versioned semantics. Three things were needed to make it work on a weightless model:

prog = torch.onnx.export(model, args, kwargs=kwargs or {}, dynamo=True, optimize=False, verbose=False)
if weights:
    prog.save(str(path))
else:
    prog.save(str(path), include_initializers=False, keep_initializers_as_inputs=True)

• Keep weights as typed inputs. include_initializers=False, keep_initializers_as_inputs=True drops the data but keeps every weight's shape. Llama-3-8B exports to an 11.3 MB file. The exporter's optimiser is off: it needs real values, and it renames the weights it pre-transposes.
• Propagate shape values. infer_shapes(model, data_prop=True). Without data_prop, Llama's reshape targets are computed, not constant, and every attention MatMul is left without dimensions.
• Fold constants, as a runtime would. A node whose inputs are all constants (a Transpose of a weight, the rotary frequencies) costs nothing per inference. A layout op on a weight is still the weight.

ONNX does not tell parameters from buffers (both are initializers), so the export returns the parameter names, including tied aliases (GPT-2's lm_head.weight is transformer.wte.weight).

Four Front Ends, One Answer

Llama-3-8B, one 2,048-token prefill, on the meta device unless stated:

Source: examples/results.md in Torch_Sim_Frontend

• Matmul FLOPs agree exactly across all four. ONNX is lower by exactly 262,144: it constant-folds the rotary-frequency matmul. Weights read agree exactly.
• Bytes moved differ, from 70 to 334 GB. ATen views are free, ONNX Transpose, Expand and Slice materialise, and each decomposition upcasts and copies in different places. The arithmetic belongs to the model; the traffic belongs to the decomposition, and a cost model must say which it assumes.
• The tests check the same agreement on small Llama and GPT-2 models with real data on a CPU, where dispatch and compile keep the fused attention kernel while export and ONNX decompose it.

Checking the Trace Against the Closed Form

Disaggregated_Inference_Sim costs a step from a closed form (InfSim 02). An operator trace is an independent count of the same thing, so each checks the other:

Source: examples/results.md in Torch_Sim_Frontend

• The weight FLOPs agree exactly, and a Hypothesis property test checks the same identity for random Llama-shaped configurations, lengths and batch sizes. (Its first failure was in the test's formula: the rotary frequencies are computed once per batch, not per sequence.)
• The trace found two errors in the closed form, since corrected. Every decode step was charged the whole embedding table, 1.05 GB or 6.4% of the step's bytes, when a lookup reads one row. And it left out the new token attending to itself, 524,288 FLOPs per sequence per step. Neither is large; both are the kind of thing only an independent count finds.
• Found, then fixed (2026-10-03). Disaggregated_Inference_Sim, its JavaScript port and its Rust port were corrected, and every number the two series quote from them was regenerated. The last column above is the closed form before the fix. The two now agree, apart from two terms the closed form does not model: the 128-FLOP rotary matmul and 532,480 bytes of RMSNorm weights. The tests that reported the discrepancy now assert the agreement. The corrected decode step at context 2,048 takes 5.70 ms, against 6.09 ms before, which equals the fused trace's ideal-fusion bound (next slide).
• Attention is a convention, not an error. The trace counts what the kernel computes (unmasked), the closed form counts useful work (causal).

Costing the Trace

The first cost model is the InfSim roofline, applied per operator: time = max(FLOPs / FLOP rate, bytes / byte rate), with the same derated H100 rates as Disaggregated_Inference_Sim (a test checks they are the same numbers). Unfused assumes every non-view operator reads its inputs from memory and writes its outputs back. Fused is the ideal-fusion bound, where only matmul, attention and gather operators touch memory.

def op_cost(self, op: Op) -> OpCost:
    if op.category in ("view", "unknown"):
        return OpCost(0.0, "free", "-")
    nbytes = op.bytes if (self.memory == "unfused" or op.category in MEMORY_OPS) else 0
    tc, tm = op.flops / self.flops_rate, nbytes / self.byte_rate
    return OpCost(max(tc, tm) + self.op_overhead, "compute" if tc >= tm else "memory", "accel")

Source: examples/results.md in Torch_Sim_Frontend

Prefill is compute-bound as a whole. Yet in the meta trace, the materialised score matrix makes more than half its time memory-bound: 126.68 ms, against 76.14 ms with fused attention and 58.53 ms from the closed form. That gap is the case for flash attention, measured from a trace. The per-operator sum is never below the whole-trace roofline, because small memory-bound operators cannot overlap with large matmuls in this model.

Interactive: Operator Coverage, Counted Three Ways

The hypothetical optical MAC part from Disaggregated_Inference_Sim (4,000 TFLOP/s peak at 40% efficiency, HBM-class memory) runs only the operator categories you choose. Everything else runs on an illustrative host CPU, and a tensor crosses the link the first time an operator on the other side reads it. Llama-3-8B, fused-attention traces. Every value is a recorded run.

Operator Coverage as an Engineering Metric

Two different coverage questions, both worth tracking release over release:

• Cost-model coverage: does every operator in the trace have a cost rule? An operator without one is costed at zero, which silently flatters the hardware. simfront's report names every one. Its first run on GPT-2's ONNX export named SplitToSequence and SequenceAt, and rules were added.
• Device coverage: which operators run on the accelerator? Count it by operators, by FLOPs and by time. They tell different stories (slide 11).

PyTorch's own FlopCounterMode shows why the first question matters:

Source: examples/results.md in Torch_Sim_Frontend

It has no formula for _scaled_dot_product_flash_attention_for_cpu, so on the fused trace it silently drops 2.22 TFLOP of attention, about 7% of the prefill. In CI, the arithmetic of reference traces is gated exactly; operator counts and bytes are reported as drift when a library upgrade changes the decomposition (deck 07).

What to Take Away

• The meta device and fake tensors trace any model without weights: Llama-3-70B on a laptop, in about a second.
• One trace format, several front ends, cross-checked. They agree exactly on arithmetic and differ on traffic, which belongs to the decomposition.
• Say which code path a trace shows. The device the tensors claim, and whether the library detected tracing, change the operators.
• An independent count finds errors: here two in the closed form it was checked against, which were then fixed and the published numbers regenerated.
• Coverage is a time metric. 99.98% of FLOPs on the device can still be under 5% of the time.

• Torch_Sim_FrontendThe front ends, cost models, coverage report, tests and examples/results.md.
• LLM Inference Simulators 09: Framework IntegrationThe integration routes, including PrivateUse1 and ONNX Runtime execution providers.
• torch.export and custom torch.compile backendsPyTorch's documentation; these APIs move quickly, so check the version you use.
• ONNX shape inferenceIncluding data_prop.
• Deck 09: Specifications, Requirements and Test PlansUses this project's requirements and traceability matrix as the worked example.