CUDA Kernel Optimizer

目录

一个智能体式 CUDA 内核优化器，通过代码生成、正确性检查、基准测试和迭代改进的自动化循环，将工作负载描述转化为 GPU 实现。该智能体由 LangGraph 驱动，会探索内核实现和启动配置，查询 GPU 属性，并可以查阅 NVIDIA 文档以获取优化指导，同时检查 Nsight Compute 计数器来为下一次实验提供依据。每次实验都会被记录，并保留经过验证的最快实现。

该模型可以同时修改内核代码和每个用例的启动配置。一个独立的 C++ 测试框架使用 NVRTC 编译内核，通过 CUDA Driver API 启动它们，并保存输出。Python 负责比较和候选方案选择。

工作原理

  1. 加载或生成签名、输入用例、参考内核和初始内核。
  2. 运行参考实现并评估初始实现。
  3. 提出一项修改，编译它，用 NumPy 比较输出，并测量内核延迟。
  4. 将结果反馈到下一次尝试中；在迭代预算内修复无效候选方案。
  5. 保存经过验证的最快候选方案、执行历史和计时热力图。

每个用例都必须通过验证。排名使用性能用例延迟的几何平均值；小型正确性用例不影响评分。计时默认使用 10 次预热启动和 100 次测量启动，基于 CUDA events。编译时间和分析器重放计时不计入排名。

安装配置

在 Windows 上使用 RTX 3060 Laptop GPU 开发。需要 Python 3.12+、NVIDIA GPU 及兼容的 CUDA Toolkit/驱动、CMake 3.24+、C++17 编译器，以及 OpenAI API key。下面的构建命令使用安装了 C++ 工具的 Visual Studio 2026。

在仓库根目录下：

python -m venv .venv
.venv\Scripts\python -m pip install -r optimizer_agent/requirements.txt
cmake -S cuda_test_harness -B cuda_test_harness/build -DCMAKE_BUILD_TYPE=Release
cmake --build cuda_test_harness/build --parallel

创建一个.env文件，放在仓库根目录下：

OPENAI_API_KEY=your-key-here

运行

.venv\Scripts\python optimizer_agent/optimizer_agent.py --description "矩形矩阵的单精度 GEMM。" --max-iterations 7

默认模型是gpt-5-mini，推理强度为中等。API 使用费用会计入你的账户。使用-h查看所有选项，或使用--config optimizer_agent/example.json查看包含的配置示例。

使用以下内容提供你自己的工作负载组件：--signature，--reference，--initial-kernel，以及--input-cases. 省略的组件会被推断或生成。要使用已保存的输入从先前的运行继续：

.venv\Scripts\python optimizer_agent/optimizer_agent.py --description "使用矩形矩阵的单精度 GEMM。" --input-cases results/run-001/input_cases.json --reference results/run-001/reference.cu --initial-kernel results/run-001/best.cu --max-iterations 7

此示例假设先前的运行生成了reference.cu；提供的参考被保存为supplied-reference.cu. 输入清单保留指向其二进制数据的路径，因此请保持这些文件可用。

可选标志：

  • --use-nsight：在每个有效候选之后对性能用例进行性能分析，并向模型暴露性能分析工具。需要 Nsight Compute 以及访问 GPU 性能计数器的权限。
  • --nvidia-research：在生成内核之前检索 NVIDIA 指导。

示例

在 RTX 3060 Laptop GPU 上进行 Float32 GEMM，启用 NVIDIA 研究和 Nsight Compute。

运行 022 运行 023
优化生成的起始内核。 使用相同的输入和参考，从运行 022 的最佳内核继续。
Run 022: timings and candidate changes Run 023: timings and candidate changes

结果与限制

每个会话在以下目录下获得一个目录results/run-NNN/其中包含内核源代码、请求、输入/输出数据、模型/工具响应、history.json，以及summary.json. 成功的运行会导出best.cu、每个用例的重放请求，以及heatmap.png/heatmap.svg. 启用性能分析时会保存 Nsight 报告。

这是一个针对单个内核的实验性优化器。通过提供的用例并不能证明一般正确性，生成的参考也不是独立的正确性判据。改进取决于工作负载；目前不包含与 cuBLAS 或其他厂商库的对比。在比较计时结果时，请保持 GPU 处于空闲状态。

生成的输入脚本会作为 Python 子进程在本地执行，没有沙箱。生成的 CUDA 内核也会在本地 GPU 上运行。

工作流图

下图由启用 Nsight 性能分析的已编译 LangGraph 工作流渲染而成。虚线边是条件路由。如果没有--use-nsight，评估会直接路由到下一次尝试或最终化；除非设置了--nvidia-research，否则会跳过 NVIDIA 研究。

LangGraph workflow showing preparation, model/tool calls, evaluation, repair, and finalization