Benchmarks against the frameworks

In short: each framework runs in its own process, confined with sched_setaffinity to as many CPUs as the thread count, and every timed result first passes a correctness gate against the framework's eager output. The latest run, a same-day sweep of every framework and rml cell on a shared 4-vCPU Xeon container with torch 2.14.1, jax 0.12.0 and tensorflow 2.21.0, puts rml ahead of JAX and TensorFlow on average at every batch and thread count, and ahead of PyTorch on one thread at batches 1 and 8. Against the best framework and mode per cell it wins 46 of 114 cells, a geometric mean of 0.96: ahead on one thread at small batch, behind at batch 64 and on four threads. Small MLP-shaped models are its strongest; the U-Net, the GCN, the LSTM, the SSM, the CNN and the RNN trail most. Four threads is slower than one in 32 of 171 rml cells, and the StableHLO route runs within 2–13% of the direct exporters.

The models are defined on Exporters and the model zoo, the runner that times rml is rml-graph-run on Executing a graph, and the optimizations each gap prompted are on Lowering and fast CPU. How CPU and GPU frameworks schedule work in general is on GPU execution.

The method

The script's docstring states the whole method (bench_frameworks.py L1–22). For each framework, model (at its bench size) and batch size it exports the model to models/graphs/<framework>_<model>_b<batch> with a golden file of the framework's outputs, then times three things at each thread count: the framework eagerly, the framework compiled, and rml running the exported graph.

Processes and threads

Each framework runs in a subprocess, so its thread pool is set before the framework starts (L252–265). The driver sets OMP_NUM_THREADS and MKL_NUM_THREADS, and for JAX an XLA_FLAGS with intra_op_parallelism_threads; inside the worker, PyTorch gets set_num_threads and one inter-op thread, TensorFlow one inter-op thread and N intra-op threads (L84–94). None of that is trusted alone: pinned(threads) confines each timed child, and each rml-graph-run, to the first N CPUs with os.sched_setaffinity, because recent JAX ignores the XLA thread flag and runs its matmuls on every core (L242–249). The report's findings open by restating that every timed process ran confined this way (frameworks.md L343). The thread counts default to 1 and every core (L403–404).

The correctness gate

Nothing is timed until it is right. In the framework worker the eager output is the reference, and the compiled output must satisfy |got − want| ≤ 10⁻⁴ + 10⁻³·|want| at every element, or the sweep stops (L53–60, L226–228). The golden file rml checks against holds the framework's own outputs, written at export (L141–168). rml-graph-run --check reports its largest absolute error, which must be at most 10⁻⁴ + 10⁻³ times the largest |want| of the first output, a bound over the whole tensor rather than per element (L277–280).

Timing

A framework cell runs its callable five times untimed, then times runs until it has at least --min-runs (10) and has spent --budget (1 s), up to --max-runs (200), waiting on the result each time so asynchronous dispatch is counted (L63–72). rml runs exactly --rml-runs (50) after the same warmup, timed with steady_clock inside the binary, so process start and graph load are outside the number (graph_run.cc L87–100). Both sides report p10, median and p90 by nearest rank: the value at index ⌈q·n⌉ − 1 of the sorted times (L75–78, graph_run.cc L68–85). The inputs are the same for every route: example_inputs(model, "bench", batch, seed=1) (L219).

Resuming and options

A full sweep is long, so every finished cell is appended to models/graphs/results.jsonl and flushed to disk at once (L41–44, L195–210). The table is written from that file, not from memory (L438–442).

Flags are parsed at L378–398. From rml's root, in the export virtualenv:

bazel build --config=opt //examples/graph_run:rml-graph-run
~/.venvs/rml-export/bin/python tools/bench_frameworks.py --out docs/bench/frameworks.md

The results

The committed report (frameworks.md L1–7) was produced on an Intel Xeon @ 2.30 GHz, 4 logical CPUs, Linux x86_64, with torch 2.14.1+cu130, jax 0.12.0 and tensorflow 2.21.0, rml built with --config=opt, at 1 and 4 threads. Every value is f32 on CPU. Each model has a table of nine columns per batch and thread count: both modes of each framework and rml from that framework's export, as median (p10–p90) milliseconds; after each rml time, the faster of that framework's two modes divided by rml's time, so above 1× rml is faster (L218). Every framework and rml cell was timed in one sweep on 2026-11-03 (L343–349).

The summary is a geometric mean over the 19 models of the faster framework mode's time divided by rml's, with the number of models rml wins; the last column takes, per cell, the fastest framework and mode against the fastest rml route (L351–362):

Taken cell by cell against the fastest framework and mode, rml wins 46 of the 114 (model, batch, threads) cells, and its geometric mean against that best is 0.96 (L364–367). The previous report claimed 80 cells and 1.31, but it mixed freshly timed rml cells with framework cells from an older sweep, some of them pathological (torch.compile on unet at four threads: 152 ms there, 0.75 ms in this sweep), so rml's notes say its summary overstated rml (L344–347). Every number on this page is from the current file; the optimizations rml has made so far are on Lowering and fast CPU.

Model by model

