论文分享-Record-Remix-Replay

如何让大模型修改复杂 GPU Kernel 源码，同时又能够快速、可靠地评估每个候选版本，并进一步为每个源码版本找到合适的编译器 Pass 和 Kernel 启动参数？

核心思路:

外层：LLM + MAP-Elites
      搜索源码实现、算法结构和代码重写

内层：贝叶斯优化
      搜索编译器 Pass、编译参数和 Launch Configuration

底层：GPU Kernel Record-Replay
      独立、快速、可重复地运行候选 Kernel

真实 GPU 应用的性能由多个层次共同决定：

算法选择
  ↓
Kernel 源码实现
  ↓
编译器优化和 Pass 顺序
  ↓
线程块大小、Grid 大小、Launch Bounds
  ↓
实际硬件执行行为

例如，一个 Kernel 是否适合使用 256 个线程，并不仅由 Kernel 的输入大小决定，还取决于：

• 源码中是否展开了循环；
• 使用了多少寄存器；
• 是否使用共享内存；
• 编译器是否完成常量传播；
• 内联后代码规模是否增大；
• 当前代码版本是否产生更多 VGPR；
• 某种编译器 Pass 是否改变了指令调度。

1 已有工作

Kernel 参数调优器

例如 CLTune、Kernel Tuner、Kernel Launcher，主要搜索：

• Block Size；
• Tile Size；
• Launch 参数；
• 模板参数。

优点是搜索空间结构化、评估成本低，但通常不能发现复杂的源码重构。

DSL 和 Auto-scheduler

例如 TVM、Ansor、Triton、Halide。

它们可以搜索：

• 调度策略；
• Tile；
• 向量化；
• 内存布局；
• 线程映射。

但是通常要求程序使用特定 DSL 或中间表示，不能直接处理任意大型 CUDA/HIP 应用中的复杂 Kernel。

编译器自动调优

例如 CompilerGym、OpenTuner，主要搜索：

• 编译器参数；
• LLVM Pass；
• Pass 顺序。

但如果每次评价都要编译并运行完整应用，成本会很高。

LLM 代码优化系统

例如 AlphaEvolve、OpenEvolve，可以修改无结构的源代码。

它们的问题是：

LLM 生成候选
    ↓
重新编译完整应用
    ↓
运行完整应用
    ↓
测量性能

2 R³介绍

2.1 Record：记录

在完整应用的一次代表性运行中，捕获某次 GPU Kernel 调用，包括：

• Kernel LLVM IR；
• Grid 和 Block 配置；
• Kernel 参数；
• 执行前 GPU 内存状态；
• 执行后 GPU 内存状态。

2.2 Remix：重组和优化

对 Kernel 进行不同层面的修改：

• LLM 修改源码；
• LLVM Pass 修改中间表示；
• 常量传播；
• 参数特化；
• Launch Bounds 插入；
• Block Size 等启动参数变化。

2.3 Replay：重放

不再运行完整应用，而是恢复捕获到的 GPU 状态，单独执行修改后的 Kernel：

恢复执行前状态
   ↓
执行候选 Kernel
   ↓
比较执行后状态
   ↓
测量运行时间

这使得候选评价从“应用级运行”变成“Kernel 级重放”。

2.4 总体结构介绍

系统采用双层搜索结构：外层负责探索不同的 Kernel 源码实现，内层负责为每个源码版本寻找最佳编译优化与启动配置。底层的 Replay Server 通过恢复已记录的 GPU 状态，对候选 Kernel 进行快速、可重复的性能评估和正确性验证。

该流程形成一个持续迭代的闭环：MAP-Elites 保留具有不同特征的高质量源码候选，LLM 基于历史候选生成新的实现；随后，贝叶斯优化针对该实现搜索 LLVM Pass、Pass 顺序和 Launch Configuration。通过正确性验证的最佳结果最终返回外层种群，参与下一轮源码搜索。

简单流程如下

LLM 修改源码
    ↓
搜索最佳编译与启动参数
    ↓
Replay Server 快速评估
    ↓
将优秀版本加入 MAP-Elites
    ↓
继续生成新的源码版本

2.5 分层搜索

2.5.1 无结构空间交给 LLM

源码修改是无结构的文本空间，例如：

