第4章 Self-Attention加速

第4章：Self-Attention加速：从QKV到Softmax的硬件实战

各位同学，今天我们来啃一块硬骨头——Self-Attention的硬件加速。说实话，我在做第一版Transformer芯片时，这块花了整整三个月才调通。你想想看，一个Attention层里既有矩阵乘法，又有softmax这种非线性操作，还有mask机制要处理，硬件上要把它们全部流水起来，确实不简单。

这一章，我会把QKV线性投影、Scaled Dot-Product Attention的并行化、Softmax的硬件实现，以及Mask机制，一个一个拆开来讲。每个部分我都会结合我踩过的坑来说，希望能帮你少走弯路。

4.1 QKV线性投影的硬件映射

先说QKV投影。说白了，就是把输入X分别乘上三个权重矩阵Wq、Wk、Wv，得到Query、Key、Value。这个操作在软件里就是三个矩阵乘法，但在硬件上，怎么映射效率最高？

我个人习惯的做法是：把三个投影合并成一个大的矩阵乘法。什么意思呢？就是把Wq、Wk、Wv拼成一个更大的矩阵，维度是[d_model, 3*d_model]。这样原来需要三次矩阵乘法，现在一次就能搞定。

核心思路： 合并QKV投影，减少数据搬运次数，提高计算单元利用率。

我在项目中遇到过一个问题：如果芯片的矩阵乘法单元（比如Systolic Array）的尺寸是固定的，比如128x128，而d_model是512，那3*d_model=1536，没法一次塞进去。怎么办？

我的做法是分块（tiling）。把大矩阵切成128x128的小块，分批计算。这里有个关键点：切分方向要沿着d_model维度，而不是batch维度。为什么？因为这样可以让输入X的数据复用，减少从DDR搬数据的次数。

// 伪代码：QKV投影的硬件映射
// 假设：d_model=512, head=8, d_k=64
// 硬件矩阵乘法单元尺寸：128x128

void qkv_projection_hw(float* X, float* W_qkv, float* Q, float* K, float* V) {
    // 1. 合并权重：W_qkv [512, 1536]
    // 2. 分块计算：沿d_model维度切分，每块128
    for (int i = 0; i < 12; i++) {  // 1536 / 128 = 12块
        for (int j = 0; j < 4; j++) {  // 512 / 128 = 4块
            // 加载W_qkv的第j行块、第i列块
            // 加载X的第j行块
            // 调用矩阵乘法单元计算
            // 结果累加到Q/K/V对应位置
        }
    }
}

小技巧： 如果芯片支持多bank的SRAM，可以把W_qkv的权重预先分布到不同bank里，这样读取时可以并行，减少访存冲突。我曾经因为没注意bank冲突，导致性能直接腰斩，后来花了三天才定位到问题。

4.2 Scaled Dot-Product Attention的并行化

接下来是Attention的核心计算：Q和K做点积，然后scale，再softmax，最后乘V。这个流程看起来是串行的，但硬件上我们可以做很多并行化的工作。

先看QK^T这一步。假设有8个head，每个head的Q和K维度是[seq_len, d_k]。如果seq_len=1024，d_k=64，那QK^T的结果是[1024, 1024]。这个矩阵乘法怎么加速？

我的做法是：把不同head的计算完全并行。8个head，每个head分配独立的计算单元。这样QK^T的计算时间就变成了原来的1/8。

但这里有个坑：softmax需要全局的max和sum，不同head之间不能完全独立。嗯，这个我们下一节再细说。

并行化策略总结：

Head间并行： 每个head分配独立计算单元
QK^T内部并行： 使用矩阵乘法单元，一次计算多行多列
流水线并行： QK^T计算完一行后，立即开始softmax，不用等全部算完

我记得有一次，团队里新来的同事把QK^T算完才做softmax，结果发现中间结果太大，片上SRAM根本放不下。后来改成逐行流水，才解决了这个问题。你想想看，1024x1024的矩阵，float16精度，那就是2MB，片上SRAM一般也就几MB，还要放权重和中间变量，确实不够用。

4.3 Softmax的硬件实现

Softmax是Attention里最麻烦的部分。它涉及指数运算、求和、除法，这些在硬件上都不好做。软件里一行代码搞定，硬件上可能要几十个周期。

常用的硬件实现方法有两种：查找表（LUT）和分段线性近似。我两种都用过，各有优劣。

4.3.1 查找表（LUT）

