KernelArc:不是单Agent死磕,而是共享记忆协同!登顶B200算子榜

KernelArc: A Multi-Agent Framework for GPU Kernel Optimization

论文原文 ↗ 论文发布 解读发布 解读:AI前沿分享

KernelArc:不是单Agent死磕,而是共享记忆协同!登顶B200算子榜 论文图示

在现代 GPU 架构上榨干硬件算力的门槛正变得空前苛刻。无论是 NVIDIA Hopper 架构引入的张量内存加速器(TMA)、异步矩阵乘累加(wgmma.mma_async),还是 Blackwell 架构支持的原生 NVFP4 精度格式与异步屏障,硬件底层的吞吐上限虽然被不断推高,但要想将这些机制转化为端到端的延迟收益,开发者必须在多维空间中进行极端精细的工程权衡。寄存器压力、数据排布、多级内存搬运、Warp 特化以及启动开销之间环环相扣,稍有不慎就会导致流水线停滞。

ArXiv URL:https://arxiv.org/abs/2608.17071v1

面对这道难题,近年来利用大语言模型(LLM)构建自动化内核研发循环(Auto-research)逐渐成为前沿方向。先前的代表性单 Agent 系统(如 AutoKernel)已经证明,让大模型通过“观察性能分析—提出假设—编辑代码—真机评测”的闭环,确实能够顺着人类预先写好的优化手册(Playbook)实现单点突破。然而,这类单 Agent 范式存在根本性的结构缺陷:它依赖单一的局部历史轨迹,很容易在固定的算法路径上陷入“过度微调”的停滞期(Plateau),误将搜索策略的瓶颈当成硬件物理极限,难以在固定计算预算内探索更多样化的算法家族或数据排布。

为了打破这种“单兵死磕”的局部最优困境,来自 IMEC 与 AILabs 的研究团队提出了名为 KernelArc 的多 Agent GPU 内核协同优化架构。该框架放弃了让单个大模型孤军奋战的做法,转而采用多个具备不同优化策略特化(Strategy-Specialized)的并行 Agent,通过“仅沉淀结论”的共享记忆机制、确定性基准哨兵以及基于性能高原触发的起草机制进行协同搜索。在 NVIDIA H100 与最新的 Blackwell B200 真实硬件上,KernelArc 不仅让 H100 上的 BF16 矩阵乘法(GEMM)达到了 766 TFLOPS、超越实测 cuBLAS 基线 3.2%,更在严格反映硬件极限制的 SOL-ExecBench 公开基准中,横跨基础算子、复杂融合块、量化算子和推理专用原语四大类别,全部斩获榜首。这项工作清晰地揭示了一个趋势:大模型在底层系统工程领域的真正威力,不在于暴力扩充单个上下文窗口去硬碰硬件细节,而在于构建一个分工明确、能敏捷共享教训与最佳实践的多 Agent 协作空间。

为什么单 Agent 优化极易撞上“性能高原”?

要理解 KernelArc 的架构动机,首先必须看清现有多模型优化方案在面对复杂工业级算子时的局限性。当前基于大模型的代码优化,多采用测量驱动的保守爬山法:设定一个初始代码种子 $k_0$,在目标硬件与软件环境 $H$ 中,针对输入工作负载集 $W$,求解满足正确性约束下的最优性能代码 $k^\star$。每当模型生成一个新的内核变体,系统就会在真实 GPU 上编译并运行基准测试。只有当候选代码既通过了正确性验证、又在真实耗时上优于当前最优解(Incumbent)时,才会被接纳为新的基准。

在单一输入形状(Shape)的简单算子场景下,这种单步递进策略依靠详尽的人类专家优化手册可以走得很深。然而,真实的深度学习算子往往需要在多个不同输入形状之间取得整体最优。以 Blackwell B200 平台上的注意力输出投影加残差(L1-030 任务)为例,基准测试要求内核在 16 个不同的输入维度下统筹评分。在面对这样多维度的评估时,单一 Agent 的搜索局限暴露无遗。

这是因为不同的矩阵形状对内存带宽和计算吞吐的要求大相径庭:某些形状极其依赖张量核心流水线的展开,某些形状则受制于访存对齐与缓存命中率,甚至不同形状在底层可能分别适合不同的 cuBLASLt 静态配置、融合策略或精度模式。单 Agent 在固定的连续上下文中,只能沿着一条狭窄的优化路线线性摸索。当这条路径的优化红利耗尽后,模型会在附近几个微小的超参数上反复打转、盲目微调,耗尽算力与 Token 预算,最终被困在性能高原上。由于模型无法跳出当下的局部历史,它很容易把搜索路线的枯竭误判为硬件算力的物理天花板。

