In short: a model is written once per framework, from shared weights. Offline Python exports it into
rml-graph, a JSON op list over mmapped safetensors weights, by two routes: each framework's
own traced graph, or the StableHLO its compiler produces. From there rml runs it in C++ with no Python.
Below the graph, runs of elementwise nodes fuse into tiled programs, and the remaining ops run on
vectorized, threaded kernels and a packed SIMD GEMM. Every model is checked across all three frameworks
and both routes, and benchmarked against each framework's eager and compiled modes.
The direction
rml's plan (docs/progress.md L135–146) extends the Python front ends to every neural model on the ML models reference. Three commitments shape it:
- Three frameworks, one model. Each model is defined in PyTorch, JAX and TensorFlow, exported, checked against the framework, and benchmarked on CPU.
- Layered like MLIR. Frameworks hand over StableHLO through their own production
exporters, which is what an MLIR pipeline consumes.
tools/lowers it to rml-graph, and inside rml, rml-graph lowers to a kernel-level plan that runs fused, threaded kernels. - Two routes, cross-checked. The direct per-framework exporters stay as a second way in, and the two routes are checked against each other.
At this commit every item of that plan is checked off except one: keeping small ops on the calling thread
(progress.md L157–175). Larger grains for exactly that have landed,
but four threads is still slower than one in 32 of 171 rml cells, so the checklist keeps it open and this
section treats it as in progress. The last two items are new examples, tiny_net and
tiny_models, which run small PyTorch, JAX and TensorFlow models from C++ and compare them bit
for bit (Exporters and the model zoo).
The levels
| Level | What it is | Status | Page |
|---|---|---|---|
| Framework model | The same network written in each framework, built from one set of seeded NumPy weights (zoo/weights.py) | 19 models | Exporters and the model zoo |
| Direct exporters | Walk torch.export's ATen graph, the jaxpr or the frozen GraphDef, keeping high-level ops such as layer_norm and gelu (graph_ir.md L10–12) | Built | Exporters and the model zoo |
| StableHLO route | from_stablehlo.py imports the StableHLO from torch-mlir, jax.export and TensorFlow's XLA StableHLO stage (graph_ir.md L13–19, progress.md L157) | Built | Exporters and the model zoo |
| rml-graph | A model as two files: a JSON op list and a safetensors file loaded by mmap (graph_ir.md L27–30) | Built | The rml-graph IR |
| Kernel-level layer | Runs of elementwise nodes with one output shape become one program that streams row tiles small enough for L1; members need not feed each other, so an LSTM cell's gates and state update form one group with two outputs (fusion.h L4–19); rml-graph-run --profile lists each step and --no-fusion turns it off (progress.md L158) | Built | Lowering and fast CPU |
| Kernels | Vectorized, threaded elementwise, softmax, norm and reduction kernels in rml/graph, buffers reused across runs through a BufferPool, and a packed SIMD GEMM with run-time dispatch (progress.md L159–161) | Built | Executing a graph, Lowering and fast CPU |
| Checks | Cross-framework tests across three frameworks and both routes, and timings against each framework's eager and compiled modes (progress.md L177–188) | Built | Benchmarks against the frameworks |
Why layers
Each level has one job, so each can be tested on its own. The framework-facing work stays in offline
Python, where the frameworks live. The format in the middle is small on purpose
(graph_ir.md L106–110): exporters lower framework ops onto it, so
JAX's dot_general becomes transpose + reshape + matmul,
and an op an exporter cannot lower fails the export with the framework's op name. There is no silent
fallback. Below the format, optimization happens in rml without touching any exporter: the fusion layer
removes the intermediate tensors of a written-out GELU or layer norm "whichever framework the graph came
from" (fusion.h L10–12).
StableHLO is the front door the plan names because each framework already produces it with its production tooling, and it is the input an MLIR pipeline expects. Both routes end in the same clean-up passes, and every fixture exists for both, so the two are independent ways into the same IR that the tests hold to the same answer (graph_ir.md L21–25).
What stays out
- Training. The format is inference only: there is no backward pass (graph_ir.md L6). Training in rml goes through Autograd.
- Classical ML and KV-cache decoding stay out of the graph stack
(progress.md L143–146). LLM decoding keeps its own
allocation-free loop in
rml/llm; the graph executor is a general model runner. - Sampling loops stay outside the graph. The exported network is one step, and noise is an input, so a diffusion sampler calls the graph once per step.
The model plan
All 19 zoo models exist in all three frameworks at this commit, each in a "tiny" size committed as a test fixture and a "bench" size that is timed (zoo/weights.py L11–19). rml's own table maps each to the /ml-models pages it covers (progress.md L190–208):
| Models | Status | Concept |
|---|---|---|
mlp | Built | Neurons, layers and the MLP, Activation functions |
cnn | Built | Convolutional networks |
transformer_block, layernorm_gelu | Built | The Transformer, Attention, Normalization and residual connections |
vit | Built | Vision Transformer |
encoder | Built | Encoder-only models, Embeddings, Positional encoding |
tiny_gpt | Built | Decoder-only LLMs, Autoregressive models |
rnn, lstm | Built | RNNs and LSTMs |
ssm | Built | State-space models |
autoencoder, vae | Built | Autoencoders, Variational autoencoders |
gan_generator | Built | Generative adversarial networks |
diffusion_mlp, unet | Built | Diffusion models, Normalizing flows and flow matching, Diffusion U-Net |
gcn | Built | Graph neural networks |
two_tower, clip | Built | Two-tower models, Multimodal models |
moe | Built | Mixture of experts |
Where it stands
The latest benchmark (frameworks.md L5–7: an Intel Xeon at 2.30GHz
with 4 logical CPUs, Linux x86_64, torch 2.14.1, jax 0.12.0, tensorflow 2.21.0, rml built with
--config=opt) times all 19 models at batch 1, 8 and 64 on 1 and 4 threads, every framework
and rml cell in one sweep on 2026-11-03. On average rml is ahead of JAX and TensorFlow in every row, and
ahead of PyTorch on one thread at batches 1 and 8. Taking the fastest framework and mode per cell, rml
wins 46 of the 114 cells, with a geometric mean of 0.96 against that best: ahead on one thread at small
batch, behind at batch 64 and on four threads (frameworks.md
L351–367). The report before it, which claimed 80 cells and 1.31, mixed fresh rml cells with older
framework cells and overstated rml (L344–347).
It is strongest on small and medium MLP-shaped models (two_tower at 2.1× the best framework) and trails most on the U-Net (0.48×, against torch.compile's convolutions), the GCN (0.6×, against PyTorch eager on four threads), the LSTM (0.64×), the SSM (0.69×), the CNN (0.75×) and the RNN (0.78×); the attention models are at 0.9–0.96×. Four threads is slower than one in 32 of 171 rml cells, almost all at batch 1. The StableHLO route runs within 2–13% of the direct exporters (frameworks.md L368–378). Benchmarks against the frameworks has the numbers, and Lowering and fast CPU what each optimization targets.
Apple Silicon numbers (NEON, Accelerate) have not been produced yet (frameworks.md L349), and GitHub CI has failed before any job starts since October 17, so these changes are verified only in rml's Linux container (progress.md L210–212).