KernelAgent：通过多智能体编排进行硬件引导的 GPU 内核优化

总结

近期，PyTorch 团队发布了 KernelAgent，这是一个开放的智能体系统，在全部 250 个 L1/L2/L3 KernelBench 任务中实现了 100% 的正确率。在本文中，我们通过在现有框架中增加一个硬件引导的优化层，进一步扩展了这项工作。基于此前侧重于正确性的流水线，KernelAgent 将 GPU 硬件性能信号集成到一个闭环的多智能体工作流中，以指导 Triton 内核的优化。

我们在全部 100 个 L1 KernelBench 任务上评估了 KernelAgent 生成的内核。总体而言，相较于早期版本生成的内核，它实现了 2.02 倍的加速。与默认的 torch.compile 相比，KernelAgent 平均实现了 1.56 倍的加速，在 100 个 KernelBench L1 任务中，有 65 个表现优于前者，并在 H100 上达到了 89% 的硬件屋顶线（roofline）效率。

优化代码库位于 KernelAgent 仓库，其中包含入门文档。我们还在 开源仓库中分享了精选的端到端 KernelAgent 优化产物。 

引言

对于现代 AI 工作负载而言，优化 GPU 内核变得愈发关键。随着模型规模的扩大和专业化程度的提高，性能瓶颈往往不再由高级算法决定，而是取决于实现这些算法的内核效率。然而，手动优化内核仍属于专家密集型工作，需要对 GPU 架构、内存层次结构和性能权衡有深入的了解。随着内核数量的增加以及每种新 GPU 架构对优化策略提出的新要求，这一挑战愈发严峻。

在实践中，经验丰富的内核工程师在优化内核时遵循系统化的工作流。他们使用 NVIDIA Nsight Compute 等工具分析内核，检查硬件性能计数器以诊断瓶颈，并迭代应用针对性的优化。是寄存器压力限制了占用率？是平铺（tiling）策略未能充分利用内存带宽？内核是否需要架构重构而不仅仅是参数调优？这一过程可能需要针对多种不同的内核架构进行推理，找出各自的瓶颈，最终收敛于一个能充分饱和硬件的设计。虽然有效，但这种迭代周期通常需要数天或数周时间。

现代编译器堆栈在自动化内核生成方面取得了重大进展。torch.compile 捕获计算图，并通过结合图变换、模式匹配和编译器启发式算法来生成 Triton 内核。TVM 和 XLA 等系统中类似的方案涵盖了许多通用内核模式，并提供了强大的开箱即用性能。然而，大多数编译器启发式算法是由静态模型引导的，而非源于真实硬件执行的直接测量。

KernelAgent 旨在通过将内核优化建立在真实的硬件信号基础上，实现这种诊断驱动的优化循环。它主要针对前向传递（推理）内核，因为延迟和吞吐量直接影响服务成本和用户体验。它建立在三个核心原则之上：

一切以硬件指标为基准。 无论是瓶颈诊断还是优化建议，都必须源自真实的分析数据。 

并行探索优化路径。 给定相同的硬件信号，可能存在多种有效的优化策略。KernelAgent 同时评估这些策略，缩短了实际优化时间，并将之前的方案综合进演进后的算法发现中。  

通过共享内存跨轮次学习。 优化智能体会反思每一轮的成功与失败，将见解总结进共享内存中，从而指导后续的迭代并避免重复陷入死胡同。

KernelAgent 优化工作流 

 

图 1： 优化工作流概览。优化过程始于将输入内核作为待优化的基准。ProfilerAgent（收集硬件信号）、JudgeAgent（诊断瓶颈）、AnalyzeAgent（提出优化建议）、Orchestrator Agent（综合知识）、Optimization Manager（通过多个优化智能体探索不同优化方向）、BenchmarkAgent（测量性能）。箭头表示智能体之间的数据流。  

KernelAgent 将经验丰富的内核工程师遵循的工作流程——分析、诊断瓶颈、提出优化方案及迭代——分解为一组协作智能体，实现了自动化。每个智能体负责优化循环中明确定义的阶段，它们共同形成了一个闭环的、硬件引导的反馈系统。

