Write Once, Run Everywhere: The Axon DSL for Shape-Safe and Framework-Agnostic LLM Architectures

arXiv:2608.198892026-08-21

A new language, Axon, lets you write an LLM once and run it on PyTorch, JAX, MLX or vLLM

The open-source language model ecosystem effectively depends on a single platform, Hugging Face, and porting models to other frameworks has meant manually rewriting them, often losing optimizations along the way. The authors built Axon, a strongly typed, Haskell-inspired language for describing model architectures once, from which a compiler automatically generates standalone code for PyTorch, Triton, JAX, MLX and vLLM. Across 467 inference benchmarks on models from 135M to 32B parameters, Axon-derived models were generally faster than reference Transformers implementations.

METAL MEDIA explanatory visual

A new language, Axon, lets you write an LLM once and run it on PyTorch, JAX, MLX or vLLM

01Problem: LLM code is locked into a specific platform's conventions, so porting to other frameworks requires manual, error-prone rewrites that often drop backend-specific optimizations
02Solution: Axon lets researchers describe a model's structure (layers, tensor shapes, parameter locations) once; the compiler lowers this into a shared Graph IR and then generates standalone code for five backends
03A strong type system checks tensor shapes ahead of time, preventing shape-mismatch errors during execution
04Results: median speedups of 7% on PyTorch, 12% on PyTorch+Triton, 91% on JAX, and 107% on MLX versus Transformers reference implementations; 58% median speedup when deployed natively on vLLM with PagedAttention and KV-cache
05Training experiments showed an Axon-derived PyTorch model matched the loss curve of the standard Transformers implementation while running about 9.6% faster per step

An explanatory diagram made by METAL MEDIA, not a figure supplied by the paper's authors.

What they did

Problem: LLM code is locked into a specific platform's conventions, so porting to other frameworks requires manual, error-prone rewrites that often drop backend-specific optimizations
Solution: Axon lets researchers describe a model's structure (layers, tensor shapes, parameter locations) once; the compiler lowers this into a shared Graph IR and then generates standalone code for five backends
A strong type system checks tensor shapes ahead of time, preventing shape-mismatch errors during execution
Results: median speedups of 7% on PyTorch, 12% on PyTorch+Triton, 91% on JAX, and 107% on MLX versus Transformers reference implementations; 58% median speedup when deployed natively on vLLM with PagedAttention and KV-cache
Training experiments showed an Axon-derived PyTorch model matched the loss curve of the standard Transformers implementation while running about 9.6% faster per step

Figure 1: Write-once, run everywhere. Axon DSL compiles axon definitions (.axon) to standalone model definitions for PyTorch, JAX, MLX and vLLM. All that is needed is a standard safetensors checkpoint.

Table 1: Per-backend comparison of Axon’s autoregressive generation performance with decoder-only models against PyTorch, Triton, and JAX. “Axon ≤1×” = count (and %) of checkpoints where Axon is faster than Transformers. “Axon >1×” = count (%) of checkpoints where Axon is slower. Median and mean report runtime ratio (lower means Axon is faster).
Backend	Checkpoints	Axon ≤1×	Axon >1×	Median ratio	Mean ratio
<4B
PyTorch	76	48 (63%)	28 (37%)	0.903	0.925
Triton	75	60 (80%)	15 (20%)	0.843	0.878
JAX	74	64 (86%)	10 (14%)	0.481	1.122
<4B total	225	172 (76%)	53 (24%)	0.804	0.974
4–32B
PyTorch	87	55 (63%)	32 (37%)	0.987	0.980
Triton	87	67 (77%)	20 (23%)	0.924	0.931
JAX	68	55 (81%)	13 (19%)	0.589	0.981
4–32B total	242	177 (73%)	65 (27%)	0.908	0.963

Figure 2: Autoregressive generation performance with decoder-only models with up to 4B parameters. Log-log plot where the diagonal line indicates equal performance. Points above the parity line indicate Axon is faster, color-fill denotes backends and border-line denotes dtype.

