【CGO’25 】CuAsmRL: Optimizing GPU SASS Schedules via Deep Reinforcement Learning

本文元信息

  • 发表会议：CGO’ 25
  • 作者：来自剑桥大学
  • 被引数：2
  • 三标论文
  • 领域：GPU、编译器、二进制优化、指令调度、强化学习、大语言模型

本文提出 CuAsmRL，一种用于 GPU SASS 调度自动优化的系统。CuAsmRL 直接在 GPU 原生汇编层进行优化，并可无缝集成至现有编译器框架，同时对 CUDA kernel 开发者保持完全透明。我们在多种常见的大语言模型（LLM） kernel 上验证了 CuAsmRL 的有效性，实验结果表明，其最高可带来 26% 的性能提升，平均提升达 9%。此外，我们还验证了其对超参数配置的稳健性，以及其发现新的优化策略的能力。

本文的叙述流畅，CGO 也是一贯如此，这篇论文从工程角度看，整合了多个底层技术模块，工程复杂度较高但是技术路线明确（也就是说好效仿，技术扎实）。它做了以下几件事：

  • 集成到 Triton 编译器：对编译流程进行修改，对编译出来的 cubin，进行反汇编（SASS 级别）再优化，然后重新组装；
  • 指令调度的动作空间管理：每条 SASS 指令都可能参与重排，但为了控制复杂度，只针对内存指令，并加入了依赖分析、动作掩码等机制，防止语义出错；
  • 强化学习训练闭环：每次动作都要跑真实的 CUDA kernel 获取反馈（执行时间），平均每个策略训练要跑上万次 GPU 任务，系统还要提供重排后的合法性检查和测试验证；
  • 微基准测试系统：为了获取 NVIDIA 没公开的固定延迟指令特性（stall count），作者专门写了低级 SASS 程序测定硬件延迟，这需要对 GPU 硬件架构非常熟悉。
  • 同时作者还对 RL 过程中的现象进行了详细分析和阐释，基本没有什么逻辑上的缺陷和不完整性。

从学术角度看，它的最大特点在于优化粒度比别人更低。现有工作大多是在 LLVM IR 或者 PTX 级别优化，像 CUTLASS、AutoTVM 之类也是自动搜索调度，但都不深入到底层的 SASS 指令。CuAsmRL 是直接在 NVIDIA 编译器输出的“最终成品”上做强化学习优化，这点是独特的。此外，它不是用传统启发式或进化搜索，而是强化学习做重排序决策，并设计了配套的状态表示、奖励函数和动作掩码逻辑，整体思路比较系统。深度方面，在工程实现和系统完整性上非常深，能跑起来本身就很不容易。但如果从强化学习方法本身来看，模型结构和策略训练没引入太多创新点（主要是用 PPO + CNN + MLP）。这不是一个方法学导向很强的论文，而是一个“用已有方法做到了别人没做到的事”的系统型工作。

摘要

大型语言模型（LLMs）因其巨大的计算需求而受到广泛关注。为了降低计算成本，研究人员开发了专门的 CUDA 算子，这些算子通常通过融合多个张量操作来最大化 GPU 资源的利用。然而，即使采用了这些专门优化的算子，仍可能无法充分发挥硬件性能。CUDA 汇编专家表明，手动优化 GPU SASS 调度能够进一步提升性能，并且目前主要依赖试错方法来寻找最优的 GPU SASS 调度方案。
在本研究中，我们提出了一种自动化方法来优化 GPU SASS 调度，该方法可以无缝集成到现有的编译器框架中。自动化优化的关键在于训练强化学习（RL）代理，以模仿人类专家进行手动调度的过程。为此，我们将调度优化问题形式化为一个“汇编博弈”（assembly game），使 RL 代理能够在其中探索最优的 GPU SASS 调度策略。
该“汇编博弈”从 -O3 优化后的 SASS 调度方案出发，RL 代理可以逐步应用一系列动作来对当前调度进行变异。如果变异后的调度能够在 GPU 上实现更高的吞吐量，则给予正向奖励。实验结果表明，CuAsmRL 方法能够透明地进一步优化现有的专门 CUDA 算子，其性能最高可提升 26%，平均提升 9%。此外，该方法还可作为工具，用于自动学习并揭示潜在的优化策略。

注：“CUDA 汇编专家” 是一种非常专业化、底层优化导向的 GPU 编程专家，专注于使用 CUDA 汇编语言（或 PTX/SASS）来挖掘 NVIDIA GPU 的极限性能。其工作内容有：分析编译生成的 PTX/SASS； 手动调整寄存器使用；优化内存访问模式；重排循环结构与指令调度；使用 SASS 编写关键算子。

注意，这个领域比算子开发还要底层，复杂度更高，对 GPU 的控制也极强。涉及的内容如下表

层级名称说明
高层CUDA C/C++常见并行程序编写语言
中层PTX (Parallel Thread Execution)NVIDIA 的虚拟汇编语言，类似 LLVM IR
底层SASS (Shader Assembly)NVIDIA GPU 的真正机器指令集（具体硬件 ISA）
工具cuobjdump / nvdisasm用于反汇编 CUDA 二进制
调试Nsight Compute / Nsight Systems性能分析、流水线瓶颈查看
编译nvcc, ptxas编译 CUDA 代码并生成 PTX/SASS

