AI芯片软硬件协同设计：从脉动阵列到Transformer矩阵乘法优化

1. 从矩阵乘法到硅片：AI芯片软硬件协同设计的核心命题

聊AI芯片的软硬件设计，绕不开一个最底层的事实：当下几乎所有主流AI加速器的算力，最终都消耗在矩阵乘法上。不管是CNN时代的卷积，还是Transformer时代的多头注意力，拆到最底层，都是形如C = A × B + bias的矩阵运算。区别只在于矩阵的维度、数据精度、稀疏性以及访存模式。所以做AI芯片设计，第一件事不是去画架构框图，而是想清楚：你要为什么形状的矩阵乘法做优化。

这个问题的答案直接决定了芯片的架构走向。如果目标负载是传统的卷积神经网络，矩阵乘法的规模相对规整，K维（规约维）通常不大，M和N维可以靠im2col展开成比较大的矩阵。但如果目标负载是Transformer系列模型，情况就完全不同了——注意力机制里的Q、K、V矩阵乘法，序列长度动辄512、1024甚至4096，K维（head dimension）却往往只有64或128。这种"瘦长型"矩阵乘法对数据复用和流水线设计提出了截然不同的要求。

我在实际参与AI加速器设计的过程中，最深的一个体会是：软硬件协同设计不是一句口号，而是从算法选型阶段就要开始介入的约束条件。算法团队想用更大的head dimension来提升模型表达能力，硬件团队就要评估SRAM带宽和PE阵列利用率能不能撑住；编译器团队想用更激进的算子融合来减少访存，硬件团队就要确认片上缓存的一致性协议是否支持。这些决策环环相扣，任何一个环节拍脑袋，最后都会在流片后的实测性能上付出代价。

这篇文章面向的是对AI芯片设计有兴趣的工程师、研究生，以及正在做相关项目的从业者。我会从脉动阵列的基本原理讲起，逐步展开到Transformer负载下的矩阵乘法优化、软硬件接口设计、编译器与硬件的配合，以及实际项目中容易踩的坑。内容会涉及一些具体的参数计算和架构取舍分析，但不会堆砌公式，重点放在"为什么这样设计"和"实际怎么做"上。

2. 脉动阵列：为什么它成了矩阵乘法的首选架构

2.1 脉动阵列的基本工作原理

脉动阵列（Systolic Array）这个概念最早由H.T. Kung在1982年提出，核心思想是让数据像心跳一样有节奏地流过处理单元阵列，每个PE只做最简单的乘加运算，数据复用靠邻居之间的传递来实现。这个设计之所以在AI芯片领域焕发第二春，是因为它完美匹配了矩阵乘法的数据复用需求。

具体来说，一个典型的N×N脉动阵列做矩阵乘法时，矩阵A的元素从左侧流入，矩阵B的元素从上方流入，每个PE（处理单元）内部有一个乘法器和一个累加器。数据每拍向右和向下流动一次，PE在每一拍完成一次乘加。以3×3阵列计算两个3×3矩阵乘法为例，总共需要7拍（2N-1）就能完成全部计算，而每个PE只需要和相邻的PE通信，不需要全局广播。

这种设计的精妙之处在于：权重矩阵B的元素在阵列中停留不动（weight stationary），输入矩阵A的元素沿水平方向流动，部分和沿垂直方向流动。这样一来，每个权重元素被复用了N次（对应N个输入），每个输入元素也被复用了N次（对应N个权重），数据复用率达到了O(N)级别。相比传统的冯诺依曼架构每次乘加都要从内存取两个操作数，脉动阵列的能效比可以高出两个数量级。

但这里有一个容易被忽略的细节：脉动阵列的利用率高度依赖于矩阵的维度是否匹配阵列大小。如果你有一个32×32的脉动阵列，但实际计算的矩阵M维只有8，那么阵列中只有8行在工作，利用率直接掉到25%。这就是为什么很多AI芯片会设计多级阵列或者可重构阵列——大阵列跑大矩阵，小阵列跑小矩阵，避免"大炮打蚊子"的浪费。

