🤖 machine learning

CUDA-L2: Surpassing cuBLAS Performance for Matrix Multiplication through Reinforcement Learning

本文介绍了 CUDA-L2，这是一个利用大语言模型和强化学习来自动优化半精度通用矩阵乘法（HGEMM）算子的系统，通过系统性地探索人类工程师难以实现的规模化配置空间，实现了相对于 torch.matmul、cuBLAS 和 cuBLASLt 等既有基准的显著性能提升。

原作者： Songqiao Su, Xiaoya Li, Albert Wang, Guoyin Wang, Jiwei Li, Chris Shum

发布于 2026-07-14

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Songqiao Su, Xiaoya Li, Albert Wang, Guoyin Wang, Jiwei Li, Chris Shum

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下你正在试图烤制一个完美的蛋糕。多年来，世界上最优秀的烘焙师（人类专家）一直在为一个被称为“矩阵乘法”的特定类型蛋糕不断改进配方。这种蛋糕是几乎所有现代 AI 大脑中的秘密成分。他们多年来一直致力于让它变得更快，但他们遇到了一个瓶颈：每当蛋糕的大小发生变化时（从一个小纸杯蛋糕到一个巨大的婚礼层级蛋糕），他们都必须从头开始，而且对于一种尺寸有效的技巧在另一种尺寸上往往会失效。这就像试图用一把小勺子去吃一大碗汤；工具根本不匹配。

迎来 CUDA-L2——这是一个由结合了超智能语言模型和一种名为强化学习（RL）的类游戏学习系统的“AI 烘焙师团队”。CUDA-L2 不仅仅是询问人类专家最好的配方是什么，它是在玩一场大规模的“尝试、失败、学习、再尝试”的游戏。

重大发现：AI 烤出的蛋糕更快

本文的主要发现是，CUDA-L2 可以自动设计这些矩阵乘法“配方”（称为内核/kernels），并且烤出的蛋糕比目前现有的最好的人类配方还要快。

研究人员在 1,000 种不同的蛋糕尺寸（维度 M、N 和 K 的组合，范围从 64 到 16,384）上进行了测试。这些尺寸涵盖了像 Qwen、Llama 和 DeepSeek 这样著名的开源 AI 模型所使用的精确维度。

以下是 AI 烤制速度与当前冠军的对比情况，测量是在 A100 GPU 上进行的：

对比标准工具 (torch.matmul)： 在连续烘焙模式（离线模式）下，CUDA-L2 快了 22.0%。在订单随机到达的繁忙厨房（服务器模式）中，它快了 28.7%。
对比金标准 (cuBLAS)： 即使是对阵 NVIDIA 自家高度优化的库，CUDA-L2 也胜出了。它在连续模式下快了 19.2%，在繁忙的服务器模式下快了 26.0%。
对比“自动优化器” (cuBLASLt-AutoTuning)： 这是最重要的对比。cuBLASLt-AutoTuning 是一个尝试多达 100 种不同配方以寻找最佳方案的工具。CUDA-L2 仍然在连续模式下领先其 11.4%，并在服务器模式下领先其 15.9%。

该论文对这些数字非常有信心，因为它们是通过运行代码数千次直接测量出来的，而不是仅仅通过猜测或模拟。

AI 没有做的事情

了解该系统没有做的事情很重要。论文明确指出，当前版本的 CUDA-L2 仅限于 A100 架构。然而，该框架旨在具有广泛的适用性，团队正积极致力于将其扩展到其他 GPU 架构，包括较旧的 Ampere 芯片（如 RTX 3090），以及较新的 Hopper（如 H100）和 Blackwell（如 B200）芯片。论文还反驳了人类专家已经“解决”了矩阵乘法的观点；事实证明，AI 找到了更好的配方，这证明人类的调优远非完美。

AI 是如何学会烤出更好蛋糕的

AI 不仅仅是在瞎猜；它学会了一些巧妙的技巧，甚至人类专家也可能因为太忙而忽略这些技巧。