这一现实催生了 KernelArc 的核心设计诉求:系统必须能够同时向不同的优化方向发散,能够跨探索路径传递验证过的成败经验,同时必须保证评估与决策过程完全确定,并在检测到陷入高原时提供明确的逃逸机制。

KernelArc 核心机制:策略特化与共享记忆的四重协同

为了将生成式大模型的发散探索能力与底层系统对严苛正确性的要求结合起来,KernelArc 在架构上做出了关键切割:将发散式的代码生成推理交给 LLM,而将评测运行、状态记录、版本回退与协同调度交由确定性的外部程序控制。具体而言,该框架由四个相互咬合的模块构成。


                  ┌──────────────────────────────────────────────┐

                  │          轻量级策略透镜 (8 种策略技能)          │

                  └───────────────┬──────────────────────────────┘

                                  │ 注入先验

               ┌──────────────────┴──────────────────┐

               ▼                                     ▼

        ┌──────────────┐                      ┌──────────────┐

        │   Agent 1    │                      │   Agent 2    │

        │ (擅长库级配置) │                      │ (擅长自定义PTX)│

        └──────┬───────┘                      └──────┬───────┘

               │ 读/写经验                            │ 读/写经验

               ▼                                     ▼

    ┌──────────────────────────────────────────────────────────────┐

    │        仅存结论的共享记忆 (Conclusions-Only Shared Memory)       │

    │        - 成功的结构与参数 (Wins)  - 失败的陷阱模式 (Traps)         │

    │        - 可配置记忆保留窗口 (Retention Horizon)                │

    └──────────────────────────────────────────────────────────────┘

               │ 提交候选代码                          │ 提交候选代码

               └──────────────────┬───────────────────┘

                                  ▼

        ┌──────────────────────────────────────────────────────────┐

        │        确定性评估哨兵 (Deterministic Benchmark Guard)       │

        │        1. 快速级联门控 (语法/依赖/越界过滤)                 │

        │        2. 真机性能评测 (SOL-ExecBench, 锁定频率与清空缓存)    │

        │        3. 动作分发: REJECT / KEEP / ACCEPT / REVERT      │

        └─────────────────────────┬────────────────────────────────┘

                                  │ 触发高原逃逸

                                  ▼

        ┌──────────────────────────────────────────────────────────┐

        │   高原触发起草 (Plateau-Triggered Drafting)                │

        │   - 连续无收益计数达标 -> 强制切换算法/DSL/数据布局           │

        │   - 只读查阅共享活动排行榜 (Read-Only Cross-Agent State)    │

        └──────────────────────────────────────────────────────────┘

第一是策略特化的并行搜索(Strategy-Specialized Parallel Search)。KernelArc 并不向每个 Agent 分配完全相同或互斥的死板任务,而是提供了 8 种“策略透镜”(Strategy Skills)。这些技能以轻量级的 YAML 配置文件和典型优化模式模板存在,涵盖了高级库自动调优、内存排布重构、深度操作符融合、低精度格式映射以及硬件原生 PTX 编写等不同层面。系统配备了一张适用性过滤表(Eligibility Map),例如硬件原生手写指令仅被允许用于 GEMM 类算子,而库参数优化则通用得多。通过给并行运行的多个 Agent 赋予不同的策略技能,系统从源头上遏制了所有 Agent 迅速退化并扎堆于同一种显而易见的传统解法。

第二是仅保留结论的共享记忆(Conclusions-Only Memory)。如果将每个 Agent 冗长的思考上下文、重试计数和完整代码全量共享,不仅会迅速挤爆大模型的有效上下文窗口,更会导致无效信息的上下文污染。KernelArc 采取了一种极度精炼的记忆策略:Agent 之间通过带文件锁的共享存储仅交换结构化的“成功经验”(Wins)与“错误陷阱”(Traps)。成功经验记录了被证实带来加速的参数或代码结构,陷阱则记录了导致编译失败、数值溢出或性能剧烈衰退的无效路径。通过设置字符上限(例如策略摘要不超过 160 字符,经验反思不超过 1200 字符)以及可配置的时间衰退窗口(Retention Horizon),共享记忆只向模型呈现紧凑而高价值的信息,而迭代计数、保底归档等元数据全部留在确定性宿主代码中。

