大数跨境

生成的Kimi注意力算子提速2.05倍!NVIDIA&陈天奇团队提出编译器与智能体协同设计范式

生成的Kimi注意力算子提速2.05倍!NVIDIA&陈天奇团队提出编译器与智能体协同设计范式 智猩猩AI
2026-09-12
4
导读:四项Kernel优化以PR形式进入FlashInfer上游~

智猩猩AI整理

编辑:BugMaker


写 GPU Kernel(算子)这件事,编程智能体现在已经能上手了。


只是大多数系统把编程环境当成一个固定的黑盒。智能体提交代码,编译器丢回一个能跑或不能跑,再跑一次基准,返回一个延迟数字。一个崩溃通常不会告诉智能体违反了哪类安全约束或硬件条件,一个延迟数字也不会解释是哪一步设计拖慢了程序。


写 Kernel 的专家不是这么干的。他们脑子里装着一张工作负载的小模型,时刻算着显存、带宽、warp 这些硬件资源,还能把上一份 Kernel 的经验带到下一份里。


编译器本身已经包含大量这类能力,包括结构化操作词表、资源模型、合法性检查、静态分析、代价模型和 lowering 规则。CAKE关注的问题,是如何把这些能力转化为面向智能体的反馈系统。


论文作者包括陈天奇,卡内基梅隆大学副教授,TVM 与 XGBoost 的共同缔造者。Luis Ceze 是华盛顿大学教授、OctoAI 联合创始人,也是 TVM 的共同作者。


这项工作由 NVIDIA 与卡内基梅隆大合作完成,本文多个优化结果最终以 FlashInfer 上游 PR 形式落地。



他们把答案叫做 CAKE,全称编译器—智能体协同设计(Compiler–Agent Co-Design)


思路很直接。让智能体改写的不是裸 CUDA,而是一套类型化、硬件显式的中间表示 Cake IR(Intermediate Representation)。编译器在代码生成之前就给出定位到具体设计决策的正确性诊断和性能诊断。


更关键的是,编译器本身也是被进化的对象,反复出现的失败会被沉淀成新的校验规则、新的原语,而不是一次性补丁。


效果立竿见影。在一个隐藏底层实现的干净启动(clean-start)实验里,80M token 预算下 CAKE 跑出的最好算子达到手工调优 FlashML 基线的 1.144 倍,而直接写 CUDA/PTX 只有 0.928 倍


针对 Kimi Delta 注意力(Kimi Delta Attention,KDA)生成的前沿算子,相对官方 FlashKDA 的几何平均加速达到 2.05 倍


01

黑盒里丢掉的,正是

专家的那点直觉  


现在的 Kernel 智能体改进的都是搜索本身,提案怎么写、变异怎么变、排序怎么排,但环境反馈只有三样,编译错误、对错结果、端到端延迟。


这些信号通常无法说明,是哪一次程序决策导致了同步失败、违反硬件契约,还是流水线停滞(pipeline stall)。而且当工作负载暴露出编译器缺的能力时,这套反馈也长不出来。


CAKE 要解决的,就是编译器那套机器变成智能体面向(agent-facing)的东西


论文给了一张直观对比,现有做法里,优化智能体对着基准环境来回试,反馈只有正确性和延迟。CAKE 把执行与反馈拆开,编译器把校验、性能模型、代码生成全部喂回给智能体,还多了一个外圈,编译器进化。



02

为智能体而生的 IR,偏偏

不碰布局代数     


CAKE 让智能体改写的,是一种类型化、硬件显式的调度表示(schedule representation)。


它把该干的事都写明白,哪些 warp 承担什么角色(warp specialization)、缓冲区 staging 到多深、哪道 barrier 卡在哪次交接、哪条指令消费哪个操作数。


分工是,调度说要发生什么,lowering 推导具体怎么做。barrier 地址、相位位、张量内存(Tensor Memory,TMEM)偏移、描述符编码,全部从声明里算出来,而不是让智能体手写。



这里有一个不同于传统 GPU DSL 的取舍。CUTLASS CuTe、Triton Linear Layout 等系统,都将布局抽象为重要的编程对象。


智能体直接写下具体承诺,一个共享内存(SMEM)视图偏移、一个 TMEM 列范围、一个 swizzle 标签,编译器负责判断这些承诺合不合法。


换句话说,它把容易导致智能体错误且难定位的布局代数操作,换成可以直接校验的显式记录。


