An arXiv preprint reports that Axon, a strongly typed specification language for large language model architectures, matched or beat Transformers in most of its tested generation comparisons. Across 467 inference experiments on models from 135 million to 32 billion parameters, the paper reports median speedups of 7% with PyTorch, 12% with PyTorch plus Triton, 91% with JAX, 107% with MLX and 58% with native vLLM.
One specification, several backends
Axon’s design centers on a single typed specification and a shared Graph IR, an intermediate representation of a model’s computation, that can produce standalone implementations for PyTorch, Triton-accelerated PyTorch, JAX, MLX and vLLM. Its contribution statement claims a compiler pipeline covering over 204 models across 60 families.
The comparison used Transformers reference implementations under identical inputs, precision and compilation settings. Generation throughput was calculated as generated tokens divided by wall time. A runtime ratio below 1 indicates that Axon took less time than the reference.
Generation was the strongest case
In the detailed autoregressive, or token-by-token, generation analysis, Axon was at or below Transformers for 76% of 225 checkpoints up to 4B parameters. The median runtime ratio was 0.804× and the mean was 0.974×, which the paper describes as a typical 24% speed advantage. Among 242 checkpoints between 4B and 32B parameters, 73% were at or below parity, with a median ratio of 0.908× and a mean of 0.963×.
Separate native vLLM tests found Axon faster than HF Transformers in 65 of 88 checkpoint-backend pairs, or 74%. At the median, generation took 63% as long as it did with Transformers, a reported 1.6× speedup. On MLX, Axon was at or below Transformers latency for 94.4% of 126 checkpoints, with a median runtime ratio of 0.483× and a mean of 0.544×. The paper reports the MLX and vLLM counts as separate analysis units, so they are not additional experiments to add to the headline total.
The results were not uniform
Speed was less consistent for forward passes, which measure a single pass through a model on an input. In the large forward analysis, JAX was above parity for 100% of checkpoints, with a median runtime ratio of 3.747×. PyTorch was at or below parity for 65%, with a median ratio of 0.990×. The analysis notes variation by architecture and weak performance for some large encoder-decoder models.
The paper also checked top-1 token parity, meaning whether Axon and Transformers chose the same highest-probability next token. In the smaller-model evaluation, nine BF16 parity failures were rerun in FP32, and the paper reports 100% parity across 303 checkpoints. The larger-model analysis was less complete: 16 BF16 pairs failed across nine unique checkpoints, 11 were fixed by FP32, while three named cases failed the criterion for every tested prompt. That is a functional equivalence check for the first token, not a complete measure of output quality.
The training test was narrower
One training comparison also favored Axon: it was reported 9.6% faster than Transformers, while both models ended with a final loss of 2.506. But the comparison covered only one 270M model and one summarization task, so it does not establish that the result generalizes to larger models or other tasks.
The evidence remains bounded
These results are descriptive medians; no confidence intervals are reported. The main benchmark set used one NVIDIA B200 configuration, while MLX used one Apple Silicon M3 Max setup with 36GB of RAM. The listed runs included one warm-up and three repetitions, with generation caps of 128 tokens for the other backends and 256 tokens for MLX. The study also reports no ablation separating AST-level from Graph IR optimization, leaving the source of the performance differences unresolved.
Paper data and sources
Original title: Write Once, Run Everywhere: The Axon DSL for Shape-Safe and Framework-Agnostic LLM Architectures
Authors: Jacob Nielsen, Danial Namazifard, Lukas Galke Poech, Peter Schneider-Kamp
Journal/Repository: arXiv
Status: Preprint, not yet peer-reviewed
First online: 2026-08-20
DOI: Not available
Original paper · Full text