MoonshotAI/FlashKDA

Source: https://github.com/MoonshotAI/FlashKDA

FlashKDA

FlashKDA: Flash Kimi Delta Attention —— 基于 CUTLASS 的高性能 KDA 内核

新闻

• 2026-04-22 — 深度解析博客：FlashKDA v1 背后的设计决策，请点击此处阅读。

要求

• SM90 及以上
• CUDA 12.9 及以上
• PyTorch 2.4 及以上

安装

git clone https://github.com/MoonshotAI/FlashKDA.git flash-kda
cd flash-kda
git submodule update --init --recursive
pip install -v --no-build-isolation .

默认情况下，构建过程会检测当前的 CUDA 设备并针对该架构进行编译。对于 wheel 或 CI 构建，可以显式编译所有支持的架构：

FLASH_KDA_CUDA_ARCHS=all pip install -v --no-build-isolation .

支持的值包括 auto（默认）、all，或以逗号分隔的架构列表，例如 90a,100a。

将 FlashKDA 用作 FLA 后端

安装完成后，FlashKDA 会被 flash-linear-attention 的 chunk_kda 自动调度。集成详情请参见 fla-org/flash-linear-attention#852 (https://github.com/fla-org/flash-linear-attention/pull/852)。

要求

1. 安装 flash-linear-attention >= 0.5.0：

   pip install -U flash-linear-attention
   

2. 在 torch.inference_mode() 下调用 chunk_kda：

   import torch
   from fla.ops.kda import chunk_kda
   
   with torch.inference_mode():
       out, final_state = chunk_kda(
           q=q,
           k=k,
           v=v,
           g=g,
           beta=beta,
           scale=scale,
           initial_state=h0,
           output_final_state=True,
           use_gate_in_kernel=True,
           use_qk_l2norm_in_kernel=True,
           use_beta_sigmoid_in_kernel=True,
           safe_gate=True,
           A_log=A_log,
           dt_bias=dt_bias,
           lower_bound=lower_bound,
           transpose_state_layout=True,
           cu_seqlens=cu_seqlens,
       )
   

退出： 设置 FLA_FLASH_KDA=0 可回退至 Triton 路径。

调试调度： 添加 logging.basicConfig(level=logging.INFO) 即可在命中时看到 [FLA Backend] kda.chunk_kda -> flashkda，未命中时看到 ... rejected:。

性能

请参见 BENCHMARK_H20.md。

测试

bash tests/test.sh

• tests/test_fwd.py — 正确性测试（与 torch 参考实现精确匹配，并与 flash-linear-attention 进行对比）。

内核 API

flash_kda.fwd

flash_kda.fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, initial_state=None, final_state=None, cu_seqlens=None)

参数：

• 当前要求 K = V = 128。
• initial_state / final_state 可接受 None（无状态）、bf16 或 fp32 张量。两者都提供时，数据类型必须一致。
• 提供 cu_seqlens 时，B 必须为 1，T 为所有序列的总长度，initial_state / final_state 的形状为 [N, H, V, K]。
• 当 cu_seqlens 为 None 时，每个批处理元素被视为独立序列，状态形状为 [B, H, V, K]。

开发

为 CUDA/C++ 源码设置智能感知（clangd），请运行：

bash setup_clangd.sh

该脚本会生成包含正确仓库路径的 .clangd 文件，并将全局 clangd config.yaml 安装到 ~/.config/clangd/。

引用

@misc{flashkda2026,
  title={FlashKDA: Flash Kimi Delta Attention},
  author={Yutian Chen, Zhiyuan Li, Yucheng Wang, Ming Wei},
  year={2026},
  publisher = {GitHub},
  howpublished = {\url{https://github.com/MoonshotAI/FlashKDA}},
}