2.2 权重固定 vs 输出固定：两种数据流的取舍

脉动阵列的数据流设计有两种主流方案：权重固定（Weight Stationary, WS）和输出固定（Output Stationary, OS）。这两种方案的选择直接影响了芯片的访存带宽需求和PE利用率。

权重固定的思路是让权重矩阵的元素预先加载到PE中并保持不动，输入数据流过阵列。这种方案的优势是权重只需要加载一次，对于推理场景（权重在多次推理之间不变）非常友好。但缺点是如果权重矩阵很大，超过阵列容量，就需要分块加载，分块之间的切换会带来额外的开销。

输出固定的思路是让每个PE负责计算输出矩阵的一个元素，输入和权重都流动。这种方案的优势是部分和不需要在PE之间传递，减少了垂直方向的数据移动功耗。但缺点是每个PE需要同时接收输入和权重，对带宽的要求更高。

我在实际项目中的经验是：对于Transformer类负载，混合方案往往更实用。具体来说，在注意力计算阶段，Q和K的矩阵乘法适合用权重固定（K矩阵作为权重），而在前馈网络（FFN）阶段，由于权重矩阵通常很大（比如4096×16384），适合用输出固定来减少部分和的移动。很多商用AI芯片（如Google TPU）采用的是权重固定为主、辅以输出固定的混合策略。

下面这个表格对比了两种方案的关键指标：

2.3 脉动阵列的边界条件与常见误解

很多人第一次接触脉动阵列时，会有一个误解：认为脉动阵列只能做方阵乘法，而且阵列越大越好。实际上，现代AI芯片中的脉动阵列设计已经非常灵活。比如，可以通过配置PE的连接方式，让同一个阵列支持不同维度的矩阵乘法；也可以通过时分复用的方式，让大阵列分时处理多个小矩阵。

另一个常见的误解是认为脉动阵列的延迟很高。确实，由于数据需要逐拍流动，脉动阵列的首次计算延迟（latency）是O(N)级别的。但对于大批量推理场景，吞吐量（throughput）才是关键指标，而脉动阵列的吞吐量可以做到每拍完成N²次乘加，这个优势足以抵消延迟上的劣势。在实际部署中，可以通过流水线并行（pipeline parallelism）来隐藏这个延迟。

还有一个需要特别注意的点：脉动阵列对数据精度的支持不是免费的。如果你需要支持INT8、FP16、BF16甚至FP8等多种精度，每个PE的乘法器设计就会变得复杂。有些芯片采用"精度可重构"的PE设计，比如一个PE内部包含多个小乘法器，通过配置可以组合成一个大乘法器（高精度）或拆分成多个小乘法器（低精度）。这种设计的面积开销大约增加30%-50%，但灵活性提升明显。

3. Transformer负载下的矩阵乘法：形状、精度与访存的三重挑战

3.1 注意力机制中的矩阵乘法形状分析

Transformer模型的核心是自注意力机制，而自注意力的核心是三次矩阵乘法：Q×K^T、softmax结果×V、以及输出投影。这三步的矩阵形状各有特点，对硬件设计的要求也不同。

以BERT-Base为例，hidden dimension是768，head数量是12，每个head的dimension是64，序列长度是512。在自注意力计算中：

• Q×K^T：形状是 (512×64) × (64×512) = 512×512。这是一个"瘦长×瘦长"的乘法，K维只有64，但M和N维都是512。
• softmax×V：形状是 (512×512) × (512×64) = 512×64。这是一个"方×瘦长"的乘法，K维是512，M维是512，N维是64。
• 输出投影：形状是 (512×768) × (768×768) = 512×768。这是一个相对规整的乘法。

