- 题目:FACT: FFN-Attention Co-optimized Transformer Architecture with Eager Correlation Prediction
- 时间:2023
- 会议:ISCA
- 研究机构:清华
主要贡献
- 算法优化:提出了Eager Prediction (EP),预测softmax前重要的部分,跳过不重要的计算,减少QKV的计算量。
- 使用基于log的加法操作进行预测,减少乘法运算。
- 通过一个乱序调度器消除预测的延迟,使得预测和计算可以完全流水线化。
- 通过一种新颖的对角线存储模式避免混合精度FFN中的内存访问冲突。
如上图所示,transformer中FFN的计算量是更大的,只有token数量足够长时attention的计算量才会超过FFN
过去的工作只是节约了S=QK的计算量,还不够,这篇论文更进一步,预测完再动态的生成QKV,从而结余QKV的计算量
Out-of-order QKV Generation Scheduler
-
输入和预测:输入token矩阵被送入系统,同时EP单元开始工作,预测注意力矩阵。
-
乱序QKV生成:
步骤1:EP单元预测第一行的注意力矩阵,确定需要计算的K、V矩阵的列。
步骤2:PE阵列立即开始为这些列生成K、V矩阵,并将结果存储在临时缓冲区中。
步骤3:当EP单元预测下一行的注意力矩阵时,PE阵列会计算额外需要的K、V矩阵,同时继续处理之前生成的Q矩阵。
并行处理:随着EP单元逐行预测注意力矩阵,PE阵列并行地生成QKV矩阵,并根据预测结果动态调整计算顺序。 -
存储和使用:生成的QKV矩阵以乱序的方式存储在临时缓冲区中,当计算某一行的注意力分数时,直接从缓冲区中读取对应的QKV矩阵。