• 改写循环；
• 重排计算顺序；
• 缓存重复加载；
• 提取公共表达式；
• 删除重复地址计算；
• 改变数据复用方式；
• 修改算法实现。

这类变化无法简单表示为几个整数参数。

传统搜索中的“随机变异”很难应用到代码：

output[i] = input[i] * scale;

如果随机修改字符，几乎一定生成非法程序。

LLM则能够生成语法上合理、语义上有一定可能正确的修改。

2.5.2 结构化空间交给贝叶斯优化

编译和启动参数通常是结构化的：

block_size ∈ {64, 128, 256, 512}
optimization_level ∈ {O1, O2, O3}
launch_bounds ∈ {128, 256, 512}
pass_A ∈ {on, off}
pass_B ∈ {on, off}

这类空间适合：

• 贝叶斯优化；
• Tree-structured Parzen Estimator；
• 随机搜索；
• 进化数值搜索。

论文认为，使用 LLM 猜测 Block Size 不如使用成熟的贝叶斯优化高效。

因此其设计原则是：

LLM 负责“创造新的实现”
BO 负责“为该实现寻找最佳配置”

这是一个比较合理的能力分工。

3 核心思想

3.1 外层搜索：LLM + MAP-Elites

3.1.1 MAP-Elites 是什么

MAP-Elites 是一种强调“多样性”的进化搜索算法。

普通进化算法可能只保留当前最快的几个候选：

Population:
A: 1.00 ms
B: 1.02 ms
C: 1.05 ms

这样容易快速收敛到某一类相似实现，失去探索能力。

MAP-Elites 不只按性能保存候选，还按特征将候选划分到不同网格单元。例如：

特征维度1：代码复杂度
特征维度2：与已有代码的差异

可能得到：

	低差异	中差异	高差异
低复杂度	Elite A	Elite B	Elite C
中复杂度	Elite D	Elite E	Elite F
高复杂度	Elite G	Elite H	Elite I

每个格子只保留该特征区域中性能最好的候选，即 Elite。

这样做的目的是同时保留：

• 简单但保守的优化；
• 大规模结构重写；
• 不同实现路线；
• 不同复杂度的代码。

因此，即使某个候选暂时不是全局最快，只要它代表了一种有价值的不同方向，也可能被保存下来，后续继续变异。

---

3.1.2 Island 模型

系统使用多个独立的 MAP-Elites 数据库，称为 islands。

每个 island 独立进行：

选择候选
→ 构造 Prompt
→ LLM 生成
→ 评价
→ 更新 Population

每隔一定迭代次数，不同 island 之间会迁移候选。

例如：

Island 1：倾向循环展开
Island 2：倾向数据复用
Island 3：倾向减少分支
Island 4：倾向向量化

周期性迁移能让优秀思想传播到其他 island，同时避免所有搜索过程过早变得完全相同。

实验设置中使用：

• 4个 islands；
• 每个 Population 大小为20；
• 每20次迭代进行迁移；
• 总共200次外层迭代。

---

3.1.3 LLM Prompt 中包含什么

R³构造的 Prompt 不只是“请优化这段代码”，而是包含：

• 当前待优化 Kernel；
• 当前任务描述；
• 历史优秀候选；
• 历史失败或较差候选；
• 优化目标；
• 代码生成约束。

其思想类似于让模型看到进化历史：

原始版本做了什么
哪些修改有效
哪些修改无效
当前候选处于哪条优化路线

模型生成的是一个新的 Kernel 实现，然后交给编译和 Replay 系统验证。

3.2 内层搜索：贝叶斯优化

每生成一个源码版本，R³不会只运行一次，而是启动一次内层调优。

实验中，论文使用并行 Tree-structured Parzen Estimator，对每个源码候选进行30次内部搜索。

每一个内部搜索点可以表示为：

$x=(CompilerPasses,LaunchConfig)$

例如：

Candidate 1:
  block_size = 128
  opt_level = O3
  constant_specialization = on
  launch_bounds = 256

Candidate 2:
  block_size = 256
  opt_level = O2
  constant_specialization = off
  launch_bounds = 512

评价一次配置时：

1. 根据配置变换 LLVM IR
2. 将 IR 编译为设备二进制
3. 恢复捕获状态
4. 运行 Kernel
5. 检查结果
6. 返回运行时间

