1. 论文基本信息

1.1. 标题

CUDA-L2: Surpassing cuBLAS Performance for Matrix Multiplication through Reinforcement Learning (CUDA-L2：通过强化学习超越 cuBLAS 矩阵乘法性能)

1.2. 作者

Songqiao Su, Xiaofei Sun, Xiaoya Li, Albert Wang, Jiwei Li and Chris Shum

1.3. 发表机构

DeepReinforce Team。论文中提到了 github.com/deepreinforce-ai/CUDA-L2，表明这是一个研究团队或机构。

1.4. 发表年份

2025 年 12 月 2 日

1.5. 摘要

本文提出了 CUDA-L2，一个结合大语言模型 (Large Language Models, LLMs) 和强化学习 (Reinforcement Learning, RL) 的系统，旨在自动优化半精度通用矩阵乘法 (Half-precision General Matrix Multiply, HGEMM) CUDA 内核。CUDA-L2 使用 CUDA 执行速度作为 RL 奖励，自动优化了 1,000 种配置下的 HGEMM 内核。该系统系统性地超越了目前主要的矩阵乘法基线，包括广泛使用的 torch.matmul 和 Nvidia 专有库 cuBLAS、cuBLASLt。在离线模式（内核连续执行无时间间隔）下，CUDA-L2 平均比 torch.matmul 提升 22.0%；比使用最优布局配置（正常-正常 NN 和转置-正常 TN）的 cuBLAS 提升 19.2%；比根据启发式建议选择算法的 cuBLASLt-heuristic 提升 16.8%；比从 cuBLASLt 建议的最多 100 个候选中选择最快算法的 cuBLASLt-AutoTuning 提升 11.4%。在服务器模式（内核随机间隔执行，模拟实时推理）下，速度提升进一步增加，分别达到 torch.matmul 的 28.7%、cuBLAS 的 26.0%、cuBLASLt-heuristic 的 22.4% 和 cuBLASLt-AutoTuning 的 15.9%。CUDA-L2 表明，即使是性能最关键、优化最彻底的 HGEMM 内核，也可以通过 LLM 引导的 RL 自动化来改进，通过系统性地探索人类难以企及的配置空间。

1.6. 原文链接

原文链接: https://arxiv.org/abs/2512.02551 PDF 链接: https://arxiv.org/pdf/2512.02551v1.pdf 发布状态：预印本 (arXiv preprint)

2. 整体概括

2.1. 研究背景与动机

2.1.1. 矩阵乘法的重要性与优化挑战

通用矩阵乘法 (General Matrix Multiply, GEMM) 是深度学习和高性能计算的核心操作。尤其是在大语言模型 (LLMs) 中，半精度通用矩阵乘法 (HGEMM) 广泛用于 attention 和 FFN 层。GEMM 的性能对整个系统的吞吐量至关重要。尽管 NVIDIA 等公司已经投入大量资源对其 cuBLAS 和 cuBLASLt 库进行了高度优化，但这些优化通常由人类专家手工完成，耗时费力，且难以在庞大的配置空间（例如不同矩阵尺寸 M, N, K）和不断演进的 GPU 架构（如 Ampere, Ada Lovelace, Hopper, Blackwell）中实现通用性。

2.1.2. 现有自动化优化的局限

当前 GPU 内核自动优化方法（如 Triton, Halide 等）虽然能够生成高效代码，但在面对 NVIDIA 专家级手动优化的 cuBLAS 库时，往往难以匹敌，尤其是在与 cuBLAS 等高度优化的库进行比较时，性能差距显著。这表明，仅仅依靠传统的编译器优化或启发式方法不足以在所有情况下达到最优。

2.1.3. 论文的切入点与创新思路

本文旨在解决 HGEMM 内核优化的自动化和超越专家级手动优化的挑战。论文的切入点是结合大语言模型 (LLM) 的代码生成能力和强化学习 (RL) 的优化搜索能力，以系统性地探索 HGEMM 的巨大配置空间。通过将 CUDA 执行速度作为 RL 的奖励信号，CUDA-L2 能够学习并生成更优的 CUDA 内核，从而超越现有最先进的、由专家优化的库。这种 LLM-guided RL 的方法，使得系统能够发现人类专家难以发现的细致优化策略。

2.2. 核心贡献/主要发现

2.2.1. 提出 CUDA-L2 系统

论文提出了 CUDA-L2，一个新颖的 LLM 和 RL 结合的系统，用于自动优化 HGEMM CUDA 内核。它通过在多样的 CUDA 代码上进行持续预训练，并结合强化学习来优化 HGEMM 内核，从而克服了 CUDA-L1 在通用性和新架构适应性方面的局限。

2.2.2. 显著超越现有基线性能

CUDA-L2 在 1,000 种 HGEMM 配置（覆盖 M, N, K 值的 $10^3$ 组合）上，系统性地超越了所有主要基线，包括 torch.matmul、cuBLAS 和 cuBLASLt。

• 离线模式 (Offline Mode)：
  • 比 torch.matmul 平均提升 22.0%。
  • 比 cuBLAS（最优布局）平均提升 19.2%。
  • 比 cuBLASLt-heuristic 平均提升 16.8%。
  • 比 cuBLASLt-AutoTuning 平均提升 11.4%。
• 服务器模式 (Server Mode)：速度提升进一步增加，分别达到 28.7%、26.0%、22.4% 和 15.9%。

2.2.3. 发现并应用高级优化技术

CUDA-L2 自动发现了多项高级 CUDA 优化技术，包括：

• 抽象选择 (Abstraction Selection)：根据矩阵大小自动选择 CUDA C/C++ 或 CuTe 抽象。
• 输入矩阵填充 (Padding the Input Matrix with Zeros)：通过填充输入矩阵来选择更高效的线程块维度 (BM)，即使这会导致少量数据冗余。
• 熟练运用 CUDA 优化技术：如共享内存银行冲突避免 (bank conflict avoidance)、双缓冲 (double-buffering)、异步内存复制 (asynchronous memory copy)、寄存器累积 (register accumulation)、块交换 (block swizzling) 和 epilogue 优化 (epilogue optimization) 等。
• 发现新颖优化变体：例如双缓冲寄存器片段 (double-buffered register fragments)、激进的寄存器级预取 (aggressive register-level prefetching)、直接寄存器到共享内存复制的 epilogue 优化 (direct register-to-shared-memory copy for epilogue) 和交错的 A-B 预取调度 (staggered/split A-B prefetch scheduling)。