图 1 展示了整体工作流。从输入内核开始，KernelAgent 会重复对内核进行分析、诊断性能瓶颈、提出架构感知的优化方案、综合优化知识、并行探索替代优化路径，并测量每个候选方案。箭头指示了跨优化轮次智能体间的信息流。 概括来说，每个优化轮次包含以下阶段：

分析（Profile） → 诊断（Diagnose） → 建议（Prescribe） → 编排（Orchestrate） → 探索（Explore） → 测量（Measure）

每个阶段都会产生结构化输出，直接馈送至下一个阶段，从而实现快速、数据驱动的迭代。

数据如何在系统中流动 

分析：收集硬件信号

优化循环始于 Profiler Agent（分析智能体），它使用 NVIDIA Nsight Compute (NCU) 检查输入内核。KernelAgent 集成了 NCU 以捕获硬件级性能指标，包括但不限于 DRAM 吞吐量和利用率、L2 缓存命中率、Warp 占用率及停顿原因、计算和张量核心利用率以及光速（Speed-of-Light, SOL）指标。这些指标为所有下游决策提供了实证基础。

输入：内核代码 + 输入规范（形状、数据类型）

输出：包含硬件指标的结构化字典

输出示例

{
   "sm__inst_executed_pipe_tensor.avg.pct_of_peak_sustained_active": 0.41,
   "smsp__warp_issue_stalled_short_scoreboard_per_warp_active.pct": 5.63,
   "gpu__compute_memory_throughput.avg.pct_of_peak_sustained_elapsed": 48.86
  ...
}

诊断：通过屋顶线分析识别瓶颈

Diagnose Agent（诊断智能体） 解析分析指标，对内核的主要性能瓶颈进行分类。它使用 SOL 指标执行屋顶线分析，并结合基于 LLM 的推理进行根本原因分析。

输入： NCU 指标 + 当前内核代码

输出： 包含以下内容的 BottleneckReport（瓶颈报告）：

• 主要瓶颈类别
• 效率百分比（计算/内存 SOL 的最大值）
• 带有证据的根本原因（引用的具体指标）

