使用Cutlass/CuTe进行HSTU注意力开发与优化

HSTU Attention 核心介绍

HSTU Attention 是 NVIDIA 开发的用于生成式推荐的注意力机制内核，其性能远超 OAI-triton，并且开源，支持多种定制化掩码，支持实验性 FP8 注意力训练，支持推理时的 Paged Attention。其优化主要从以下几个方面进行：

• 内核融合: 将多个计算步骤融合成一个内核，减少内核调用次数，提高效率。
• 块拆分: 将注意力矩阵拆分成更小的块进行计算，减少内存访问压力，提高缓存利用率。
• 形状优化: 优化计算顺序和同步时间，隐藏计算和内存访问的延迟。
• Warp 拆分: 在线程块内部对 Warp 进行特殊化处理，提高并行性。
• 负载均衡: 在整个数据范围内对线程块进行特殊化处理，平衡 SM 的负载。
• 代数优化: 使用更高效的数学公式替换原有的计算公式，减少指令数量。
• 掩码设计: 将掩码函数几何化，减少不必要的计算。

关键数据

• HSTU Attention 的内存占用约为 0.98 GB。
• HSTU Attention 的计算量约为 9.510.912 个指令。
• HSTU Attention 的性能提升显著，例如在 AMPERE 架构上性能提升 3-7%。

研究结论

HSTU Attention 通过多种优化手段，显著提高了生成式推荐中注意力机制的计算效率和性能，为生成式推荐的应用提供了强大的技术支持。