2.2.4. 揭示超参数选择模式

论文分析了 CUDA-L2 学习到的超参数选择模式，例如：

• BM (Block M) 和 BN (Block N) 值与 M, N 维度强相关，且 BM 和 BN 之间高度相关。
• 多阶段流水线中的阶段数 (n_stage) 与 $K$ 维度正相关。
• 块交换 (Block Swizzling) 的使用与问题规模强相关，对于大问题几乎总是使用，且其步长 (stride) 值与问题规模呈现正相关。

3. 预备知识与相关工作

3.1. 基础概念

3.1.1. 半精度通用矩阵乘法 (HGEMM)

概念定义： HGEMM 指的是使用 16 位浮点数（半精度，FP16）执行的通用矩阵乘法操作。它是深度学习模型中计算密集型任务（如神经网络的 attention 机制和全连接层 FFN）的核心组成部分，因为它能在保证一定精度的前提下，显著提高计算速度和降低内存占用。 数学公式： HGEMM 计算的是矩阵乘积 $C = \alpha AB + \beta C$，其中 $A \in \mathbb{R}^{M \times K}$， $B \in \mathbb{R}^{K \times N}$， $C \in \mathbb{R}^{M \times N}$。 $$C_{ij} = \alpha \sum_{p=1}^{K} A_{ip} B_{pj} + \beta C_{ij}$$ 符号解释：

• A, B, C: 输入矩阵和输出矩阵。
• M, N, K: 矩阵的维度。$A$ 是 $M \times K$ 矩阵，$B$ 是 $K \times N$ 矩阵，$C$ 是 $M \times N$ 矩阵。
• $\alpha, \beta$: 标量系数，通常情况下 HGEMM 关注的是 $\alpha=1, \beta=0$，即 $C = AB$。
• $A_{ip}$: 矩阵 $A$ 的第 $i$ 行第 $p$ 列的元素。
• $B_{pj}$: 矩阵 $B$ 的第 $p$ 行第 $j$ 列的元素。
• $C_{ij}$: 矩阵 $C$ 的第 $i$ 行第 $j$ 列的元素。

3.1.2. CUDA (Compute Unified Device Architecture)

概念定义： CUDA 是 NVIDIA 推出的一种并行计算平台和编程模型，允许开发者使用 $C/C++$ 等高级语言直接访问 GPU 的计算资源。它提供了一套编程接口 (API) 和运行时环境，使得编写在 GPU 上高效运行的并行程序成为可能。CUDA 内核 (kernel) 是在 GPU 上并行执行的函数。

3.1.3. 强化学习 (Reinforcement Learning, RL)

概念定义： RL 是一种机器学习范式，智能体 (agent) 通过与环境 (environment) 交互来学习如何做出决策，以最大化累积奖励 (reward)。在本文中，LLM 作为代码生成器，其生成的 CUDA 内核的执行速度被用作 RL 的奖励信号，引导 LLM 生成更优的代码。

3.1.4. 大语言模型 (Large Language Models, LLMs)

概念定义： LLMs 是基于深度学习的语言模型，拥有大量参数（通常数十亿到数千亿），并在海量文本数据上进行预训练。它们能够理解、生成人类语言，并通过指令微调 (instruction tuning) 或上下文学习 (in-context learning) 来执行各种任务，包括代码生成和代码优化。

3.1.5. cuBLAS 和 cuBLASLt

概念定义： cuBLAS 是 NVIDIA 针对 GPU 提供的 BLAS (Basic Linear Algebra Subprograms) 库，它包含了高度优化的矩阵乘法等线性代数操作。cuBLASLt 是 cuBLAS 的轻量级 (Lightweight) 版本，提供了更底层的 API，允许用户对 GEMM 算法选择、参数配置有更大的控制权，通常用于需要极致性能调优的场景。cuBLASLt 包含启发式算法选择 (heuristic) 和自动调优 (AutoTuning) 机制。

3.2. 前人工作

3.2.1. 自动化 GPU 内核优化

早期的 GPU 内核优化主要依靠人工调优，但随着 GPU 架构复杂性的增加和应用场景的多样化，自动化优化工具变得越来越重要。

• TVM / Apache TVM: 一个开源的深度学习编译器栈，旨在将深度学习模型部署到各种硬件后端。它通过自动优化张量表达式来生成高效的内核。
• Halide: 一种领域特定语言 (DSL)，用于图像处理和计算，它通过将算法和调度分离来帮助生成高效代码。
• Triton: OpenAI 开发的用于编写高效 GPU 内核的 DSL，其目标是简化 GPU 编程，并能生成接近专家手写代码的性能。
• LLM 辅助的代码生成和优化： 近年来，随着 LLM 的发展，研究开始探索 LLM 在代码生成、代码优化甚至 GPU 内核生成方面的潜力。例如 KernelBench [10]、SkyPilot [6] 和 CUDA-L1 [7] 都尝试利用 LLM 来生成和优化 CUDA 内核。

3.2.2. CUDA-L1 (前作)

CUDA-L2 是 CUDA-L1 的扩展。CUDA-L1 是一个 LLM-RL 系统，用于优化 KernelBench [10] 中描述的 CUDA 内核。它专注于优化特定类型的内核，但可能在通用性、适应新 GPU 架构和处理更复杂 CUDA 构造方面存在局限。

3.3. 技术演进

从手动专家优化到 DSL 编译，再到 LLM 结合 RL 的自动化优化，GPU 内核优化技术不断演进。