全文翻译

1. 引言

LLMs 是基于 Transformer 的深度神经网络，包含多层自注意力机制和线性投影结构。自其问世以来，在图像生成和自然语言处理等多个领域均取得了最先进的性能表现。据统计，OpenAI 每日生成的文本已超过千亿词汇。因此，LLMs 已成为深度学习社区中的关键计算负载，受到广泛关注。然而，由于其通常包含由数十亿参数构成的多层 Transformer 主干结构，训练和推理过程均伴随巨大的计算开销。为此，研究人员开发了专用的 CUDA 算子以加速 LLM 计算流程，替代了高层语言生成 CUDA 算子的方式。例如，融合注意力机制（flash-attention）通过优化 NVIDIA GPU 的共享内存使用效率提升了注意力机制的计算性能；而融合前馈模块则在 LLAMA 中融合了多个算子以提升整体效率；均方根层归一化则是 Transformer 中广泛使用的归一化操作。这些工作通常采用手工编写的、面向硬件高效执行的 CUDA 算子实现，借助硬件厂商提供的编程模型，兼顾了灵活性与性能表现。

在本研究中，我们进一步探讨通过更底层的硬件原生汇编级优化手段提升手工编写 CUDA 算子性能的可能性，具体聚焦于 NVIDIA CUDA 算子。相较于传统 CUDA 优化，底层优化能够进一步提升现有专用 CUDA 算子的效率，已有相关研究表明，GPU 原生命令（SASS）调度的人工优化确实可以带来显著性能提升。然而，当前的方法依赖于繁琐的试错过程，由专家手动探索最优的 SASS 调度策略，不仅费时费力，而且难以适应深度学习算子快速演化的趋势。此外，手动优化方式难以集成进现有编译流程。

为此，我们提出 CuAsmRL —— 一个用于自动优化 NVIDIA GPU SASS 调度的优化器。该方法通过训练一个强化学习（RL）代理，模拟人类专家进行调度优化的过程，自动搜索最优的 SASS 调度策略。据我们所知，这是首次将 SASS 调度优化问题形式化为汇编游戏。自动优化 SASS 调度的能力使我们能够将 CuAsmRL 集成进 OpenAI Triton 编译框架中，Triton 是一个基于 MLIR 的 GPU 算子开发平台。具体而言，我们首先通过自动调优器搜索最优算子配置，随后重用 Triton 的编译流程，在其生成的 cubin 阶段进行拦截与解析，提取 SASS 指令后进行优化，并重新汇编生成优化后的 cubin。

在针对典型 LLM 算子的评估中，我们发现 CuAsmRL 能够自动发现优于 -O3 优化级别的调度策略，在几何平均吞吐量上获得了 1.09 倍的性能提升。由于该优化发生在 CUDA 算子之下的底层，对于 CUDA 开发者是完全透明的。考虑到 LLM 的训练与推理过程可能会消耗数百万小时的 GPU 计算资源，此类算子级优化预计将带来显著的收益。

注：计算机里的透明一般是指看不见！被隐藏

综上所述，本文的主要贡献包括：

  • 我们将 SASS 调度优化问题形式化为一个汇编游戏，并实现了 CuAsmRL，一个用于自动优化 NVIDIA GPU SASS 调度的优化器；
  • 我们将 CuAsmRL 集成进现有编译框架 OpenAI Triton 中，作为一个 SASS 到 SASS 的优化器，且对 CUDA 算子开发者保持透明；
  • 实验结果表明，代表性的 LLM 专用算子在 Ampere 架构 GPU 上通过本方法可获得最高 26%、平均 9% 的性能加速；
  • 我们展示了 CuAsmRL 可作为自动发现优化策略的工具，揭示强化学习所学调度操作，为 SASS 指令级优化带来新的启示。

2 背景与动机

2.1 GPU 编程与 CUDA 算子编译

GPU 是高性能的并行计算加速器，能够高效执行张量操作。为在 GPU 上编程，开发者需遵循 CUDA 所提供的编程模型。概念上，CUDA 算子由多个线程块（thread blocks）组成，这些线程块可并行运行；每个线程块内部包含多个线程，且每个线程块映射到一个 GPU 流多处理器（streaming multiprocessor），以独立方式执行。

CUDA 算子开发通常使用高级编程语言，如 C++ 或 Python，随后由编译器将算子代码编译为设备代码。在 C++ 情况下，编译任务由 NVIDIA 的编译器 NVCC 完成；若采用 Python，则可借助 Triton 编译框架。整个编译流程包含多个阶段：首先，源代码被编译为 PTX——一种与 GPU 架构无关的中间语言。值得注意的是，开发者亦可在高级语言中直接嵌入 PTX 代码。接下来，PTX 会被进一步编译为 SASS，这一过程只能通过 NVIDIA 专有的编译器 ptxas 实现。SASS 是目标 GPU 架构的原生汇编语言，具有高度架构依赖性。本研究限定于 Ampere 架构的 GPU。

