inclusionAI/cuLA

解决的问题

cuLA 解决了长上下文 LLM 工作负载中线性注意力机制的性能瓶颈。虽然线性注意力将标准注意力的二次复杂度降低为线性时间更新，但在现代 GPU 上高效实现这些机制需要深度的硬件级优化，以充分利用最新 NVIDIA 架构的计算能力。

工作原理

cuLA 提供了使用 CuTe DSL 和 CUTLASS C++ 编写的手动调优 CUDA 内核。它实现了线性注意力的高性能变体，包括 GLA、KDA、GDN 和 Lightning Attention。这些内核专门针对 NVIDIA Blackwell (SM10X) 和 Hopper (SM90) GPU 进行了优化，利用 WGMMA 数据流和持久打包调度（persistent packed scheduling）等特性，在 prefill 和 decode 阶段均实现吞吐量最大化。

适用对象

专为处理长上下文 LLM 且需要在 NVIDIA Hopper 或 Blackwell 硬件上获得最大推理和训练性能的 AI 研究人员和工程师设计。它还旨在作为 flash-linear-attention (FLA) 库的高性能后端。

亮点

• 硬件优化：专门针对 NVIDIA Blackwell (SM100X) 和 Hopper (SM90) 架构进行了调优。
• 显著加速：展示了相较于 FLA Triton 实现的巨大性能提升，包括在 Hopper 上 FlashKDA prefill 实现高达 7.56 倍的加速。
• 广泛的算法支持：实现了多种线性注意力变体，如 KDA (Kimi Delta Attention) 和 Lightning Attention。
• 无缝集成：设计为只需简单的 import 更改即可直接替换 FLA 内核。