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

TL;DR - Axon is a strongly typed, shape-safe DSL for defining LLM architectures once and compiling them into standalone implementations for PyTorch, Triton, JAX, MLX, and vLLM. It aims to reduce framework lock-in while improving model portability, auditability, and performance.

• Uses concise Haskell-like specifications to support specialized architectures across training and inference frameworks.
• Compiles models into native framework implementations rather than relying on a single compatibility layer.
• Across 467 inference benchmarks on models from 135M to 32B parameters, median speedups ranged from 7% on PyTorch to 107% on MLX versus Transformers reference implementations.
• Native vLLM deployments using PagedAttention and KV caching achieved a 58% median speedup over Transformers implementations.