从硬件设计的角度看，Q×K^T的K维只有64，这意味着如果脉动阵列的深度是256，那么利用率只有25%。这是一个非常现实的挑战。很多AI芯片在处理这类负载时，会采用"K维拼接"的策略——把多个head的K维拼在一起，比如把4个head的64维拼成256维，这样就能填满阵列深度。但这样做的前提是硬件支持灵活的SRAM寻址和PE配置。

另一个挑战来自序列长度的动态变化。在实际推理中，序列长度可能是128、256、512、1024不等。如果硬件设计只针对固定序列长度优化，那么在其他长度下性能会大幅下降。我在项目中见过一种做法：把序列长度维度映射到阵列的M维和N维，通过分块（tiling）来处理任意长度。比如512×512的矩阵乘法，如果阵列是128×128，就分成4×4=16个块来计算。这种分块策略的代价是需要额外的片上缓存来存储中间结果。

3.2 精度选择：FP16、BF16还是INT8

Transformer模型对精度的敏感度比CNN更高，尤其是注意力分数经过softmax之后，数值范围会变得很窄，低精度表示容易导致精度损失。这就给AI芯片的精度设计带来了两难：用FP16能保证精度，但算力密度和能效比不如INT8；用INT8能提升吞吐，但可能需要量化感知训练（QAT）来补偿精度。

从硬件设计的角度，支持多种精度意味着PE的乘法器和累加器需要更灵活的设计。一个常见的方案是采用"分时复用"的乘法器：在FP16模式下，一个乘法器完成一次FP16乘法；在INT8模式下，同一个乘法器拆成两个INT8乘法器，算力翻倍。这种设计的面积开销大约增加20%-30%，但能同时满足精度和性能需求。

我在实际测试中对比过不同精度下的Transformer推理精度和性能：

需要特别注意的是，INT8的精度损失在注意力机制中会被放大。因为softmax对输入的微小变化很敏感，如果Q和K的量化误差较大，注意力分数的分布就会偏移，导致模型关注到错误的位置。实践中，通常会对Q和K保持较高精度（FP16），而对V和输出投影使用INT8。

3.3 访存瓶颈：为什么算力利用率总是上不去

AI芯片设计中最让人头疼的问题不是算力不够，而是算力利用率（MFU）上不去。一个标称256 TFLOPS的芯片，在实际跑Transformer时可能只能达到30%-50%的利用率。原因通常出在访存上。

以Q×K^T为例，计算512×512的输出，需要读取512×64的Q矩阵和512×64的K矩阵，总共64K个元素。如果每个元素是FP16（2字节），总访存量是128KB。而计算量是512×512×64=16.7M次乘加。计算访存比（Compute-to-Memory Ratio）大约是130 ops/byte。如果芯片的片上带宽是1TB/s，算力是256 TFLOPS，那么带宽能支撑的算力是130×1T=130 TFLOPS，只有标称算力的一半。

这就是为什么片上缓存（SRAM）的设计比PE阵列的设计更考验功力。好的AI芯片会把Q、K、V矩阵分块加载到SRAM中，让PE阵列在计算当前块时，下一块的数据已经在预取。这种"双缓冲"（double buffering）策略可以把访存延迟隐藏起来，把利用率提升到70%以上。

我在项目中的一个实际经验是：SRAM的容量和带宽需要根据目标模型的算子形状来反推。比如，如果目标模型的最大head dimension是128，序列长度是2048，那么单个head的Q矩阵大小是2048×128×2字节=512KB。如果芯片要同时处理4个head，就需要2MB的SRAM专门用于Q矩阵缓存。这个数字直接影响了芯片的面积和成本。

4. 软硬件接口设计：编译器如何与脉动阵列对话

4.1 指令集设计：从微码到粗粒度指令

AI芯片的指令集设计是一个容易被低估的环节。传统的CPU指令集是细粒度的（比如ADD、MUL），但AI芯片如果也采用细粒度指令，编译器生成的代码会非常冗长，取指功耗会吃掉很大一部分能效预算。所以现代AI芯片普遍采用粗粒度指令（Coarse-Grained Instruction），一条指令就能触发一个完整的矩阵乘法块。