查找表说白了就是提前算好指数函数的值，存到ROM里，用的时候直接查。比如输入范围是[-16, 0]，我们以0.001为步长，那需要16000个表项。每个表项用16bit存，那就是32KB。这个大小在芯片上是可以接受的。

但有个问题：输入范围怎么定？ 如果输入超出范围，直接截断还是做特殊处理？我在项目中遇到过输入跑到-20的情况，直接截断到-16，结果精度损失很大。后来我加了一个边界检测，超出范围的就走一个简单的线性近似。

// 查找表实现softmax中的exp
// 输入范围：[-16, 0]，步长0.001
// 表项数量：16001

float lut_exp(float x) {
    if (x > 0) return 1.0;  // 理论上x不会大于0，因为减去了max
    if (x < -16) return 0.0;  // 截断
    
    int idx = (int)((-x) * 1000);  // 映射到表索引
    return exp_lut[idx];
}

注意： 查找表的精度和表大小是trade-off。步长越小，精度越高，但ROM也越大。我建议用0.001的步长，配合线性插值，精度可以做到1e-5以内。

4.3.2 分段线性近似

另一种方法是分段线性近似。把指数函数分成若干段，每段用一条直线来拟合。比如把[-16, 0]分成16段，每段宽度1.0，每段用两个参数（斜率和截距）来拟合。

这个方法的好处是：不需要大ROM，只需要存16组参数。计算时先判断输入落在哪一段，然后做一次乘加运算就行。缺点是精度不如查找表高，但很多场景下够用了。

方法	存储开销	计算延迟	精度
查找表（步长0.001）	32KB	1周期（查表）	高（1e-5）
分段线性（16段）	128字节	2周期（判断+乘加）	中（1e-3）
分段线性（64段）	512字节	2周期	较高（1e-4）

我个人习惯在精度要求高的场景用查找表，在面积受限的场景用分段线性。如果你拿不准，可以先在软件里仿真一下两种方法的精度损失，再做决定。

4.4 Mask机制的硬件支持

最后说说Mask。在Decoder的Self-Attention里，我们需要一个上三角mask，让每个token只能看到它前面的token。这个在软件里很简单，但在硬件上怎么高效实现？

我的做法是：在QK^T计算完成后，直接对结果矩阵做mask操作。具体来说，就是用一个mask矩阵（也是上三角）和QK^T的结果做element-wise的与操作。mask矩阵里，需要mask掉的位置是0，不需要mask的位置是1（或者一个很大的负数，比如-1e9）。

但这里有个性能问题：mask矩阵也需要存储和搬运。如果seq_len=1024，mask矩阵就是1024x1024，1M个元素。虽然可以用bit来存（每个位置1bit），但还是要128KB。

优化技巧： 其实不需要显式存储mask矩阵。因为mask是固定的上三角模式，硬件上可以做一个简单的地址判断：如果行索引小于列索引，就把结果置为-INF。这样连存储都省了。

我曾经在项目里犯过一个错误：把mask操作放在softmax之后做。结果发现softmax已经算完了，再mask掉某些位置，概率分布就不归一化了。后来才意识到，mask必须在softmax之前做，把需要mask的位置设成-INF，这样softmax之后这些位置的概率就是0。

嗯，这里还要注意一点：padding mask和causal mask要分开处理。padding mask是对所有head都一样的，而causal mask是每个head都一样。我建议把两种mask合并成一个，在硬件上统一处理，减少控制逻辑的复杂度。

本章小结

这一章我们聊了Self-Attention的四个关键加速点：

QKV投影： 合并权重矩阵，分块计算，注意bank冲突
Scaled Dot-Product Attention： head间并行，逐行流水，减少中间结果存储
Softmax： 查找表和分段线性近似两种方案，根据精度和面积需求选择
Mask： 用地址判断代替显式存储，注意mask必须在softmax之前做

这些内容看起来多，但核心就一句话：把计算和数据搬运用好，把流水线排好。下一章我们会继续深入，看看多头注意力（Multi-Head Attention）的硬件实现，以及怎么把多个head的计算进一步优化。

避坑指南： 我曾经在流片前一周发现，因为QK^T的中间结果太大，导致片上SRAM溢出。后来紧急改成逐行流水，才赶上了流片窗口。所以，一定要在架构设计阶段就估算好中间结果的存储需求，不要等到RTL写完了才发现问题。

📚 课程目录