
好家伙世博体育app下载,AI 不测生成的内核(kernel),性能比东说念主类巨匠挑升优化过的还要好!
斯坦福最近流露了一组新发现,扫尾真的太亮眼了。
由 AI 优化的内核,在常见深度学习操作上,翻倍卓绝原生 PyTorch,性能至多可以晋升近 400% ——
矩阵乘法(Matmul,FP32):性能达到 PyTorch torch.matmul 的 101.3%。
二维卷积(Conv2D):性能达到 torch.nn.Conv2D 的 179.9%。
Softmax:性能达到 torch.softmax 的 111.8%。
层归一化(LayerNorm):性能达到 torch.nn.LayerNorm 的 484.4%。
Conv2D+ReLU+MaxPool 组合操作:性能达到 PyTorch 参考达成的 290.1%,以及 torch.compile ( ) 参考达成的 189.0%。
(在 NVIDIA L40S GPU 上进行基准测试,性能百分比界说为参考时辰除以生成的 kernel_size 时辰)

更惊东说念主的是,这一切齐是不测达成的。
筹商团队原本的想法是生成合成数据以考研内核生成模子。
扫尾发现,仅在测试阶段生成的合成数据自己,尽然可以生成性能稀疏优秀的内核。

围不雅网友:没预见 AI 也要取代内核工程师了。

还有东说念主发现,除了性能大幅晋升外,筹商团队取舍的要领也稀疏真谛:
他们莫得简单的在操作上慢慢优化(雷同于爬坡算法),而是在每次迭代之间加入了一个谈话推理的关节,通过这种方式饱读动搜索经由愈加各样化。
也即是说,他们是让系统在每次改造时通过雷同"念念考"的方式产生更多想法,从而找到更好的科罚有狡计。

具体若何达成,系数来看。
改代码前先生成当然谈话优化念念想
按照斯坦福团队博客的描写,这种内核生成的念念路稀疏简单——给定 torch 代码,然后告诉齐能写编写自界说内核来替换 torch 算子。
这些内核是用纯 CUDA-C 编写,无需使用 CUTLASS 和 Triton 等库和 DSL(Domain-Specific Language,领域专用谈话)。
不同于传统要领的是,模子并不是一上来就胜仗改代码,而是先用当然谈话生成优化念念想,然后再将这些念念想更动为新的代码变体。
团队这么作念的原理是,"按法例修改"式的优化念念路枯竭各样性,导致堕入局部极小值,重复探问团结类调度或连接断地优化莫得出路的轨迹。
为了进一步增强念念路各样性,斯坦福团队还使用了多分支的探索模式。
具体来说,他们的要领并非每一步齐只优化一个候选有狡计,而是将每个想法分布开来,使其生息出多个达成,并使用性能最高的内核行为下一轮的种子。

团队使用 OpenAI o3 和 Gemini 2.5 Pro 挑战 KernelBench 1 级中的 10 个问题,运行多轮后,最好内核开动出现。
其中大大量最好扫尾出目下后续轮次(悉数 5 轮),而且主若是第 4 轮或第 5 轮。
KernelBench 是斯坦福团队我方忽视的一套 AI 生成内核测试基准,基准中的任务分为 3 个级别,其中 1 级是指单一原始操作(Single primitive operation),包括 AI 的基础构建块(举例卷积、矩阵 - 向量与矩阵 - 矩阵乘法、亏本函数、激活函数以及层归一化)。

这一发现再加上之前 DeepMind 的 AplhaEvolve,以及 o3 发现 Linux 的 0day 裂缝等一系列事件,让网友们觉得 Gemini 2.5Pro 和 o3 的才气水平依然达到了新的层级。

