技术进展

字节清华AIR发布CUDA Agent加速GPU内核

Heooo 08月18日12时01分 26 阅读

「字节跳动Seed团队与清华AIR联合推出CUDA Agent,通过强化学习训练大模型编写高性能GPU内核。在KernelBench基准上,通过率提升至98.8%,96.8%生成内核超越torch.compile。开放数据集与训练配方,助力AI基础设施等领域。」

字节跳动Seed团队与清华大学智能产业研究院(AIR)近日联合发布了全新系统——CUDA Agent。该系统专门用于训练大语言模型编写高性能GPU内核,标志着AI在底层代码优化领域取得了重要进展。

在过去的技术探索中,前沿大模型虽然能够生成语法正确的CUDA代码,但生成的代码效率往往较低,在实际运行速度上难以超越传统编译器。以KernelBench基准测试为例,基础模型Seed1.6的任务通过率虽已达到74.0%,但在运行速度上超越torch.compile的比例仅为27.2%,平均几何速度反而落后于编译器本身。这一现象揭示了大模型在理解底层硬件特性和性能优化方面仍存在明显短板。

为了突破这一瓶颈,CUDA Agent将大模型置于一个真实的CUDA开发环境中,并为其配备了性能分析工具、正确性校验机制以及受权限保护的安全沙盒。系统通过近端策略优化(PPO)算法开展长达150步的深度训练,使模型能够在迭代中学习如何针对具体任务生成更高效的并行计算代码。最终,在包含250个任务的基准测试中,模型整体通过率飙升至98.8%,且96.8%的生成内核运行速度超越了传统的torch.compile,展现出接近乃至超过人工优化水平的性能表现。

在实际应用与开源规划方面,尽管该系统的完整模型权重暂未公开,但研究团队已面向公众开放了包含6000个样本的CUDA-Agent-Ops-6K数据集、相关的SKILL.md规范说明以及奖励和热身训练配方。这些资源为学术界和工业界进一步研究提供了宝贵的基础,也使得其他团队能够复现并改进相关工作。

这一成果未来有望广泛应用于AI基础设施、大模型推理服务、自动驾驶以及量化交易等对延迟高度敏感的行业。通过让大模型学会自我优化底层计算内核,AI系统能够在硬件层面获得更大的性能优势,从而推动整个技术栈的加速迭代。

# CUDA Agent # GPU内核 # 强化学习 # 字节跳动 # 清华AIR

来源:Heooo AI工具导航