谈谈AI的软硬件交付界面

TL;DR

聊这个话题的出发点有几个：

0.1 知识屏障

大多数搞并行计算开发的业务方，其实是数学、物理、生物统计、材料这些专业出身的算法博士。他们对数学算法本身理解很深，但对GPU的微架构了解相对有限，写代码的能力也偏弱。当然，少数OI竞赛和数学/物理竞赛双修的选手不是没有，但毕竟是少数。再往上，能从代码一路通到芯片架构层面的人就更少了。这几年工作学习搬了一些砖，感触尤深。

说到底，交付界面的核心还是那三个问题：

。这事儿值得好好捋一捋。

0.2 异构计算的同构表达

这词儿听着有点绕：异构意味着专用，可人类写程序的心智偏好，还是希望一套系统能尽量通用、泛化能力强。本质上，我们是在异构器件上追求某种程度的同构表达。前段时间翻了翻Jim Keller在61DAC的Keynote^[1]，再结合今年HotChips上各种AI翻跟斗百花齐放、国产GPU层出不穷的形势，讨论这个话题更有意思了。

1. The Good Old days

Jim Keller提到，过去

，只要一套标准的指令集加内存交付界面就够了。

有了这套标准界面，软件工程师可以尽情优化算法，硬件工程师只管闷头提升性能。但从标量时代切换到向量时代后，SIMD自动向量化这条路一直走得不太顺。Jim Keller在那场Keynote里说

——确实如此。它用大量线程把向量化操作掰成了人类心智能接受的标量编程，这个点子太妙了。

2. AI时代的问题

计算规模一上来，深度学习如今已经离不开大规模张量运算。而CUDA在张量时代也面临着一系列抽象的难题。

2.1 指令集扩展不可行

国内还有不少人单纯地认为：从标量指令集扩展到向量指令集，再延伸到张量指令集，就能搞出一套通用深度学习指令集。是不是弄个TensorCore，塞几个MMA指令就万事大吉了？但现实是，这条泛化路径已经走不通了。一方面，矩阵计算本身的计算访存比很高，必须大量处理数据局部性。数据搬运需要专门的流程，还要对算子做时空切分，在算法层面上做各种时空折中。

根子其实在于缓存策略。标量的预取和分支预测已经很成熟了，SIMT靠更大的寄存器文件和Warp调度来隐藏延迟，也能把吞吐率做上去。可到了张量时代，矩阵乘法、TensorCore相关的东西就复杂了：访存冲突、寄存器/SMEM复用、后续的Element-Wise操作（比如Softmax、LayerNorm等算子融合）——这些问题全都牵涉到GPU微架构的取舍。

2.2 标准算子作为交付界面

那能不能拿一些标准算子来做交付界面呢？国产GPU生态里，很多专家都在呼吁赶紧做算子库。特别是部分人还抱着一个错误认知：大模型架构基于Transformer，算子好像已经收敛了。从算法角度看，Transformer确实优秀，效率极高，几乎榨干了GPU的计算和访存能力。加上Flash-Attention这类与硬件深度绑定的算子优化，已经快到极致了。CUDA的壁垒其实变相成了一种“人工”智能——写一个算子的Kernel，基本要花两三周时间。

另一方面，从代数的视角看，模型的算法并没有收敛。Transformer的效率问题仍然待解。再从范畴论的角度看，Transformer本质上是在构建一系列态射并构成预层范畴，那能不能最终形成一个稀疏表示来降低计算复杂度？MoE某种意义上就在做这种尝试，但还有很多问题没解决。

这条路也在快速变化中，交付界面依然不确定。

2.3 互联架构和模型并行策略的变化

从局部的GEMM访存优化再往外扩一层，自然就走到多芯片协同的分布式集群训练了。无论是DP、TP还是CP并行策略，本质都是在做矩阵拆分——拆数据维度、拆模型矩阵维度、拆序列长度维度，逻辑是相通的。只是在处理前向和反向时，还要考虑通信与计算的Overlap策略。而这些灵活的策略，又得跨越两套不同的网络：ScaleUP（NvLink）和ScaleOut（RDMA）。可惜的是，这两套网络各自都有不少短板。

2.4 从系统的视角：金字塔结构

Jim Keller从指令数的角度切入这个问题，非常有启发性。

这其实回到了一个更宽泛的问题：如何在分布式系统中调度算子、分发指令，同时保证更好的数据局部性。从系统角度去看待这件事，方向就对了。

3. Contract

答案是什么？我也不知道，因为不懂的领域实在太多了。

开个无厘头的玩笑：既然没有满足软硬件交付界面的合约，那就去找足够满足合约的人不就行了？这倒是挺“人工”智能的。当然，这条路也走不通。本质上，我们可能得回到张量代数层面来考虑这件事，还有太多东西值得探索。

现阶段，有两个方向值得讨论：

3.1 SM微架构

其实从之前《Tensor-003 TensorCore架构》的分析就能看到，TensorCore的计算密度太高，已经出问题了。更大的脉动阵列构成的TensorCore，对某些扁平矩阵（比如M远大于K）也有问题。英伟达在Hopper上临时搞了个WGMMA“胶布”凑合。然后从GMEM到SMEM再到RF，全链路异步化，这导致CUDA的编程难度相比十年前大幅提高。

我估计在B200上，这个难度还会继续走高。原因很简单：要在维持CUDA编程框架的前提下继续提升TensorCore算力，只能从几个方向着手——进一步扩大分布式共享内存的范围，覆盖更多SM。这必然导致SM互联网络的变化，2D-Mesh或Ring这样的局部拓扑会出现。同时，SM之间的异步操作可能还会引入更多L1.5缓存/SMEM。

但有没有更好的微架构？CUDA兼容还是不兼容？走DSA还是GPGPU？最终还得看算法上的代数抽象。

3.2 互联的架构

Jim Keller后来还讲了Tenstorrent的架构：从成本角度抛弃HBM，从互联角度用以太网做胶布——这些思路可能都对。

但问题在哪？在拓扑上。

2D-Mesh拓扑会带来编译和算子放置上的一系列麻烦。怎么解？这是个很困难的问题。说不定也是个很简单的问题——毕竟对称才是美。