第三是外置的确定性基准哨兵(External Deterministic Guard)。在内核优化中,大模型绝对不能既当选手又当裁判。KernelArc 将评估权限完全收拢到一个独立进程中。当 Agent 提交新的代码候选时,首先必须穿过一个级联门控(Cascade Gate),廉价而快速地筛除语法错误、非法系统调用或越界脚本,避免宝贵的 GPU 测时资源被浪费。随后,合规代码被送入受控的硬件环境进行多轮冷缓存基准测试。哨兵只向 Agent 返回四种高度明确的操作指示:

第四是基于高原触发的起草机制(Plateau-Triggered Drafting)与只读对标。每个 Agent 内部都维护着一个停滞计数器 $r$。只有在经历了真实 GPU 评测但未带来性能提升时(即收到 ACCEPTREVERT),计数器 $r$ 才会递增;一旦命中 KEEP 则立即清零。当 $r$ 累加达到阈值 $r_{\mathrm{draft}}$ 时,系统判定当前方向陷入局部最优,此时强制触发逃逸指令:要求该 Agent 彻底放弃当前代码的局部微调,转向完全不同的算法实现、换用另一种领域特定语言(DSL)或更改输入张量的数据排布。在这个过程中,陷入瓶颈的 Agent 拥有只读权限去检索内部活动排行榜上其他 Agent 沉淀出的优秀快照,以此为灵感草稿重新出发,实现多路 Agent 之间的优势重组。

极端深度的单点验证:Hopper H100 上的 GEMM 演进

为了验证该优化循环究竟能把单个硬件原语的算力压榨到何种地步,作者首先构建了一个受控的极端边界实验:在配备 132 个流式多处理器(SM)、HBM3 显存带宽 3.35 TB/s、理论 BF16 算力约为 989 TFLOPS 的 NVIDIA H100 SXM 硬件上,优化一个固定的 $4096 \times 4096$ 矩阵乘法。

在这项实验中,单 Agent 被赋予了一份结构完备、由浅入深的 Hopper 专属优化课程手册(Curriculum),包含从硬件先验、剖析器指标到典型报错模式的完整指引,但完全不提供最终代码或固定的调度参数。Agent 从一个教科书级别、几乎没有做优化的原生 CUDA 内核出发,在 8 小时的物理墙上时钟预算内自主迭代。整个过程没有任何人类工程师插手修改代码。


[4.6 TFLOPS]   初始朴素 CUDA 内核

      │

      ▼ (引入张量核心基元、WGMMA、异步拷贝与多级流水线)

[541 TFLOPS]   初具规模的硬件专用实现

      │

      ▼ (生产者-消费者 Warp 特化: 1个 Warp 组负责 TMA 搬运, 2个负责 WGMMA 计算)

[617 TFLOPS]   彻底解耦搬运与计算吞吐

      │

      ▼ (异步屏障 mbarrier 细粒度同步、Epilogue 尾部暂存与等待深度放宽)

[720 TFLOPS]   逼近供应商高度优化的底层实现

      │

      ▼ (Descriptor 缓存、K 轴循环栈开销消除、256B 扇区缓存提示)

[766 TFLOPS]   最终产出 (超越 cuBLAS 实测基线 3.2%, 达到硬件峰值 77%)

这一连贯演进轨迹充分展示了该循环对底层硬件特性的消化能力。从最初的 4.6 TFLOPS 起步,内核在经历张量核心转换与软件流水线化后冲上 541 TFLOPS;随后,模型成功实现了极具挑战性的生产者-消费者 Warp 特化(Warp Specialization),将一个 Warp 组专门分配给 TMA 数据加载,另外两个 Warp 组专注作为 WGMMA 计算单元,通过异步屏障 mbarrier 的相位追踪协调节奏,瞬间将吞吐推至 617 TFLOPS。经过对尾部写回机制和等待深度的进一步松弛,吞吐达到 720 TFLOPS。

在最后的冲刺阶段,Agent 甚至自主推导并探索出了一些优化手册中未作显式规定的底层技巧,例如在多次重复调用的评测 harness 下对张量描述符(Descriptor)及转置矩阵进行缓存、彻底消除 K 循环的栈展开开销、将 L2 缓存提示配置为 256 字节扇区等。最终,该内核锁定了 766 TFLOPS 的有效算力,超越了在同一评测基准下实测的 cuBLAS 最佳基线(742 TFLOPS)达 3.2%,达到了 H100 理论硬件峰值的 77%。