如果出现以下情况，则该点被判定无效：

• 编译失败；
• Kernel 崩溃；
• 输出错误；
• 执行后内存不匹配。

最终，内层 BO 返回：

{
  best_runtime,
  best_compiler_configuration,
  best_launch_configuration,
  correctness
}

外层 MAP-Elites 用 best_runtime 评价该源码版本。

3.3 Record-Replay 引擎

3.3.1 Instrumentation：插桩

R³扩展了 Proteus JIT 基础设施。

在编译应用时，系统会：

1. 识别 GPU Kernel；
2. 找到 Kernel 调用到的所有传递依赖；
3. 提取 Kernel 及依赖的 LLVM IR；
4. 将 LLVM IR 嵌入应用二进制；
5. 将 Kernel Launch 重定向到 Proteus Runtime。

所谓传递依赖闭包，包括 Kernel 调用的：

• 内联设备函数；
• 模板实例；
• 辅助函数；
• 所依赖的全局定义。

系统不能只提取 Kernel 函数本体，否则重放时会缺少依赖。

---

3.3.2 Recording：记录

在一次代表性的完整应用运行中，预加载库会拦截：

• Kernel Launch；
• GPU 内存分配；
• GPU 内存复制；
• 相关设备内存操作。

对于一次 Kernel 调用，记录以下内容：

① Kernel LLVM IR

后续可以重新编译和施加 LLVM Pass。

② Launch Configuration

例如：

gridDim
blockDim
dynamic shared memory size
stream

③ Kernel Arguments

例如：

kernel<<<grid, block>>>(A, B, C, n, alpha);

需要记录：

• 指针参数；
• 标量参数；
• 结构体参数；
• 参数布局。

④ 执行前设备状态

论文称为 prologue snapshot。

它代表 Kernel 开始执行前所看到的 GPU 内存状态。

⑤ 执行后设备状态

论文称为 epilogue snapshot。

它代表原始 Kernel 执行完毕后的参考结果，用于正确性检查。

一个 Replay Unit 可以概括为：

Replay Unit =
{
    LLVM IR,
    launch config,
    arguments,
    prologue state,
    epilogue state
}

---

3.3.3 Replay：重放

重放时，系统首先恢复 prologue：

GPU Memory ← recorded pre-execution state

然后加载候选 Kernel：

候选 LLVM IR
→ 应用编译变换
→ 编译设备二进制
→ 按候选 Launch Config 启动

最后：

实际执行后状态
       vs.
记录的 epilogue 状态

如果一致，则候选在该捕获输入下正确。

3.4 LLVM IR

R³把 LLVM IR 作为源码搜索和底层 Replay 的接口。

1. 避免重新编译整个应用

传统流程可能是：

修改一个头文件
→ 许多 Translation Unit 重新编译
→ 链接整个应用
→ 执行应用

而R³只需要：

修改 Kernel
→ 生成该 Kernel 的 LLVM IR
→ 发送给 Replay Server

不需要重新构建整个大型项目。

在 QUDA 中，论文提到完整重新编译大约需要30分钟，即使使用增量编译，成本也依然很高。

---

2.可以直接进行编译器变换

在 LLVM IR 上，Replay 引擎可以执行：

• 常量传播；
• Kernel 参数特化；
• Thread/Block Dimension 特化；
• 插入 Launch Bounds；
• LLVM 标准优化级别；
• 自定义 Pass Pipeline；
• 不同 Pass 顺序。

4 优化

4.1 Persistent Replay Server

普通 Record-Replay 如果每次候选都重新启动进程，仍会有大量开销：

• 初始化 GPU；
• 加载 Replay 数据；
• 分配设备内存；
• 复制 prologue；
• 创建校验缓冲区；
• 读取中间文件；
• 将结果传回 CPU。

R³提出了持久化 Replay Server。

工作方式：

Server 启动时：

加载所有捕获的 Kernel
预分配 GPU Buffer
加载 Prologue/Epilogue
将 Replay Unit 映射到不同 GPU
保持 Worker 常驻

评价候选时只需要发送：

Kernel ID
候选 LLVM IR
编译配置
Launch 配置

Server 返回：

runtime
correct / incorrect

因为 Worker 和 GPU 状态一直存在，所以初始化成本只支付一次。