虽然开发者可以利用 CUDA 提供的二进制工具获取 CUDA 算子所对应的 SASS 代码，但该指令集缺乏详细文档，导致该阶段的代码下沉（lowering）与优化过程基本处于不可见状态。最终，SASS 指令被汇编为可在 GPU 上直接执行的二进制文件 cubin。该完整的编译流程如图 1 所示。
image-20250801091833.webp

CUDA 编译流程如 NVIDIA 官方文档所示：C++/Python 以及 PTX 被标注为绿色，表示它们是常用的编程接口；SASS 作为 GPU 的原生汇编语言，被标注为红色，意味着其未被官方文档充分公开；而 Cubin 作为可执行的二进制文件，则以灰色标识。

2.2 GPU SASS 指令优化

虽然已有大量研究致力于在 C++/Python 层面对 CUDA 算子进行优化（例如内存访问优化与负载均衡优化），但针对 GPU SASS 调度优化的研究相对稀缺。这主要由于 SASS 属于闭源格式，且缺乏官方汇编器。然而，正因为 SASS 处于编译流程的底层，若在该层进行优化，不仅对开发者完全透明，还可惠及所有已存在的专用 CUDA 算子。

近年来，开源社区已开发出定制的汇编器，使得 SASS 层的优化成为可能。例如，MaxAs 是首个能够解析 CUDA 二进制并重组早期 GPU 架构 SASS 代码的工作。随后，针对更现代 GPU 架构，如 TuringAs 与 Cuasm 的汇编器相继被提出，从而进一步推动了 GPU SASS 指令的优化研究。接下来的章节将依次讨论 SASS 指令的结构以及相关工作的优化方法论。

2.3 SASS 指令解析

典型的 SASS 指令形式如下，包含若干字段：控制码（control code）、操作码（opcode）以及操作数（operands）：

[B------:R-:W2:Y:S02] LDG.E R0, [R2.64];

控制码被方括号包裹，并由多个冒号分隔的字段组成。第一个字段是等待屏障掩码（wait barrier mask），若其中某一位被设置，当前指令将在该位清除前保持等待状态。第二与第三字段分别为读屏障（read barrier）与写屏障（write barrier）掩码；上述例子中，设置写屏障为 2，意味着后续若有指令使用寄存器 R0，其执行将被阻塞，直至 R0 准备就绪。第四字段是 yield 标志，据信用于负载均衡。最后一个字段是停顿周期计数（stall count），指示当前指令在发出下一条指令前需等待的周期数。操作码目前在官方文档中的说明极为有限。以该例为例，LDG 表示从全局内存中加载数据。操作数由寄存器与内存地址构成。若需更系统地理解 SASS 指令的解码过程，建议参考相关研究工作。

固定延迟与可变延迟指令。
SASS 指令可分为固定延迟指令与可变延迟指令两类。前者如 IADD3 与 FFMA 等数学运算，具有确定的执行周期；而后者如 LDG.E 等全局内存加载指令，其执行周期取决于 GPU 内存系统的访问路径（包括 L1 缓存、L2 缓存与全局内存）而呈不确定性。因此，在程序编译时，无法预知其所需周期数。此外，自 Kepler 架构起，GPU 指令执行模式为静态调度，编译器需通过设置控制码主动规避数据相关性问题。

例如上述 LDG.E 指令具有可变延迟，其控制码便设置了第 2 写屏障，确保后续读取 R0 的操作在数据未就绪前不会提前执行。

2.4 延迟隐藏

在解决数据依赖所引起的执行停顿时，会引入额外延迟。为缓解此类延迟，GPU 通常采用两种机制实现延迟隐藏：其一，当某一 warp 执行长延迟操作时，硬件可切换至下一个具备执行资格的 warp；其二，则是调度下一条独立指令执行。上述两种机制分别称为线程级并行（Thread-Level Parallelism, TLP）与指令级并行（Instruction-Level Parallelism, ILP）。

已有研究表明，通过手动重排 SASS 指令，可实现内存访问延迟的隐藏，尽可能实现内存读写指令与计算指令的并发执行，从而提升 ILP，减少执行管线的停顿。尽管 GPU 能够通过 TLP 方式切换到其他 warp，但能调度的 warp 数量受到算法结构、算子配置（如 tile 大小与寄存器使用）的制约，故此策略在某些情况下难以持续生效，执行停顿仍可能成为瓶颈。

因此，研究人员尝试通过手动交错（interleaving）内存读写与计算指令来隐藏延迟。在 MaxAs 中采用试错式策略，而 TuringAs 则引入基于性能分析的调度策略。

注：LLVM 后端有支持 Interleave 操作的 Pass。

2.5 强化学习

强化学习（Reinforcement Learning, RL）是一类用于解决序列决策问题的算法，其核心思想是通过在环境中反复试探性地行动并从反馈中学习，从而优化策略。应用 RL 的关键在于将待优化问题建模为马尔可夫决策过程（Markov Decision Process, MDP），该过程包含状态空间、动作空间以及奖励函数。

RL 旨在最大化长期回报，其最优策略可表示为：

π=arg⁡max⁡πE[∑t=0∞γtrt∣s0]\pi = \arg\max_\pi \mathbb{E} \left[ \sum_{t=0}^{\infty} \gamma^t r_t \mid s_0 \right]