Table 2: Per-backend comparison of Axon’s Forward performance with Encoder-Only and Encoder-Decoder Models against PyTorch, Triton, and JAX backends. “Axon ≤1×” = count (and %) of checkpoints where Axon is faster than Transformers. “Axon >1×” = count (%) of checkpoints where Axon is slower. Median and mean report runtime ratio (lower means Axon is faster).
Backend	Checkpoints	Axon ≤1×	Axon >1×	Median ratio	Mean ratio
≤4B
PyTorch	26	15 (58%)	11 (42%)	0.883	1.786
Triton	26	9 (35%)	17 (65%)	1.332	2.169
JAX	26	7 (27%)	19 (73%)	1.079	1.260
≤4B total	78	31 (40%)	47 (60%)	1.084	1.738
4–32B
PyTorch	20	13 (65%)	7 (35%)	0.990	1.506
Triton	20	8 (40%)	12 (60%)	1.133	1.546
JAX	14	0 (0%)	14 (100%)	3.747	3.356
4–32B total	54	21 (39%)	33 (61%)	1.175	2.001

Figure 3: Benchmarking Autoregressive models between 4B and 32B models. Log-log plot where the diagonal parity line indicates equal performance. Generally, Axon produces comparable or superior throughput on the bigger models.

Table 3: Autoregressive generation comparison of Axon (vLLM native) against Transformers. “Axon ≤1×” = number (and %) of checkpoints where Axon is faster than Transformers. “Axon >1×” = number (and %) of checkpoints where Axon is slower than Transformers. Median and mean report Axon-to-Transformers runtime ratio (lower is faster).
Model size	Checkpoints	Axon ≤1×	Axon >1×	Median ratio	Mean ratio
Small (≤4B)	54	38 (70%)	16 (30%)	0.656	1.430
Large (4B–32B)	34	27 (79%)	7 (21%)	0.609	2.889
Total	88	65 (74%)	23 (26%)	0.631	1.994

Figure 4: vLLM: Axon (vLLM native) vs. Transformers generation throughput on 88 checkpoints. Points above the parity line indicate Axon is faster. 74% fall above.

Table 4: Breakdown of Axon performance by precision, model type, and sequence length. “Axon ≤1×” = number (and %) of checkpoints where Axon is faster than Transformers. “Axon >1×” = number (and %) of checkpoints where Axon is slower than Transformers. Median and mean report Axon-to-Transformers runtime ratio (lower = Axon is faster).
Group	Points	Axon ≤1×	Axon >1×	Median ratio	Mean ratio
BF16	63	61 (97%)	2 (3%)	0.423	0.500
FP32	63	58 (92%)	5 (8%)	0.503	0.588
causal_lm	84	82 (98%)	2 (2%)	0.390	0.437
seq2seq_lm	42	37 (88%)	5 (12%)	0.574	0.752
len=64	42	42 (100%)	0 (0%)	0.417	0.448
len=128	42	39 (93%)	3 (7%)	0.483	0.525
len=256	42	38 (90%)	4 (10%)	0.494	0.652
Total	126	119 (95%)	7 (5%)	0.483	0.544

Figure 5: MLX. Benchmarking conventional HF models against Axon derived standalone model definitions. Axon yields some considerable speed-ups across the board.

Table 5: Axon Compiler Phase and Stage with Invariants
Phase	Representation	Primary invariant
Parse	one-file AST	Syntactic structure and explicit MAIN pragma insertion
Load	loaded AST set	Imports and builtins located without rewriting semantics
Materialize	one-file AST	Optional checkpoint/config specialization for generic models
Resolve/validate-closed	closed AST	No unresolved imports or names; unreachable definitions pruned from MAIN
Normalize	normalized AST	Call syntax, pipes, path sugar, and zero-arg call/name distinctions made explicit
Elaborate/validate-elaborated	elaborated AST	Default arguments filled and call arguments positionalized
Flatten/validate-flat	flat AST	Explicit evaluation order; flat calls and binds accepted by typecheck and Graph IR lowering
Typecheck/validate-typed	Typed flat AST	expression types, arities, dimensions, and primitive rules applied to a fixpoint
Optimize-ast	typed flat AST	Optional conservative AST cleanup with retype/validation
Graph lowering/validation	Graph IR	Typed graph modules, multi-output nodes, structured paths, constraints, and metadata
Optimize-graph/validation	Graph IR	Optional graph cleanup, specialization, backend-neutral rewrites, and opt-in backend intrinsics
Backend	generated/runtime code	Executable tensor program consuming the validated Graph IR contract

Figure 12: Benchmarking Encoder-only and Encoder–Decoder up to 4B models. Log-log plot where the diagonal parity line indicates equal performance. Here, we see the majority of the points in the vicinity below the parity line.