这对进化搜索非常重要。外层要生成大量候选，每个候选内部还要运行多次 BO 配置。如果每次都重新初始化，Record-Replay 的优势会被系统开销抵消。

4.2 Prefix-cache Aware Prompting

Record-Replay 将执行评价加速后，系统出现了一个有趣变化：

原来瓶颈：编译和运行候选
后来瓶颈：LLM 生成候选

因此论文又优化了 LLM 推理。

现代 LLM 服务通常支持 Prefix Cache：

Prompt A:
[相同静态指令][代码1][代码2]

Prompt B:
[相同静态指令][代码1][代码3]

如果前缀相同，模型可以复用前缀阶段的计算结果。

但普通 OpenEvolve 风格的 Prompt 往往将动态内容放在较前位置：

[随机候选][随机示例][静态任务说明]

每次 Prompt 很早就发生变化，导致 Prefix Cache 命中率低。

R³采用两种方法。

1. 静态内容放前面

[固定系统说明]
[固定优化目标]
[固定格式要求]
[动态代码候选]
[动态历史示例]

从而让较长的静态前缀被缓存。

2. 重排 Inspiration Samples

假设本次选择的历史示例是：

E = [e3, e1, e5]

最近缓存的某个 Prompt 前缀是：

H = [e1, e5, e8]

那么R³可能将本次示例重排为：

[e1, e5, e3]

使 [e1, e5] 命中已有缓存。

它只改变示例顺序，不改变选中了哪些样本，因此作者认为不会改变 MAP-Elites 的基本搜索语义。

5 实验设置

5.1 硬件环境

AMD 实验：

GPU：AMD MI300A
软件：ROCm 7.0.2

NVIDIA 实验：

GPU：NVIDIA H100
软件：CUDA 12

本地 gpt-oss 模型通过 vLLM 部署。

---

5.2 测试应用

论文选择四个科学计算应用。

1）LULESH

• 冲击流体力学 Proxy Application；
• 非结构六面体网格；
• 15个 GPU Kernel。

2）MiniFE

• 有限元 Proxy Application；
• 包含稀疏矩阵生成、装配和迭代求解；
• 7个 GPU Kernel。

3）S3D

• 湍流燃烧 DNS；
• 包含详细化学反应和分子输运；
• 54个 GPU Kernel。

4）miniWeather

• 干燥可压缩非静力天气模拟；
• 总计9个 Kernel；
• 实验数据路径涉及7个 Kernel。

5.3 对比方法

Baseline 1：Record-Replay + BO

只调：

• 编译参数；
• 编译器配置；
• Launch 参数。

不修改源码。

共搜索200次。

Baseline 2：OpenEvolve

使用 LLM + MAP-Elites 修改源码，但通过较传统的方式编译和运行应用评价。

R³

同时搜索：

• 源码；
• 编译器 Pass；
• Launch 参数。

并使用 Kernel Replay 快速评价。

6 结果

1. 编译和评价开销显著下降

从 OpenEvolve 迁移到 Record-Replay 后，在 QUDA 上：

• 编译时间降低约86%；
• 评价时间降低约88%。

OpenEvolve 原本是 evaluation-bound：

LLM生成时间 < 编译和运行时间

R³之后变成 inference-bound：

编译和运行时间 < LLM生成时间

说明 Record-Replay 已经把评价降得足够低，新的瓶颈变成了 LLM。

---

2. Kernel Speedup

在 MI300A 和 H100 上，Figure 6和Figure 7均显示：

• R³的 Kernel 中位加速比高于两个基线；
• R³的最大加速比更高；
• OpenEvolve 通常优于只做 BO；
• 说明源码重写确实提供了编译参数调优无法覆盖的优化空间；
• R³又优于 OpenEvolve，说明源码和底层配置联合调优有价值。

有些 Kernel 的加速超过3倍甚至接近4倍，但这些是分布中的高值或离群点，不能理解为所有 Kernel 都获得这种提升。

---

3. miniWeather 的逐 Kernel 对比

Figure 8 将R³与两个基线逐 Kernel 比较。

坐标含义是：

x轴：基线优化后的Kernel时间
y轴：R³优化后的Kernel时间

如果点位于：

$y<x$

则表示R³更快。

