.@dMatrix_AI had matrix multiplication running in Mojo on their Corsair accelerator 6 days after getting access to the Mojo compiler.

At ModCon 2026, our Chief Scientist Abdul Dakkak walked through why this was possible with the Modular stack: https://t.co/GQd3sPwPIA

TL;DR：Modular 首席科学家 Abdul 以 Trainium、TPU 等案例说明，只要实现一个小型 C 接口并补充算子，Mojo/Max 技术栈就能以极短时间、几乎不改模型与服务层的方式扩展到全新硬件。

Modular 如何把 Max 拓展到 GPU 之外

Modular 首席科学家 Abdul 的这场分享，围绕一个核心问题展开：当 AI 加速器日益碎片化时，如何让同一套软件、同一份模型代码，稳定地跑在越来越多的硬件上？他给出的答案是 Modular 的模块化技术栈，以及一套用于引入新硬件的验证方法论。

问题背景：AI 软件的碎片化

如今的 AI 软件非常碎片化，每一种加速器都有自己的一套编译器、算子库和运行时，以及各自独立的开发路径。硬件厂商反复与不同框架集成，贡献模式往往最终演变成各种分叉版本。

最终承受这种混乱的是用户：他们不得不把这些零散的组件拼接在一起，用胶带强行缝合起来部署自己的模型。这绝不是一种愉快的体验。

Modular 的使命，是提供一个统一的计算层——足够稳定，让用户可以一致地部署；同时足够灵活，能够吸收新的模型、新的模态和新的硬件。

技术栈三层：Mojo、Max、Endpoints

Modular 平台建立在模块化组件之上，整个技术栈可以提炼为三层：

• Mojo——可移植的编程语言。今天运行的所有 Max 算子都是用 Mojo 编写的。
• Max——模型服务和框架层，即编写模型和提供模型服务的方式。
• Endpoints——把下面所有内容以 API 端点的形式暴露出来，可部署在 Modular 云端或用户自己的环境中。

一个重要的分离原则是：新硬件的引入主要应当影响技术栈的底层，技术栈的其余部分不受影响。 这意味着模型和服务层不需要为每一种加速器重新构建。

引入新硬件的方法论

1. Max Driver：让 Max 看得见硬件

要使用任何设备，Max 需要知道如何与它通信，而通信方式是通过 Max Driver。它提供连接层，处理诸如内存传输之类的事情——数据如何从 CPU 复制到加速器、如何从加速器复制回 CPU，以及函数启动、同步操作等。

引入新硬件的团队需要做的，就是实现那个小型的 C 接口（通过 C API 暴露）。一旦实现，Max 就可以将其作为插件加载进来。

2. 可选的设备驱动模块抽象

到这一步，Max 已经能看到并通信了，但还没有编程模型和性能。在驱动之上，设备驱动模块提供了可以跨硬件工作的共享工具和库抽象。

这些抽象是可选的——Modular 希望你使用它们，但不强制。因为存在逃生通道的需求：当你想充分利用硬件的最佳能力时，需要暴露目标硬件特定的能力。虽然这种情况与抽象理念相悖且很少见，但遇到时你仍希望停留在同一种语言中，而不需要切换到另一种语言。

此外，所有这些抽象都放在库中，而不是硬编码在编译器里。这使 Modular 能够快速演进，同时仍能写出与厂商实现相媲美的竞争性算子。

3. 验证方式：三个逐步增加难度的案例

Modular 用三种方式验证这套方法论：

1. 自己的团队在 Trainium 上；
2. 外部团队在 TPU 上；
3. 芯片厂商接入自己的技术栈。

在每种情况下都追问：什么改变了？什么是共享的？多快达到目标？迭代经验再反馈回技术栈，使其更加通用。

验证目标是端到端运行 Gemma 4 31B dense 模型（Google 开发，演示中省略了视觉编码器和投机解码器）。关键在于：模型只编写了一次，就部署到了本次演示的每一个目标上，也运行在主题演讲提到的另外六种硬件上。模型是现代架构，包含滑动窗口注意力（sliding window attention）、全连接块、采样操作等。

案例一：Trainium

Amazon Trainium 是一个矩阵乘法引擎，每颗芯片吞吐量达 630 TFLOPs（BF16），配备约 96 GB 高速 HBM 内存，带宽约 3 TB/s，每个节点有 16 颗芯片。

Trainium 暴露了三个需要技术栈泛化的层面：

1. 静态形状：Max 同时支持动态和静态形状（LLM 中批处理大小和序列长度是动态的），但此前没有在整个算子实现中保留静态信息。修改是确保静态形状信息在全过程中被保留。
2. Tile（分块）：Trainium 以 tile 方式运算，而此前的抽象基于向量（SIMD vectors）。技术栈被扩展为同时支持一维 SIMD 向量和二维 tile 张量。
3. 关键算子优化：模型跑通后，剩下的是对注意力（attention）和矩阵乘法这两个性能关键算子的精确优化。

模型服务和模型定义本身完全共享——它已在 B200 上验证过，因此移植到 Trainium 时可以确认它是在忠实实现这个模型。演示中的 MatMul 代码拥有与 CPU、GPU 上相同的抽象，但针对 Trainium 做了专门化：共享结构，同时不隐藏硬件细节。

移植时间线

