用一种语言定义大模型,跨框架高效运行,解决部署锁定问题。
Write Once, Run Everywhere: The Axon DSL for Shape-Safe and Framework-Agnostic LLM Architectures

- 基于类似Haskell的强类型语法,统一定义大模型架构。
- 在多种框架上实现平均7%至107%的推理加速,最高提升107%。
- 适合需要跨平台部署、避免框架依赖的研究者和开发者。
开源大模型生态高度依赖单一平台,若该平台停用将导致严重后果。当前模型实现与跨训练/推理环境迁移成本高昂,限制了效率与可移植性。本文提出Axon,一种类Haskell的强类型领域专用语言,支持‘写一次,到处运行’的大模型架构设计。通过以语言规范为基础协作,而非特定框架,Axon促进开放合作,使研究者可在不牺牲优化能力的前提下构建专用架构,避免部署锁定。Axon支持简洁可审计的模型描述,并能自动编译为PyTorch、PyTorch+Triton、JAX、MLX和vLLM等主流框架的独立实现。在467次推理基准测试中(模型规模135M至32B参数),相比Transformers参考实现,其在PyTorch上平均提速7%,PyTorch+Triton上12%,JAX上91%,MLX上107%。当以原生vLLM架构部署并启用PagedAttention与KV缓存时,平均速度提升达58%。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。