一次编写，到处运行：用于形状安全且框架无关的 LLM 架构的 Axon DSL

arXiv 论文提出 Axon，一种具有 Haskell 风格语法的强类型领域特定语言，支持 LLM 架构的“一次编写，到处运行”。Axon 可自动编译为 PyTorch、PyTorch with Triton、JAX、MLX 和 vLLM 的独立实现。在 467 次推理基准测试中，相比 Transformers 参考实现，Axon 在 PyTorch 上中位加速 7%，Triton 上 12%，JAX 上 91%，MLX 上 107%，作为 vLLM 原生架构时加速 58%。

当前开源语言模型生态几乎依赖单一平台，存在关闭风险且移植成本高。Axon 通过语言规范而非特定框架来实现协作，避免供应商锁定，同时提供显著性能提升，可能改变 LLM 架构的开发和部署方式，降低跨框架迁移成本。

Axon 是一种强类型 DSL，利用类型系统保证形状安全，通过自动编译生成各框架的原生代码，从而优化推理。其性能提升可能源于针对不同框架的代码生成优化，如利用 vLLM 的 PagedAttention 和 KV-cache，以及 JAX 和 MLX 的硬件加速特性，使得模型定义与实现解耦，允许跨框架高效部署。

产品团队可采用 Axon 定义模型架构，一次开发即可部署到多种推理后端，提升部署灵活性并降低成本。用户可能获得更快的推理速度，尤其在使用 JAX、MLX 或 vLLM 的场景下，性能提升明显。研发流程中模型审计和框架迁移工作简化。

• 评估在推理服务中引入 Axon 编译的 vLLM 实现，以利用 58% 的中位加速，降低算力成本。
• 面向多框架部署需求的企业用户，提供基于 Axon 的模型架构咨询和迁移服务，降低锁定风险。
• 开发基于 Axon 的模型架构市场，促进社区共享可移植的 LLM 架构定义。
• 结合 Axon 的类型安全特性，设计更严格的模型验证流程，减少生产环境中的形状错误。

• Axon 的编译是否完全自动化，是否需要人工干预或调整？
• Axon 是否支持训练场景，还是仅限推理？
• 性能基准是否覆盖所有常见模型和硬件组合，是否具有代表性？
• Axon 的语法学习曲线和调试工具是否成熟，社区支持如何？