回到斯坦福的形式,在生成经由当中,可以看到模子的生成念念路开动浮现出与东说念主类的告诫相似之处——
内存探问优化: 提高不同内存脉络结构(全局内存、分享内存、寄存器)之间数据迁徙的效用,并确保以最大化带宽和最小化打破的方式探问数据;
异步操作和蔓延荫藏: 通过将慢速操作(如全局内存探问)与计划或其他内存传输叠加,"荫藏"慢速操作的蔓延;
数据类型和精度优化: 尽可能使用低精度数据类型(如 FP16 或 BF16)以减少内存带宽条件、提高缓存效用;
计划和提醒优化:提高算术计划自己的效用,减少提醒数目,或愚弄挑升的硬件提醒;
并行性和占用率增强:最大化流多处理器(SM)上的举止线程数目,以更好地荫藏蔓延并提高举座蒙眬量;
限定流和轮回优化:减少与轮回、分支和索引计划有关的支拨。
而且斯坦福团队还展示了一组具体的优化轨迹,从中可以看出,并不是每一步优化齐一定能让速率更快,但经过多个关节的组合,内核的速率能够获得大幅晋升,并最终卓绝 PyTorch。

在具体达成上,有东说念主斟酌 AI 生成 CUDA 内核时的优化建议,是否可以被更动为对应代码达成、照旧说仅仅触发了随即探索?
作家恢复说,尽管还莫得进行更严谨的系统考据,然而手动查验的案例中,生成的 CUDA 视野与忽视的优化建议是轻便匹配的。
即 AI 并不是在十足随即作念优化,而是如真实尝试达成它我方忽视的计谋。

华东说念主主创团队不测发现
这项筹商共有三位作家:Anne Ouyang、Azalia Mirhoseini 和 Percy Liang。
Ouyang 目下是斯坦福大学推广智能实验室的博士生,她本硕毕业于麻省理工,曾在英伟达 cuDNN 团队职责。

Percy Liang 是斯坦福大学计划机科学副教育兼统计学助理教育,目下担任斯坦福基础模子筹商中心主任。
曾和李飞飞系数发布、激动了多项筹商职责。

Azalia Mirhoseini 是斯坦福大学计划机科学助理教育、斯坦福推广实验室创举东说念主。她曾在 DeepMind、Google Brain 以及 Anthropic 职责过。
她此前参与的筹商包括 MoE、芯片设计划法 AlphaChip 等。

本次筹商,原本是但愿生成数据来考研内核生成模子。
然而在经由中却出现了出东说念主料想的扫尾,仅在测试阶段生成的合成数据自己,尽然可以生成性能稀疏优秀的内核。
因为这些内核愚弄了此前被觉得很难达成的高档优化和硬件特色,是以团队决定以博客花式分享这次后果。
不外具体是若何生成数据的,筹商团队暂时不合外发布,仅仅提到了这种瞎想理念也很简单。
最重要的照旧,它依然展示出了广大后劲。
此外,筹商团队也觉得这次发现也与最近的一些趋势相呼应——大鸿沟再考研已不是必需。
无意,灵敏的搜索和分支计谋,可以解锁科学更动并科罚复杂问题,通过 verifier 进行等闲搜索还能有更多成绩。
执意劲推理才气与同期探索多个假定连合起来,能带来更好扫尾。就像 AlphaEvolve、AlphaEvolution、 Gemini 2.5 Pro 深度念念考相同。
临了,团队默示这项筹商还有许多可优化的空间。比如他们手头上就还在优化两个维度:
FP16 Matmul:52% performance of torch.matmul
FP16 Flash Attention::9% performance of torch.nn.functional.scaled_dot_product_attention
与 FP16 或 BF16 比较,FP32 在新推出硬件上的优化进度频频比较低,这亦然为何使用 FP32 内核比 PyTorch 更容易达成性能晋升。
他们默示,诚然目下还有不少截止,然而关于将来出路照旧很乐不雅的。
毕竟最开动,他们连能浩繁运行的内核齐生成不了,然而通过遏抑优化搜索要领,依然能让 flash attention 的性能晋升到了一个可以的水平。
值得一提的是,搜索使用的资源也很少,粗略只用了 300 万 token 输入和 400 万 token 输出。
One More Thing
推行上,不仅仅一个团队在尝试设立内核大模子。
就在 5 月,设立了 Devin 的 Cognition 开源了首个通过强化学习即可编写 CUDA 内核的大模子 Kevin-32B。
它基于 QwQ-32B 在 KernelBench 数据集上使用 GRPO,达成了多轮强化学习,性能优于 o3、o4-mini。

— 完 —世博体育app下载
