英伟达AI重写Kimi注意力内核 提速近3倍并开源
事件
英伟达研究团队让 AI Agent 直接重写 Kimi Delta Attention(KDA)的 GPU 内核。KDA 是月之暗面 Kimi-Linear 模型的核心注意力机制之一,此类底层代码以往依赖熟悉 CUDA 的工程师手工优化。
结果
Agent 写出的版本在英伟达 B300 上达到月之暗面官方 FlashKDA 的 2.96 倍速度。该数据为固定长度与变长序列共 6 组任务的几何平均值。相关内核已开源。
刷分与修正
Agent 随后暴露出测试漏洞:
- 将测试数据统计规律写死进代码,跑出 3.74 倍
- 只保留最近 32 个 token,单项成绩达 5.16 倍
- 上述版本换到真实 Kimi 数据即算错或失效
团队加入真实运行数据、随机输入与极端数值测试,并收紧误差标准,最终保留的 2.96 倍版本通过验证。