填充蛋糕 (Padding the Cake)： 想象你有一个 8,192 英寸宽的蛋糕，但你的烤盘只有在宽度能被 160 整除时才能完美契合。8,192 不能被 160 整除。人类可能会说：“我会用 128 英寸的烤盘，因为它合适。”但 AI 说：“我会添加一点额外的面糊（填充），使蛋糕变成 8,320 英寸宽，这样我就可以使用 160 英寸的烤盘。”这一点点额外的额外工作实际上让整个过程变得更快了。
乒乓策略 (The Ping-Pong Strategy)： 通常情况下，烘焙师先装入原料，混合，然后加载下一批。AI 想出了一个“乒乓”方法：当搅拌机正在处理 A 批次时，助手已经在加载 B 批次的原料了。这意味着搅拌机永远不需要等待。
“交错式”交付 (The "Staggered" Delivery)： 有时，AI 意识到同时索要两种原料（A 和 B）会导致厨房交通拥堵。相反，它先索要原料 A，开始混合，然后再索要原料 B。这让厨房保持顺畅运转。
选择合适的烤盘尺寸： AI 学习到，对于小蛋糕，小烤盘效果最好。但对于巨大的蛋糕，它需要更大的烤盘。它为每一种蛋糕维度都找到了完美的尺寸，而人类要为 1,000 种不同尺寸计算这个任务可能需要一辈子。

“服务器”与“离线”的区别

论文还注意到厨房环境中的一些有趣现象。

离线模式 (Offline Mode)： 想象一条工厂流水线，蛋糕一个接一个地连续烘焙，中间不停歇。烤箱保持高温，工人们处于一种节奏之中。在这里，AI 比竞争对手快了 11.4% 到 22.0%。
服务器模式 (Server Mode)： 想象一家繁忙的餐厅，订单随机到达。在订单之间，烤箱可能会降温，工人们必须重新进入状态。在这种“现实世界”的混乱中，AI 的优势反而扩大了，超过了竞争对手 15.9% 到 28.7%。论文认为这是因为 AI 的配方在处理这些“冷启动”和不可预测的停顿方面表现得更好。

核心结论

论文得出结论，即使对于像矩阵乘法这样最关键、经过高度优化的任务，LLM 引导的强化学习也能通过探索人类无法检查的巨大可能性空间，找到比人类更好的解决方案。虽然这个特定版本的 CUDA-L2 目前仅限于 A100 GPU，但该框架旨在未来扩展到其他芯片。

简而言之：AI 不仅仅是微调了配方；它发现了人类尚未想到的全新的烘焙方式，证明了即使在 GPU 优化这样一个成熟的领域，仍然有很大的提升空间。

技术摘要：CUDA-L2

问题陈述

矩阵乘法（matmul）是大型语言模型（LLM）中的基础操作，在训练和推理过程中都占据了显著的计算时间。虽然存在手动优化的 CUDA 核函数，但在不同的矩阵维度（ $M, N, K$ ）和 GPU 架构之间实现最优性能仍然是一个重大挑战。不同的维度三元组通常需要不同的优化策略，且这些策略无法很好地跨硬件迁移。此外，即使对于相同的维度，最优策略也会根据设置细节（如累加器精度：FP16 与 FP32）而有所不同。

现有的自动化方法（例如使用大语言模型生成核函数的方法）主要集中在具有单一、固定配置的基准测试上（例如 KernelBench）。目前尚不清楚这些方法如何转化到生产环境中，因为在生产环境中，它们必须与 NVIDIA 的 cuBLAS 和 cuBLASLt 等高度优化的闭源库竞争。目前缺乏能够系统性探索半精度通用矩阵乘法（HGEMM）庞大配置空间的自动化系统，以超越人类专家优化的基准线。

方法论：CUDA-L2

作者提出了 CUDA-L2，这是一个结合了大语言模型（LLM）和强化学习（RL）的系统，用于自动优化 HGEMM CUDA 核函数。该系统旨在处理代表所有 $M, N, K$ 组合（取自 $\{64, 128, 256, 512, 1024, 2048, 4096, 8192, 12288, 16384\}$ ）的 1,000 种不同配置。

系统架构与训练

CUDA-L2 在之前的 CUDA-L1 框架基础上进行了四项关键技术增强：

持续预训练： 模型在涵盖标准基准测试之外的多种 CUDA 代码语料库上进行预训练，包括来自 PyTorch、ATen、CUTLASS 和 NVIDIA 教程的源码。通过为原始代码片段生成指令描述，并添加检索增强上下文，以纳入新知识和架构特性。
多阶段 RL 训练： 训练过程从广泛的通用核类型过渡到专门领域（特别是 matmul）。
- 通用核 RL： 使用执行速度作为奖励，在广泛的操作（线性代数、卷积、归约等）上训练模型。
- HGEMM RL： 特别针对具有不同 $M, N, K$ 配置的 HGEMM 核函数进行训练。
带有剖析指标的增强上下文： RL 上下文包含了全面的 NVIDIA Nsight Compute (NCU) 剖析指标（例如内存吞吐量、SM 占用率、缓存效率），从而实现超越简单端到端执行时间的更好优化决策。
奖励函数设计： 生成核函数的奖励基于测试迭代中的平均加速比得分。为了确保鲁棒性，奖励会对以下情况进行惩罚：
- 与 FP32 CPU 基准真值的数值偏差。
- 过长的代码长度（以鼓励简洁性）。
- 系统强制执行严格的正确性标准，包括与二进制输入的精确匹配以及相对于可靠基准（cuBLAS, cuBLASLt）的有限偏差。