近年来，深度强化学习（Deep RL）使用深度神经网络对策略进行建模，从而增强了对复杂和动态环境的处理能力。其优势在于可在几乎无需人工干预的前提下学习复杂的决策问题，且具备长期优化能力，即能够容忍短期损失以换取长期收益。因此，深度强化学习已广泛应用于诸多领域，包括电子游戏、机器人控制、数据中心能耗管理与设备布局优化等。

2.6 研究动机

我们观察到，当前针对 SASS 调度的优化工作严重依赖人工，工作量庞大且易出错。首先，每个 CUDA 算子通常包含数千行 SASS 指令，需由人工识别潜在的优化模式。其次，SASS 指令之间的依赖关系必须精确维护。更为复杂的是，一旦输入数据类型或算子配置发生变化，对应的 SASS 指令将发生根本性改变，需重新优化。最后，人工调度难以集成进现有编译器框架，除非实现自动化。

为此，我们提出使用强化学习来弥合这一差距。原因在于，指令交错问题可建模为一个离散优化问题，而 RL 可通过学习一系列操作序列来最大化长期收益。此外，基于 RL 的优化具备自动化能力，便于与现有编译器框架集成。因此，借助此方法，经过编译器生成的 CUDA 算子可由 RL 智能体自动进行调度优化，最大限度地减少人工干预。

CuAsmRL 架构

3.1 分层搜索空间

本节将对 CuAsmRL 的整体结构进行概述。随后我们将分别介绍其分层搜索空间以及如何将 SASS 调度问题形式化为强化学习问题。图 2 展示了 CuAsmRL 的整体工作流程。CuAsmRL 以面向 Triton 编程接口的源代码作为输入，首先通过自动调优器（autotuner）枚举并寻找最优的算子配置，然后将该最优配置输入至 Triton 编译流程，生成中间表示并最终编译成 cubin。之后，系统利用官方工具对 cubin 反汇编得到 SASS 指令，并由强化学习智能体通过“汇编游戏”对 SASS 调度进行优化，最终输出优化后的 cubin 可执行文件。
image-20250801094241.webp

CuAsmRL 的整体工作流程如下：系统首先以面向 Triton 编程接口的源代码作为输入，接着通过自动调优器枚举并搜索最优的算子配置；随后，利用 Triton 的编译流程基于该最优配置对代码进行编译。最后，通过训练强化学习智能体参与“汇编博弈”，对 SASS 指令调度进行优化，并输出最终的优化后 cubin 可执行文件。

得益于与 Triton 的集成，CuAsmRL 能够以 Triton 编程接口为目标的 CUDA 算子源码为输入，并借助其内置的自动调优管线，搜索最优算子配置。完成配置选择后，算子将通过 Triton 编译流程被编译生成 cubin，再由工具反汇编为 SASS 指令。此后，强化学习代理被训练以进行汇编级调度优化，输出最终优化结果。

自动调优器在该流程中至关重要，因为如 tile 大小等算子配置可能导致高达 2 倍的吞吐量差异，并生成完全不同的 SASS 指令，从而影响调度结构。因此，CuAsmRL 采用分层搜索策略：首先搜索最优算子配置，然后在此基础上优化对应的 SASS 调度。

自动调优器采用类似网格搜索（grid search）的策略：对用户提供的算子配置逐一枚举，使用目标 GPU 编译并运行相应算子，测量吞吐率，并通过贪心策略选取最优配置并缓存。性能测量采用平均 100 次重复执行结果，并在前置执行 100 次热身迭代后进行。

3.2 汇编博弈前(Pre-Game)静态分析

CuAsmRL 配备了解析器用于解析 SASS 指令。该解析器不仅将每条指令拆分为控制码、操作码及操作数等组成部分并存储为数据结构，还对操作数进行了扩展处理。由于通用寄存器为 32 位，若操作数后缀带有 .64，则表明该操作涉及相邻寄存器。通过构造微基准程序，故意污染相邻寄存器并对比输出与预期值，可以验证这一推断。此模式在内存操作中极为常见，为此我们扩展了操作数结构以覆盖相邻寄存器，从而准确提取依赖关系。相邻寄存器的计算方式如下：

base=reg 号2,mod=reg 号 mod 2,flip=1−mod,adj.reg=base×2+flip\text{base} = \frac{\text{reg 号}}{2} \quad,\quad \text{mod} = \text{reg 号} \bmod 2 \quad,\quad \text{flip} = 1 - \text{mod} \quad,\quad \text{adj.reg} = \text{base} \times 2 + \text{flip}

base：计算得到寄存器的基底，mod：计算寄存器奇偶数，flip：取反操作（0 变 1，1 变 0），adj. reg：最终计算得到相邻寄存器编号。它是与当前寄存器配对的寄存器，用于 .64 类型的双寄存器操作

