Setup
aaronnet is the lab’s in-house autograd engine: written from scratch, CPU-only, and accelerated with Numba. It is the engine our transformer work runs on, so we wanted a baseline. On a workload small enough to run end to end in seconds, how far off is it from PyTorch?
- Model: two Linear layers with a ReLU between them, dim to 2 times dim and back to dim.
- Operation: one forward plus one backward pass. Gradients are zeroed before each iteration.
- Batch size: 32. Dtype: float32 on both sides.
- Device: CPU. Both sides pinned to one thread (OMP, MKL, OpenBLAS, Accelerate, Numba, and torch).
- Warmup: one forward plus backward on each side before timing, matching how we exercise aaronnet in practice.
- Timing: 100 iterations per repeat, 7 repeats per cell. We report median and standard deviation across repeats, in ms per iteration.
- Environment: Apple M5, macOS 26.5.1, Python 3.14.3, PyTorch 2.10.0, NumPy 2.3.5, Numba 0.63.1, aaronnet at commit
28b81ba.
Results
| Input size | AaronNet (ms) | PyTorch (ms) | Ratio |
|---|---|---|---|
| 128 | 0.096 ± 0.040 | 0.059 ± 0.001 | 1.63x |
| 256 | 0.191 ± 0.002 | 0.103 ± 0.001 | 1.85x |
| 512 | 0.558 ± 0.012 | 0.251 ± 0.003 | 2.22x |
| 1024 | 2.425 ± 0.041 | 0.966 ± 0.011 | 2.51x |
Cells show median ± standard deviation across 7 repeats of 100 iterations each. Ratio is AaronNet median divided by PyTorch median. Lower is better; above 1.0x means AaronNet is slower.
Interpretation
Across every size we tested, aaronnet is slower than PyTorch, and the gap widens with the input dimension: 1.63x at dim 128, up to 2.51x at dim 1024. The dispersion within a cell is tight (standard deviation under one percent of the median at the larger sizes), so the ratios are not an artifact of noise.
The direction of the trend is consistent with where most of the work lives: as the matrices get larger, the backward matmul becomes the dominant cost, and that is the operation where our engine has the most ground to make up on an optimized library like PyTorch.
Limits of this measurement
- A single model shape and a single batch size.
- CPU only, on one machine.
- AaronNet’s matmul backward unconditionally computes the input-gradient array. We create the input with gradient tracking off so the result is discarded, but the compute cost is not skipped. PyTorch, with the same input configuration, skips that work entirely. This mismatch counts against aaronnet in every cell.
Script and environment
The benchmark script and the full per-repeat results, including the environment used, are published alongside this note.
Previous note: taking the transformer apart