如何将 Mojo 语言和 MAX 引擎统一到一个代码库中

任何在生产环境中尝试优化重型神经网络推理的人都知道这种痛苦。首先，数据科学家用 Python 和 PyTorch 编写原型，一切运行得漂亮清晰。然后代码遇到性能瓶颈，在生产环境中他们开始用 C++、CUDA 或 Triton 重写关键部分。这造成了高级代码便利性和硬件之间的经典鸿沟。

Modular 团队由 Chris Lattner（LLVM、Clang 和 Swift 的创造者）领导，致力于从根本上解决这个问题。在他们的代码库 modular/modular 中，他们整合了平台的开放组件：Mojo 编程语言和 MAX 计算框架。

让我们弄清楚这个代码库里有什么，以及如何在实践中应用它。

单体仓库里有什么

长期以来，这个项目部分以封闭形式开发，但现在作者们正在逐步发布更多源代码。Modular 单体仓库包含几个相互关联的层：

• Mojo 编译器源代码（目录 KGEN）和语言标准库（mojo/stdlib）。
• 一组优化的计算内核 max/kernels，用于加速 GPU 和 CPU 上的矩阵运算。
• 一个开箱即用的推理服务器 max/python/max/serve，兼容 OpenAI API。
• 用于构建计算图的模型管道 max/python/max/pipelines。

有趣的是贡献权限的划分。社区的拉取请求被欣然接受到 Mojo 标准库、MAX 内核和架构管道中。然而，编译器源代码本身仍由 Modular 团队直接管理，尽管他们的代码是开放的。

没有魔法的 Mojo 语言

Mojo 被设想为 Python 思想的直接演进，但具有 C 级别的性能和严格的类型控制。语法对任何 Python 开发者来说都很熟悉，而该语言通过 MLIR 和 LLVM 直接编译成机器代码。

在代码中，你可以将 Python 的动态类型与刚性结构、手动内存管理和 SIMD 向量指令结合起来。

下面是一个用 Mojo 编写的简单向量加法函数示例：

from algorithm import vectorize
from sys.info import simdwidthof

fn add_vectors[type: DType, size: Int](
    a: DTypePointer[type],
    b: DTypePointer[type],
    result: DTypePointer[type]
):
    alias width = simdwidthof[type]()

    @parameter
    fn vector_add[simd_width: Int](idx: Int):
        let val_a = a.load[width=simd_width](idx)
        let val_b = b.load[width=simd_width](idx)
        result.store[width=simd_width](idx, val_a + val_b)

    vectorize[vector_add, width](size)

注意用 fn 而不是 def。这个关键字在编译时启用严格的静态类型检查和变量生命周期检查，几乎像 Rust 一样。同时，在单个项目中，如果某个模块的最大速度还不是关键要求，你可以自由导入 NumPy 或 SciPy 等标准 Python 库。

MAX 框架和模型运行

语言本身没有运行时环境就没什么用处。对于神经网络工作，Modular 开发了 MAX（Modular Accelerated Execution）。这是一个引擎，接收现成的计算图，优化张量内存放置，并在可用硬件上运行它们。

代码库中最有用的工具之一是现成的推理服务器。它启动一个完全复制 OpenAI /v1/chat/completions 接口的 HTTP 端点。这意味着你可以用它替换 vLLM 或 TGI，而无需重写服务中的客户端代码。

在底层，MAX 使用自定义的 Mojo 内核进行注意力机制和矩阵乘法计算。因为内核是用 Mojo 而不是纯 CUDA 汇编编写的，所以开发者更容易为新芯片和专用加速器调整架构。

实践中你会遇到什么

尽管这个项目引起了很多关注，但重要的是要冷静评估其当前状态。

许可模式是混合的。代码库中的大部分开源代码采用 Apache License v2.0（带 LLVM 例外）分发。但是，MAX 二进制构建的使用受单独的 Modular Community License 管辖。如果你计划在封闭的商业环境中部署该解决方案，你应该仔细阅读他们网站上的法律条款。

第二个问题是成熟度。纯 Mojo 的第三方库生态系统仍在形成中。是的，Python 互操作性运行顺畅，但调用 Python 代码时会失去速度优势，回到 CPython 运行时开销。

谁应该关注这个项目

如果你正在部署大语言模型并遇到推理延迟瓶颈，MAX 和现成管道的组合绝对值得在你的基准测试中进行测试。推理服务器部署快速，提供了良好的基准测试起点。

对于底层工程师和自定义 CUDA 内核作者，该项目提供了编写快速代码的能力，而无需在 C++ 中使用冗长的模板样板。在 Mojo 中编写 SIMD 优化的入门门槛明显低于传统工具。

Modular 团队正在有条不紊地开放其平台的组件。你可以探索编译器结构、摆弄标准库，并直接从代码库运行本地推理服务器。看来 Python 和 C++ 的组合在 ML 基础设施领域终于有了一个真正的竞争对手。