诊断示例

        "category": "memory",
        "summary": "Kernel is memory-bound at 70.3% DRAM throughput with significant long scoreboard stalls from memory latency",
        "reasoning": "The roofline analysis shows Memory SOL at 70.3% while Compute SOL is only 45.2%...",
        "root_causes": [
            { "cause": "High memory latency stalls due to long scoreboard waits blocking warp execution",
              "evidence": [
               {"metric": "smsp__warp_issue_stalled_long_scoreboard_per_warp_active.pct", "value": 37.69, "interpretation": "37.7% of warp stalls are due to waiting for memory operations (global/shared memory loads), indicating memory latency is a significant bottleneck"},
                    {"metric": "sm__warps_active.avg.pct_of_peak_sustained_active", "value": 30.08, "interpretation": "Only 30% warp occupancy suggests insufficient parallelism to hide memory latency"}
                ]...
            },

建议修复：架构感知的推荐方案

针对诊断出的瓶颈，Analyzer (Prescriber) Agent（分析与建议智能体） 生成具体的、架构感知的优化建议。它结合了瓶颈分类、GPU 规格（例如 A100 与 H100 的差异）以及从精选数据库中检索到的优化模式。这使得 KernelAgent 能够根据目标硬件量身定制建议。

输入： BottleneckReport + GPU 规格 + 优化数据库 + 内核代码

输出： 包含优化建议及其理由的列表

建议示例

"recommended_fixes": [
         {"fix": "Increase pipeline depth with more stages (num_stages=4-5) and reduce register pressure by using smaller BLOCK_K or enabling register spilling to shared memory", 
    "rationale": "More pipeline stages help hide memory latency by overlapping loads with computation. Reducing register usage from 91 per thread would allow more concurrent warps to better hide the 37.7% long scoreboard stalls and improve the 30% warp occupancy"}
...
        ]

编排：将分析转化为搜索策略

Orchestrator Agent（编排智能体） 综合当前诊断结果与历史优化数据，为下一轮制定具体的搜索策略。它汇总先前的诊断、建议和结果，整合搜索策略（如束搜索、贪婪搜索等），并确定接下来要探索哪些修复方案。 

在每一轮之后，KernelAgent 会生成结构化的自我分析：诊断是否正确？修复方案是否解决了根本原因？什么有效，为什么？这些信息实现了推理时学习（inference-time learning）

输入： 建议修复方案 + 尝试历史 + 反思（Reflexion） 

输出： 最终确定的优化提示词（prompt） 

反思示例

"was_diagnosis_correct": true,
    "was_fix_effective": false,
    "expected_outcome": "...should reduce memory latency stalls by allowing more in-flight memory operations, improving memory throughput and reducing warp stalls",
    "actual_outcome": "Performance degraded significantly by 37.4% (1.0910ms → 1.4996ms)....
    "reasoning": "The fix backfired because: 1) Doubling BLOCK_N (128→256) and BLOCK_K (32→64) dramatically increased shared memory and register usage per block, likely reducing occupancy significantly....",
    "lessons": [
        "Increasing BLOCK_N and BLOCK_K together with num_stages creates compound pressure on shared memory and registers",
        ...
    ],
    "avoid_patterns": [
        "Simultaneously increasing multiple tile dimensions (BLOCK_N, BLOCK_K) along with pipeline stages",
        "...
    ],
    "try_patterns": [
        "Try smaller BLOCK_K (16 or 32) with increased num_stages to reduce register pressure while improving pipelining",
      ...

探索：并行优化 

Optimization Manager（优化管理器） 执行探索阶段。它维护 Top-K 性能表现的内核，并为每个内核生成多个优化工作线程以并行探索不同的修复方案。如果一条优化路径降低了性能，探索其他修复方案的工作线程可能会成功，从而防止搜索陷入局部最小值。每个工作线程应用不同的优化，编译内核，并将其传递给 Measure Agent（测量智能体）。

输入： 内核候选者 + 不同的优化计划

输出：用于评估的已编译优化内核 

结果示例

BeamSearch initialized: 2 kernels × 2 bottlenecks = 4 workers
---------------------------------------
Round 1: 4/4 workers succeeded
--------------------------------
Round 2: 3/4 workers succeeded
...

测量正确性与性能 

Benchmarking Agent（基准测试智能体） 为探索过程中产生的每个内核变体验证正确性并测量实际性能。对于每个候选内核，该智能体会首先对照受信任的参考实现进行正确性检查。仅通过验证的内核才会被纳入基准测试。性能测量在受控的基准测试协议下进行，以确保稳定性和可复现性。

性能测量

• 预热迭代（默认：25 次），以排除冷启动效应
• 重复迭代（默认：100 次），以获得稳定的测量值
• 共享基准锁防止工作线程之间的 GPU 资源竞争

输入： 已编译内核变体，参考实现（pytorch），测试形状输入

输出： 正确性判定结果。测得的内核运行时间

结果示例

Round 1: 4 successful, best new: 7.8000ms
Round 2: 4 successful, best new: 4.0457m
Round 3: 4 successful, best new: 3.1118ms
...

性能总结

我们使用 triton.testing.do_bench 获取一致的性能测量结果，针对 H100 上的每个内核变体报告 100 次重复运行（含 >1s 预热）的平均运行时间。具体而言，我们将 KernelAgent 与以下对象进行比较：

• 仅使用正确性循环生成的 KernelAgent 内核（我们之前的基准），
• 开箱即用的 torch.compile，指 PyTorch Inductor 默认模式，使用静态形状（未启用动态形状支持），且禁用了 CUDA 图。 

在 100 个 L1 KernelBench 问题中，相对于开箱即用的 torch.compile，KernelAgent 在 65 个任务中表现更好。总体而言，KernelAgent 实现了较早期仅关注正确性的基准 2.02 倍的几何平均加速比，相较于默认 torch.compile 实现了 1.56 倍的加速。它还在 Nvidia H100 上达到了 89% 的硬件屋顶线效率（屋顶线效率通过 Nvidia Nsight Compute 得出的计算 SOL 和内存 SOL 的最大值计算得出，即流式多处理器或内存吞吐量占硬件峰值的百分比）。 

我们在开源仓库中分享了精选的端到端 KernelAgent 优化产物。我们还在不同输入形状下测试了每个类别中的部分内核。在 12 个内核/144 个形状的测试中，我们观察到了相似的加速效果。 

关于测试时间缩放的影响： 

图 2：KernelAgent 的性能随优化轮次增加而演进。 

虽然大部分性能提升是在第一轮中实现的，这反映了硬件引导的诊断和粗粒度修复的有效性，但系统通过额外的轮次持续取得稳定进展。随着轮次增加，KernelAgent 能够通过爬山法（hill climb）实现超越初始改进的成果，优化早期方案并探索在解决主要瓶颈后才显现的次要瓶颈。这种行为突显了迭代式、反馈驱动优化的重要性。 

下面，我们提供一个端到端的案例研究，以便更好地理解 KernelAgent 在不同轮次中正在学习和应用哪些优化技术。 

案例研究：矩阵-向量乘法 (A @ x)

操作： C = A @ x
形状： M=2048, K=1,048,576
数据类型： BF16 输入, FP32 累加, BF16 输出
硬件： H100  

结果概览

PyTorch Compile 基准： 2.09 ms 

仅有正确性流水线的 KernelAgent： 9.52 ms 

LLM 基准：无硬件反馈的直接提示词。 每一轮将上一轮的输出作为输入：（顺序探索，8 轮，opus-4.5）。最佳结果：3.1985 ms  

带有优化层的 KernelAgent（4 个工作线程，8 轮，opus-4.5）。最佳结果：1.95 ms： 

图 3：逐轮内核性能：带有优化层的 KernelAgent 与无硬件反馈直接提示的矩阵-向量乘法对比

关键见解： 

1. LLM 中的启发式优化知识是有效的，例如“更大的块可以提高带宽”。然而，如果没有性能反馈，这些启发式方法会将内核带入局部最小值，且随着 LLM 无法感知它所处的性能权衡曲线，这些方法会失效。
2. 如果没有结构化的探索，LLM 会被束缚在种子内核的轨迹中。它从未考虑过将 split-K 切换为更简单的每个线程处理一行的设计，也无法超越 eager 模式的性能。
3. KernelAgent 的多工作线程探索、基于分析的方法以及反思性的知识共享机制，使其能够探索不同的替代方案并找到优化路径。 

为何基准方案缓慢： 初始的 Triton 内核使用了带有向量累加器的 2D 平铺。分析显示，该内核主要受限于寄存器导致的占用率瓶颈，因此无法发出足够的并发内存请求来掩盖 DRAM 延迟。

KernelAgent 识别出的第一个改进： 

• 瓶颈：由于寄存器压力限制了占用率，导致 SM 利用率不足。
• 建议：将大型向量累加器替换为标量累加器，每个程序处理较少行数，并增加网格并行度
• 性能： 9.52 ms → 6.80 ms。占用率增加了 8 倍，内存 SOL 从 18.5% 上升至 25.8%。
• 反思：在任何其他优化生效之前，减少寄存器状态是必要的。

# NUM_ROWS=4: four scalar accumulators instead of a vector
acc0 = 0.0
acc1 = 0.0
acc2 = 0.0
acc3 = 0.0

for k0 in range(0, K, BLOCK_K):
    # Load B vector tile once [BLOCK_K]
    b = tl.load(b_ptrs, mask=k_mask, other=0.0).to(tl.float32)

    # Process each row individually with its own 1D load
    if row_start + 0 < M:
        a0 = tl.load(a_ptr + (row_start + 0) * stride_am + offs_k * stride_ak,
                      mask=k_mask, other=0.0).to(tl.float32)
        acc0 += tl.sum(a0 * b)
    if row_start + 1 < M:
        a1 = tl.load(a_ptr + (row_start + 1) * stride_am + offs_k * stride_ak,
                      mask=k_mask, other=0.0).to(tl.float32)
        acc1 += tl.sum(a1 * b)
    # ... (acc2, acc3 similar)

# Launch config: BLOCK_K=512, NUM_ROWS=4, num_warps=4, num_stages=4
# Grid: (cdiv(M, 4),) = (512,)

KernelAgent 识别出的第二个改进： 

• 瓶颈：仍然主要受内存延迟限制；改进趋于平稳。
• 建议：为向量 x 引入有限的缓存/复用，减少冗余的全局内存流量。避免增加 num_stages，因为这先前增加了寄存器压力。
• 性能： 6.80 ms → 6.20 ms。通过对 B 向量的共享内存缓存获得了小幅提升，减少了冗余的全局内存访问。
• 反思：矩阵-向量乘法与 GEMM 的行为大不相同；平铺策略不能直接迁移。

KernelAgent 识别出的第三个改进

• 瓶颈：在降低寄存器压力后，性能受限于低效的内存事务，而非 Warp 数量不足。
• 建议：返回至 向量化 2D 加载 以获得更好的合并访问，但需谨慎控制寄存器：更小的平铺（BLOCK_M=32），更大的 K 平铺（BLOCK_K=512），以及 num_stages=1 以消除流水线寄存器开销
• 性能： 6.20 ms → 4.03 ms。 
• 反思：在任何其他优化生效之前，减少寄存器状态是必要的。

 

### Before (Step 1 approach):
```python
# Sequential scalar accumulators, NUM_ROWS=4
acc0 = 0.0; acc1 = 0.0; acc2 = 0.0; acc3 = 0.0
# ...process rows one at a time with branching...
```

@triton.jit
def matvec_kernel(A_ptr, x_ptr, C_ptr, M, K, stride_am, stride_ak,
                  BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr):
    pid_m = tl.program_id(0)
    row_start = pid_m * BLOCK_SIZE_M
    row_offsets = row_start + tl.arange(0, BLOCK_SIZE_M)
    row_mask = row_offsets < M

    # Back to vector accumulator, but only 32 elements (not 128)
    acc = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)

    for k_start in range(0, K, BLOCK_SIZE_K):
        k_offsets = k_start + tl.arange(0, BLOCK_SIZE_K)
        k_mask = k_offsets < K

        x_vals = tl.load(x_ptr + k_offsets, mask=k_mask, other=0.0)
        a_ptrs = A_ptr + row_offsets[:, None] * stride_am + k_offsets[None, :] * stride_ak
        a_vals = tl.load(a_ptrs, mask=row_mask[:, None] & k_mask[None, :], other=0.0)
        acc += tl.sum(a_vals.to(tl.float32) * x_vals.to(tl.float32)[None, :], axis=1)

    tl.store(C_ptr + row_offsets, acc.to(tl.bfloat16), mask=row_mask)

# Launch: BLOCK_SIZE_M=32, BLOCK_SIZE_K=512, num_stages=1, num_warps=4
# Grid: (cdiv(M, 32),) = (64,)

最终改进：每个程序一行（架构转变）以饱和带宽

• 瓶颈：由于寄存器压力限制了占用率，导致 SM 利用率不足。
• 建议：进行 架构变更：为每个程序分配一行。标量累加器（最少寄存器），巨大的网格并行度（2048 个程序），纯 1D 流式加载，以及大的 BLOCK_K 以摊销循环开销
• 性能： 4.03 ms → 1.95 ms。Warp 活跃度 ~95%
• 反思：该工作负载本质上是 内存带宽绑定（memory-bandwidth bound） 的。最大化占用率和并行度比平铺的优美性更重要。有时需要进行架构变更来跳出局部最优。

@triton.jit
def matvec_kernel(A_ptr, x_ptr, C_ptr, M, K, stride_am, stride_ak,
                  BLOCK_SIZE_K: tl.constexpr):
    pid_m = tl.program_id(0)
    if pid_m >= M:
        return

    # Scalar accumulator --- minimal register usage
    acc = 0.0
    a_row_ptr = A_ptr + pid_m * stride_am

    num_k_blocks = tl.cdiv(K, BLOCK_SIZE_K)
    for k_block in range(num_k_blocks):
        k_start = k_block * BLOCK_SIZE_K
        k_offsets = k_start + tl.arange(0, BLOCK_SIZE_K)
        k_mask = k_offsets < K

        x_vals = tl.load(x_ptr + k_offsets, mask=k_mask, other=0.0)
        a_vals = tl.load(a_row_ptr + k_offsets * stride_ak, mask=k_mask, other=0.0)

        prod = a_vals.to(tl.float32) * x_vals.to(tl.float32)
        block_sum = tl.sum(prod, axis=0)  # Scalar reduction
        acc += block_sum

    tl.store(C_ptr + pid_m, acc.to(tl.bfloat16))

# Launch: BLOCK_SIZE_K=1024, grid=(M,) = (2048,)
# No explicit num_warps or num_stages (defaults)

图 4：通过 KernelAgent 的改进所实现的内存吞吐量（GB/s 及 % SOL）

经验教训

我们愿意分享我们在编排多个智能体解决复杂内核工程问题时学到的经验。

问：如何在没有人工监督的情况下保持智能体不偏离轨道？ 

关键在于设定硬性、可验证的约束。在 KernelAgent 中，正确性和性能是通过门控评估强制执行的。每个内核变体必须通过数值验证，性能通过真实的硬件基准进行测量。当进度由可执行、可衡量的结果定义时，智能体就能保持在轨道上。

问：如何构建工作流，使多个智能体能够并行取得进展，同时共享上下文，以便未来的迭代能够建立在共享内存之上？ 

仅有并行性是不够的；没有协调，智能体会迅速重复工作或探索冗余路径。在每一轮内，优化工作线程独立且并行地操作，探索不同的优化策略。轮次结束后，它们的成果（无论成败）被总结到一个共享的、结构化的上下文中，捕捉了尝试的内容、成功的内容及原因。随后，该共享内存会在后续轮次中广播给所有智能体。

问：如何防止智能体陷入局部最小值，并清楚知道何时停止？  

避免局部最小值需要探索的多样性和明确的终止标准。KernelAgent 维护着一组表现最好的内核束，而非单一的胜出者。并行探索进一步降低了早期次优决策主导搜索的风险。 

具体到 GPU 优化，它可能会陷入顺序参数优化的陷阱。虽然优化 A 没奏效，优化 B 也没奏效，但当两者结合时，可能会出现性能突破。因此，KernelAgent 的目标是最大化对可提出想法的探索。 

KernelAgent 会监控性能增量和硬件利用率指标。当连续多轮未能对屋顶线效率或运行时间产生有意义的提升时，系统会判定进一步优化不太可能带来回报。

结论

KernelAgent 展示了从之前关注正确性的循环中得出的深度智能体原则（包括扎实的工具使用、并行探索、确定性控制）能够自然地扩展到性能优化领域。通过在循环中增加硬件分析和工作记忆，允许多智能体学习并探索不同的优化路径，我们可以推动已验证的内核从“正确”迈向“正确且快速”。

亲自尝试一下。 KernelAgent 是一个正在积极开发中的开源项目。我们欢迎社区的反馈、贡献和新的用例，希望这项工作能有助于推进 PyTorch 生态系统内实用、可扩展的内核优化。

致谢

我们也感谢以下人员提供的反馈：Paulius Micikevicius, Yang Wang, Lu Fang, Jie Liu, Zacharias Fisches, Alec Hammond, Richard Li, Chris Gottbrath, Davide Italiano, Joe Spisak 和 John Myles White。