这一单点实验有力地证明了观察-编辑-评测闭环的深度上限,但它同样揭示了大模型优化的现实:这样的深度是在极其详尽的人类手册引导、极具针对性的输入形状以及连续 8 小时的单一算子打磨下才实现的。当面对跨多种形状、跨多种算子架构的复杂全景时,单一纵深路径就会寸步难行,必须引入多 Agent 的广度协作。

与此同时,研究团队还考察了不同底层大模型作为优化核心时的性价比表现。横向对比 7 种模型基座后可以发现:Claude Opus 5 凭借极强的逻辑推理与复杂的 PTX/CUDA 代码驾驭能力,最终跑出了最高的 766 TFLOPS 成绩,但其调用 Token 成本也最为高昂;而国产大模型 Kimi K3 则展现出了极佳的性价比,在消耗仅仅几美元的成本下就迅速冲破了 700 TFLOPS 大关,并在 718 TFLOPS 附近平稳收敛。这表明在实际系统落地时,完全可以通过策略性地混合不同档次的大模型来平衡研发预算与性能收益。

广度协作突破:Blackwell B200 与 SOL-ExecBench 全品类登顶

为了在真正具有挑战性的广谱工作负载下检验多 Agent 架构的实用价值,研究团队将 KernelArc 部署到了由 NVIDIA 主导的权威评测基准 SOL-ExecBench 上。该基准在锁定 SM 时钟频率、严格冲刷 L2 缓存并借助 CUPTI 工具精准采样的 Blackwell B200 GPU 上运行,包含 235 个深度学习核心算子。其评分机制(Speed-of-Light, SOL)以理论硬件物理上限为 1.0,根据内核在多个动态输入形状下的综合耗时进行加权汇总,对算子的通用泛化性能提出了极高要求。

SOL-ExecBench 被划分为四大难度区间:代表单算子基础构件的 L1(如 GQA、RMSNorm、SwiGLU)、复杂度高出 3 到 10 倍的端到端融合模块 L2、显式引入 FP8 块级缩放和原生 NVFP4 低精度的 Quantization(Q)分类,以及面向推理服务核心原语的 FlashInfer-Bench(FI)。

在 L1 类别的经典任务 L1-030(带有残差相加的注意力输出投影,需在 16 个各异的形状上综合评分)上,单 Agent 路径的弊端彻底暴露。由于不同形状之间的最佳配置完全互斥,单 Agent 很快在 SOL 评分为 0.441 处彻底陷入停滞,反复调整局部参数均无法寸进。而启用了双 Agent 协同的 KernelArc 迅速破局:一个 Agent 探索库级参数,另一个 Agent 尝试操作符内存融合。最终系统提交的优化方案极为精妙——它并没有死磕手写 GEMM 主循环,而是直接将输出张量与残差缓冲区进行了内存别名化(Aliasing),使 cuBLASLt 能够顺理成章地将残差加法完全折叠进 GEMM 的尾部处理(Epilogue)中;同时,根据 16 种不同形状的 $M$ 维度切分,建立了一张静态的 cuBLASLt Expert-API 配置查询表,并在未命中时平滑回退。这种“内存级融合 + 动态配置分表”的复合设计,让 KernelArc 直接将该任务的分数提升至 0.481(在后续迭代中进一步达到了 0.4905,耗时压缩至 0.036 毫秒)。

跨类别的实测数据印证了多 Agent 体系在多样化算子上的强悍迁移能力。在 2026 年 7 月 30 日记录的公开排行榜截屏中,KernelArc 产出的实现在所有四个代表性任务上均拿下了第一名:

  1. L1-030(GEMM + 残差):由 Claude Opus 4.6 驱动,通过残差别名折叠与静态配置表,在 16 个输入形状下显著打破单 Agent 僵局;

  2. L2-025(MoE 反向传播融合):由 Kimi K3 驱动,成功重组了专家混合架构反向传播时稀疏梯度的写回与带宽绑定,登顶复杂融合类别;

  3. Q-031(NVFP4 分组查询注意力):由 Kimi K3 驱动,精准利用 Blackwell 的 SM100 指令集实现了原生 NVFP4 块级缩放张量核心的高效利用,登顶低精度量化类别;

  4. FI-014(分页 Causal GQA Prefill 算子):由 Kimi K3 驱动,针对超长上下文 Prefill 阶段的不规则内存访问,设计了专门的分页访存对齐与共享内存缓冲,登顶推理评测类别。