这套词表也不是拍脑袋定的,而是从生产 Kernel 的语料里自底向上长出来的。智能体从 FlashInfer、CUTLASS、DeepGEMM 这些库里识别反复出现的调度模式或缺失能力,改 IR、移植、验证,再进入下一轮。



03

编译器是会进化的队友    


编译器进化走两条互补的路。


一条是主动补课。智能体翻生产 Kernel 和硬件文档,找出 Blackwell 缺失的指令形态、资源类型、同步惯用法,形成编译器改动提案。


另一条是被动总结。智能体从失败的候选里分析 sanitizer 报告、对错不匹配、调试日志,把高频失败蒸馏成新分析。


一次不透明的运行时崩溃变成一条校验规则,一个反复出错的非法 lowering 模式变成一条静态检查,一次系统性预测偏差变成一个校准目标。


两条路是耦合的。新原语向编译器暴露更多硬件事实,让分析更强,新分析反过来约束未来原语的设计空间。所有编译器改动都要过语料测试这道关,因为原语和分析必须一起进化。



04

在 clean-start 实验中直接超过

手工调优基线     


最能说明问题的,是那场 Flash-KMeans 干净启动实验。


任务被设成隐藏底层实现,智能体只能看到数学规格、评测契约和正确性 oracle,看不到任何低层的 CUDA、PTX 或 SASS,因为这些引用里已经编码了要让智能体自己发现的调度决策。


三条 B200 跑下来,80M token 预算时 CAKE 组 3/3 达到平台期判据,最好算子达到手工调优 FlashML 基线的中位 1.144 倍


对照组直接写 CUDA/PTX,0/3 达到平台期,中位只有 0.928 倍中位进化时间也要 3.73 小时,比 CAKE 组的 1.89 小时慢了近一倍


这是论文中最能体现 CAKE价值的实验之一,也是差距最大的结果之一。



换到前沿算子上,差距更直观。针对 Kimi Delta 注意力的 prefill,CAKE 生成的实现相对官方 FlashKDA 达到 2.05 倍几何平均加速,并在 SGLang 的端到端 Kimi-K3 服务里验证通过,生成的 CUDA 已进入 FlashInfer 上游 PR。


在 Alpha-MoE 上,API级调用速度最高提升6.204倍和4.025倍,而单看GPU执行时间,提升分别为1.215倍和1.170倍


KNN 和 KMeans 两个家族在 400 多个 shape 上整体提升 1.42 到 2.12 倍。在已知算子复现实验中,11个固定对比结果里有10个达到或超过对应基线,包括 TensorRT-LLM、CUTLASS 和 DeepGEMM 等。


05

从一个 shape 到能上库,

还有诚实的边界   


论文单独拎出一步。很多人以为把内层循环多跑几个 shape 就行,研究团队说这是两回事。


单 shape 优化可以激进特化,泛化阶段要换目标、换排序信号、换失败模式。他们把测过的种子分桶,生成专用或共享变体,排好 guard 和显式回退,再验证边界、尾部和重叠情形。


GB200 上,泛化后的 KNN build、KNN search、KMeans 分别拿到 1.418 倍、2.116 倍、1.803 倍的几何平均,KNN 召回率 1.0,零错误输出。



研究团队也把边界写得很诚实。这套东西目前只测了 NVIDIA 从 Ampere 到 Blackwell 的目标,非 NVIDIA 平台没有结果。


大部分性能证据在 B200,时延模型只校准了 B200 和 H100。静态分析和性能模型故意留得不完整,GPU 实测仍是最终裁判,编译器进化的 merge gate 也还是人在把守。


CAKE 真正改变的不是让智能体更快地搜 Kernel,而是把编译器从黑盒改造成一个会进化的协作者,让编译器知识在智能体循环里不断积累。四个改动已经以 FlashInfer 上游 PR 的形式交了回去,这大概是最实在的注脚。

END


关注+星标,获取AI前沿进展与开源一线动态

【声明】内容源于网络
0
0
智猩猩AI
智猩猩旗下AI技术内容账号,关注AI大模型掀起的范式革命,追踪AI大时代涌现的开源项目。
内容 2350
粉丝 0
智猩猩AI 智猩猩旗下AI技术内容账号,关注AI大模型掀起的范式革命,追踪AI大时代涌现的开源项目。
总阅读3.2k
粉丝0
内容2.4k