• 手动优化阶段： 早期和目前 cuBLAS 等库的大部分优化都依赖于经验丰富的 GPU 专家手动编写和调优 CUDA 代码。这能达到非常高的性能，但成本高昂，且难以扩展。
• 编译器/ DSL 优化阶段： TVM、Halide、Triton 等工具通过提供高级抽象和自动调度机制，使得开发者能更容易地生成高效的 GPU 代码，减少手动调优的需求。然而，这些工具生成的代码在某些情况下仍难以匹敌专家手写的 cuBLAS 等库。
• LLM 结合 RL 阶段 (本文工作)： 本文将 LLM 的语义理解和代码生成能力与 RL 的试错和优化能力结合，形成一个强大的自动化优化框架。LLM 负责生成初始 CUDA 内核代码，RL 则通过实际执行反馈（即 CUDA 执行速度）来指导 LLM 迭代优化。这种方法旨在突破 DSL 优化的局限，实现对现有最先进库的超越，并通过持续预训练和适应新架构来提高通用性。

3.4. 差异化分析

CUDA-L2 与现有工作的主要区别和创新点在于：

• 超越 cuBLAS 性能： 现有 LLM 驱动的内核优化工作（如 KernelBench、CUDA-L1）通常难以在性能上超越 NVIDIA 高度优化的 cuBLAS 库。CUDA-L2 首次展示了通过 LLM-guided RL 自动化，能够系统性地超越包括 cuBLASLt-AutoTuning 在内的 NVIDIA 最先进库。
• 更广泛的通用性和适应性： CUDA-L2 在 CUDA-L1 的基础上，通过在更广泛、更多样化的 CUDA 代码（包括 CUTLASS、CuTe 等）上进行持续预训练，增强了 LLM 的通用 CUDA 优化能力。它还结合了 contextual memory 机制，使其能够根据新的知识或架构特性进行调整，从而更好地适应新的 GPU 架构（如 Ampere、Ada Lovelace、Hopper、Blackwell）。
• 发现深层优化技术： CUDA-L2 不仅仅是简单地搜索参数，而是能够生成包含高级 CUDA 优化技术（如填充、异步内存复制、双缓冲、块交换等）的内核，甚至能够发现这些技术的变体和组合，这在以往的自动化方法中是罕见的。这种能力得益于 LLM 的代码生成和 RL 的深度探索。
• 系统性探索复杂配置空间： 针对 1,000 种不同的 HGEMM 配置进行优化，这表明 CUDA-L2 能够系统性地探索人类专家难以手动穷举的巨大配置空间。

4. 方法论

4.1. 方法原理

CUDA-L2 的核心思想是结合大语言模型 (LLM) 的代码生成能力和强化学习 (RL) 的优化搜索能力，以自动化且超越人类专家水平的方式优化 HGEMM CUDA 内核。 其背后的理论基础在于：

1. LLM 的代码生成潜力： 经过大量代码训练的 LLM 能够理解 CUDA 编程范式，并根据描述生成功能正确的代码。
2. RL 的优化探索能力： 通过将 CUDA 内核的实际执行速度作为奖励信号，RL 能够引导 LLM 探索代码空间，学习哪些代码模式和参数组合能够带来更高的性能。
3. 迭代优化： LLM 生成的内核经过编译和执行，其性能反馈给 RL 模块，RL 模块再调整 LLM 的生成策略，形成一个闭环的迭代优化过程。
4. 持续学习和适应性： 通过在更多样化的 CUDA 代码上进行持续预训练，并结合 contextual memory，CUDA-L2 能够不断吸收新知识，并适应新的 GPU 架构和更复杂的 CUDA 优化技术。

4.2. 核心方法详解 (逐层深入)

CUDA-L2 建立在 CUDA-L1 的基础上，并通过以下关键技术增强其能力：

4.2.1. 持续预训练 (Continued Pretraining)

为了实现更通用化的 CUDA 优化，CUDA-L2 扩展了预训练数据集。

• 多样化 CUDA 代码： 不仅限于 KernelBench，还包括来自 GitHub 上广泛的 CUDA 代码库。
• 整合 CUTLASS、CuTe 等库： 这些库代表了 NVIDIA 官方的最新 GPU 编程抽象和优化策略，将它们纳入预训练可以使 LLM 更好地理解和生成高性能的 CUDA 内核。
• 更大的 LLM 基础模型： 使用了如 DeepSeek 671B 等更大的基础模型，这使得 LLM 能够获得更强的通用 CUDA 优化能力和代码生成质量。

4.2.2. 通用内核强化学习 (General Kernel RL)

CUDA-L2 使用 RL 来优化 LLM 生成的 CUDA 内核。这个过程涉及：

• LLM 作为策略网络： LLM 接收任务描述 (例如 HGEMM 的 M, N, K 维度) 和当前环境状态 (例如之前的优化尝试)，并生成 CUDA 内核代码。
• 环境执行和反馈： 生成的 CUDA 代码会被编译并在 GPU 上执行。执行时间被测量并作为 RL 的奖励信号。
• RL 算法： CUDA-L2 采用了一种 RL 策略，该策略能够学习生成包含 CUTLASS 模板、CuTe 库和 NVIDIA 内联 PTX 汇编指令的高性能 CUDA 内核。

4.2.3. HGEMM 强化学习 (HGEMM RL)

针对 HGEMM 任务，CUDA-L2 采用了一个特定的 RL 策略，其奖励函数设计如下：

$$r(\mathrm{custom}) = \frac{1}{N} \sum_{i=1}^{N} \left[ \frac{t_{\mathrm{ref}}^i}{t_{\mathrm{custom}}^i} - \alpha \cdot \mathrm{diff}^i \right] - \beta \mathrm{L}(\mathrm{custom})$$

符号解释：

• $r(\mathrm{custom})$: LLM 生成的自定义 CUDA 内核的奖励值。

• $N$: 评估该内核的测试用例数量。

• $t_{\mathrm{ref}}^i$: 第 $i$ 个测试用例中参考内核的执行时间。

• $t_{\mathrm{custom}}^i$: 第 $i$ 个测试用例中自定义内核的执行时间。

• $\frac{t_{\mathrm{ref}}^i}{t_{\mathrm{custom}}^i}$: 代表自定义内核相对于参考内核的速度提升。这个比值越大，说明自定义内核越快，奖励越高。