在进入“汇编游戏”之前，CuAsmRL 会对汇编文件执行若干分析流程以实现静态分析：

  • 依赖分析：记录每条内存指令是否依赖于同一基本块中某个固定延迟指令的结果。对于每条内存指令，该分析过程会向前扫描其前序指令，查找其操作数寄存器的赋值源头。一旦遇到标签（label）即中止扫描，并将当前指令加入黑名单（denylist）；否则，记录从定义到使用之间累计的停顿周期数。如果某个固定延迟指令的停顿周期在微基准测试中已知，或已由先前分析推断得到，则取其中较小值作为保守估计。该分析限定在基本块内进行，因为我们不允许在标签间进行指令重排（见 §3.5）。在实践中，这一分析手段效果显著。例如，在某个算子上运行该分析过程即可推断出 IADD3.X 指令的停顿周期为 5，与微基准测试结果仅相差 1 个周期。即使存在轻微高估，也不影响调度正确性；因原始调度本身为合法调度，推导值始终是保守或准确的。未来，我们可替代人工微基准测试，在大规模 SASS 算子代码上运行该分析过程，自动构建一个停顿周期查找表。以 CUDA 工具链的每次更新为契机，可将共享库中的大量算子（如 libcu*.so）进行反汇编与分析。
  • 嵌入分析：为后续嵌入过程（见 §3.4）做准备，例如构建寄存器到整数的映射表。此外，考虑到 SASS 指令的操作数数量不固定，我们在分析时记录汇编文件中操作数的最大数目，后续在嵌入过程中，对操作数不足的指令用虚拟值（−1）进行补齐。
  • 动作空间分析：统计汇编文件中除黑名单外的内存指令数量，该信息用于定义强化学习过程中的动作空间，详细见 §3.5。

3.3 强化学习

在完成对反汇编后的 SASS 指令的分析后，系统训练一个强化学习（Reinforcement Learning, RL）智能体，通过参与“汇编游戏”（assembly game）以优化 SASS 调度策略。汇编游戏是一个迭代过程：在每一轮迭代中，RL 智能体感知当前的 SASS 调度状态（即状态），并选择一个动作对其进行修改，从而生成新的调度。该变异后的 SASS 文件将被汇编并部署至 GPU 执行，GPU 的实际运行结果将产生一个奖励值反馈给智能体。该过程如图 3 所示。为构建汇编游戏的形式化问题，我们分别定义了动作空间、状态空间以及奖励函数，具体内容将在后续小节中详细阐述。
image-20250801095807.webp

图 3 展示了汇编游戏的工作机制：在每一轮迭代ii 中，SASS 文件被嵌入为状态表示SiS_i，并输入至以深度神经网络为结构的 RL 智能体。智能体输出一个动作AiA_i，对 SASS 文件进行修改，产生变异后的 SASS 代码。此代码随后被汇编并部署至目标 GPU 执行，从而获得奖励值RiR_i，该值反馈至智能体。新的调度状态记为Si+1S_{i+1}。

3.4 状态空间

为了使强化学习（RL）智能体能够处理 SASS 指令调度，我们对 SASS 指令进行了嵌入。回忆一下，一个典型的 SASS 指令由控制码（control code）、操作码（opcode）和操作数（operands）组成，如 §2.3 所示，我们分别对每个字段进行嵌入，然后将它们的嵌入结果拼接起来。例如，读/写屏障可以取 0 到 5 之间的任意整数，其嵌入也相应地映射为这些整数；如果屏障不存在，则用 −1 填充。对于操作码，我们只区分其是否为内存操作指令。预处理分析过程已经从 SASS 文件中提取出了内存指令；如果是非内存指令，则用 −1 表示。对于操作数，我们将内存位置转换为其在内存表中的索引（内存表由预处理分析过程建立），然后将索引除以总内存位置数进行归一化。如果操作数数量不足 SASS 文件中操作数数量的最大值，则用 −1 补齐，因为 SASS 指令的操作数数量是可变的。图 4 展示了一个对 SASS 指令进行嵌入的示例。
image-20250801100212.webp
因此，嵌入完成后，一条 SASS 指令的状态表示就是一个向量，而整个汇编文件则通过将这些指令向量按行拼接，变成了一个矩阵。

3.5 动作空间

在定义了状态空间之后，我们接着定义动作空间。鉴于专家在优化过程中常通过交错计算指令与内存指令以实现延迟隐藏，我们希望强化学习（RL）智能体也具备类似的灵活性。因此，我们允许智能体选择一条指令，并将其与上一条或下一条指令进行交换。我们认为，这种操作方式与专家进行指令交错的策略较为相似，如下面所示所示：

IMAD.WIDE R14, R84, R8, c[0x0][0x160];
LDGSTS.E.BYPASS.LTC128B.128 [R74], desc[UR18][R18.64], P4;

（Listing 1：重排序前）

LDGSTS.E.BYPASS.LTC128B.128 [R74], desc[UR18][R18.64], P4;
IMAD.WIDE R14, R84, R8, c[0x0][0x160];

（Listing 2：重排序后）

然而，允许任意指令重排序会导致极大的动作空间，因为一个 kernel 可能包含成千上万条 SASS 指令。考虑到隐藏延迟的关键在于将内存加载/存储指令放置于更优位置，我们选择仅探索动作空间的一个子集，从而实现裁剪。具体来说，我们仅允许智能体选择内存加载/存储类指令（如 LDG、LDGSTS、STG），其索引由预处理分析阶段记录，并在每次迭代时动态更新。RL 智能体输出一个离散值，该值映射到某条指令的索引及其重排序方向（向上或向下）。在网络结构方面，智能体使用卷积神经网络（CNN）对状态进行编码，随后通过多层感知机（MLP）输出各动作的概率分布。