论文报告 miniWeather 的7个 Kernel 全部位于对角线下方，即R³在每个 Kernel 上都优于 Record-Replay + BO 和 OpenEvolve。

7 案例分析

7.1 QUDA 是什么

QUDA 是一个格点量子色动力学库，包含：

• 大量源文件；
• 数万行代码；
• 复杂模板；
• 大规模迭代求解器；
• 多重网格预条件；
• 很长的编译时间。

它比 MiniFE 这类 Proxy Application 更接近真实生产级 HPC 应用。

目标 Kernel 是：

CoarseDslash

它在 QUDA 多重网格预条件器中应用粗网格 Dirac 算子。

每个线程大致负责一个格点，执行：

• 邻居访问；
• Gauge 相关计算；
• Spinor 操作；
• Clover 贡献；
• Color/Spin 维度循环。

Kernel 本身加上内联设备函数达到数百行，是一个复杂真实 Kernel。

---

7.2 LLM 找到的源码优化

最终版本主要进行了以下改动。

① 提高输入数据复用

原实现可能在不同输出 Color Row 计算中重复加载同一个：

• 输入 Spinor；
• Halo 数据；
• Gauge 数据。

优化后：

某个Spin/Color-column对应的数据只加载一次
    ↓
在线程处理多个输出Color Row时重复使用

减少了：

• 全局内存访问；
• 字段对象访问；
• 重复加载。

② 将循环不变量移出热点循环

例如：

• Gauge Field 引用；
• Clover Field 引用；
• Halo Buffer；
• Source Spinor；
• Row Base Index。

从内层循环移到外层：

for (...) {
    auto base = compute_base(...); // 原来可能重复计算
    ...
}

变成类似：

auto base = compute_base(...);
for (...) {
    ...
}

减少：

• 地址计算；
• 索引计算；
• 对象解引用；
• 重复指令。

③ 暴露编译期常量

LLM还重构了部分代码，使下列信息更容易被编译器识别为常量：

• Dslash/Clover 选择；
• Warp-fission 索引；
• 条件分支；
• 固定模式参数。

然后内层 Replay Engine 的常量特化和编译器优化进一步删除无效路径。

这里体现了分层优化的价值：

LLM重构源码
    ↓
暴露常量和结构
    ↓
LLVM常量传播
    ↓
删除分支和死代码
    ↓
获得最终性能提升

只使用 LLM 或只使用 Compiler Pass 都不一定能发现完整优化链。

---

3. QUDA 最终结果

CoarseDslash Kernel：

$Speedup=1.33×$

将其放回完整 QUDA 应用，在64张 MI300A GPU 的代表性工作负载上：

$Compute Time Reduction=28.4$

作者排除了数据加载和清理时间，只统计 Compute 阶段。

搜索时间方面：

R³：108分钟
OpenEvolve预测：约42小时

由于成本太高，作者没有真正完整运行 OpenEvolve，而是根据单次评价时间进行估算。因此42小时是预测值，不是直接实测的端到端 OpenEvolve 实验结果。

8 论文创新点

创新点1：分层优化架构

将搜索空间按性质拆分：

无结构源码空间 → LLM + MAP-Elites
结构化参数空间 → Bayesian Optimization

每个源码候选都由内层调优后再评价。

---

创新点2：用 Record-Replay 支撑大规模代码进化

传统 LLM Evolution 的最大问题是评价太慢。

R³将：

完整应用构建和运行

替换为：

LLVM IR 编译 + 独立Kernel重放

使在真实大型应用上搜索成为可能。

---

创新点3：面向紧密搜索循环的 Persistent Replay Server

Record-Replay 并非论文首次提出，但R³针对数千次候选评价进行了工程优化：

• 常驻 Replay Worker；
• 内存中处理 LLVM IR；
• 状态驻留 GPU；
• 预分配缓冲区；
• 并行 BO；
• 快速正确性验证。

这是从“离线重放工具”到“自动搜索内层执行引擎”的转变。

---

创新点4：搜索吞吐优化

论文进一步针对新瓶颈设计：

• Prefix-cache-aware Prompt；
• Runtime-aware LLM Scheduling；
• 多 Kernel 并行；
• 多 GPU Replay；
• 避免过多并发造成 Population Staleness。

因此它不仅关注“找到什么优化”，也关注“如何让整个 Agent 搜索系统规模化运行”。