• 5 月：两名工程师负责基础模块启用，大约一周内向量加法就跑起来了。
• 随后扩大团队构建所需 Max 操作，到 7 月底 Gemma 功能性可运行，并加入了集合通信（collectives）——因为该模型跨多颗芯片，实际使用了度数为 4 的张量并行。
• Max Serve 几乎不需要改动，耗时极少。
• 最后阶段是性能优化。

性能结果

矩阵乘法：用那页幻灯片大小的 Mojo MatMul 代码与现有最佳公开实现比较，在 Gemma 4 的形状上始终更快——至少持平，某些情况下比最先进水平接近快 1.6 倍。超越通常用汇编编写的厂商实现非常困难。

优化速度：以模型功能可用（Day 0）时的性能为 1 倍基准，到 Day 5：

• 首 token 延迟（time to first token）提升约 50 倍
• 每个输出 token 延迟提升约 7 倍

重点不在数字本身，而在速度（velocity）：模型跑通后，共享的基础设施（工具链、分析能力、相关技能）能让团队快速优化、识别瓶颈并改进。

什么变了，什么没变

改变的：设备驱动中新增了与 Trainium 通信的能力，并为矩阵乘法和注意力实现了关键算子。

不变的：模型定义、Max Serve 层、从 chunk prefill 到投机解码（spec decoders）的一切——全部完全一样。不需要更改任何与算子无关的内容。Modular 表示将继续与 AWS 合作，把 Trainium 支持引入 Max。

案例二：外部团队接入 TPU（Htech）

Trainium 证明了平台可以扩展到根本不同的架构，但那是 Modular 自己设计的。更严峻的考验是：外部团队能否做到同样的事情？

Modular 与 Htech 合作验证这一点。Htech 编译器工程师 Mihilo 介绍了外部合作伙伴视角下的体验——尤其是在只能从 Modular 工程团队获得有限支持的情况下。Htech 是一家拥有超过 250 名员工、在引入新硬件平台方面有成熟记录的工程公司，几个月前与 Modular 建立合作，目标是为 Mojo Max 带来 TPU 支持。

TPU 架构特点

TPU 是围绕大型矩阵乘法单元构建的领域专用加速器，是一种顺序执行机器，类似 VLIW CPU，但拥有非常宽的寄存器。所用 v6e TPU 包含：

• 单个张量核心：两个 256×256 矩阵乘法脉动阵列（MXU）
• 一个 8×128 向量单元（VPU）
• 单个标量核心，负责编排任务（DMA 调度、地址计算）

计算单元只能访问 scratchpad 中的数据——VM（Vector Memory）和 SPM（Scalar Memory），完全由软件管理。数据必须在 scratchpad 和 HBM 之间显式 DMA 传输，而这里的“用户“实际指的是编译器。

与 GPU 的对比：GPU 是高度并行的机器，拥有数百个 SM、数千线程，通过硬件动态隐藏内存延迟；TPU 是只有少量计算单元的顺序机器，但在矩阵运算上表现出色，且有不同的内存层次结构。

软件栈差异带来的核心挑战

GPU 有专用的 LLVM 后端，但 TPU 没有。它们使用领域专用编译器 XLA 和运行时 PGRT，例如完全禁止指针算术，并限制常规控制流构造的使用。而 Mojo 目前只会降级到 LLVM IR——当目标硬件没有 LLVM 后端时，这就是重大问题。

作为 LLVM IR 的替代，XLA 编译器的入口是 Mosaic——一个 ML 方言（ML dialect）集合，包含 arith、func、mem 等上游方言，加上一个用于表达 TPU 硬件能力的 TPU 专用方言。Mosaic 实现的 MatMul 算子中，参数由带形状的内存（shaped memref）表示，使用向量方言的加载操作加载，计算本身由 TPU 方言操作表示。

一个值得注意的趋势：领域专用加速器中，ML 正在成为一种统一层，取代 LLVM IR。 Mojo 开发者意识到了这一点，通过暴露 **ML 互操作（ML interop）**工具提供了集成途径。

解决方案：Mosaic 降级流水线

Mojo 通过 ML 扩展（ML extensions）提供 ML 互操作，让算子内部可以直接使用 TPU 操作。但这类算子无法降级到 LLVM IR，因此需要修改 Mojo 编译器实现，偏离默认降级路径，创建一条新的 Mosaic 降级流水线，从而把算子编译成 Mosaic。

Mihilo 还指出一个尚未解决的问题：Mojo 中宿主端与设备端张量实现之间的契约只是一个简单指针，不携带形状信息，这与 Mosaic 所需的带形状内存（shaped memref）存在张力——原始讲解在这一点上中断，未给出最终结论。

核心启示

Modular 的方法论可以概括为一条边界清晰的路径：实现小型 C 接口让 Max 看见硬件 → 用可选的库抽象覆盖跨硬件能力 → 只在注意力与矩阵乘法这类关键算子上做硬件专门化。模型定义、服务层和其余基础设施保持不变，因此同一份 Mojo/Max 代码可以跨越 NVIDIA、AMD、Qualcomm、Trainium 等多种硬件，用相同的命令运行。Trainium 案例表明，功能可用之后，共享的工具链与分析能力能让团队在数天内完成大幅性能优化。

Source: https://www.youtube.com/watch?v=hYjKbTAGOo0