AI 内核的可移植性在结构上是不可能的，因为 TPU 的 Pallas、NVIDIA 的 CuTile 和 CUTLASS、AWS 的 NKI、AMD 的 FlyDSL 以及 Tenstorrent 的 tt-Metalium 各自暴露了硬件特定的概念，没有任何通用 DSL 能够统一这些概念。证据是：MaxText 的 MoE 分组矩阵乘法在 TPU 上以 282 行 Pallas 代码实现，而 flashinfer 在 Blackwell SM100 上等效的实现需要 400 万行生成的 CUDA，两者共享代码为零，因为算法本身在不同硬件上就存在分歧。

可移植性是个神话：为什么最好的 AI 栈永远不会与硬件无关

在 LinkedIn 上联系我： https://www.linkedin.com/in/patrick-toulme-150b041a5/

在 X 上关注我： https://x.com/PatrickToulme

动机

免责声明： 这篇博客比我之前的博文更带有个人观点。

AI 行业说它想要可移植的内核。但它一直在构建硬件特定的内核。

TPU 有 Pallas。Trainium 有 NKI。NVIDIA 有 CUDA C、CUTLASS、Triton、CuTile、CuTe —— 已经有五个 DSL 了，而且还在增加。AMD 采用了 Triton，这是最接近“可移植”内核 DSL 的东西——然后还是构建了 FlyDSL。Tenstorrent 有 tt-Metalium。Mojo 的目标是 NVIDIA 和 AMD。

如果可移植性行得通，我们早就有一个通用 DSL 了。但我们有很多个——而且数量还在增长。

在 TPU 上是 282 行的 Pallas Python。在 Blackwell 上是 400 万行生成的 CUDA C++。同一个操作——MoE 分组矩阵乘法。共享代码为零。这不是因为谁没能实现可移植性，而是因为硬件差异足够大，导致最优算法、分块策略、内存分阶段和同步方式都完全不同。在性能层面，可移植性从来就不可能实现。

然而团队们仍在追求它。初创公司耗费数月构建“硬件无关”的训练栈，到处都能跑，但跑得快的地方一个都没有。硬件供应商扭曲自己的芯片去适应别人的编程模型。 工程师们编写可移植的内核，在两个平台上达到 30% 的模型利用率（MFU），而不是在一个平台上达到 90%——在前沿规模下，这个差距相当于数百万美元的计算浪费。

这篇文章认为，未来最好的 AI 栈将不是可移植的——而且这是一件好事。 它阐述了为什么内核层面的可移植性是个神话，为什么每个硬件供应商最终都会构建自己的 DSL，以及这些 DSL 应该是什么样子才能最大化性能和可用性。

AI 栈的三个层次

每个 AI 训练或推理栈都有三个层次。理解哪些层次是可移植的——哪些不是——是理解为什么行业不断从头重建基础设施的关键。

第一层：数学（可移植）

# 这段代码在 TPU、GPU、CPU、Trainium 上都能运行。数学计算到处都一样。
y = jnp.dot(x, w)                    # 矩阵乘法
y = jax.nn.softmax(logits)            # softmax
loss = optax.softmax_cross_entropy(logits, labels)

PyTorch 和 JAX 在这一层是可移植的。torch.matmul 和 jnp.dot 表达了一个数学操作——两个矩阵相乘。框架没有指定怎么做。这是真正的可移植性。它行得通，而且很重要。

但 jnp.dot 并不是真正在硬件上运行的代码。真正运行的是编译器输出，它做出了数千个硬件特定的决策，涉及分块、内存放置、指令选择、同步和调度。数学是一样的。其他一切都不一样。

第二层：编译器（可移植性的幻象）

XLA、torch.compile、Triton 的编译器——它们位于你的可移植数学代码和硬件之间。它们把 jnp.dot 转换成 TPU VLIW 指令包、或 CUDA PTX、或 Trainium NeuronCore 指令。API 是可移植的。输出是不可移植的。

这就是可移植性幻象所在。**当有人说“JAX 在 TPU 和 GPU 上都能跑”时，他们的意思是：JAX 的追踪和 HLO 生成是共享的，但 XLA-TPU 和 XLA-GPU 是不同的编译器，它们为本质不同的硬件生成本质不同的代码。**它们共享一个 IR 格式。它们不共享后端。