Table 6: Wall-clock time and GPU idle for autoregressive generation (112 tokens for Qwen 0.5B, 96 for Pleias 3.8B). “—” indicates the metric is not directly measurable: JAX has no torch.profiler equivalent for per-kernel timing, and GPU active/idle is a synchronous-execution concept that does not apply to JAX’s async dispatch model.
	Qwen2.5-0.5B (0.67B)	Pleias-3b-Preview (3.8B)
Metric	Transformers	Axon-Torch	Axon-JAX	Transformers	Axon-Torch	Axon-JAX
Wall-clock (ms)	1380	1185	541	1091	1071	1160
Speedup vs Transformers	1.0×	0.86×	0.39×	1.0×	0.98×	1.07×
GPU active (ms)	373	370	—	557	561	—
GPU idle (ms)	1007	815	—	534	510	—
GPU idle %	73.0%	68.8%	0.0%	48.9%	47.6%	0.0%
CUDA kernels	129,189	120,987	—	101,935	104,123	—

Figure 13: Benchmarking Encoder-only and Encoder-Decoder between 4B and 32B models. Log-log plot where the diagonal parity line indicates equal performance.

Table 7: Python function call counts and cProfile cumulative time. “—” indicates the category does not apply: JAX fuses all per-op dispatches into a single XLA compilation per step, so individual F.linear and SDPA calls do not exist as Python-level dispatches.
	Qwen2.5-0.5B (0.67B)	Pleias-3b-Preview (3.8B)
Metric	Transformers	Axon-Torch	Axon-JAX	Transformers	Axon-Torch	Axon-JAX
Total Python calls	455,769	508,390	137,805	359,909	330,929	82,210
Calls vs Transformers	1.0×	1.12×	0.30×	1.0×	0.92×	0.23×
cProfile time (s)	1.13	0.95	0.54	0.87	0.78	1.16
Per-op dispatch (Python calls per generate pass):
nn.Module.__call__	35,616	0	0	28,032	0	0
F.linear	18,928	18,928	—	14,880	14,880	—
SDPA	2,688	2,688	—	2,112	2,112	—
rope_apply	0	5,376	—	0	4,224	—
forward (jit dispatch)	—	—	112	—	—	96

Figure 14: Training of Gemma 3 270M on summarization task for 2000 steps, demonstrating identical training behaviour between conventional Transformers-definition and Axon derived PyTorch model. Note the Axon and Transformers loss curves are on top of each other together.

Why it matters

Today's LLM tooling concentrates around a handful of platforms, creating a single point of failure for the whole ecosystem; Axon proposes a shared language specification instead of a shared framework to reduce that dependency. This matters for anyone who wants fast, portable models without being locked into one company's optimization stack.

Terms in this paper

DSL (domain-specific language) · a programming language built for one narrow purpose, here describing neural network architectures
strongly typed · the compiler strictly checks the kind and shape of values before running the program
Graph IR · a shared intermediate representation that all backend code generators consume
PagedAttention / KV-cache · vLLM's memory management technique that speeds up serving by efficiently handling cached tokens during generation
top-1 token parity · a correctness check confirming different implementations predict the same next token at every step

Original abstract (English)

The entire ecosystem of open-source language models effectively relies on a single platform. What if this platform was forced to shut down tomorrow? Implementing and maintaining efficient model definitions and translating them between different training and inference regimes is a resource-heavy task that severely limits model efficiency and portability, hindering both scaling and deployment. Here, we present Axon, a strongly typed domain-specific language with Haskell-like syntax, that enables a write-once, run everywhere paradigm for LLM architectures. By basing collaboration on a language specification rather than a specific framework's vision, Axon fosters open cooperation and empowers researchers to implement highly specialized architectures without giving up optimization infrastructure or accepting deployment lock-in. Axon allows for concise, auditable specifications that can be automatically compiled to standalone implementations for leading frameworks: PyTorch, PyTorch with Triton, JAX, MLX and vLLM. In 467 inference benchmarking experiments on models ranging from 135M to 32B parameters, we demonstrate median speedups of 7% on PyTorch, 12% on PyTorch with Triton, 91% on JAX, and 107% on MLX, compared to the reference implementations from Transformers. When deployed as native vLLM architectures with PagedAttention and KV-cache, Axon models achieve a 58% median speedup over Transformers implementations.

Authors · Jacob Nielsen, Danial Namazifard, Lukas Galke Poech, Peter Schneider-Kamp

Read on arXiv

Latest papers

All papers →

Latest from METAL MEDIA

NCT 127 drop BLINGY, their first full album in two years

Figures: Jacob Nielsen et al., arXiv:2608.19889, arxiv-nonexclusive