Axon DSL Compiles One Model Definition to PyTorch, JAX, MLX and vLLM, with a 107% Median Speedup on MLX

arXiv 2608.19889·high signalXBlueskyLinkedInCopy link

Axon is a strongly typed, Haskell-syntax DSL that lets you write an LLM architecture once and compile it to standalone PyTorch, PyTorch+Triton, JAX, MLX and native vLLM implementations, explicitly as a hedge against the open-model ecosystem's dependence on a single platform. Across 467 inference benchmarking experiments on models from 135M to 32B parameters, Axon shows median speedups of 7% on PyTorch, 12% on PyTorch with Triton, 91% on JAX and 107% on MLX versus HuggingFace Transformers reference implementations. Deployed as native vLLM architectures with PagedAttention and KV-cache, Axon models hit a 58% median speedup over Transformers.