AI 推論與開發工具
Axon Generates Five LLM Backends from a Strongly Typed DSL, Cutting MLX Median Latency to 48.3% of Transformers
Axon expresses model architectures as functional specifications with symbolic tensor shapes, then compiles them to PyTorch, Triton, JAX, MLX, and native vLLM. Across 467 inference tests, it accelerated most generative models, though some BF16, encoder, and sequence-to-sequence models still showed accuracy or performance regressions.
The newly proposed Axon aims to elevate LLM architectures from framework-specific Python classes into a strongly typed domain-specific language that can be compiled across backends. Developers use a Haskell-like syntax to describe attention, MoE, state-space modules, parameter paths, and KV cache behavior; types such as `Tensor[B,S,D]` allow the compiler to validate symbolic shapes before generating code. The compilation pipeline first expands functions and libraries, normalizes them into a typed Graph IR, and then lowers the representation to PyTorch, Triton, JAX, MLX, or vLLM, without requiring the Axon runtime in the deployment environment.
The [paper](https://arxiv.org/abs/2608.19889) reports 467 inference comparisons spanning 47 model definitions and models ranging from 135M to 32B parameters. Relative to Hugging Face Transformers, the reported median speedups for PyTorch, Triton, JAX, and MLX were 7%, 12%, 91%, and 107%, respectively. Across 126 MLX tests, median execution time was 48.3% of the baseline.
When compiled into native vLLM architectures supporting PagedAttention and KV cache, the models achieved a median speedup of 58%, rather than being loaded through the slower Transformers-compatible path. This distinction aligns with the differences between native implementations and the Transformers backend described in [vLLM’s model implementation documentation](https://docs.vllm.ai/en/latest/models/supported_models/).
The technical value lies in enabling a single auditable specification to serve training, desktop deployment, and server inference simultaneously, reducing the cost of rewriting every new architecture five times. However, the current results come solely from the authors’ tests: short forward-pass workloads such as T5 and mT5 showed significant slow outliers; MLX BF16 outputs for five models failed to meet numerical-consistency criteria; and the paper does not include a public code repository.
Engineering teams should next watch for the release of the compiler and full benchmark suite, support for dynamic control flow and quantized models, and whether the speedups persist under continuous batching and multi-user serving workloads.