这些生成的最终代码形态各异:有高度精炼的底层 PTX 嵌入式汇编,有纯粹通过重塑生命周期的操作符融合,也有针对驱动层 API 的精细参数编排。这充分证明了 KernelArc 并不是一个死记硬背某种特定算子模板的脚本工具,而是一个真正具备在多条技术路线上自主权衡并合成复杂工程解法的系统框架。

协同机制消融:为什么共享教训能撬动更高的优化上限?

多 Agent 系统的设计往往容易伴随争议:算力收益究竟来自于 Agent 数量增多带来的“算力堆叠”,还是真正来自于架构层面的协同设计?为了剥离混杂因素,研究团队针对具有代表性且算力优化空间最充沛的推理任务 FI-014(分页 Prefill GQA)展开了严格的等预算消融实验。

实验将所有方案的探索尝试次数严格限制在 100 次候选提交之内。这意味着,无论是单 Agent 还是两个 Agent 并行,耗费的物理尝试预算完全对等(双 Agent 配置下每个 Agent 仅分得 50 次尝试机会,编译崩溃或门控拦截同样扣除计数)。研究人员在此约束下对比了三组系统变体,每组重复运行 5 条独立轨迹,并采用几何平均值和 95% 置信区间消除随机波动:

在私有记忆的单 Agent 对照组(Single)中,轨迹很快就在前 30-40 次尝试后进入缓慢的饱和区,最终获得的端到端加速比显著受限。

而在两个引入了共享记忆的多 Agent 组中,搜索效率呈现出了系统性的跨越。其中,Multi-bounded 组采用上限为 16 条记录的受限共享记忆并开启活动排行榜只读查阅,Multi-unbounded 组则不设记忆上限。实验统计数据表明,无论是受限还是不受限的共享多 Agent 架构,其几何平均加速比均大幅领先于单 Agent 方案。

更深入的轨迹分析揭示了协同机制起作用的关键时机:共享机制并不是在搜索的每分每秒都以均等的力量发挥作用,其核心价值集中在早期探索高原逃逸两个关键节点。在初期,一个 Agent 踩中的编译错误或无效参数布局(Traps)会被迅速沉淀,另一个 Agent 检索后会直接将该类方案从可选掩码中剔除,从而避免了两个人掉进同一个坑里消耗宝贵的尝试预算。在中后期,当某一个 Agent 连续多次尝试未果、触发高原起草(Drafting)时,它能够借助只读机制立刻吸纳同伴探索出的较优代码结构作为起点,在完全不同的策略透镜下展开突变。这种“一人遇挫,全员换道”的机制,使得系统在固定的候选预算内,能够真正覆盖更多本质不同的算法家族,最终触达更具统治力的性能前沿。

值得注意的是,在进一步的探索中,研究团队发现当在固定的 100 次尝试预算下将 Agent 数量从 2 个继续扩充至 4 个时,并没有在 FI-014 任务上观察到稳定的额外收益。这一结果给后续研究敲响了警钟:多 Agent 协同的核心在于信息互补与精准逃逸,盲目增加 Agent 并行度如果缺乏更深层次的策略正交性,反而可能因为分摊了有限的尝试预算而削弱单条路径的深挖质量。

从手写内核到自主协同的演进分水岭

回顾整个 GPU 性能工程的发展脉络,早期的 TVM、Ansor 等传统编译器主要在预设的调度树或中间表示(IR)空间内进行基于规则或简单统计模型的自动调优,无法自由进行任意源码级的架构重写与硬件专属 PTX 逻辑编排。而后来的单 Agent 自动化系统虽然引入了大模型的代码生成灵性,却又深陷在狭窄历史视野与频繁陷入停滞的局部泥潭中。

KernelArc 的探索给高性能计算与 AI 系统的交汇点提供了一个极具参考价值的样本:底层性能优化不必在“僵硬的编译器自动化”与“昂贵脆弱的人工调优”之间做非此即彼的二选一。 通过将发散生成赋予大模型,将确定性约束、基准防线与版本控制交由严格代码,并辅以轻量透明的经验共享机制,大模型完全有能力在极其复杂的现代 GPU 架构上,自主寻找出比肩甚至超越人类专家与商业供应商库的系统级实现。

当底层芯片架构的迭代速度越来越快、编程模型复杂度呈指数级上升时,这种由多样化策略驱动、具备自我反思并能集体逃逸局部最优的多 Agent 框架,或许正是未来高性能计算自动化破局的关键通路。