• $\mathrm{diff}^i$: 第 $i$ 个测试用例中最大元素级绝对误差，定义为 $\max_{j} |\mathrm{out}_{\mathrm{FP32}}^i[j] - \mathrm{out}_{\mathrm{custom}}^i[j]|$。这用于衡量生成内核的正确性。

• $\alpha$: 正的惩罚系数，用于惩罚不正确的内核。如果误差 $\mathrm{diff}^i$ 很大，则奖励会显著降低。

• $\mathrm{L}(\mathrm{custom})$: 生成代码的长度。

• $\beta$: 正的惩罚系数，用于惩罚过长的代码，鼓励生成简洁高效的代码。

  这个奖励函数的目标是：

1. 最大化速度： 通过 $\frac{t_{\mathrm{ref}}^i}{t_{\mathrm{custom}}^i}$ 项，鼓励生成执行速度更快的内核。

2. 保证正确性： 通过 $-\alpha \cdot \mathrm{diff}^i$ 项，惩罚生成不正确结果的内核。

3. 鼓励简洁性： 通过 $-\beta \mathrm{L}(\mathrm{custom})$ 项，惩罚生成冗长代码的内核。

   此外，CUDA-L2 还融入了：

• Contextual Memory： 这是一个关键的增强。它允许 LLM 在生成代码时，不仅考虑当前任务，还能从一个动态更新的知识库中获取信息，包含最新的 NVIDIA 优化实践、新的 GPU 架构特性等。这使得 LLM 能够适应未在基础模型中涵盖的新知识或架构特性。
• 支持 CUDA C/C++、CuTe、内联 PTX 汇编、CUDA 内联函数和 CUTLASS 模板： 这意味着 LLM 生成的代码可以利用 NVIDIA GPU 提供的各种底层优化手段，而不仅仅是高级语言抽象。

4.2.4. 内核正确性检查 (Kernel Correctness)

一个自定义 HGEMM 内核被认为是成功的，必须同时满足可执行性和正确性。

• 可执行性 (Executability)： 使用 compute-sanitizer --tool memcheck 工具检查内存访问冲突，确保生成的内核不会导致 GPU 崩溃或运行时错误。
• 正确性 (Correctness)： 由于浮点运算的非结合性，精确匹配 FP32 CPU 结果通常不可行。CUDA-L2 采用两种实际标准：
  1. 二进制输入精确匹配 (Exact Match with binary Inputs)：
     • 生成元素为 $\{0, 1\}$ 的二进制输入矩阵 A, B。
     • 使用 FP32 在 CPU 上计算参考结果 $C^{\mathrm{ref}}$，这将是精确的整数。
     • 计算 FP16 GPU 内核的输出 $C^{\mathrm{custom}}$。
     • 对于 $C^{\mathrm{ref}}_{ij} < 2048$ 的元素，要求 $C^{\mathrm{custom}}_{ij} = C^{\mathrm{ref}}_{ij}$。因为在半精度浮点数中，[0, 2048) 范围内的整数可以精确表示。
     • 如果 $C^{\mathrm{ref}}_{ij} \geq 2048$，则允许一定的误差范围。
     • 如果所有元素 $a_{ik} \cdot b_{kj} \in \{0, 1\}$，且部分和单调非递减，那么当 $C^{\mathrm{ref}}_{ij} < 2048$ 时，如果 $C^{\mathrm{custom}}_{ij} \neq C^{\mathrm{ref}}_{ij}$，则内核被认为不正确。
  2. 基线容差 (Baseline Tolerance)：
     • 使用 cuBLASLt-AutoTuning 作为基线，计算其相对于 FP32 CPU 参考结果的最大元素级绝对误差 $\max_{j} |\mathrm{out}_{\mathrm{FP32}}^i[j] - \mathrm{out}_{\mathrm{cuBLASLt-AutoTuning}}^i[j]|$。
     • 自定义内核的误差不得超过这个基线值。

4.2.5. 优化技术发现 (Optimization Techniques Discovered)

CUDA-L2 通过 LLM-RL 循环，自动发现了并应用了多种高级优化技术：

• 抽象选择 (Abstraction Selection)： 根据 M, N, K 自动选择 CUDA C/C++ (对于小矩阵，注重少流水线、低同步) 或 CuTe (对于大矩阵，使用复杂 tiled MMA 操作)。

• 零填充输入矩阵 (Padding the Input Matrix with Zeros)：

  • 原理： HGEMM 通常采用分块 (tiled) 方式，每个线程块计算一个 $\mathrm{BM} \times \mathrm{BN}$ 的输出块。为了避免边界访问问题，矩阵维度 $M$ 必须能被 $\mathrm{BM}$ 整除，$N$ 必须能被 $\mathrm{BN}$ 整除。
  • 发现： CUDA-L2 发现，即使 $M$ 不能被某个 $\mathrm{BM}$ 整除，通过将 $M$ 填充到能被 $\mathrm{BM}$ 整除的更大值，然后选择一个更优的 $\mathrm{BM}$，可以带来性能提升。
  • 示例： 对于 $M=8192, N=512, K=2048$，CUDA-L2 选择了 $\mathrm{BM}=160$，将 $M$ 填充到 8320 (1.6% 开销)，这比传统选择 $\mathrm{BM}=128$ 性能更好。

• 熟练运用 CUDA 优化技术：

  • 共享内存银行冲突避免 (Shared memory with bank conflict avoidance)： 通过 swizzle pattern 重新组织数据布局，防止并发访问同一银行 (bank) 导致的冲突。
  • 双缓冲 (Double-buffering)： 在 GPU 内核中，尤其是在内存密集型操作中，使用两个缓冲区交替进行数据加载和计算，以隐藏内存访问延迟。这通常涉及 n_stage 来控制流水线阶段数量。
  • 异步内存复制 (Asynchronous memory copy)： 允许非阻塞的内存传输（例如从全局内存到共享内存），使得计算和数据传输可以并行进行。
  • 寄存器累积 (Register accumulation)： 将部分计算结果存储在寄存器文件 (register files) 中，以最小化内存流量。
  • 块交换 (Block swizzling)： 优化线程块 (thread block) 访问全局内存的模式。
  • Epilogue 优化 (Epilogue optimization)： 有效地将寄存器中的计算结果写回全局内存，通常涉及 user-defined shape 和 1-bit predicate 来减少内存流量。
  • 循环展开 (Loop unrolling)： 编译器优化技术，减少循环开销，增加指令级并行性。

