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

Source

• Source name: arXiv
• Source kind: paper
• Authors / organization: Jacob Nielsen, Danial Namazifard, Lukas Galke Poech, Peter Schneider-Kamp
• Venue: arXiv: cs.AI
• Published: 2026-08-20
• Collected: 2026-08-21
• URL: http://arxiv.org/abs/2608.19889v1
• PDF: https://arxiv.org/pdf/2608.19889v1
• DOI: UNKNOWN
• Registry group: academic_apis_and_preprints
• Evidence role: canonical_academic_index
• Language: en
• Country / region: UNKNOWN

Plain-Language Brief

Researchers introduce Axon, a strongly typed domain-specific language with Haskell-like syntax that aims to enable a write-once, run everywhere paradigm for Language Model architectures. This allows for efficient model definitions and seamless transition between different training and inference regimes, promoting scalability and deployment flexibility. The study shows speedups of up to 107% compared to reference implementations from Transformers across various frameworks.

What the Document Says

• A new language called Axon is presented with Haskell-like syntax for defining Language Model architectures (LLM).
• Axon promotes a write-once, run everywhere paradigm to reduce resource usage and improve portability.
• By basing collaboration on a language specification instead of specific frameworks, Axon encourages open cooperation and supports highly specialized architectures without compromising optimization infrastructure or deployment lock-in.
• Axon allows for automatic compilation of concise, auditable specifications into standalone implementations for leading frameworks such as PyTorch, PyTorch with Triton, JAX, MLX, and vLLM.
• The study conducted 467 inference benchmarking experiments on models ranging from 135M to 32B parameters, showing speedups of 7% on PyTorch, 12% on PyTorch with Triton, 91% on JAX, and 107% on MLX compared to reference implementations from Transformers.
• When deployed as native vLLM architectures with PagedAttention and KV-cache, Axon models achieve a median speedup of 58% over Transformers implementations.

Key Concepts Mentioned

• Axon
• Language Model Architecture (LLM)
• Haskell-like syntax
• write-once, run everywhere paradigm
• compilation
• benchmarking experiments
• speedups

Claims Made by the Document

• Axon enables a write-once, run everywhere paradigm for LLM architectures.
• Axon allows automatic compilation of specifications into leading framework implementations.
• The study shows speedups ranging from 7% to 107% compared to reference Transformers implementations.

Evidence Used

• Introduction of Axon and its purpose for LLM architectures.
• Description of the benefits of Axon's syntax and how it fosters collaboration and specialized architecture implementation.
• Explanation of Axon's ability to compile specifications into multiple framework implementations.
• Results from the 467 inference benchmarking experiments showing speedups across various models and frameworks.

Methods / Approach

UNKNOWN

Limitations Stated by the Authors

UNKNOWN

Limitations Noted by the Agent

Agent interpretation — not from the source.

UNKNOWN

Open Questions

UNKNOWN

Notes for Future Q&A

May include interpretation. Labeled where it does.

UNKNOWN

Raw Abstract / Excerpt

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.