在重排序过程中，保持数据依赖关系至关重要，否则可能导致程序语义错误。为此，我们采用动作掩码（action masking）机制来过滤掉潜在的非法操作。以下是需要考虑的依赖关系：

  • 寄存器依赖：使用某寄存器的指令不得出现在该寄存器赋值指令之前；
  • 屏障依赖：读写屏障指令不得被重排序至其设置者之前。例如，若一条指令需等待第 2 个屏障完成，则不能重排序至设置第 2 个屏障的指令之前。该依赖通过比较相邻指令的控制码实现；
  • 停顿周期依赖：固定延迟指令通过延迟若干周期以解决依赖关系，该延迟周期数即为“停顿计数”。由于原始 SASS 指令由 NVIDIA 的专有编译器调度，所有依赖关系均已满足。尽管官方未公开停顿计数，我们通过微基准 (§4.3) 或分析流程 (§3.2) 获取其值。若某内存指令使用了由具有未知停顿计数的固定延迟指令赋值的寄存器，则该指令将被加入拒绝列表，始终被掩码。否则，我们检查其前后邻接指令是否会因重排序而造成潜在依赖冲突。例如，算法 1 展示了判断将一条内存指令向上移动是否满足停顿周期的掩码算法：该算法累加停顿计数并与最小要求进行比较，若不满足，则将该动作掩码；
  • 其他依赖：我们还识别了一些硬件层面上的隐式依赖。例如，一组 LDGSTS 指令以寄存器为偏移连续写入全局内存的相邻地址时，任意指令的重排序都可能导致错误。此类依赖可能与 Ampere 架构中从全局内存向共享内存的数据搬运机制有关，鉴于缺乏公开文档，我们只能手动识别此类模式。我们还禁止指令跨越标签或任何屏障/同步指令进行重排序，因而所有调度仅限于基本块内部。官方文档中列出了所有屏障与同步指令。此外，部分依赖关系以启发式规则形式表达并硬编码实现，任何违反这些规则的操作均被掩码。在 §5 中针对若干 LLM kernel 的评估显示，当前规则集在该领域内已足够完备。在 §5.7 中，我们还对优化后的 kernel 逐步进行人工验证，确保其正确性。
    image-20250801102759.webp

综上所述，我们在每次迭代中为每个动作动态生成动作掩码。若某动作可能违反任何依赖关系，则其掩码值为 0，对应的概率为零，表示该动作不可选。若当前无可行动作，则立即终止该训练回合。

3.6 奖励函数

获取反馈信号是强化学习过程中的关键环节，因为它直接引导智能体探索更优的调度策略。在本工作中，我们主要关注优化后的 CUDA kernel 的运行时间，因此必须在每次动作应用之后对运行时间进行测量。具体而言，我们采用 CUDA events 来测量 kernel 的执行时间。按照标准流程，首先对 GPU 进行 100 次热身运行，然后再执行 100 次迭代以测量平均耗时。在每次迭代之间清除 L2 缓存，以确保测量结果的准确性。最终返回的平均执行时间即作为反馈信号。我们观测到两次独立测量之间的标准差通常在 1% 以内，说明测量具有较高稳定性。奖励值的计算公式如下：

Ri=Ti−1−TiT0×100R_i = \frac{T_{i-1} - T_i}{T_0} \times 100

其中，T0T_0 表示初始运行时间，TiT_i 表示当前动作应用后的运行时间，Ti−1T_{i-1} 表示前一次动作应用后的运行时间，系数 100 用于数值缩放。该公式的直观含义是：如果动作使运行时间下降，则获得正向奖励；反之，则为负向奖励。依据优化目标函数，RL 智能体将学习一个策略网络，用以最大化累积奖励，最终目标是减少 kernel 的总体执行时间。此外，该目标函数也鼓励智能体容忍短期的性能下降，只要该动作可能带来长期收益。

3.7 强化学习算法

CuAsmRL 默认采用了 Proximal Policy Optimization（PPO）算法的参考实现。我们在所有实验中使用统一的超参数设置，因为针对具体场景微调强化学习超参数在计算上代价高昂。所使用的默认超参数来自一项横跨多个领域的大规模实证研究，该研究总结了一组在实际中表现良好的通用超参数。

在实现上，我们对算法进行了定制：使用卷积神经网络（CNN）对汇编文件的嵌入表示进行编码，并采用基于 actor-critic 架构的策略梯度方法以学习最优策略。重排序过程被封装在环境转移逻辑中，并遵循标准的 Gym 接口，这一设计为将来替换或扩展其他 RL 算法提供了便利。

训练过程中，我们记录包括每轮的奖励值（episodic rewards）与智能体损失函数在内的统计信息，并定期保存智能体的权重检查点。

4. 实现

4.1 与 Triton 的集成

