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

Jacob Nielsen; Danial Namazifard; Lukas Galke Poech; Peter Schneider-Kamp

cs.AIAug 20, 2026

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

Organizations: University of Southern Denmark Campusvej 55, 5230 Odense, Denmark

Abstract

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.

Figures & tables

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.

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.

<4B
Backend	Checkpoints	Axon $\leq 1\times$	Axon $>1\times$	Median ratio	Mean ratio
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

Table 1: Per-backend comparison of Axon’s autoregressive generation performance with decoder-only models against PyTorch, Triton, and JAX. “Axon

\leq 1\times

” = count (and %) of checkpoints where Axon is faster than Transformers. “Axon

>1\times

” = count (%) of checkpoints where Axon is slower. Median and mean report runtime ratio (lower means Axon is faster).

$\leq$ 4B
Backend	Checkpoints	Axon $\leq 1\times$	Axon $>1\times$	Median ratio	Mean ratio
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
$\leq$ 4B total	78	31 (40%)	47 (60%)	1.084	1.738
4–32B

Table 2: Per-backend comparison of Axon’s Forward performance with Encoder-Only and Encoder-Decoder Models against PyTorch, Triton, and JAX backends. “Axon

\leq 1\times

” = count (and %) of checkpoints where Axon is faster than Transformers. “Axon

>1\times

” = count (%) of checkpoints where Axon is slower. Median and mean report runtime ratio (lower means Axon is faster).

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.

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.

Model size	Checkpoints	Axon $\leq 1\times$	Axon $>1\times$	Median ratio	Mean ratio
Small ( $\leq$ 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

Table 3: Autoregressive generation comparison of Axon (vLLM native) against Transformers. “Axon

\leq 1\times

” = number (and %) of checkpoints where Axon is faster than Transformers. “Axon

>1\times

” = number (and %) of checkpoints where Axon is slower than Transformers. Median and mean report Axon-to-Transformers runtime ratio (lower is faster).

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

Group	Points	Axon $\leq 1\times$	Axon $>1\times$	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

Table 4: Breakdown of Axon performance by precision, model type, and sequence length. “Axon

\leq 1\times

” = number (and %) of checkpoints where Axon is faster than Transformers. “Axon

>1\times

” = number (and %) of checkpoints where Axon is slower than Transformers. Median and mean report Axon-to-Transformers runtime ratio (lower = Axon is faster).

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

Table 5: Axon Compiler Phase and Stage with Invariants

	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 $\times$	0.86 $\times$	0.39 $\times$	1.0 $\times$	0.98 $\times$	1.07 $\times$
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%

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.

Per-op dispatch (Python calls per generate pass):
	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 $\times$	1.12 $\times$	0.30 $\times$	1.0 $\times$	0.92 $\times$	0.23 $\times$
cProfile time (s)	1.13	0.95	0.54	0.87	0.78	1.16
nn . Module . call	35,616	0	0	28,032	0	0

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.

Panel of Figure 6 — Figure 6: Graph IR view of the running example. Blue denotes typed graph values, orange operation nodes, and green structured operands and attributes. Paths, literals, and shape lists are structured operands, not strings that backends must parse.

Figure 8: 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.

Figure 9: 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.

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

74\%

fall above.

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

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.

Figure 13: Benchmarking Encoder-only and Encoder-Decoder between

4B

and

32B

models. Log-log plot where the diagonal parity line indicates equal performance.

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.

Efficient Language Model Inference

Efficient language model inference reduces the compute, memory, latency, and energy needed to generate useful outputs.

See all 318 papers

IronLLM: Forging Compact Edge-Native Language Models for Real-Time Embodied Intelligence
Sep 29, 2026Changdi Yang, Fengquan Jiao, Haochih Lin +14Multi-Token Prediction On-Device Language Model Inference
Robotics Foundation Model Team, Xpeng Inc.
Unlocking Lossless Speedups in LLMs via Discrete Diffusion
Sep 3, 2026Subham Sekhar Sahoo, Lingjie Chen, Khiem Pham +14Diffusion Language Model Inference LLM Inference Acceleration
Institue of Foundation Models · University of Illinois Urbana-Champaign · Cornell Tech +2
RoofLang: Enabling AI-Driven Architecting of LLM Inference Systems
Sep 11, 2026Ziyue Yang, Yuting Jiang, Lei Qu +1LLM Inference Large Language Model-Guided Optimization
Shanghai Xingyunzhili Artificial Intelligence Institute, Shanghai, China · Microsoft Research, Redmond, WA, USA

Sep 29, 2026J/K move · Enter open · S save

IronLLM: Forging Compact Edge-Native Language Models for Real-Time Embodied Intelligence

Changdi Yang, Fengquan Jiao, Haochih Lin, Haoran Yang+13

Robotics Foundation Model Team, Xpeng Inc.

Multi-Token Prediction On-Device Language Model Inference Small Language Models Speculative Decoding Efficient Language Model Inference

Open paper AskPDF arXiv

We present IronLLM-0.6B, a 654M-parameter language model designed for efficient on-device inference. IronLLM-0.6B combines a hybrid attention architecture with X-MTP, a lightweight shared-KV multi-token prediction design that eliminates per-depth KV-cache replay and employs a lightweight verification head for rollback-free drafting, achieving a 1.48x decoding speedup. The model is pretrained on approximately 6.2 trillion tokens using a quality-oriented data pipeline and is further post-trained with Multi-Domain On-Policy Distillation to integrate capabilities from domain-specialized teachers. To better meet the low-latency requirements of on-device scenarios, IronLLM-0.6B adopts an Instruct-Only design. Evaluations show that IronLLM-0.6B achieves competitive performance relative to larger models such as Qwen3.5-0.8B and MiniCPM5-1B, while producing more concise responses on many tasks. We further present IronLLM-0.6B-Light, which replaces RMSNorm with Dynamic Tanh and simplifies several computationally expensive components to improve inference and quantization efficiency. Together, the IronLLM models provide an effective performance-efficiency trade-off for resource-constrained deployment.

Abstract

Efficient Language Model Inference

Explore similar work

IronLLM: Forging Compact Edge-Native Language Models for Real-Time Embodied Intelligence

Unlocking Lossless Speedups in LLMs via Discrete Diffusion

RoofLang: Enabling AI-Driven Architecting of LLM Inference Systems