• 新颖优化变体： CUDA-L2 还能发现标准实现之外的变体。

  • 双缓冲寄存器片段与乒乓执行 (Double-Buffered Register Fragments with Ping-Pong Execution)：

    • 标准： 单缓冲区 (Single-buffer)，每次只能加载或使用一个数据片段。

    • 优化： 双缓冲区 (Double-buffer)，在当前数据片段计算时，可以预取下一个片段，通过乒乓 (ping-pong) 机制，在两个缓冲区之间切换，实现内存访问和计算的重叠。 下图（原文 Listing 2）展示了两种实现方式：

      该图像是代码片段，展示了用于优化半精度通用矩阵乘法（HGEMM）CUDA内核的部分实现细节。代码涉及到矩阵的分区和复制，以及使用条件语句选择不同的缓冲区处理数据。这些细节是通过自动化和强化学习策略改进计算性能的关键部分，反映了CUDA-L2系统的操作机制和效率提升策略。

  • 激进的寄存器级预取 (Aggressive Register-Level Prefetching)：

    • 标准： 单步预取 (Single-step Prefetch)，只预取下一个迭代所需的数据。

    • 优化： 多步预取 (Multi-step Prefetch)，在迭代次数高且寄存器空间充足时，预取多个未来迭代所需的数据，进一步隐藏内存延迟。 下图（原文 Listing 3）展示了两种实现方式：

      该图像是示意图，展示了CUDA-L2算法中半精度通用矩阵乘法（HGEMM）CUDA内核的优化代码。左侧代码段为K+1预取策略，右侧代码段展示了K+0、K+1和K+2的管道预取过程。两段代码通过使用#pragma unroll和合适的内存预取策略，以提高HGEMM运算的执行效率。这些优化策略是实现性能提升的关键因素。

  • 带有直接寄存器到共享内存复制的 Epilogue 优化 (Epilogue Optimization with Direct Register-to-Shared-Memory Copy)：

    • 标准： 两步 (Two-step) 复制，通常 CUTLASS/CuTe 会创建一个中间张量来重组数据，然后从中间张量复制到共享内存。

    • 优化： 直接宽复制 (Direct wide copy)，直接从寄存器向共享内存进行宽位宽（如 128 位）复制，避免了中间张量，减少了数据移动和延迟。 下图（原文 Listing 4）展示了两种实现方式：

      该图像是代码片段，展示了两种在 CUDA 上执行张量拷贝的不同方法。左侧代码通过中间张量进行拷贝，右侧则直接执行 R2S（Row to Sparse）拷贝操作。这些代码涉及张量的初始化以及使用类库进行数据复制的操作，以优化矩阵乘法的性能。

  • 交错/分步 A-B 预取调度 (Staggered/Split A-B Prefetch Scheduling)：

    • 标准： 连续预取 (Consecutive Prefetch)，先预取矩阵 $A$ 的数据，紧接着预取矩阵 $B$ 的数据。这可能导致内存和计算流水线未能完全重叠。
    • 优化： 交错预取 (Staggered Prefetch)，将 $B$ 的预取操作插入到 $A$ 的 MMA (Matrix Multiply Accumulate) 执行之后，但下一个 $A$ 的预取之前，从而更好地填充 GPU 流水线中的空闲时间，实现计算和内存操作的最大并行化。 下图（原文 Listing 5）展示了两种实现方式：

    Listing 5: A-B prefetch scheduling: consecutive (left) vs. staggered (right).
    

    # Standard: Consecutive Prefetch
    # cute::copy(s2r_tiled_copy_a, # Prefetch A
    #            tAsA(_, _, ik_next, ismem_read),
    #            tCrA_view(_, _, ik_next));
    # cute::copy(s2r_tiled_copy_b, # Prefetch B
    #            tBsB(_, _, ik_next, ismem_read),
    #            tCrB_view(_, _, ik_next));
    # cute::gemm(tiled_mma, tCrD, # MMA compute
    #            tCrA(_, =, ik), tCrB(_, _, ik), tCrD);
    
    # Optimized: Staggered Prefetch
    # cute::copy(s2r_tiled_copy_a, # Prefetch A
    #            tAsA(_, _, ik_next, ismem_read),
    #            tCrA_view(_, _, ik_next));
    # cute::gemm(tiled_mma, tCrD, # MMA compute
    #            tCrA(_, =, ik), tCrB(_, _, ik), tCrD);
    # cute::copy(s2r_tiled_copy_b, # Prefetch B (staggered)
    #            tBsB(_, _, ik_next, ismem_read),
    #            tCrB_view(_, _, ik_next));
    

5. 实验设置

5.1. 数据集

实验使用了 1,000 种 HGEMM 配置。

• 维度组合： 这些配置代表了 M, N, K 值从 $\{64, 128, 256, 512, 1024, 2048, 4096, 8192, 12288, 16384\}$ 中所有 $10^3$ 种组合。
• 覆盖范围： 这些配置已经涵盖了在广泛开源模型（如 Qwen、Llama 和 DeepSeek）的 attention 和 FFN 层中使用的维度。
• 数据类型： 所有矩阵使用半精度浮点数 (FP16)。
• 目标平台： 实验在 NVIDIA A100 GPU 上进行。

5.2. 评估指标

论文主要使用加速比 (Speedup) 作为评估指标，并结合正确性检查。

5.2.1. 加速比 (Speedup)

概念定义： 加速比衡量了自定义内核相对于参考内核的性能提升。它量化了自定义内核在完成相同计算任务时，执行时间减少的百分比。 数学公式： $$s(\mathrm{custom}) = \frac{t_{\mathrm{ref}}}{t_{\mathrm{custom}}} - 1$$ 符号解释：