PyTorch 也是如此。torch.compile 配合 Inductor 后端针对 NVIDIA GPU。torch.compile 配合 XLA 后端针对 TPU。torch.matmul 调用是一样的。它下面的所有东西都分道扬镳。

第三层：硬件原生代码（不可移植）

这一层是 MFU 赢或输的地方。

在 TPU 上，一个融合注意力内核是一个 Pallas 程序，它显式管理 VMEM，在 8 个子通道 × 128 通道上进行分块，发出用于 HBM↔VMEM 传输的 DMA 操作，并生成将 MXU 矩阵乘法、VPU 元素操作和 XLU 跨通道操作打包到单周期中的 VLIW 指令包。

在 Blackwell 上，同一个逻辑操作是一个 CuTile 程序，它分配带有原子竞争处理的 TMEM 列，用 elect.sync 选举一个 warp 领导者，发出 tcgen05 MMA 指令，管理 20 个 mbarrier 对象用于异步流水线，并在 TMEM 分配失败时用 NANOSLEEP 退避。

在 Trainium 上，那个操作是一个 NKI 内核，它将工作划分到各个 NeuronCore 上，通过 SBUF（片上 SRAM）和 PSUM（累加内存）分阶段处理数据，并用显式的分区维度分块发出张量引擎指令。

这些不是同一个算法的不同实现。**它们是由不同硬件约束驱动的不同算法。**你不能把一个移植到另一个。你得从头重写。

证据：相同的数学，零共享代码

以下是这个论证最具体的版本。混合专家分组矩阵乘法——同一个操作为两个不同的硬件目标编译。

在 TPU 上：MaxText MegaBlox（Pallas）

从我上一篇文章中：MaxText 的 MoE 层编译成 29 个 GMM Pallas 内核调用。每个都是一个 tpu_custom_call，它：

• 在一个 3D 网格（tiles_n， num_active_tiles， tiles_k）上运行
• 从 HBM 到 VMEM DMA 获取 [512， 1024] 大小的已排序 token 分块
• 通过 group_ids[grid_id] 选择正确的专家权重分块
• 在 VMEM f32 临时缓冲区中累积 dot(lhs_tile, rhs_tile)
• 应用组边界掩码并将结果存储回 HBM

反向传播使用两个不同的内核：gmm 带 transpose_rhs=True 用于输入梯度，以及 tgmm——一个结构上不同的内核——用于权重梯度。分块大小为 (512， 1024， 1024)，针对 TPU MXU 的 256×256 脉动阵列进行了调优。

手写内核代码总数：282 行 Pallas Python，针对 TPU 内存层次结构和 MXU。

在 NVIDIA 上：flashinfer MoE（CUDA）

flashinfer PR #2917：300 个文件，为同一个逻辑操作——MoE 批量 GEMM——生成的 CUDA 内核。每个文件针对 SM100（Blackwell）并直接在文件名中编码了优化参数：

batched_gemm_e2_sm100_s128x128x128_...bf16_tma_warpspecialized_cooperative_align.cu
batched_gemm_e2_sm100_s128x256x128_...fp8_tma_warpspecialized_pingpong_align.cu

分块大小（128×128×128， 128×256×128）。数据类型（bf16， fp8）。内存访问模式（TMA）。调度策略（cooperative， pingpong）。架构目标（SM100）。全都硬编码在内。该 PR 增加了大约 400 万行生成的 CUDA——不是因为开发者喜欢写 CUDA，而是因为在 SM100 上达到峰值性能需要 SM100 特定的代码。

为什么可移植的 DSL 无法解决这个问题

显而易见的回应是：“如果我们有一个能抽象掉这些差异的 DSL 呢？”

这就是 Triton 的承诺，在某种程度上也是 MLIR 生态系统的承诺。把你的内核写在一个硬件无关的 DSL 中，编译器将它降低到 TPU、GPU 或 Trainium。

**问题在于，一个可移植的 DSL 只有在编译器足够复杂以弥合抽象操作和硬件特定指令之间的鸿沟时才能工作。**而那个编译器需要做一些从来没有编译器做到过的事情。

可移植编译器需要做什么