我们选择将 CuAsmRL 集成至 OpenAI Triton，这是一个用于编写 GPU kernel 的编译器。Triton 允许用户使用 Python 语法编写 kernel 程序，并通过即时编译（JIT）将其编译为可在 NVIDIA 或 AMD GPU 上运行的目标代码。此外，Triton 也是深度学习框架 PyTorch 的默认后端之一，而 PyTorch 是当前最流行的深度学习平台之一。因此，通过与 Triton 的集成，我们希望本工作能够直接服务于更广泛的深度学习社区。

Triton 中 kernel 编写语法的示例如下所示（见代码片段 3）：

@triton.jit
def matmul(x_ptr, y_ptr, out_ptr):
   ...

CuAsmRL 复用了 Triton 的编译流程，但对自动调优器进行了扩展，并拦截了编译生成的 cubin 文件。随后，它将 cubin 反汇编为 SASS 指令，提取出包含调度信息的 kernel 部分，同时保留其他元信息。这一点尤为关键，因为如符号表、ELF 文件格式等元信息必须被完整保留。之后，系统对该 kernel 部分训练 RL 智能体进行优化，并将优化后的 SASS 部分重新嵌入 cubin 中替换原始 kernel。为了应用 CuAsmRL 的优化，用户仅需在 Triton 源代码中更改一行代码，如代码片段 4 所示：

@cuasmrl.jit(ret_ptr=1)
def matmul(x_ptr, y_ptr, out_ptr):
    ...

其中，ret_ptr 表示输出缓冲区的索引，可用于概率性测试（probabilistic testing）。概率性测试会生成随机输入和参考输出，并与程序执行结果进行比对。我们使用该机制作为有效性校验手段，并在 §5.7 中对优化后的 kernel 的每一步骤进行了人工验证。由于 SASS 缺乏官方语义定义，因此无法使用形式化验证方法；此外，按位穷举输入空间在实际中计算成本极高，因为 kernel 通常处理大规模输入数据。可选地，用户还可通过添加参数指定 RL 智能体的超参数，例如学习率、训练批大小等。

4.2 工作流程

由于训练 RL 智能体是一个耗时过程，我们建议用户采用“离线搜索 + 部署时查找”的工作流。这不仅避免了运行时开销，同时更高的训练预算也能使 RL 智能体更充分地探索动作空间，从而获得更优的调度结果。代码片段 5 展示了如何调用 CuAsmRL 进行优化以及如何加载优化后的 cubin：

“离线搜索 + 部署时查找”（offline search + deploy-time lookup）是一种在强化学习（RL）场景下常见的优化策略，其核心思想是：在程序运行之前花时间训练和优化模型，运行时则直接使用最优结果，避免一切实时计算的开销。

# 调用优化
matmul(x_ptr, y_ptr, out_ptr)
# 部署时加载已优化版本
matmul(x_ptr, y_ptr, out_ptr, load_dir='path-to-cubin')

用户在完成 kernel 编写后，应调用 CuAsmRL 以执行分层优化（hierarchical optimization）。整个强化学习过程中的最优 cubin 文件将被保存至文件系统，并使用 GPU 类型、工作负载类型等作为前缀组成查找键。在部署阶段，系统将接收该查找键并直接执行 cubin 的加载，而无需重新训练，从而避免了运行时开销，仅保留离线搜索时间的代价。我们实测 RL 智能体的训练时间通常不超过 5 小时。相对于大规模语言模型的训练与推理过程通常消耗数百万 GPU 小时而言，这一成本可以视为一次性且可忽略不计。

4.3 停顿周期表（Stall Count Table）

CuAsmRL 内置了一张映射表，用于将常见的固定延迟指令（fixed-latency instructions）名称映射到其对应的 stall count（停顿周期）。该映射表通过微基准测试（microbenchmarking）获得，并被用于第 §3.5 节所述的动作掩码逻辑。表 1 展示了该映射关系，涵盖了常见的整数操作指令，因为这些指令在地址计算中使用频繁，其输出通常被后续内存访问指令所依赖。

image-20250801104816.webp

下面介绍微基准测试的具体方法。与已有研究在 Ampere 架构上使用 PTX 层进行微基准测试不同，我们直接采用 SASS 指令进行编程，这使得我们可以构造出明确的使用-定义（use-definition）指令对，从而精确测定固定延迟指令的 stall count。这种方法在以往对 Volta 与 Turing 架构的研究中也曾被使用。

我们的测试过程如下：首先编写一个简单的 CUDA kernel，编译并导出其 SASS 指令，然后据此手动编写 SASS 程序。例如，代码片段 6 展示了用于测试 MOV 指令的微基准程序：

[B - - - - - -:R -: W -: -: S04 ] MOV R15, 0x1;
[B - - - - - -:R -: W -: -: S04 ] STG.E desc[UR4][R4.64], R15;

在该示例中，第 2 行指令使用了 MOV 指令（第 1 行）的输出，并将其写入全局内存。我们通过逐步减少 MOV 指令的 stall count 值，观察其输出是否仍与预期一致。一旦输出结果开始偏离，说明该值不足以满足依赖关系；此时的最小 stall count 即为 MOV 指令所需的延迟周期数。

