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 内核，以降低服务延迟。