SuperAPI AI 工具集

flashkda-delta-attention

通用与其它

查看项目

更新于 10-15 21:05

技能介绍

用途:FlashKDA 提供 Kimi Delta Attention(KDA)的高性能 CUDA 内核,基于 CUTLASS 实现,面向 SM90+ GPU(H100、H20 或更新),可作为 flash-linear-attention 中 chunk_kda 操作的即插即用后端,用于高效循环状态推理。

主要能力:提供底层内核 flash_kda.fwd,支持 q/k/v/g/beta 等张量输入、可选初始与最终循环状态、以及 cu_seqlens 变长打包批处理;支持多轮/流式推理中跨块传递状态。K 与 V 必须为 128。

使用方式:安装需克隆仓库并初始化子模块后 pip install -v .;可选安装 flash-linear-attention(>=0.5.0)以启用自动分发。推荐通过 FLA 的 chunk_kda 调用,也可直接调用 flash_kda.fwd。环境变量 FLA_FLASH_KDA=0 可强制回退 Triton 路径。测试可运行 bash tests/test.sh。

适用场景:需要 KDA 循环状态推理的 GPU 推理任务,包括变长序列打包批处理与多轮/流式状态维护。

限制:要求 SM90+ GPU、CUDA 12.9+、PyTorch 2.4+、Python 3.8;K 与 V 必须为 128;变长模式下 B 必须为 1;initial_state 与 final_state 同时提供时 dtype 必须一致;lower_bound 有效范围为 [-5.0, 0];FLA 分发需使用 torch.inference_mode()。其余收费、价格、平台等信息暂未提供。