• $s(\mathrm{custom})$: 自定义内核相对于参考内核的加速分数。
• $t_{\mathrm{ref}}$: 参考内核的执行时间。
• $t_{\mathrm{custom}}$: 自定义内核的执行时间。
• 每个评估分数是多次运行 (runs) 的平均速度分数。

5.2.2. 内核正确性 (Kernel Correctness)

概念定义： 确保 LLM 生成的 CUDA 内核不仅运行快，而且能够产生正确的结果。这通过检查可执行性和结果的数值正确性来保证。 可执行性： 使用 compute-sanitizer --tool memcheck 检查内存访问违规。 数值正确性：

• 精确匹配二进制输入： 对于由 $\{0, 1\}$ 元素构成的矩阵乘法，当输出元素 $C^{\mathrm{ref}}_{ij} < 2048$ 时，要求 $C^{\mathrm{custom}}_{ij} = C^{\mathrm{ref}}_{ij}$。因为 FP16 可以精确表示 [0, 2048) 范围内的整数。
• 基线容差： 自定义内核的最大元素级绝对误差不得超过 cuBLASLt-AutoTuning 基线相对于 FP32 CPU 参考结果的最大误差。

5.2.3. 测量模式 (Timing Measurement)

为了避免 Python 的惰性求值和确保精确计时，采用了标准的 CUDA 事件计时机制：

torch.cuda.synchronize()
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
start_event.record()
kernel(a, b, b_col_major, out)
end_event.record()
torch.cuda.synchronize()
elapsed_time_ms = start_event.elapsed_time(end_event)

说明：

• torch.cuda.synchronize(): 确保所有之前的 CUDA 操作完成。
• start_event.record() / end_event.record(): 记录 CUDA 事件发生的时间戳。
• kernel(a, b, b_col_major, out): 执行自定义 CUDA 内核。
• start_event.elapsed_time(end_event): 计算两个事件之间经过的时间。

5.2.4. 离线模式 (Offline Mode) 与 服务器模式 (Server Mode)

• 离线模式： 模拟批处理推理，内核以最快速度连续执行，GPU 缓存保持“热”状态。
• 服务器模式： 模拟实时推理，内核以随机时间间隔执行，GPU 缓存可能“冷却”并从冷启动状态开始。服务器模式下的速度提升通常更能体现实际部署环境中的优势。

5.3. 对比基线

CUDA-L2 与多个广泛使用的和最先进的矩阵乘法库进行了比较。

5.3.1. torch.matmul

• 概念定义： PyTorch 深度学习框架中提供的默认矩阵乘法操作。对于大多数用户而言，这是最常用且最易于使用的基线，因为它处理了底层的硬件优化细节。

5.3.2. cuBLAS

• 概念定义： NVIDIA 提供的 GPU BLAS 库，包含高度优化的矩阵乘法实现。本文评估了两种最常用的矩阵布局：
  • cuBLAS-NN (Normal-Normal)：两个输入矩阵都采用行主序 (row-major)。
  • cuBLAS-TN (Transposed-Normal)：第一个输入矩阵转置，第二个输入矩阵采用行主序。
• cublas_max： 对于每个 (M, N, K) 配置，选择 cuBLAS-NN 和 cuBLAS-TN 中性能最优的布局作为 cuBLAS-max 的基线。
• 操作： 使用 cublasGemmEx 函数，指定 CUBLAS_GEMM_DEFAULT_TENSOR_OP 操作以启用 Ampere FP16 Tensor Cores。

5.3.3. cuBLASLt

• 概念定义： cuBLAS 的轻量级版本，提供更底层的 API，允许更细粒度的控制和算法选择。
  • cuBLASLt-heuristic： 查询 cuBLASLt 库，并根据其内置的启发式算法建议来选择最佳算法。选定的算法会被缓存以消除评估时的开销。
  • cuBLASLt-AutoTuning： 这是最具竞争力的基线。它会查询 cuBLASLt 库，获得最多 100 个候选算法，然后对这些算法进行基准测试，选择实际执行最快的算法。这个过程通常在程序启动时进行一次性设置，以避免运行时开销。
• cuBLASLt_max： 类似 cuBLAS-max，对于每个 (M, N, K) 配置，选择 cuBLASLt-heuristic-NN 和 cuBLASLt-heuristic-TN (以及 cuBLASLt-AutoTuning-NN 和 cuBLASLt-AutoTuning-TN) 中性能最优的布局作为 cuBLASLt-heuristic-max (以及 cuBLASLt-AutoTuning-max) 的基线。

6. 实验结果与分析

6.1. 核心结果分析

6.1.1. 总体性能对比

CUDA-L2 在所有基线上均持续表现出优异的性能，无论是在离线模式还是服务器模式。

• 布局影响： CUDA-L2 在 cuBLAS 和 cuBLASLt 基线上，相对于 TN (转置-正常) 布局的性能提升略高于 NN (正常-正常) 布局。例如，离线模式下，相对于 cuBLAS-NN 提升 20.0%，而相对于 cuBLAS-TN 提升 21.4%。这可能表明 CUDA-L2 在处理转置矩阵方面有独特的优化能力。
• 基线强度：

  • torch.matmul 是最弱的基线，CUDA-L2 对其提升最大。

  • cuBLAS-max 通过选择最优布局，性能优于单一布局的 cuBLAS。

  • cuBLASLt-heuristic 通过算法选择进一步提升。

  • cuBLASLt-AutoTuning 是最强的基线，因为它为每个配置穷尽测试多达 100 个内核候选。即便如此，CUDA-L2 仍能超越。

    以下是原文 Table 1 的结果，展示了 CUDA-L2 相对于各基线在离线和服务器模式下的平均加速比： 以下是原文 Table 1 的结果：

• 离线模式 (Offline Mode)：
  • torch.matmul：平均加速比 22.0%。
  • cuBLAS-max：平均加速比 19.2%。
  • cuBLASLt-heuristic-max：平均加速比 16.8%。
  • cuBLASLt-AutoTuning-max：平均加速比 11.4%。
• 服务器模式 (Server Mode)： 性能提升更加显著，因为 CUDA-L2 能够更好地处理冷启动情况，有效利用 GPU 缓存。
  • torch.matmul：平均加速比 28.7%。
  • cuBLAS-max：平均加速比 26.0%。
  • cuBLASLt-heuristic-max：平均加速比 22.4%。
  • cuBLASLt-AutoTuning-max：平均加速比 15.9%。
