MoonshotAI/FlashKDA

FlashKDA: high-performance Kimi Delta Attention kernels

解決的問題

FlashKDA 提供 Kimi Delta Attention (KDA) 的高效率 CUDA 內核。KDA 是一種線性注意力機制，旨在解決使用此特定注意力變體的模型在速度與記憶體效率方面對高效、硬體加速計算的需求。

工作原理

本專案基於 CUTLASS 庫實作 KDA 內核，專門針對 NVIDIA SM90（H100）及更高架構。它作為 flash-linear-attention (FLA) 庫的後端整合，讓使用者可將標準的 Triton 基礎實作，替換為這些經過優化的 C++ 內核，以獲得更佳效能。

適用對象

專為從事線性注意力模型的 AI 研究人員與工程師設計，也適用於使用 flash-linear-attention 框架，在現代 NVIDIA GPU 上優化模型推論與訓練效能的使用者。

主要亮點

  • 基於 CUTLASS 的高效率內核，專為 SM90+ GPU 設計。
  • 透過自動分發機制，與 flash-linear-attention 庫無縫整合。
  • 支援使用累積序列長度（cu_seqlens）的可變長度批次處理。
  • 支援無狀態與有狀態（初始/最終狀態）的遞迴運算。

相關

  • 專案
    inclusionAI/cuLA

    針對 KDA 與 Lightning Attention 等線性注意力變體的高性能 CUDA 核心庫，專門針對 NVIDIA Blackwell 與 Hopper GPU 進行了優化。

  • Dispatch
    Kimi Linear 以減少 75% 記憶體佔用與 6 倍解碼速度超越全注意力機制

    Kimi Linear 的混合線性注意力架構在各項任務中表現優於全注意力機制，同時將 KV-cache 記憶體減少高達 75%，並在 1 M-token 上下文中實現高達 6 倍的解碼加速。

  • 專案
    NVIDIA/cutlass

    A collection of CUDA C++ and Python DSL abstractions for implementing high-performance matrix-matrix multiplication (GEMM) and the related linear algebra operations on NVIDIA GPUs.

  • Dispatch
    Kimi Delta Attention：推導、實作與社群回應

    Kimi Delta Attention（KDA）透過將標量保留替換為每通道的對角矩陣，實現對 key 通道的獨立遺忘，同時保留 delta‑rule 寫入，且可使用融合遞迴或分塊的 Triton 核心高效執行。

  • 專案
    Tencent/hpc-ops

    針對 LLM 推論的生產級運算子庫，提供針對注意力、MoE 與取樣高度優化的 CUDA 內核，以降低服務延遲。