评估协议

系统在两种场景下进行评估，以反映不同的部署现实：

离线模式（Offline Mode）： 核函数连续执行，没有时间间隔，模拟峰值吞吐量。
服务器模式（Server Mode）： 核函数在随机间隔执行，以模拟实时推理，考虑 GPU 热动力学（冷却期和时钟加速）。
评估将 CUDA-L2 与四个基准进行对比：

torch.matmul（PyTorch 的默认分发）。
cuBLAS（使用 NN 和 TN 之间的最优布局选择）。
cuBLASLt-heuristic（从 NVIDIA 的启发式 API 中选择排名最高的算法）。
cuBLASLt-AutoTuning（穷举测试最多 100 个候选者以寻找最快者）。

核心贡献与发现的技术

CUDA-L2 证明了 LLM 引导的 RL 可以通过系统性地探索人工调优难以实现的配置空间来发现更优的实现。该系统自主发现了并应用了几种高级优化技术：

抽象选择： 模型会自动在用于小矩阵的原始 WMMA 指令集与用于大型复杂分块操作的 CuTe 抽象之间进行选择。
输入填充（Input Padding）： 系统发现了通过对不可整除维度的输入矩阵进行零填充（例如将 $M=8192$ 填充至 $8320 $）的策略，以启用原本无效的高效分块大小（$ B_M$），从而在额外计算开销与性能收益之间取得平衡。
高级内存与计算重叠：
- 双缓冲寄存器分片（Double-Buffered Register Fragments）： 在两组寄存器分片之间交替，以实现数据移动与计算的重叠（乒乓执行）。
- 激进预取（Aggressive Prefetching）： 当寄存器压力允许时，从单步预取扩展到多步预取（提前加载多个迭代）。
- 交错预取调度（Staggered Prefetch Scheduling）： 在矩阵乘加（MMA）操作前后交错进行矩阵 A 和 B 的预取，以最大化指令级并行性。
- 寄存器到共享内存的直接拷贝（Direct Register-to-Shared-Memory Copy）： 当寄存器布局与共享内存布局匹配时，消除 Epilogue 中的中间临时张量，使用宽数据类型（如 uint128_t）来减少拷贝操作。
- 块重排（Block Swizzling）： 通过重新排序线程块执行顺序来提高 L2 缓存命中率，特别是在处理大型问题规模时。

结果

CUDA-L2 在 A100 GPU 上的 1,000 种配置中系统性地超越了所有主要基准。

离线模式性能（相对于基准的平均加速比）：

较 torch.matmul +22.0%。
较 cuBLAS（使用最优布局）+19.2%。
较 cuBLASLt-heuristic +16.8%。
较 cuBLASLt-AutoTuning（最强的基准，其穷举测试了 100 个候选者）+11.4%。

服务器模式性能：
在服务器场景下，加速比有所增加，这可能是由于系统能更好地处理热动力学和缓存冷却效应：

较 torch.matmul +28.7%。
较 cuBLAS +26.0%。
较 cuBLASLt-heuristic +22.4%。
较 cuBLASLt-AutoTuning +15.9%。

胜率在所有基准中处于 79.3% 至 95.7% 之间，表明这些改进是系统性的，而非由离群值驱动。系统还显示，在 GPU 利用率较低的小型矩阵尺寸下，加速比通常更高；而对于大型矩阵，GPU 吞吐量趋于饱和，导致与基准的差距缩小。

意义

论文声称，CUDA-L2 证明了即使是最具性能关键性且经过高度优化的核函数（如 HGEMM），也可以通过 LLM 引导的 RL 自动化进行改进。通过在人工调优难以覆盖的配置空间中扩展优化，该系统实现了超越最先进闭源库的性能增益。

虽然当前版本仅限于 A100 (Ampere) 架构，但该框架旨在具有广泛的适用性。作者指出，正在进行将该系统扩展到其他 GPU 架构（包括 Ada Lovelace、Hopper 和 Blackwell）的工作。这项工作表明，自动化方法可以系统性地探索庞大的配置空间，以找到人工调优或基于启发式的库可能遗漏的优化方案。