• 胜率 (Win Rates)： CUDA-L2 在 79.3% 到 95.7% 的配置上超越了基线，这表明其改进是系统性的，而非由少数异常值驱动。

6.1.2. max(CUDA-L2, baseline) 性能

如果允许在 CUDA-L2 和基线之间为每个配置选择更快的内核，则可以获得额外的性能提升。这模拟了实际部署中，用户可以动态选择最优内核的情况。

以下是原文 Table 2 的结果，展示了 max(CUDA-L2, baseline) 相对于基线的平均加速比： 以下是原文 Table 2 的结果：

• 在离线模式下，cuBLASLt-AutoTuning-max 的加速比从 11.4% 增加到 13.2%。
• 在服务器模式下，cuBLASLt-AutoTuning-max 的加速比从 15.9% 增加到 18.1%。 这进一步验证了 CUDA-L2 的有效性，并且在实际应用中，可以通过这种选择机制进一步提升性能。

6.1.3. 速度提升与问题规模的关系

下图（原文 Table 3）展示了 CUDA-L2 相对于 cuBLASLt-AutoTuning-max 的加速比随矩阵大小的变化趋势。

该图像是图表，展示了相对于不同矩阵大小的相对加速比，包含三个子图(a)、(b)和(c)。(a)显示了相对加速比与$ext{log}_2( ext{M} imes ext{N} imes ext{K})$的关系；(b)表示相对加速比与平均维度的关系；(c)则展示了相对加速比与最大维度的关系，提供了统计分布信息。

• 趋势： 随着问题规模的增大（即 $\log_2(M \times N \times K)$ 增大），CUDA-L2 的加速比呈现下降趋势。
• 小问题： 对于较小的问题（例如 $\log_2(M \times N \times K) \approx 18-20$），CUDA-L2 能够实现约 $1.3-1.4 \times$ 的高加速比。这表明 CUDA-L2 在优化小矩阵乘法方面具有显著优势。
• 大问题： 对于非常大的问题（例如 $\log_2(M \times N \times K) > 38$），加速比逐渐接近 $1.0 \times$，即性能与 cuBLASLt-AutoTuning-max 相当。
• 原因分析： 论文认为，对于大矩阵，由于 GPU 的流水线和内存带宽瓶颈被充分利用，进一步优化的空间变小。而小矩阵由于其特有的调度和资源利用模式，CUDA-L2 能发现更多细粒度的优化，从而获得更高的相对加速比。这意味着 CUDA-L2 对于 attention 机制和 FFN 层中的中小型矩阵乘法，可以带来显著的性能提升。

6.2. 消融实验/参数分析

6.2.1. BM, BN, BK 的选择模式

CUDA-L2 学习到了针对不同维度 M, N, K 的最优分块大小 (BM, BN, BK) 选择模式。 下图（原文 Table 4）展示了超参数选择模式。

该图像是一个六部分的箱线图，展示了不同超参数对优化CUDA矩阵乘法内核的影响。部分(a)和(b)分析了随着M和N的增加，BM的变化情况，分别显示了相关系数ρ=0.652和ρ=0.705。部分(c)揭示BK与K之间的弱相关性（ρ=0.256），部分(d)显示BM和BN趋于相似（ρ=0.695）。部分(e)比较了随着K的增加阶段数的变化（ρ=0.473），而部分(f)展示了块交换的影响，表明问题规模与性能关系的复杂性。

• BM 与 $M$ 的关系 (图 4a)： BM (线程块 M 维度) 与 $M$ (矩阵 M 维度) 呈正相关 ($\rho = 0.652$)。当 $M$ 从小 ($\leq 256$) 增加到大 ($> 4K$) 时，BM 从约 60 增加到 160。这表明更大的 $M$ 维度需要更大的 M-tiles 以保持效率。
• BN 与 $N$ 的关系 (图 4b)： BN (线程块 N 维度) 与 $N$ (矩阵 N 维度) 呈更强的正相关 ($\rho = 0.705$)。
• BK 与 $K$ 的关系 (图 4c)： BK (线程块 K 维度) 与 $K$ (矩阵 K 维度) 仅呈弱相关 ($\rho = 0.256$)。这可能是因为 BK 的选择更多受到 GPU 寄存器限制和流水线阶段数等其他因素的约束。
• BM 与 BN 的关系 (图 4d)： BM 和 BN 之间高度相关 ($\rho = 0.695$)。这表明最优配置倾向于平衡的 tile 大小，以匹配 GPU 的计算单元指令格式（如 $16 \times 8 \times 16$），并减少资源不平衡导致的瓶颈。

6.2.2. 多阶段流水线 (n_stage) 中的阶段数选择

• n_stage 与 $K$ 的关系 (图 4e)： n_stage (流水线阶段数) 与 $K$ 维度呈正相关 ($\rho = 0.473$)。
• 趋势： 随着 $K$ 的增加，流水线阶段数应增加以达到最佳性能。
  • 小 $K$ ($\leq 128$)：需要 2-3 个阶段来隐藏足够的延迟。
  • 大 $K$ ($> 8K$)：需要 6 个或更多阶段，通过同时处理多个数据加载来保持高吞吐量。

6.2.3. 块交换 (Block Swizzling) 的使用时机和方式

• Block Swizzling 与问题规模的关系 (图 4f)： Block Swizzling 用于优化线程块访问全局内存的方式，其使用与问题规模 ($M \times N \times K$) 强相关。
• 使用时机：
  • 小问题 (小于 $2^{27}$ 或约 1.34 亿次操作)：Block Swizzling 是可选的，仅在 44% 的配置中使用，因为其开销可能抵消收益。
  • 中等问题 ($2^{27}$ 到 $2^{33}$)：使用率增加到 73%-80%。
  • 超大问题 (大于 $2^{36}$ 或 640 亿次操作)：几乎总是使用 (99%)，因为它对于复杂的内存访问模式至关重要。
