明敏 克雷西 发自 凹非寺
量子位 | 公众号 QbitAI
哎呀,这AI竟然能随机生成一个内核,其性能竟然超过了我们这些专家精心调优的结果!
斯坦福最近披露了一组新发现,结果真的太亮眼了。
AI优化的核心,在执行深度学习常规操作时,其表现超越了原始的PyTorch,性能提升幅度最高可达近400%。
在NVIDIA L40S GPU上进行的基准测试中,性能百分比是通过将参考时间与生成kernel_size所需时间进行对比计算得出的。
更惊人的是,这一切都是意外实现的。
研究团队本来的目标是生成合成数据以训练内核生成模型。
研究发现,在测试环节所生成的模拟数据,居然能够产生性能极为卓越的核心模块。
围观网友:没想到AI也要取代内核工程师了。
还有观察者指出,除了显著的性能增强之外,该研究团队所采纳的策略同样颇具趣味:
他们并非仅仅在操作层面进行逐步改进(比如采用爬坡算法),而是在每次迭代周期中穿插了一个语言推理环节,以此促使搜索过程变得更加丰富多元。
也就是说,他们促使系统在每次升级过程中,以类似“思考”的方式不断涌现新思路,进而探寻出更优的解决途径。
具体如何实现,一起来看。
改代码前先生成自然语言优化思想
根据斯坦福团队在博客中的阐述,该内核构建的方法相当简便——首先提供torch代码,随后指导大型模型创作定制的内核,以替代原有的torch运算符。
这些内核完全采用CUDA-C进行编写,无需借助CUTLASS和Triton等库,亦无需使用领域专用语言(DSL)。
与常规做法不同,模型并非直接着手修改代码,而是首先运用自然语言构建优化理念,随后再将这些理念转换成全新的代码版本。
团队之所以采取这种方式,是因为“逐个调整”的改进策略缺少变化,容易陷入局部最优解,反复探索同一种转换或是无止境地优化无望的路径。
为了提升思维的丰富性,斯坦福研究团队额外采用了多路径的探索策略。
具体而言,他们的做法并非每一步仅仅针对一个候选方案进行优化,相反,他们会将每一个创意进行拆分,从而衍生出多种可能的实现路径,随后选取其中性能最为卓越的核心部分,作为下一轮迭代的基础。
该团队运用OpenAI的o3模型与Gemini 2.5 Pro设备,对KernelBench 1级测试中的10个问题进行了多次挑战,经过多轮比拼,最终最佳的内核版本逐渐显现。
大多数杰出的成果在后续的几轮中涌现,尤其是第4轮和第5轮。
KernelBench是由斯坦福团队自主研发的一套针对AI生成内核的测试标准。这套基准将任务划分为三个不同等级,其中第一级涵盖了单一的原始操作,这些操作构成了AI的基础构件,例如卷积、矩阵与向量的乘法、矩阵与矩阵的乘法、损失函数、激活函数以及层归一化等。
这一新发现,与DeepMind的AlphaEvolve项目、o3揭露Linux系统0day漏洞等事件相结合,使得众多网友普遍认为,Gemini 2.5Pro和o3的技术实力已经跃升至一个新的高度。
回到斯坦福的项目中,我们能够观察到,在生成过程中,模型的思维模式逐渐展现出与人类经验相吻合的特征——
斯坦福团队还呈现了一系列具体的改进路径,这些路径表明,并非每一次的改进都能确保速度的提升,然而,通过多个步骤的累积,核心的速度可以显著增加,甚至有可能超越PyTorch。
在具体实施过程中,有人对AI在生成CUDA内核时的优化策略提出了疑问,他们想知道这些建议是否能够直接转化为相应的代码执行,抑或是仅仅启动了随机的搜索过程?
作者在回答中指出,尽管尚未进行严格系统的验证,但在手动审查的案例里,所生成的CUDA代码与所提出的优化方案基本相符。
AI并非纯粹进行随机优化,它实际上是在努力实施自己设定的策略。
华人主创团队意外发现
本研究的合作者包括三位学者,他们分别是Anne Ouyang、Azalia Mirhoseini以及Percy Liang。
欧阳目前就读于斯坦福大学扩展智能实验室,担任博士生一职;她的本科和硕士学业均是在麻省理工学院完成的;此外,她曾在英伟达的cuDNN团队中担任过职务。
李佩慈教授在斯坦福大学担任计算机科学副教授和统计学助理教授的双重职务,并且现正领导着斯坦福基础模型研究中心。
曾和李飞飞一起发布、推进了多项研究工作。
Azalia Mirhoseini担任斯坦福大学计算机科学领域的助理教授,同时也是斯坦福扩展实验室的创始人。她的职业生涯中,曾在DeepMind、Google Brain和Anthropic等知名机构任职。
她之前参与的科研项目涵盖了MoE(模块化计算架构)以及AlphaChip等芯片设计算法领域。
本次研究,本来是希望生成数据来训练内核生成模型。
在实施过程中,却意外地产生了未曾预料的成效;特别是在测试环节,所生成的模拟数据本身,竟然展现出了极为出色的内核性能。
由于这些核心技术实现了之前认为难以达成的高级优化和硬件功能,因此该团队选择了通过撰写博客的方式来公布他们的研究成就。
不过,关于数据生成的具体方法,研究团队目前并未对外公布,但它们只是简单提及,这种设计理念其实相当简便。
最关键的还是,它已经展示出了巨大潜力。
除此之外,研究团队还指出,这一新发现与近期出现的某些动向不谋而合——大规模的再训练行为已不再成为必要条件。
偶尔,运用高明的搜索与分叉策略,能够开启科学创新的门径,并攻克繁难问题;借助verifier进行深入的搜索,还能获得更多的成果。
将卓越的推理技巧与并行考察众多假设的方法相结合,往往能够产生更优的成效。此类方法正如AlphaEvolve、AlphaEvolution、Gemini 2.5 Pro等深度学习系统所展现的那样。
最终,该团队指出,该研究尚存在诸多改进的余地。例如,他们目前正致力于对两个具体方面进行优化。
相较于FP16或BF16,新硬件上FP32的优化水平通常不高,这便是为何采用FP32架构相较于PyTorch,更易于实现性能的显著提升。
他们指出,尽管目前仍存在诸多约束,但对未来的发展前景仍抱有积极态度。
起初,他们甚至无法构建出一个能够正常运作的核心,然而,经过不懈的改进搜索策略,他们已成功将flash attention的表现力提升至相当可观的层次。
值得注意的是,在搜索过程中所依赖的资源相当有限,大约仅涉及300万token的输入以及400万token的输出。
One More Thing
实际上,不只是一个团队在尝试开发内核大模型。
在五月,成功研发了Devin的认知模型,并开源了名为Kevin-32B的大模型,该模型能够通过强化学习技术来编写CUDA内核。
该系统依托于QwQ-32B模型,在KernelBench数据集上运用了GRPO技术,经过多轮强化学习训练,其性能显著超越了o3和o4-mini。
斯坦福大学的研究团队在2025年5月28日发布的文章中,探讨了快速核函数的相关内容,详细介绍了该领域的最新研究成果和技术进展。
该链接指向的是安妮·杨在社交媒体平台发布的特定状态,其编号为1928124885567467768。
认知实验室在社交媒体平台发布的最新动态,可在此链接查看:https://x.com/cognition_labs/status/1919835720493236295。