以脉动阵列为例，一条典型的矩阵乘法指令可能包含以下字段：

• 操作码：标识这是矩阵乘法
• 输入地址：矩阵A在SRAM中的起始地址
• 权重地址：矩阵B在SRAM中的起始地址
• 输出地址：结果矩阵C在SRAM中的起始地址
• 维度信息：M、N、K的值
• 精度模式：FP16/INT8/FP8
• 转置标志：是否需要对某个矩阵转置

这种粗粒度指令的好处是取指带宽需求低，编译器优化空间大。但缺点是灵活性受限，如果遇到不规则的算子（比如稀疏矩阵乘法），就需要额外的指令来辅助。

我在设计指令集时踩过的一个坑是：没有预留足够的维度字段。最初设计时只预留了8位来表示M、N、K，最大支持255。结果后来遇到序列长度2048的模型，直接溢出。后来改成16位才解决问题。这个教训是：AI模型的维度增长很快，指令集设计要留足余量。

4.2 数据布局：NCHW、NHWC还是自定义

数据布局对AI芯片的性能影响巨大。CNN时代，NCHW和NHWC之争持续了很久。到了Transformer时代，数据布局变得更加复杂，因为注意力机制涉及大量的reshape和transpose操作。

一个典型的例子是Q×K^T的计算。在内存中，Q矩阵通常以(sequence_length, hidden_dim)的形式存储，但脉动阵列需要的是(sequence_length, head_dim)的分块形式。如果硬件不支持灵活的stride寻址，就需要软件先做一次数据重排，这个开销可能占到总计算时间的20%以上。

我在项目中的做法是：在SRAM中维护多种数据布局的副本，通过硬件DMA在布局之间快速切换。比如，Q矩阵在SRAM中同时以"按序列"和"按head"两种布局存储，计算时根据算子需求选择对应的布局。这种做法的代价是SRAM容量翻倍，但换来了计算效率的提升。对于SRAM资源紧张的芯片，也可以只维护一种布局，通过硬件地址生成器（Address Generator）来实现stride寻址。

4.3 编译器与硬件的协同优化

编译器在AI芯片中的作用不仅仅是把算子映射到硬件，更重要的是做算子融合和内存规划。以Transformer的一个注意力层为例，原始的计算图包含Q投影、K投影、V投影、Q×K^T、softmax、×V、输出投影等多个算子。如果每个算子都单独执行，中间结果需要反复写入和读取SRAM，访存开销巨大。

编译器可以通过算子融合把多个算子合并成一个"超级算子"，让中间结果留在PE阵列或寄存器中，不写回SRAM。比如，Q×K^T的结果可以直接送入softmax单元，softmax的输出直接送入下一个矩阵乘法单元。这种融合可以把访存开销降低50%以上。

但算子融合不是免费的。它要求硬件支持灵活的PE阵列配置和可编程的数据通路。有些芯片采用"硬连线"的数据通路，融合能力有限；有些芯片采用"可重构"的数据通路，融合能力强但面积开销大。我在项目中的经验是：对于Transformer负载，至少需要支持"QK^T+softmax"和"softmax×V+输出投影"两个融合模式，这两个融合能覆盖注意力层80%的计算量。

5. 从RTL到流片：实际项目中的工程化考量

5.1 时钟频率与功耗的平衡

AI芯片的时钟频率选择是一个典型的工程取舍。频率越高，算力越大，但功耗也越高。对于一个16nm工艺的AI芯片，脉动阵列的频率通常在800MHz到1.2GHz之间。超过1.2GHz后，功耗会急剧上升，能效比反而下降。

