论文
一次编写，处处运行：面向形状安全、框架无关LLM架构的Axon DSL
Write Once, Run Everywhere: The Axon DSL for Shape-Safe and Framework-Agnostic LLM Architectures
摘要
整个开源语言模型生态系统实际上依赖于单一平台。如果这个平台明天被迫关闭会怎样？实现和维护高效的模型定义，并在不同训练与推理范式之间进行转换，是一项资源密集型任务，严重限制了模型的效率与可移植性，阻碍了规模化与部署。在本文中，我们提出 Axon，一种具有类Haskell语法的强类型领域专用语言（DSL），为 LLM 架构实现“一次编写、处处运行”的范式。通过将协作建立在一门语言规范之上，而非某个特定框架的愿景之上，Axon 促进开放合作，使研究者能够实现高度专业化的架构，而无需放弃优化基础设施或接受部署锁定。Axon 支持简洁、可审计的规范，并可自动编译为面向主流框架的独立实现：PyTorch、PyTorch with Triton、JAX、MLX 和 vLLM。在针对参数量从 135M 到 32B 模型的 467 项推理基准实验中，相对于 Transformers 参考实现，我们展示了如下中位加速：PyTorch 上 7%、PyTorch with Triton 上 12%、JAX 上 91%、MLX 上 107%。当作为带有 PagedAttention 和 KV 缓存的原生 vLLM 架构部署时，Axon 模型相对 Transformers 实现取得 58% 的中位加速。