假设在一个抽象 DSL 中有一个 tile_matmul(a, b)。编译器必须决定：

对于 TPU：

• 为两个操作数分配 VMEM 缓冲区
• 发出 HBM→VMEM 传输的 DMA copy-start/copy-done 对
• 将分块大小设为 (512， 1024, 1024) 以匹配 256×256 MXU 脉动阵列
• 将 MXU 矩阵乘法 + VPU 元素操作 + XLU 跨通道混洗打包进 VLIW 指令包
• 在循环迭代中对 DMA 和计算进行双缓冲

对于 Blackwell：

• 使用原子 UTCATOMSWS.FIND_AND_SET 分配 TMEM 列
• 每次 MMA 之前用 elect.sync 选举一个 warp 领导者
• 仅通过领导者线程发出 tcgen05 MMA 指令
• 管理 20 个基于奇偶校验重用的 mbarrier 对象用于异步流水线
• 在 TMEM 分配竞争时用 100ns 的 NANOSLEEP 退避
• 在 10ms 屏障等待超时后等待结果

对于 Trainium：

• 通过 SBUF（片上 SRAM）配合显式的分区维度分块来分阶段处理数据
• 在 NeuronCore 分区上发出张量引擎指令
• 在 PSUM 内存空间中累积
• 管理 NeuronCore 级别的并行性

这些不是同一个算法的不同指令选择。它们是不同的算法。TPU 上的双缓冲 DMA 方法，Blackwell 上的领导者选举异步 MMA，以及 Trainium 上的分区分块张量引擎，是结构上不同的程序，通过完全不同的机制解决同一个数学问题。

并集问题

一个可移植的 DSL 如果要暴露足够的硬件细节来编写高性能内核，就需要包含所有硬件概念的并集：

• TPU： VMEM， HBM， DMA 引擎， MXU， VPU， XLU， 子通道/通道， VLIW 指令包
• Blackwell： SMEM， TMEM， L2 缓存， 张量核心， tcgen05， mbarrier， TMA， warp 特化
• Trainium： SBUF， PSUM， HBM， 张量引擎， NeuronCore， 分区维度

那不是 DSL——那是三件 DSL 穿了一件风衣。里面写的每个内核都需要有 if target == TPU: ... elif target == NVIDIA: ... elif target == TRAINIUM: ... 的转义出口，这破坏了共享抽象的目的。

或者，一个隐藏这些差异的可移植 DSL 会产生泛化代码，不利用其中任何一个差异——然后你又回到了编译器层，希望编译器足够聪明，能恢复你丢掉的性能。

DSL 的激增就是证明

如果内核层面的可移植性可以实现，市场现在应该已经收敛到一个 DSL 上了。相反，每个硬件供应商都独立构建了自己的 DSL：

仅 NVIDIA 就有四个用于编写内核的 DSL——CUDA C、CUTLASS/CuTe DSL、Triton 和 CuTile——从线程级 C 到分块级 Python。针对一个供应商的硬件就有四个抽象级别，因为即使在一个架构系列内，一个 DSL 也不够用。

**AMD 的轨迹尤其能说明问题。**他们采用了 Triton——这个行业里最接近“可移植”内核 DSL 的东西——并为其做了 ROCm 分支。但采用 Triton 还不够。AMD 仍然构建了 FlyDSL：一个 Python DSL，带有显式的寄存器级布局、通道级操作和 CuTe 风格的分块代数，通过他们自己的 MLIR Fly 方言降低到 ROCDL 机器码。即使你采用了一个“可移植”的 DSL，你最终还是得构建一个硬件原生的 DSL 来获得可移植 DSL 无法提供的性能。

这不是工程失败。而是证据表明硬件特定的内核编程是本质上硬件特定的。

为了最高 MFU：用 Python 写 ISA

以下是所有成功 DSL 的统一原则：性能最高的内核 DSL 是硬件指令集的薄 Python 包装器。

Pallas 并没有抽象掉 MXU——它给了你 dot()，直接映射到 MXU 矩阵乘法指令；load()/store() 映射到 DMA 操作；以及对应 VMEM/HBM/SMEM 内存空间的 Ref 类型。

