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 (SM10X) 與 Hopper (SM90) 架構進行了調優。
• 顯著加速：展示了相較於 FLA Triton 實作的巨大效能提升，包括在 Hopper 上 FlashKDA prefill 實現高達 7.56 倍的加速。
• 廣泛的演算法支援：實現了多種線性注意力變體，如 KDA (Kimi Delta Attention) 與 Lightning Attention。
• 無縫整合：設計為只需簡單的 import 更改即可直接替換 FLA 核心。