确定 MOV 指令 stall count 后，我们可以控制寄存器的值并构造类似的微基准测试来测量其他指令的 stall count。对于延迟更高的指令，则在使用指令与其使用者之间插入适量的 NOP 指令，直至输出结果正确，从而推断出其所需的最小 stall count。这些测得的 stall count 数值将被硬编码进 CuAsmRL 系统中。

我们发现基于数据依赖关系的微基准方法相比基于时钟的方式更加精确。已有研究中采用基于时钟的方法会低估 stall count，其基本形式如代码片段 7 所示（控制码已省略）：

CS2R R2, SR_CLOCKLO;         // t1
// 一系列 IADD3 指令
CS2R R6, SR_CLOCKLO;         // t2
IADD3 R6, P0, -R2, R6, RZ;   // t2 - t1

在该方法中，IADD3 指令的平均 stall count 被测为 2.6 个周期，显著低于表 1 所列值。我们认为其原因在于：第二个时钟采样点（t2）并不能保证之前所有 IADD3 指令已经执行完成，从而造成了低估。为解决这一问题，必须为 IADD3 序列与最终计时指令之间构造明确的读/写依赖关系。这进一步强调了构造依赖路径在测量 SASS 指令真实 stall count 中的重要性。

5. 评估

本研究旨在回答以下三个核心问题：

  • CuAsmRL 能在多大程度上透明地提升 Triton 及其他基线系统的性能？
  • CuAsmRL 对其超参数配置是否敏感？
  • 为什么必须在 SASS 层进行优化？RL 智能体采用了哪些调度动作来优化 SASS 指令？

5.1 实验设置

  • NVIDIA A100 80GB PCIe GPU、NVIDIA ptxas 12.2、Triton v2.1.0

Baseline：

  • PyTorch 基线（v2.1.2）：通过组合多个 PyTorch 运算构建。其“eager 模式”将计算任务分发至 NVIDIA 高性能计算库 CuBLAS（v12.1），但该库对运算融合的自定义能力有限。
  • Cutlass 基线（v3.5）：用于评估融合 GEMM 与 LeakyReLU 的场景。
  • Flash-Attention 基线（v2.3.3）：用于评估自注意力计算中的优化情况。

为评估 kernel 吞吐量，我们对每个 kernel 进行 5 次重复测量。每次测量通过 CUDA event 记录 kernel 执行时间，采用 100 次热身迭代后再执行 100 次正式测量迭代的方式。为了研究更细粒度的 kernel 性能指标（详见 §5.4），我们在训练完成后将优化得到的 cubin 文件写入文件系统，并使用 Nsight Compute 工具（NVIDIA 提供的 kernel 级性能分析器）对比分析 CuAsmRL 与 Triton 所生成的 kernel。Nsight Compute 可访问 GPU 硬件性能计数器，提取精细化的性能统计数据。我们选取了一系列具有代表性的 LLM kernel 对 CuAsmRL 进行评估。

  • 计算密集型 kernel 包括：
    • 融合 GEMM 与后处理阶段（如 LeakyReLU）；
    • 融合前馈网络；
    • 批量矩阵乘法；
    • Flash-Attention。
  • 内存带宽受限型 kernel 包括：
    • Rmsnorm；
    • Softmax。

上述融合 kernel 来自 Triton 官方仓库与 Kernel 项目。所有实验均使用常见 kernel 尺寸与配置，数据类型统一采用 float16。所评估的 kernel 总览列于表 2。
image-20250801110332.webp

其他实验部分（略）

  1. 指令延迟
  2. 算子吞吐量
  3. 加速比
  4. 训练数据
  5. 该优化的必要性
  6. 自动分析 RL 的优化偏好和趋势

局限性和未来工作

将 CuAsmRL 应用于其他领域的 kernel 优化可能需要额外的依赖约束，超出 §3.5 中所列内容，这是由于缺乏公开的底层硬件与指令数据所致。因此，用户需参考 §5.7 中的方法对优化结果进行人工验证，以确保正确性。

CuAsmRL 的另一项局限在于，其奖励函数的计算依赖于实际在 GPU 上执行 kernel 以获取运行时反馈信号。每个训练步骤中需执行约 200 次 kernel，而训练出一个有效策略通常需约 1.5 万个步骤（详见图 8），这导致训练成本较高。若能构建一个可近似 kernel 执行时间的成本模型，将显著降低训练开销。然而，由于 SASS 指令集的底层行为和数据尚未公开，这一成本模型的构建充满挑战。

从本工作的指令重排序问题建模角度来看，也可尝试使用其他搜索算法进行优化，例如进化搜索（evolutionary search）。进化搜索无需训练过程，但容易陷入局部最优解，从而导致性能下降。我们之所以选择强化学习方法，是基于其在多个领域中已展现出的先进性能，并具备在未见过的 SASS 调度中泛化的潜力。

然而，为了实现泛化能力，未来工作需考虑在多种 CUDA kernel 的 SASS 调度数据上对 RL 智能体进行预训练。一旦预训练完成，RL 智能体即可作为常规编译器优化阶段的一部分，无需为每个 kernel 单独耗费数小时重新训练，从而实现更高效的编译过程。

第 8 章 结论