1. AI芯片软硬件协同设计的核心逻辑1.1 为什么软硬件必须一起设计做AI芯片这行的人都有一个共识芯片架构和算法模型是互相塑造的关系。你设计一款芯片如果只盯着硬件指标——峰值算力多少TOPS、功耗多少瓦、面积多大——而不去考虑它要跑什么模型、模型的计算模式长什么样那这颗芯片大概率会沦为“跑分好看、落地拉胯”的典型。我见过不少团队踩这个坑。硬件团队闷头把MAC阵列堆到几千个结果软件团队拿到芯片后发现模型里的reshape、transpose、softmax这些操作根本没法高效映射上去数据在片上来回搬运的开销比计算本身还大。最后实际推理性能只有理论峰值的百分之十几。所以这一节我想先把一个核心观点讲透AI芯片的设计起点不是晶体管而是计算模式。你得先搞清楚目标负载的计算特征再倒推硬件应该长什么样最后再设计软件栈怎么把模型高效映射到这套硬件上。这个顺序不能反。1.2 从Transformer的计算模式说起现在做AI芯片绕不开Transformer架构。不管你接的是视觉任务还是NLP任务Transformer已经成了事实上的通用骨架。Vision Transformer把图片切成patch当token处理Swin Transformer引入窗口注意力降低计算量各种轻量Transformer也在移动端铺开。所以芯片设计必须把Transformer的计算模式吃透。Transformer的核心计算是什么拆开看就三块QKV投影输入序列乘以三个权重矩阵得到Query、Key、Value。本质是大规模矩阵乘法。注意力计算Q乘以K的转置经过softmax再乘以V。这里涉及矩阵乘、逐元素操作和归一化。前馈网络两层全连接加激活函数还是矩阵乘法为主。你会发现矩阵乘法占了整个计算量的绝大部分。这就给硬件设计指明了方向把矩阵乘法的吞吐做上去同时把数据复用做到极致减少片外访存。这也是为什么脉动阵列会成为TPU系列的核心计算单元。1.3 软硬件协同的分层思路我习惯把AI芯片的软硬件设计分成四层来看层级硬件侧软件侧协同关键点算法层无模型结构、算子定义确定计算模式编译层指令集定义图优化、算子融合、调度映射策略运行时层驱动、DMA引擎内存分配、流水线编排数据搬运效率计算层MAC阵列、缓存算子库、微内核计算密度这四层里编译层是软硬件协同的咽喉。编译器要把模型图拆成硬件能执行的指令序列决定哪些算子融合、哪些数据放片上缓存、哪些走DMA搬运。编译器做得好不好直接决定芯片的实际利用率。实操心得很多团队把编译器当成“翻译工具”觉得只要把模型转成指令就行。实际上编译器是性能优化的主战场。同样的硬件编译器优化前后性能差距可以到3到5倍。建议在芯片定义阶段就让编译器团队深度参与而不是等流片回来再补。2. 脉动阵列的基本原理与设计取舍2.1 脉动阵列到底在解决什么问题脉动阵列这个概念不新上世纪七八十年代就有了但它在AI芯片领域焕发第二春核心原因是它完美匹配了矩阵乘法的数据复用需求。先想一个最朴素的问题做矩阵乘法C A × B如果用一个PE处理单元算一个输出元素那每个输出元素需要读取A的一整行和B的一整列。数据复用率极低访存带宽会成为瓶颈。脉动阵列的思路是让数据在PE阵列中有节奏地流动每个数据进入阵列后能被多个PE复用。就像心脏泵血一样数据一波一波地“脉动”穿过计算阵列每个PE在数据经过时完成一次乘累加。具体来说权重矩阵B的元素预先加载到各个PE中并保持不动输入矩阵A的元素从左侧依次流入部分和从上方或下方依次流出。每个时钟周期数据向前推进一格。这样A的每个元素会与一整行PE交互B的每个元素会与一整列PE交互数据复用率大幅提升。2.2 脉动阵列的三种数据流实际设计中脉动阵列有三种经典数据流方式选择哪种直接决定了硬件结构和编译器策略权重固定Weight Stationary权重预加载到PE中不动输入和部分和在阵列中流动。适合权重矩阵较大的场景因为权重加载一次可以复用很多次。TPUv1采用的就是这种模式。输出固定Output Stationary每个PE负责一个输出元素累加结果留在PE中输入和权重都流动。适合输出矩阵较小的场景。行固定Row Stationary折中方案每行PE固定一组权重输入在行内流动。灵活性更高但控制逻辑更复杂。我个人的经验是做推理芯片优先考虑权重固定。因为推理场景下权重是固定的输入是变化的权重固定能让权重加载的开销摊薄到极致。训练芯片则要更灵活一些因为权重也在变。2.3 阵列尺寸怎么定这是每个做AI芯片的人都会纠结的问题脉动阵列到底做多大128×128256×256还是更大先说结论阵列尺寸不是越大越好要和目标模型的矩阵维度匹配。假设你的目标模型里最常见的矩阵乘法维度是M×K乘以K×N那阵列做M×N的大小最理想这样一次就能算完一个完整的矩阵乘。如果阵列比矩阵大多出来的PE就浪费了如果阵列比矩阵小就需要多次累加增加控制开销。但现实是Transformer里不同层的矩阵维度不一样。注意力层的维度可能是512×512前馈网络可能是512×2048。你不可能为每一层都做一个完美匹配的阵列。所以实际设计中的做法是选一个主流维度作为阵列尺寸然后用tiling分块策略处理其他维度。比如TPUv1用了256×256的阵列对于更大的矩阵就切成256×256的块来算对于更小的矩阵就做padding或者用部分阵列。注意事项阵列尺寸增大带来的收益是递减的。256×256的阵列有65536个PE功耗和面积已经很可观了。再往上做布线延迟、时钟树功耗、良率问题都会急剧恶化。我见过一些团队盲目追求大阵列最后芯片面积爆了功耗也压不住。2.4 数据位宽的选择脉动阵列的PE里做的是乘累加运算数据位宽直接影响计算精度、面积和功耗。训练场景通常需要FP32或BF16推理场景可以用INT8甚至INT4。但这里有个坑不是所有层都能低精度。注意力计算里的softmax对精度很敏感用INT8可能会溢出或者精度损失严重。所以实际芯片往往支持混合精度——矩阵乘用INT8softmax和归一化用FP16或FP32。TPUv1的设计就很典型矩阵乘单元用8位整数但累加器用32位防止溢出。这个设计在当时是非常超前的因为8位整数的推理精度已经能满足大部分场景而32位累加器保证了数值稳定性。3. Transformer在AI芯片上的映射实操3.1 从模型图到硬件指令的完整流程把一个Transformer模型映射到AI芯片上大致要经过这几个步骤模型导入与图解析把训练好的模型通常是ONNX或PyTorch格式转换成芯片编译器能识别的中间表示。这一步要处理算子版本兼容、动态shape等问题。图优化做算子融合比如把MatMulBiasReLU融合成一个算子、常量折叠、死代码消除。算子融合能显著减少访存次数是性能优化的关键。算子调度与分块根据片上缓存大小把大矩阵切成小块决定计算顺序和数据搬运策略。这一步要平衡计算和访存尽量让DMA搬运和计算重叠。指令生成把调度方案翻译成芯片的指令序列包括DMA指令、计算指令、同步指令。运行时执行驱动加载指令配置DMA和计算单元启动流水线。这个流程里第3步是最考验编译器能力的。分块策略的好坏直接决定了片上缓存的命中率和DMA的利用率。3.2 注意力层的映射细节注意力计算是Transformer里最复杂的部分也是映射到硬件上最容易出问题的部分。以标准的自注意力为例计算过程是Q XW_qK XW_kV XW_v然后Attention softmax(QK^T/√d)V。前三步是标准的矩阵乘法映射到脉动阵列上比较直接。麻烦的是后面QK^T这是两个矩阵相乘但K需要转置。如果硬件不支持转置操作就需要在搬运时做处理或者用额外的转置单元。softmax涉及求最大值、指数运算、求和、除法。这些是逐元素操作不适合脉动阵列。通常需要单独的向量处理单元来加速。softmax结果乘以V又是一个矩阵乘法但左矩阵是softmax的输出右矩阵是V。我的实操经验是把注意力层拆成“矩阵乘-向量处理-矩阵乘”三段流水线矩阵乘走脉动阵列向量处理走SIMD单元中间结果放片上缓存。这样能最大化利用不同类型的计算单元。3.3 前馈网络的映射策略前馈网络相对简单就是两个全连接层加一个激活函数。但它的参数量很大通常占整个模型参数量的三分之二。映射时的核心问题是权重矩阵太大放不下片上缓存怎么办。以BERT-base为例前馈网络的中间维度是3072输入输出维度是768。权重矩阵是768×3072和3072×768每个都是两百多万个参数。如果用INT8存储每个矩阵大约2MB多。两个矩阵加起来接近5MB。而很多推理芯片的片上缓存总共才几MB。所以必须做权重分块加载。具体做法是把权重矩阵按行或按列切成小块每次加载一块到片上算完再加载下一块。这里的关键是让权重加载和计算重叠用双缓冲double buffering技术一块在算的时候另一块已经在搬了。实操心得双缓冲的缓冲区大小要仔细设计。太小了搬运次数多DMA启动开销大太大了片上缓存放不下。一般建议缓冲区大小至少能覆盖一次DMA搬运的延迟通常几KB到几十KB比较合适。3.4 一个具体的分块计算示例假设我们要在一个256×256的脉动阵列上计算一个512×512的矩阵乘法C A × B其中A是512×512B是512×512。分块策略如下把A按行切成2块每块256×512把B按列切成2块每块512×256这样得到4个子矩阵乘每个是256×512乘以512×256但脉动阵列的K维度只有256所以每个子矩阵乘还需要沿K维度再切一次每个子矩阵乘切成2次每次K256所以总共需要2×2×28次阵列计算。每次计算前需要加载对应的A块和B块到片上缓存。这个例子里分块顺序很重要。如果先遍历A的行再遍历B的列那B的每一列块会被重复加载2次。如果调整循环顺序让B的列块在片上缓存中停留更久就能减少重复加载。编译器通常会做循环重排loop reordering来优化这个问题。手动写kernel的话这个顺序需要根据具体的缓存大小和DMA带宽来调。4. 常见问题与排查技巧实录4.1 性能不达预期的排查思路芯片流片回来跑模型发现性能只有理论峰值的20%这是最常见的问题。排查思路应该从大到小第一步确认瓶颈在计算还是访存。用性能计数器看计算单元的利用率和DMA的带宽占用率。如果计算单元利用率低但DMA带宽跑满了说明是访存瓶颈如果两者都低说明是调度问题或者同步开销太大。第二步如果是访存瓶颈检查数据复用是否充分。常见原因是分块策略不合理导致同一块数据被反复从片外搬运。解决办法是调整分块大小和循环顺序增大片上缓存的数据复用窗口。第三步如果是调度问题检查DMA和计算是否重叠。很多芯片支持DMA和计算并行但如果指令序列里DMA和计算是串行排列的那就白白浪费了并行能力。需要编译器做软件流水software pipelining。第四步如果是同步开销检查同步指令的频率。过多的同步会导致流水线停顿。可以尝试合并同步点或者用更细粒度的同步机制。4.2 精度问题的排查低精度推理时出现精度下降通常有这几个原因累加器溢出INT8乘法的结果用INT8累加会溢出必须用更宽的累加器。TPUv1用32位累加器就是这个道理。softmax数值不稳定softmax需要先减最大值再取指数如果实现时漏了减最大值这一步指数运算会溢出。量化误差累积每一层的量化误差会逐层累积到后面层可能偏差很大。解决办法是对敏感层保持高精度或者用量化感知训练来补偿。我踩过的一个坑是LayerNorm层的量化。LayerNorm涉及均值和方差的计算对精度很敏感。用INT8量化后方差计算误差很大导致归一化结果偏差明显。后来改成FP16计算LayerNorm精度问题就解决了。4.3 常见问题速查表问题现象可能原因排查方法解决思路计算单元利用率低访存瓶颈看DMA带宽占用率优化分块策略增大数据复用性能波动大调度不稳定看指令序列和同步点软件流水减少同步精度下降明显累加器溢出检查累加器位宽加宽累加器混合精度功耗超标数据搬运过多看片外访存次数增大片上缓存减少搬运编译时间过长搜索空间太大看编译器的调度算法限制搜索空间用启发式算法4.4 几个容易忽略的细节DMA的描述符开销每次DMA搬运都需要配置描述符如果搬运的数据块太小描述符配置的开销可能比搬运本身还大。建议单次DMA搬运至少几KB。片上缓存的bank冲突脉动阵列需要高带宽的数据供给如果缓存bank设计不合理会出现bank冲突实际带宽远低于理论值。设计时要根据访问模式来规划bank数量和宽度。时钟域 crossing如果计算单元和DMA引擎在不同的时钟域跨时钟域同步会引入延迟。尽量让关键路径在同一个时钟域内。温度对性能的影响芯片温度升高会导致时序变差可能需要降频。散热设计要在芯片设计阶段就考虑进去不能等封装完了再补。5. 从TPUv1看AI芯片设计的工程取舍5.1 TPUv1的设计哲学TPUv1是Google在2015年部署的推理芯片虽然现在看来规格不算高但它的设计思路至今仍有参考价值。TPUv1的核心是一个256×256的脉动阵列支持INT8乘法32位累加。片上缓存24MB用来存放权重和中间结果。指令集很简单只有几条指令矩阵乘、卷积、激活、归一化等。它的设计哲学很明确不做通用计算只做神经网络推理需要的那几件事把这几件事做到极致。这种专用化的思路让TPUv1的能效比远超同时期的GPU。5.2 工程取舍的启示TPUv1有几个取舍值得思考取舍一只支持INT8推理。这限制了它的应用范围但换来了更高的能效比和更小的面积。对于推理场景这个取舍是合理的。取舍二权重预加载。TPUv1的权重需要预先加载到片上缓存这意味着它不适合权重频繁变化的场景。但推理场景下权重是固定的所以这个限制可以接受。取舍三简化的指令集。TPUv1的指令集很简单编译器不需要做复杂的指令调度。这降低了编译器的开发难度但也限制了硬件的灵活性。这些取舍背后的逻辑是一致的明确目标场景围绕场景做优化不为通用性牺牲效率。这个思路对现在做AI芯片的团队仍然适用。5.3 对当前设计的参考价值虽然TPUv1已经过去好几年了但它的一些设计思路仍然值得借鉴脉动阵列仍然是矩阵乘法的最高效实现之一尤其是在推理场景下。片上缓存的容量和带宽是决定实际性能的关键比峰值算力更重要。编译器的质量直接决定芯片的易用性和实际性能不能重硬件轻软件。专用化设计在特定场景下仍然有优势不必盲目追求通用性。当然现在的模型和场景比当年复杂得多。Transformer的注意力机制、动态shape、多模态融合等需求都对芯片设计提出了新的挑战。但底层的方法论没有变理解计算模式匹配硬件结构优化数据流动。我个人在实际项目中的体会是AI芯片设计最难的不是某个单点技术而是在算力、功耗、面积、灵活性、开发成本之间找到平衡点。每个决策都会影响其他维度没有完美的方案只有适合当前场景的方案。踩过几次坑之后我越来越倾向于“先做减法再做加法”——先把核心场景做透再考虑扩展。贪多求全往往什么都做不好。