• 步长 (Stride) 选择： 随着问题规模的增大，最优的步长值也随之增大，对于大问题通常选择 512-16384。这表明 Block Swizzling 不仅在何时使用，而且如何使用（即选择何种步长）都对性能有显著影响，并且 CUDA-L2 能够发现这些模式。

7. 总结与思考

7.1. 结论总结

CUDA-L2 系统成功地将大语言模型 (LLMs) 与强化学习 (RL) 相结合，实现了半精度通用矩阵乘法 (HGEMM) CUDA 内核的自动化优化。通过将 CUDA 执行速度作为 RL 奖励，并对 1,000 种矩阵配置进行系统性探索，CUDA-L2 显著超越了包括 torch.matmul、cuBLAS 和 cuBLASLt-AutoTuning 在内的所有主要基线，在离线模式下，相对于 cuBLASLt-AutoTuning 实现了 11.4% 的加速，在服务器模式下更是达到了 15.9%。对于更常用的基线，加速比甚至更高。此外，CUDA-L2 还自动发现了并应用了多种高级 CUDA 优化技术，例如输入矩阵填充、双缓冲寄存器片段、激进寄存器级预取以及交错预取调度等，甚至能够发现这些技术的创新变体。这项工作证明了即使是最为性能关键且经过高度优化的 HGEMM 内核，LLM 引导的 RL 自动化也能通过探索人类难以企及的配置空间，发现并实现超越现有最先进解决方案的性能提升。

7.2. 局限性与未来工作

论文虽然展示了 CUDA-L2 的强大能力，但仍存在一些潜在的局限性，并为未来工作提供了方向：

• 计算成本： cuBLASLt-AutoTuning 需要对多达 100 个候选算法进行基准测试，这本身就具有一定的计算开销。CUDA-L2 训练过程中的 LLM 代码生成、编译和 RL 迭代可能会产生更高的计算资源和时间成本，尤其是在面对新的 GPU 架构或更广泛的内核类型时。论文并未详细说明 CUDA-L2 的训练时间。
• 泛化能力： 尽管 CUDA-L2 经过了持续预训练并使用了 contextual memory，但在面对全新的 GPU 架构或完全不同的内核类型时，其泛化能力仍需进一步验证。LLM 的能力边界和 RL 探索空间的效率是关键。
• 可解释性： LLM 生成的 CUDA 代码，尤其当它包含内联 PTX 汇编或复杂 CuTe 模板时，可能难以被人类理解和调试。其发现的某些优化技巧，虽然有效，但其背后的深层原因和适用范围可能不完全透明。
• 错误处理和鲁棒性： LLM 生成的代码可能包含语法错误或逻辑错误，尽管有正确性检查，但如何更高效地从这些错误中恢复并学习，以提高生成代码的鲁棒性，是一个挑战。
• 对 LLM 性能的依赖： CUDA-L2 的成功在很大程度上依赖于底层 LLM 的代码生成质量和对 CUDA 编程范式的理解。LLM 模型的进步将直接影响 CUDA-L2 的上限。

7.3. 个人启发与批判

7.3.1. 个人启发

• LLM 作为自动化优化引擎的巨大潜力： 这项工作极大地拓宽了 LLM 的应用场景，将其从自然语言处理扩展到复杂的系统级代码优化。LLM 的代码生成能力与 RL 的性能反馈相结合，形成了一个强大的自动化发现和优化循环，为未来在其他计算密集型领域（如科学计算、数据库查询优化、操作系统内核优化）的应用提供了模板。
• 超越人类专家瓶颈： 传统的手动优化严重依赖专家知识和经验，且效率低下。CUDA-L2 展示了机器智能在系统性探索优化空间方面的优势，甚至能够发现人类专家难以察觉的细微但高效的优化策略。这对于释放硬件潜力、应对快速迭代的硬件架构具有重要意义。
• “黑盒”优化到“白盒”洞察： 尽管 LLM-RL 的生成过程可能相对“黑盒”，但论文对 CUDA-L2 发现的优化技术（如填充、双缓冲、预取策略）和超参数选择模式（如 $BM/BN/BK$ 与 $M/N/K$ 的关系）进行了详细分析。这种从自动化优化中提取可解释性洞察的能力，对于指导未来的手动优化和工具开发非常有价值。
• 跨领域融合的典范： LLM (AI)、RL (AI)、GPU 编程 (系统/硬件) 的深度融合，是典型的跨学科研究典范，预示着 AI 赋能其他领域的强大趋势。

7.3.2. 批判

• 训练成本与可访问性： 训练一个像 CUDA-L2 这样基于大型 LLM 和 RL 的系统，需要巨大的计算资源和时间投入。这使得其技术可能难以被小型研究团队或个人开发者复制和应用。如何降低这种系统化的训练和部署成本，是其走向更广泛应用的关键。
• “超越 cuBLAS”的语境： 尽管 CUDA-L2 在 1000 种配置上超越了 cuBLAS，但 cuBLAS 本身涵盖了各种数据类型、操作模式和 GPU 架构。CUDA-L2 专注于 HGEMM 在 A100 上的性能。cuBLAS 作为一个通用的、经过多年优化和验证的库，其全面性和稳定性仍然是 CUDA-L2 这样的特定优化器短期内难以完全替代的。未来的研究需要验证 CUDA-L2 在更广泛的 GEMM 类型（如 FP32, INT8）、更多 GPU 架构以及更复杂场景下的表现。
• 代码可维护性： LLM 生成的包含 CuTe 模板、内联 PTX 汇编的复杂 CUDA 代码，虽然性能优异，但其可读性、可调试性和可维护性可能是一个挑战。在工业界部署时，代码的这些非功能性属性同样重要。如何平衡性能与代码质量，是此类自动化生成系统需要考虑的问题。
• 对 LLM 的“信仰”问题： 整个系统的成功高度依赖于 LLM 的代码生成能力。如果 LLM 犯下严重的逻辑错误或产生无法修复的漏洞，RL 优化循环可能无法有效地收敛到高性能解，甚至可能陷入局部最优。对 LLM 内部机制的更深层理解和控制，对于提高这类系统的可靠性至关重要。