我在项目中的做法是：把脉动阵列和SRAM放在不同的时钟域。脉动阵列跑在1GHz，SRAM跑在500MHz，通过异步FIFO进行跨时钟域通信。这样做的好处是SRAM的功耗可以降低30%左右，而脉动阵列的性能不受影响。代价是异步FIFO的面积和验证复杂度增加。

另一个降低功耗的手段是电压频率缩放（DVFS）。在推理负载较轻时，降低电压和频率，把功耗控制在预算内。但DVFS的切换需要时间（通常几十微秒），如果负载变化频繁，切换开销可能抵消节能收益。所以DVFS策略需要根据实际负载特征来设计。

5.2 验证策略：如何确保脉动阵列的正确性

脉动阵列的验证是一个大工程。一个32×32的阵列有1024个PE，每个PE都有乘法器、累加器、寄存器和多路选择器。如果只靠定向测试，覆盖率很难上去。我在项目中采用的是形式验证+随机测试+FPGA原型验证的三层策略。

形式验证用于证明PE阵列的数据通路满足设计规范，比如"第i行第j列的PE在第t拍应该接收到正确的数据"。随机测试用于覆盖各种边界条件，比如矩阵维度不是阵列大小的整数倍、精度模式切换、转置操作等。FPGA原型验证用于在真实负载下测试性能和功耗。

这里有一个容易忽略的点：脉动阵列的复位和初始化序列需要特别验证。因为阵列中有大量的寄存器和FIFO，如果复位不彻底，残留的数据可能导致计算错误。我在项目中遇到过一个问题：在连续执行多个矩阵乘法时，前一个乘法的部分和没有完全清空，导致后一个乘法的结果出现偏差。后来在每轮计算前增加了一个"清空周期"才解决。

5.3 流片后的调试：性能不达预期怎么办

流片回来之后，最怕的就是性能不达预期。我在项目中遇到过一种情况：芯片的标称算力是128 TFLOPS，但实际跑Transformer时只能达到40 TFLOPS。排查过程如下：

第一步，检查时钟频率是否达到设计值。用示波器测量时钟输出，发现实际频率只有800MHz，而设计值是1GHz。原因是时钟树上的缓冲器插入过多，导致延迟增加。解决方案是优化时钟树，减少缓冲器级数。

第二步，检查SRAM带宽是否成为瓶颈。用性能计数器统计SRAM的读写次数，发现SRAM的利用率只有60%。原因是数据布局不合理，导致大量的bank冲突。解决方案是调整数据布局，把连续访问的数据映射到不同的bank。

第三步，检查PE阵列的利用率。用硬件性能计数器统计PE的活跃周期，发现只有50%的PE在大部分时间处于活跃状态。原因是矩阵维度不匹配，很多PE在"空转"。解决方案是优化编译器的分块策略，把多个小矩阵拼成一个大矩阵来计算。

经过这三步优化，最终性能提升到了90 TFLOPS，达到了标称值的70%。这个案例说明：AI芯片的性能优化是一个系统工程，需要从时钟、访存、计算三个维度同时入手。

6. 踩坑实录：那些文档里不会写的经验教训

6.1 脉动阵列的"边缘效应"：为什么第一拍和最后一拍总是浪费

脉动阵列有一个固有的问题：在计算的开始和结束阶段，阵列中只有部分PE在工作。以N×N阵列计算N×N矩阵为例，第一拍只有1个PE在工作，第二拍有2个，直到第N拍才有N²个PE全部工作。计算结束时也是类似，最后一拍只有1个PE在工作。这个"边缘效应"导致的实际利用率是：

实际利用率 = N² / (2N-1) ≈ N/2

对于N=32的阵列，利用率只有约94%。看起来还可以接受，但如果矩阵维度小于阵列大小，利用率会进一步下降。比如用32×32阵列计算8×8矩阵，利用率只有8²/(2×32-1)≈12.7%。

