表格基础模型注意力机制在不同硬件与形状下的基准评测

01 摘要

论文指出表格上下文学习模型依赖行列交替注意力，其行注意力序列长、列注意力序列短，且表格数据内存布局分散，当前模型隐藏维度较小，而现有高效注意力研究多集中在一维序列。作者在 A100、H100、B200 上对多种注意力后端进行吞吐基准测试，发现最优后端在列注意力和行注意力之间不同，并随硬件与模型细节变化；针对各代 GPU 优化的 FlashAttention 整体表现最佳，但在长序列列注意力下有时被 cuDNN 超越，交叉点取决于头维度。作者认为该可复现基准为表格原生注意力改进奠定基础（属作者观点）。

02 关键点

1. 表格注意力与语言模型一维注意力在序列长度和内存布局上差异显著，但此前缺乏二维表格场景的高效注意力基准。
2. 在 A100/H100/B200 上测试 Torch SDPA、FlashAttention-2/3/4、vLLM 和 SageAttention 的前向与反向吞吐。
3. 最优后端取决于注意力类型（列/行）、序列长度、头维度与硬件；FlashAttention 整体最好，cuDNN 在长序列列注意力中有时更优，SageAttention 在行注意力超过 16k 行时表现良好。