Write One Model, Run It Anywhere: The Axon DSL

The entire open-source LLM ecosystem depends on one codebase: HuggingFace Transformers. What if it disappeared tomorrow? (It won't, but the question exposes a fragility most people ignore.)

We all pin our architectures to PyTorch implementations. Porting a model from Transformers to MLX, JAX, or a production vLLM deployment means rewriting it — differently — for each target. The rewrites introduce bugs, diverge in behavior, and bleed engineering cycles that should go into research.

"Write Once, Run Everywhere: The Axon DSL for Shape-Safe and Framework-Agnostic LLM Architectures" by Nielsen, Namazifard, Poech, and Schneider-Kamp at University of Southern Denmark is a direct shot at this problem. They built a domain-specific language that compiles LLM architecture definitions into native implementations for five different backends — and it's faster than hand-written code on every single one of them.

---

The Problem: Framework Lock-In

Right now, if you want to deploy a novel architecture (say, a hybrid attention-MoE transformer with custom gating), you write it in PyTorch. Then if you want Apple Silicon inference, you port it to MLX. If you want TPU training, you port it to JAX. If you want production serving with PagedAttention, you port it to vLLM. Each port is a separate engineering project with its own surface for bugs and performance regressions.

The result: researchers stick with whatever framework they started in. Novel architectures rarely escape their original ecosystem. The community loses the portability that makes open collaboration work.

What Axon Is

Axon is a strongly typed, Haskell-like DSL. You describe your model architecture — shapes, layers, attention patterns, normalization — in a single specification file. The Axon compiler then generates standalone, optimized implementations for:

PyTorch	Reference implementation, +7% median speedup over Transformers
PyTorch + Triton	Kernel-fused variant, +12% median speedup
JAX	TPU-friendly, +91% median speedup
MLX	Apple Silicon native, +107% median speedup
vLLM	Production serving with PagedAttention & KV-cache, +58% median speedup

Those numbers come from 467 inference benchmarks across models from 135M to 32B parameters. Every backend beats its equivalent Transformers reference — often by a wide margin.

How Does It Get Faster?

The key insight: by separating architecture specification from implementation optimization, Axon lets each backend apply its best compilation strategy without the spec author having to know anything about it. The compiler knows the exact tensor shapes and dataflow from the strongly typed spec — no dynamic shape inference, no guesswork. That information feeds directly into:

• Triton kernel fusion for PyTorch: the compiler identifies fusion opportunities at the graph level
• XLA compilation for JAX: static shapes enable aggressive HLO optimizations
• Metal Performance Shaders for MLX: the compiler maps operations to Apple's accelerator primitives
• PagedAttention integration for vLLM: native KV-cache management without adapter layers

The result is a model that's correct by construction across all targets and faster than hand-written code on each.

mindmap
  root((Axon DSL))
    Spec
      Strongly typed
      Haskell-like syntax
      Shape-safe
    Compiler
      Graph-level optimization
      Backend codegen
    Backends
      PyTorch +7%
      Triton +12%
      JAX +91%
      MLX +107%
      vLLM +58%
    Benefits
      Write once
      No framework lock-in
      Auditable specs
      Research → production

What This Unlocks

The implication isn't just "my benchmarks are faster." It's that the friction between research and production collapses. A researcher can prototype a new attention mechanism in Axon, validate it on PyTorch, and deploy it on vLLM — all from the same spec. No porting. No "this feature doesn't exist in MLX." No production bugs that weren't in the research code.

For the ecosystem, it means specialized architectures (linear attention, state space models, hybrid MoE) can circulate as Axon specs rather than as PyTorch checkpoints that only work on CUDA. Apple Silicon users, TPU users, and GPU users all get the same architecture without anyone maintaining three codebases.

Limitations

• New operations require extending the compiler. If your architecture uses a custom CUDA kernel that has no Axon primitive, you're writing compiler support first.
• Training coverage is not yet benchmarked. The 467 experiments cover inference only. Training pipelines (gradient computation, sharding, FSDP) require additional compiler work.
• Haskell syntax is a barrier for researchers more comfortable with Python DSLs. The paper acknowledges this but argues the type safety justifies the learning curve.
• Ecosystem adoption is the real open question. Axon competes with network effects — Transformers has millions of users and thousands of pre-built models.

The Takeaway

Axon is one of those ideas that feels obvious after you hear it. Of course model architectures should be framework-agnostic. Of course the compiler should handle backend optimization. The fact that it took until 2026 for someone to ship this tells you how much implicit inertia the HuggingFace monoculture has created.

For anyone deploying LLMs across diverse hardware — and especially for anyone who wants their research to actually reach production — this is worth watching closely. The benchmarks are real. The speedups are real. The question is whether the DSL can build enough momentum to break the monoculture.