CuTile 并没有抽象掉 tcgen05——它给了你 tile_matmul()，映射到 tcgen05 MMA；tile_load() 映射到 TMA 描述符；以及自动的 TMEM 分配，生成 UTCATOMSWS.FIND_AND_SET 指令。

NKI 并没有抽象掉张量引擎——它给了你 nisa.nc_matmul()，映射到张量引擎指令，以及显式的 SBUF/PSUM 内存分阶段。当你写 NKI 时，你本质上是在用 Python 写 Trainium 的 ISA。

这些 DSL 的价值不在于抽象——而在于人体工程学。**它们让你用 Python 而不是 C 或汇编来写硬件特定的代码。**关键是它们没有隐藏硬件。它们是通过熟悉的语法暴露了硬件。

这是正确的权衡。C 级别的 ISA 访问（CUDA C++、原始 PTX、汇编）给你最大的控制，但迭代速度和可读性很差。一个抽象的、可移植的 DSL 给了你人体工程学，但剥夺了决定 MFU 的硬件控制。一个基于硬件 ISA 的 Python DSL 两者兼得：硬件控制加上 Python 的迭代速度。

flashinfer 的 PR 是 400 万行生成的 CUDA C++。MaxText 的 MegaBlox 内核是 282 行 Pallas Python。两者都在各自的硬件上达到了高 MFU。开发者体验的差异是巨大的——但 Pallas 之所以工作，不是因为它比 CUDA 更抽象。而是因为它直接将 Python 映射到 TPU 硬件概念，所以这 282 行表达了与生成的 CUDA 相同级别的硬件特定意图，只是更简洁。

给硬件供应商的建议

如果你在构建 AI 加速器硬件，并希望开发者在你的芯片上达到峰值性能，以下是来自每个成功 AI 硬件平台的证据告诉我们的：

1. 构建你自己的 Python DSL

不要等待一个开源 DSL 来支持你的硬件。 Google 构建了 Pallas。NVIDIA 构建了 CuTile。AWS 构建了 NKI。在每种情况下，供应商构建的 DSL 都优于第三方替代方案，因为 DSL 设计者对其硬件的性能模型有深刻理解。

你的 DSL 应该让开发者能够用 Python 编写映射到硬件执行模型的内核。不要用 C。不要用自定义语言。用 Python——因为 ML 生态系统存在于 Python 中，从研究者洞察到高性能内核的最快路径就是一个直接表达你硬件概念的 Python 文件。

为芯片构建 DSL。不要为 DSL 构建芯片。

2. 暴露 ISA，而不仅仅是高级抽象

你的 DSL 应该在两个级别上运行：

低级：虚拟 ISA 访问。 让开发者能够表达硬件的指令集——你的矩阵乘法单元的分块操作、内存层次结构的显式分阶段、同步原语。这是峰值 MFU 的来源。理解你硬件的开发者应该能够用 Python 编写接近 ISA 级别的代码，并确切知道会生成哪些指令。

高级：编译器降级操作。 并非所有用户都想编写 ISA——这没问题。对于不需要峰值性能或正在原型开发的开发者，提供更高级别的操作（如 tile_matmul、tile_load），由编译器降级到 ISA。这就是 Pallas 的 dot() 和 CuTile 的 tile_matmul() 所在的层级——高级到足以提供人体工程学，低级到专家能够预测生成的代码。

关键是，这两个级别必须在同一内核中可混合使用。 Pallas GPU 在这方面做对了：你可以编写一个内核，在编译器足够好的部分使用高级 jnp 操作，而在需要直接硬件控制的部分降级到显式的 tcgen05 指令——在同一个函数、同一个 Python 文件中。开发者可以逐操作（而非逐内核）选择抽象级别。

这是扩展采用率的方法。新用户从高级操作开始，获得合理的性能。专家用户对热路径降级到 ISA 级别控制。一个只提供高级接口的 DSL 会让专家用户留下性能遗憾。一个只提供 ISA 访问的 DSL 学习曲线太陡，难以被采用。最终胜出的 DSL 将是那些让开发者能够在这两个级别之间无缝滑动的 DSL。

3. 不要为了开源 DSL 而约束你的硬件

这是最反直觉的建议，但证据非常清楚。

Triton 是一个好的 DSL，用于非 Blackwell 的 NVIDIA GPU。