Search papers, labs, and topics across Lattice.
This paper introduces Axon, a strongly typed domain-specific language designed to facilitate the development of shape-safe and framework-agnostic architectures for large language models (LLMs). By enabling a write-once, run everywhere approach, Axon allows researchers to create specialized model definitions that can be efficiently compiled across multiple frameworks, significantly enhancing model portability and efficiency. The evaluation reveals substantial performance improvements, with median speedups of up to 107% on JAX and 58% on native vLLM architectures compared to traditional Transformers implementations.
Axon achieves up to 107% speedup on JAX, revolutionizing how LLMs can be efficiently deployed across different frameworks without sacrificing optimization.
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.
