2、AI芯片基础回顾：GPU与TPU架构对比、AI芯片的关键指标（算力、带宽、延迟）、片上互联与内存层次结构
好，咱们正式开始聊MoE之前，得先把底层的硬件基础捋一遍。我见过不少同学，模型结构讲得头头是道，一问到芯片怎么跑的就卡壳了。说白了，MoE的很多设计决策，本质上都是在跟硬件特性做博弈。你不懂芯片，就理解不了为什么MoE要那样切分专家、为什么通信会成为瓶颈。
2.1 GPU vs TPU：两种不同的哲学
先聊聊GPU和TPU。这两者都是当前AI计算的主力，但设计思路完全不同。
GPU（以NVIDIA为例），本质上是为图形渲染设计的。图形渲染的特点是啥？大量并行的、相对独立的计算任务。所以GPU的核心思路是：堆大量的CUDA Core，每个Core能力一般，但数量多到吓人。我早期做训练时，用的还是V100，那时候就觉得这卡真猛。后来用上A100、H100，你会发现它的设计一直在强化一个点——通用性。GPU什么都能算，FP32、FP16、INT8，甚至BF16，全支持。代价呢？控制逻辑相对复杂，片上缓存（L1/L2）占了不少面积。
TPU（Google的定制芯片），则是从出生就只为深度学习服务的。它最核心的单元叫MXU（矩阵乘法单元），是一个巨大的脉动阵列。TPU的思路很纯粹：我就把矩阵乘法做到极致。你想想看，Transformer里95%以上的计算量都是矩阵乘法，那我专门为这个设计一个超级高效的单元，是不是很合理？TPU v4开始支持BF16，但它的灵活性远不如GPU。我曾在Google的论文里看到，TPU的MXU利用率能做到80%甚至90%以上，这在GPU上很难达到。
核心差异一句话总结：
- GPU：通用计算平台，灵活，生态好，适合各种模型结构。
- TPU：专用计算引擎，极致效率，但需要模型设计去适配它。
我个人习惯，做MoE推理时，如果追求极致性价比，会优先考虑TPU。但做训练和实验迭代，还是离不开GPU的生态。嗯，这里要注意，TPU对动态shape的支持比较弱，MoE的专家路由恰好是动态的，这是个不小的坑。
2.2 AI芯片的关键指标：算力、带宽、延迟
这三个指标，是衡量一块AI芯片好不好的核心。我建议你把它刻在脑子里。
| 指标 | 定义 | 对MoE的影响 |
|---|---|---|
| 算力（Compute） | 单位时间内能完成的浮点运算次数（TFLOPS） | 决定了专家计算的快慢。算力越高，单个专家处理得越快。 |
| 带宽（Bandwidth） | 单位时间内能搬运的数据量（GB/s） | 决定了参数和激活值加载的速度。MoE的All-to-All通信极度依赖带宽。 |
| 延迟（Latency） | 完成一次操作所需的时间（ms） | 决定了路由决策和通信的响应速度。高延迟会导致GPU空转等待。 |
我曾经在一个MoE推理项目中，只盯着算力看，选了个算力很高的卡。结果跑起来发现，大部分时间都花在跨卡通信上了，算力根本喂不饱。说白了，MoE模型对带宽的敏感度远高于算力。你想想看，一个token只激活少数几个专家，计算量不大，但要把token从这张卡路由到那张卡，数据搬运量可不小。
避坑指南： 我曾经犯过一个错误，只关注峰值算力，忽略了持续算力。很多芯片峰值算力很漂亮，但实际跑起来因为散热或功耗限制，根本达不到。看芯片参数时，多看看它的TDP（热设计功耗）和实际跑分。
2.3 片上互联与内存层次结构
这部分是MoE能否高效运行的关键。我把它分成两个层面来讲。
2.3.1 内存层次结构
从离计算单元最近到最远，依次是：
- 寄存器（Register）：最快，但容量极小，KB级别。
- L1/L2缓存（Cache）：速度很快，容量MB级别。GPU的L1和共享内存通常是同一块物理空间，可以灵活配置。
- HBM（高带宽内存）：显存，容量GB级别，带宽极高（H100可达3.35TB/s）。这是模型参数和中间激活的主要存放地。
- 主存（CPU内存）：容量最大，但带宽最低，延迟最高。一般只做数据预处理和模型加载用。
MoE模型的一个特点是，专家参数很大，但每个token只访问其中一小部分。这就导致了一个问题：访存模式是稀疏且随机的。GPU的缓存机制对这种随机访问很不友好，很容易出现Cache Miss。我建议在设计MoE时，尽量让同一个batch的token路由到相同的专家，这样能提高缓存命中率。
2.3.2 片上互联
当MoE模型大到一张卡放不下时，就需要多卡互联了。这里主要看两种互联方式：
- NVLink（NVIDIA）：GPU之间的高速互联。H100的NVLink 4.0单向带宽可达900GB/s。延迟极低，适合做All-to-All通信。
- InfiniBand / RoCE：跨节点（服务器之间）的互联。带宽通常为200Gbps或400Gbps，延迟比NVLink高一个数量级。
我做过一个实验，把MoE的专家分布在8张A100上，用NVLink做All-to-All通信，通信开销只占整个推理时间的15%左右。但如果换成跨节点的InfiniBand，这个比例直接飙升到40%以上。所以，MoE的专家尽量放在同一台机器内，用NVLink互联，这是性能的关键。
注意： 不要忽视PCIe的瓶颈。即使有NVLink，很多数据还是要走PCIe总线。比如CPU下发指令、DMA传输等。PCIe 4.0 x16的单向带宽只有32GB/s，跟NVLink比差了一个数量级。我曾经遇到过，因为PCIe带宽打满，导致GPU空转等待数据的情况。
2.4 知识体系总览
下面这张图，是我自己梳理的本章知识结构。你可以把它当作一个思维导图来看。
嗯，这张图基本把本章的核心逻辑串起来了。你从中心出发，左边是架构对比，中间是量化指标，右边是硬件结构。这三者共同决定了MoE在芯片上的实际表现。
最后说一句，理解这些硬件细节，不是为了让你去造芯片，而是为了让你在设计MoE时，能做出更聪明的取舍。比如，知道带宽是瓶颈，你就会想办法减少通信量；知道缓存层次结构，你就会优化数据局部性。这些，都是后面章节要深入展开的内容。
公众号:蓝海资料掘金营，微信deep3321