我在项目中的解决方案是**"矩阵拼接"**：把多个小矩阵拼成一个大矩阵来计算。比如，把4个8×8矩阵拼成16×16矩阵，利用率提升到16²/(2×32-1)≈50%。但拼接需要额外的数据重排逻辑，而且要求这些小矩阵之间没有数据依赖。对于Transformer中的多头注意力，不同head之间的计算是独立的，正好适合拼接。

6.2 精度切换的"暗坑"：累加器的位宽设计

在支持多精度的AI芯片中，累加器的位宽设计是一个容易被忽视的问题。以INT8乘法为例，两个8位整数相乘得到16位结果，但如果累加多次（比如K=512），结果可能达到16+log2(512)=25位。如果累加器只有24位，就会溢出。

我在项目中遇到过这个问题：在跑一个K=1024的矩阵乘法时，INT8模式的输出结果出现周期性错误。排查后发现是累加器溢出。解决方案是把累加器位宽从24位增加到32位。这个改动的面积开销大约是5%，但避免了精度问题。

对于FP16模式，累加器的位宽需求更高。因为FP16的尾数只有10位，累加多次后舍入误差会累积。通常的做法是用FP32累加器来累加FP16的乘积，这样可以把精度损失控制在可接受范围内。但FP32累加器的面积是FP16的4倍左右，需要在精度和面积之间做取舍。

6.3 编译器与硬件的"接口错位"：一个真实的调试案例

我在项目中遇到过这样一个问题：编译器生成的指令序列在仿真器上运行正常，但在FPGA原型上跑出来的结果不对。排查过程如下：

首先，对比仿真器和FPGA的波形，发现指令的取指地址不一致。仿真器中指令从地址0x1000开始，FPGA中从0x2000开始。原因是编译器的链接脚本和FPGA的存储器映射不一致。

修复链接脚本后，结果仍然不对。继续排查，发现数据地址的对齐方式不同。编译器假设SRAM支持字节寻址，但硬件只支持16字节对齐的寻址。导致某些数据访问被硬件忽略。

修复对齐问题后，结果基本正确，但偶尔出现单个元素的错误。进一步排查，发现是跨时钟域FIFO的亚稳态问题。在高速运行时，FIFO的读指针和写指针偶尔会出现竞争。解决方案是在FIFO的读写指针上增加格雷码编码和同步器。

这个案例说明：软硬件协同设计中的"接口"不仅仅是逻辑上的接口，还包括地址映射、对齐规则、时钟域 crossing 等物理层面的接口。这些细节在架构设计阶段很容易被忽略，但到了调试阶段就会变成大问题。

7. 写给正在做AI芯片设计的你：一些个人体会

做AI芯片设计这些年，我最大的感受是：这个领域的门槛不在单点技术上，而在系统思维上。你可能对脉动阵列的原理了如指掌，也可能对Transformer的每个算子都烂熟于心，但如果不能把算法、架构、编译器、验证、物理实现串起来看，做出来的芯片大概率是"纸面参数漂亮，实际跑分拉胯"。

我见过太多项目在架构设计阶段只盯着峰值算力，忽略了访存带宽和编译器成熟度，结果流片后才发现实际利用率只有30%。也见过项目为了追求极致的能效比，把SRAM容量砍得太狠，导致大模型根本跑不起来。这些教训归结起来就是一句话：AI芯片的设计目标不是"最高算力"，而是"在目标负载下的最高有效算力"。

如果你正在做相关项目，我的建议是：在架构设计阶段就拉上算法团队和编译器团队一起评审，把目标模型的实际算子形状、访存模式、精度需求都摸清楚。在验证阶段不要只跑合成测试，要跑真实的模型负载。在流片前留足时间做FPGA原型验证，把软硬件接口的问题尽早暴露出来。

最后分享一个我在项目中总结的"三问法则"，每次做架构决策时问自己：这个设计对目标负载的利用率是多少？这个设计的访存带宽够不够？这个设计的编译器能不能高效映射？如果这三个问题都有明确答案，那这个决策大概率是靠谱的。