The factor for each framework at the two corners of the sweep, copied from each model's table: batch 1 on one thread, where per-op overhead dominates, and batch 64 on four threads, where arithmetic and threading do.

The two corners tell different stories. At batch 1 on one thread rml beats PyTorch and TensorFlow on 17 of the 19 models each, but JAX on only 9. At batch 64 on four threads, where large GEMMs and convolutions dominate, it beats PyTorch on 4 of 19, against 16 for both JAX and TensorFlow (L357–362).

Where rml wins

• Small and medium MLP-shaped models, as a geometric mean over batches and threads against the best framework: two_tower 2.1×, diffusion_mlp 1.7×, mlp 1.3×, gan_generator and vae 1.2×, autoencoder 1.06× (L368–370). These are a handful of matmuls and elementwise activations, the shapes the few-row GEMM path and fused programs on Lowering and fast CPU were built for.
• Against JAX and TensorFlow on average in every row of the summary (L364–365), and against TensorFlow by large factors at small batch; diffusion_mlp at batch 1 on one thread is 0.0232 ms in rml from TensorFlow's export against 0.947 ms TF eager and 0.652 ms TF XLA (L156).
• Graphs that arrive decomposed. layernorm_gelu (1.15× overall) at batch 64 on four threads runs in 20.9 ms from JAX's export against 37.7 ms for jax.jit (1.8×), and 19.9 ms from TensorFlow's against 27.3 ms for TF XLA (1.4×) (L51). JAX and TensorFlow hand this model over with layer norm and GELU written out, and rml's export passes fuse both back into single ops before rml runs it.

Where it trails, and why

The findings name six models furthest behind the best framework, as a geometric mean over batches and threads, and put the attention models (vit, transformer_block, moe, tiny_gpt, encoder) at 0.9–0.96× (L371–374):

The gcn rows show the second case plainly: at batch 64 on four threads PyTorch eager takes 4.2 ms and rml from PyTorch's export 11.4 ms (0.37×), while at batch 64 on one thread the gap is 15 against 21.3 ms (L182–183).

Worked example: the LSTM

The lstm rows at batch 64 (L105–106), median milliseconds:

Reading it: rml scales from 38.8 to 12.7 ms with four threads on PyTorch's export, and on four threads it is ahead of jax.jit and TF XLA, but PyTorch's fused ATen LSTM runs in about half rml's time at either thread count. The previous report had PyTorch eager at 8.81 ms on four threads from the older framework sweep; timed on the same day it is 6.9 ms. rml's LSTM work so far is constant weights packed once at compile time, each step's gates and cell update fused into one program reading the gate slices in place, and GEMM column chunks balanced across threads (Lowering and fast CPU). rml also tried folding the recurrence's bias and input-projection adds into the step GEMM and dropped it: the per-step GEMM already runs at 245–290 GFLOP/s on four threads, about the machine's peak, so the remaining gap is the GEMM itself (kernels.md L258–273).

This sweep was taken after two changes that rml's checklist first measured only as interleaved A/B ratios. The export pass hoist_step_projections gives the JAX and TensorFlow LSTM one input projection instead of one per step, and Im2ColPositions fills its row segments with plain loops: jax/tf lstm_b64 about 1.15–1.2× and unet_b64 about 1.25–1.5× faster on 4 threads (progress.md L169–170, Lowering and fast CPU).

Four threads slower than one

Four threads is slower than one in 32 of rml's 171 cells (19 models × 3 batches × 3 exports), almost all at batch 1 with sub-millisecond times (L375–376): mlp from PyTorch at batch 1 is 0.0561 ms on one thread and 0.0735 on four, and lstm goes from 1.23 to 1.53 (L13–14, L101–102). Waking a pool costs more than such work saves; the grains that keep small ops on the calling thread are on Lowering and fast CPU.

Direct against StableHLO

The report also times rml on the graph each framework's StableHLO lowered to, against the direct export of the same model (L220–339). The factor is direct / StableHLO, so above 1× the StableHLO graph runs faster. The geometric means are 0.96 for PyTorch, 0.87 for JAX and 0.98 for TensorFlow: the StableHLO route is within 2–13% of the direct one (L377–378). Single cells swing both ways; mlp at batch 1 on one thread is 0.0561 ms direct and 0.0616 ms through StableHLO from PyTorch, but 0.0879 direct and 0.0632 through StableHLO from TensorFlow (L226). Both routes are described on Exporters and the model zoo.

Reading the numbers

• A shared machine. The container is a shared 4-vCPU cloud VM where single cells vary by up to about 30% between runs, so rml's notes say to compare the means (L347–348).
• One sweep. Every framework and rml cell in this report was timed on the same day. The report before it kept framework cells from an older sweep, some far off their neighbours (PyTorch on vit at batch 1 on four threads took 344 ms there; it is 1.91 ms here, L58), which is why its summary read so much better for rml (L344–347).
• Not the target hardware. Apple Silicon numbers, with NEON and Accelerate, have not been produced (L349).

Checks

The script has no unit tests of its own; its correctness gate is the check, and the sweep stops at the first failure. The agreement it relies on is tested elsewhere: