中文摘要

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

Introduction

  • 开篇有点冗余，老生常谈的内容
  • 首个NVIDIA GPU SASS 调度器
    • 集成于 Triton
    • 和普通 CUDA 编译一致，生成 CUBIN 时反汇编然后走优化流程
  • 贡献
    • SASS 调度优化问题形式化为汇编博弈
    • 集成于 Triton，实现 SASS-to-SASS 优化
    • 可以提供加速

      背景与动机

  • ①GPU 和 CUDAd 的编译流程
    • 线程块、NVCC、Triton、SM、PTX、SASS、CUBIN
    • image-20250307083841185.webp
  • ②GPU SASS 指令优化
    • 大量研究致力于==访存优化==和==负载均衡优化==，对于 SASS 的研究少
      • 主要是闭源，且未有汇编工具
    • 其他工作关于 SASS 的：MaxAs、TuringAs、Cuasm
  • ③SASS 指令介绍（略）
    • 示例: [B------:R-:W2:Y:S02] LDG.E R0, [R2.64];
    • 固定延迟和可变延迟指令
      • 固定延迟如 IADD3，FFMA 等计算指令
      • 可变延迟如 LDG. E 受限于 Cache，不确定性
    • 自 Kepler 以来，GPU 指令为静态调度
  • ④延迟隐藏
      - TLP：当一个 warp 执行长延迟操作时，GPU 可切换到下一个可执行的 warp 以隐藏延迟。
  - ILP：在 warp 内部，GPU 可调度独立指令以填补空闲周期，从而减少流水线暂停（stall）。
    • 先前研究是手动排序（专家）SASS 指令
  • ⑤强化学习（略）
    • 将优化问题建模为马尔可夫决策过程（Markov Decision Process, MDP），该过程包括 状态空间（state space）、动作空间（action space） 和 奖励函数（reward function）

      主架构

      image-20250307090322803.webp

      主架构图：吐槽一下，这是小学生画的吗？甚至连白底都不想去

首先还是大篇幅背景介绍 AI 编译的自动调优过程（这个和作者无关？为啥要讲？）

介绍了一个通用寄存器的背景，给出了确定相邻寄存器的算法

静态分析

然后第二小节终于开始讲主要内容了（可能因为本质就是个 RL 一会就讲完了吧？）。在初始化 汇编博弈（assembly game） 之前，CuAsmRL 需要对 汇编文件 运行多个 静态分析（static analysis） 过程，以便优化 SASS 调度策略。
作者的静态分析流程略，其主要目的还是分析 stall 。

静态分析 Pass 的核心目标是记录所有 内存指令（memory instructions），如果其操作数来自于同一基本块（basic block）内的固定延迟指令（fixed latency instruction），则进行标记和分析。
对于每一条 内存指令，分析 Pass 通过扫描前序指令来查找其操作数寄存器的赋值来源。如果在扫描过程中首先遇到了代码标签（label），则分析过程终止，并将当前指令加入拒绝列表（denylist），防止误判。否则，分析 Pass 记录该内存指令与其依赖指令之间的累计 stall 计数，即 use-definition 指令对（use-definition instruction pair） 之间的执行停滞周期数。
如果某条固定延迟指令的 stall 计数已被记录（例如通过 微基准测试（§4.3） 或之前的推测值获得），则选择最小的 stall 计数，以优化调度策略。
该分析 Pass 仅在同一基本块内执行，因为指令重排序不会跨越 代码标签（§3.5）。实践表明，该分析 Pass 在优化中非常有效。例如，在某个 CUDA 内核上运行该 Pass 后，我们推测 IADD3.X 指令的 stall 计数为 5，这一结果与微基准测试测得的值仅相差 1 个周期。这种轻微的过估（overestimation）是安全的，因为原始调度始终是可执行的，因此推测值要么是准确的，要么是略微保守的。
在未来，我们可以用 大规模 SASS 内核代码 运行此分析 Pass，并自动构建 stall 计数查找表（stall count look-up table），从而避免手动微基准测试。例如，每次 CUDA 工具包发布时，我们都可以提取共享库（libcu*.so）中的大量 CUDA 内核，进行分析，并自动生成 stall 计数数据。

强化学习

然后即可通过强化学习了：汇编博弈的优化过程是 迭代进行的（iterative）：

  • -RL 代理感知当前的 SASS 调度状态（state）。
  • RL 代理选择一个动作（action），==即对 SASS 调度进行某种修改==。
  • 修改后的 SASS 文件重新汇编，并在 GPU 上执行，计算新的执行吞吐量。（这是潜在的缺陷）
  • GPU 返回执行结果作为奖励（reward），RL 代理基于奖励反馈调整策略。
    （这个图就不放了，很容易理解）
    状态空间
  • 读写屏障：可以取值 0~5
  • 操作码：内存指令与非内存指令
  • 操作数：仅对内存指令操作，转为归一化后的内存表索引（内存表由静态分析得到）
  • 可变操作数用-1 补齐占位
    动作空间
    在人工优化 SASS 时，专家通常采用交错计算指令和内存指令的方法。作者允许 RL 选择一条指令和上一条或者下一条交换。但是一个 Kernel 可能有千万条指令，所以为了剪枝 (prune) 作者仅仅以隐藏延迟为目的，来优化内存读写指令。
    依赖关系约束和动作屏蔽
  • 寄存器依赖
  • 屏障依赖
  • 停滞计数依赖（微基准和静态分析得到）
  • 还有其他依赖不一一列举
    奖励函数
    具体而言，我们使用 CUDA 事件（CUDA events） 来测量 CUDA 内核的执行时间。测量流程遵循标准方法：
  1. 首先进行 100 次 GPU 预热（warming up），以确保 GPU 进入稳定运行状态。
  2. 随后执行 100 次重复运行，并记录 总执行时间（elapsed time） 。
  3. 在每次迭代之间清空 L2 缓存（L2 caches），确保测量结果的准确性。
  4. 计算平均执行时间，并将其作为 RL 代理的反馈信号（feedback signal）。
    实验表明，两次独立测量的标准偏差（standard deviation）通常在 1% 以内，表明测量方法具有稳定性和可靠性。
    奖励值（很简单的一个函数）：
    $Ri=\frac{T{i-1}-T_i}{T_0}*100$
    强化学习算法
    在默认情况下，CuAsmRL 采用 近端策略优化算法（Proximal Policy Optimization, PPO） 作为 RL 训练框架。
  • 不针对特定任务进行超参数微调
  • 基于 “演员-评论家”（Actor-Critic）的策略梯度方法（policy gradient algorithm）

    缺陷和未来工作

    泛化性不强：应用于其他领域的 CUDA内核优化可能需要额外的依赖约束。
    训练时计算成本高。
    理论也可以用其他迭代类型的算法来解决问题。