KDA²驱动Delta Attention CUDA内核优化：从串行递推到2.4倍加速

1. 先说清楚问题：Delta Attention 的“delta”到底增量在哪

1.1 从线性注意力的递推说起

Kimi 的 Delta Attention 不是一个新概念，但它确实把“增量”这两个字刻在了骨子里。传统 Softmax Attention 的思路是每次解码都拿当前 Query 去和全部历史 Key 做点积，再对 Value 做加权求和——这是一个典型的“全量重算”模式。而 Delta Attention 走的是线性注意力的递推路线：维护一个隐状态矩阵，每来一个新 token，就基于一个带衰减的更新规则把该 token 的信息“叠”进去，后续计算 Query 时只需要读取这个隐状态，而不需要再扫描整个 KV 历史。

听起来很美，对吧？但这个东西一旦落到 CUDA kernel 层面，就立刻暴露出一个尴尬的事实：递推本身是高度串行的。每一步更新都依赖上一步的隐状态，你没法像 FlashAttention 那样把一个长序列切成若干独立的 block，分给不同 SM 并行算完再合并。这是 Delta Attention 与生俱来的“硬伤”，也是它让内核设计变得不愉快的根源。

我在做这个项目之前，先在 PyTorch 里写了一个朴素版本：一个 for 循环，每次迭代更新一个二维隐状态矩阵。功能完全正确，数值也对得上，但性能惨不忍睹——2000 长度的序列在 H100 上要跑 1 毫秒以上。当时我就意识到，真正的性能瓶颈不是算法本身，而是这个递推循环的 kernel 实现方式。

1.2 增量版本为什么反而很难写内核

有人会问：既然递推是串行的，那把它写成 kernel 是不是很简单？不就是把 for 循环搬到 GPU 上吗？问题就在这里——把 for 循环搬到 GPU 上恰恰是最大的陷阱。

CPU 上的 for 循环，每一步只处理一个 token 的矩阵运算，循环次数即序列长度。GPU 上如果照搬这种结构，你会发现：

• SM 利用率极低。每一步只启动一个小规模矩阵运算，大量线程在等待状态同步，GPU 的并行度完全发挥不出来。
• 隐状态存哪里？它必须在整个递推过程中保持可见。如果放全局内存，每一步读写都伴随几百纳秒延迟；如果放共享内存，每个 SM 只能处理一个序列片段，多头并行时又需要跨 block 通信。
• 数值累积问题。隐状态矩阵是不断被“更新”的，每一步都有乘加和归一化操作。浮点误差会随着序列长度递增而累积，如果不做补偿，最终结果和参考实现会越差越远。

这些问题的组合，使得 Delta Attention 没有一个现成的 kernel 模板可以套。FlashAttention 的经典分块策略在这里直接失效，因为块与块之间存在严格的时序依赖。CUTLASS 的 GEMM 模板也不能直接用，因为它的设计前提是数据可以被自由切分并行。

1.3 我们用的评测基线：H100 上的 PyTorch 与 FlashAttention 对比

在引入 KDA² 之前，我需要一个可靠的性能基线。当时我手上的环境是 H100 SXM 80GB，CUDA 12.4，PyTorch 2.3。我对比了两个参照实现：

第一个是 PyTorch 原生实现——就是我们上面说的 for 循环版本。第二个是“伪 FlashAttention 化”的版本：把序列切块，每一块内部用 FlashAttention 逻辑并行计算，块之间保留递推依赖。后者写起来更复杂，但至少能部分利用 GPU 并行能力。

我在四种配置下测了一组基准，输入维度统一为 8 个注意力头、每头 128 维，序列长度分别是 512、1024、2048、4096：

这个结果让我很困惑：理论加速比稳定在 1.8 倍左右，但无论怎么调 block 大小、换 tile 策略，都突破不了 2 倍。原因在于——伪 FlashAttention 版本只是把“块内并行”做出来了，块间的递推步骤仍然是串行的，全局内存读写更是被反复执行了多次。换句话说，瓶颈已经从“计算”转移到了“数据搬运”。

这时候我才下定决心，开始研究如何用 Kernel Design Agents 来设计真正适合这个问题的内核。

2. KDA² 是怎么拆问题的：四层设计代理的分工

2.1 Planner：把大 kernel 拆成可验证的功能块

KDA² 给我的第一个启发是：不要让一个 Agent 试图一口气生成整个 CUDA kernel。一个 Delta Attention kernel 涉及访存模式、递归状态管理、数值补偿、block 间同步等多个逻辑层面，任何单一模型在长上下文中都容易丢失细节，生成出来的代码往往是“能编译但跑不对”。

KDA² 的做法是引入一个 Planner Agent，它不写任何 CUDA 代码，只负责拆解问题。它拿到一个需求描述，会输出一个功能块分解树：

• 阶段 A：KV 输入管线。决定历史 Key/Value 怎么从全局内存搬到共享内存，用什么数据布局。
• 阶段 B：隐状态递推。决定每个线程块如何处理递归依赖，是否引入 block 间协作。
• 阶段 C：Query 输出计算。把当前 Query 与最新隐状态做矩阵乘，得到注意力输出。
• 阶段 D：数值纠偏。针对累积误差设计补偿逻辑。

每个功能块被赋予明确的输入输出接口、边界条件和可验证的单元测试。Planner 的输出不是一行 kernel 代码，而是一份“设计规格书”。

实际使用中我发现，这一步对最终质量的贡献甚至超过了代码生成本身。没有清晰的功能边界，后面所有 Agent 都会陷入无休止的“改一处坏了另一处”的循环。有了规格书，每个功能块可以被独立生成、独立测试、独立替换。

2.2 Kernel Writer：CUDA 与 PTX 的双轨生成

Kernel Writer 是真正写代码的 Agent。它的输入是 Planner 输出的功能规格书，输出是 CUDA C++ 代码，但在关键逻辑处会同时生成一段 PTX 汇编作为参考实现。

这里有一个值得注意的设计细节：KDA² 并不是只用 CUDA 编译器去优化。它会让 Writer 先生成一份“教科书式”的 CUDA 版本，这份版本的代码清晰但性能一般；然后再让 Writer 针对性能热点生成 PTX 级优化版本。两份代码都会被编译、跑同样的正确性测试。

为什么这么做？因为 CUDA C++ 的编译器优化有时会“好心办坏事”——比如把某些原子操作重排、把循环展开成超长指令序列。当你怀疑编译器做了非预期的事情时，有 PTX 参考版本对照，可以快速定位问题出在源代码层面还是编译层面。

在我们的场景里，PTX 版本主要用来验证两个关键操作：隐状态更新的乘加融合（FMA 是否真的被融合了），以及共享内存 bank conflict 的规避是否在指令层面真实生效。这些是纯 CUDA 代码难以感知的底层行为。

2.3 Verifier：仿真器与 Nsight Compute 的自动复核

Verifier Agent 负责“挑刺”。它干三件事：

第一，跑正确性测试。用随机生成的输入数据对比 Agent 版本和 PyTorch 参考实现，设置atol=1e-2, rtol=1e-2的宽松边界——注意，是宽松边界，不是严格边界。Delta Attention 的数值特性决定了它不可能做到和 FP32 逐位一致，所以我们一开始就用误差允许范围来约束，而不是咬死精度。

第二，跑 Nsight Compute 的 profile 分析。Verifier 会自动解析ncu输出的性能指标，重点看共享内存 bank conflicts、寄存器溢出（local memory spill）、warp 占用率、全局内存吞吐这几项。如果发现 bank conflicts 次数超过阈值，Verifier 会生成一条结构化反馈：哪个循环、哪条访存指令、涉及哪些数组，然后返回给 Writer 修改。

第三，跑指令级回归。这一步是我觉得最聪明的部分。Verifier 会把编译后的 SASS（机器汇编）反汇编出来，检查关键循环中是否存在非预期的分支指令或额外的内存访问。有一次我们发现 CUDA 编译器把我们的共享内存读操作优化成了“先读全局再算地址”的方式，性能完全不正常，靠 SASS 级检查才定位到问题。

2.4 为什么不让一个 Agent 直接写完整内核

我最初的想法是：直接把需求扔给一个 Agent，让它输出整个 kernel 文件，我负责编译测试和修改。试了一轮之后我放弃了这个想法——不是因为生成不了正确代码，而是因为迭代成本太高。

一个完整 Delta Attention kernel 大约 400-600 行 CUDA 代码。任何一个小改动，比如把共享内存 tile 从 32x32 改成 64x16，都可能引发连锁反应：访存模式变了、bank conflicts 变了、寄存器分配变了、甚至 block 调度都可能受影响。如果只靠一个 Agent 在“生成代码→报错→修改→再报错”的循环里打转，它很快就会迷失方向。

KDA² 的四层分工解决了这个问题。Planner 稳定了整体结构，Writer 专注于局部实现，Verifier 提供自动反馈，最后还有一个 Tuner Agent 负责调整参数组合。每个环节的职责清晰，修改可以局部化，迭代路径是可控的。

3. 三版内核的演进：从“能跑”到“能看”

3.1 V1：初版直接生成，KV 流水线并行

V1 是 KDA² 生成的第一个可用版本。它的思想是：在递推循环内部做流水线并行——每个 block 负责一段连续的 token 序列，前一个 block 计算出的隐状态通过全局内存传递给下一个 block。

这个版本的优点是好写、好调试，逻辑结构与 PyTorch 参考实现一一对应。但性能只比“伪 FlashAttention”版本好了 15% 左右，无限接近 2 倍的加速比天花板。Verifier 的 profile 报告给出了清晰的瓶颈定位：

• 全局内存读写带宽利用率只有 42%，大量时间花在等待隐状态数据同步。
• 每个 block 处理 128 个 token，递推步内有大量空转的 warp——因为隐状态矩阵的更新依赖是串行的，同一时刻只有一个 warp 在活跃计算。

V1 的意义不在于性能，而在于它验证了 KDA² 工作流的正确性。从问题描述到可运行的 kernel，Agent 只用了不到 20 分钟，正确性测试一次通过。这让我确认了整体方案可行——接下来的两个版本都在 V1 的骨架上做深度优化。

3.2 V2：增量 softmax 统计量的归一化处理

V2 的改进聚焦在数值精度上。Delta Attention 的隐状态更新中包含一个归一化因子——你可以类比成 softmax 中的累计 exp 和。每一次递推，累计量被更新一次。问题在于，当序列长度超过 1000 时，累计量的数值范围可能跨越十几个数量级，FP32 的精度开始不够用。

V2 中 KDA² 引入了“增量归一化”策略：不再直接维护大规模的累计量，而是把累计量拆成“当前主值”和“补偿项”两个部分。每次更新时，先把补偿项叠加进主值，再基于主值做归一化。这个技巧在数值计算领域很成熟，但在 CUDA kernel 里实现时有一个隐含的坑：FMA 融合可能导致补偿项被“吞掉”，让补偿逻辑白干。

Kernel Writer 在这里做了一次漂亮的改动——通过内联 PTX 把补偿项的计算强制拆分成独立的乘法和加法，阻止编译器把它融合进 FMA。改完之后，2048 序列长度的最大误差从 1.8e-2 降到了 4.2e-3，误差增幅原本随序列长度线性上升，现在变成了对数增长。

V2 的性能也比 V1 提升了一截：递推步骤内的空转减少了，因为归一化逻辑不再需要大量的额外寄存器，Warp 占用率从 56% 提高到了 72%。

3.3 V3：融合输出缓冲与跨 block 的原子更新

到了 V3，我们开始动更核心的刀。V2 仍然是“分块递推，块间传递隐状态”的结构，全局内存隐状态的读写依然是性能大敌。V3 的目标是把递推过程中的隐状态尽量留在片上——共享内存或寄存器中——避免每步都访问全局内存。

KDA² 的 Planner 在这个阶段做了一个很精巧的拆解：它把隐状态矩阵按“行”切分成两个部分：一部分是当前 block 独立维护的行，这些行可以完全留在共享内存里；另一部分是需要跨 block 协作的行，只在必要的时点做全局同步。

进一步地，V3 引入了一个全局原子计数器，用来协调跨 block 的隐状态合并。核心逻辑是：每个 block 计算完自己负责的片段后，将结果以原子方式累加进全局隐状态，并判断“所有 block 是否都完成了”？如果是，则触发一次全局隐状态的同步刷写。这个计数器不是每步都更新，而是每若干个递推步骤更新一次——我们管这个叫“延迟刷写”。

V3 的性能表现终于有了质的飞跃：2048 序列长度的端到端 kernel 时间从 V2 的 472 μs 降到了 386 μs，相比于 PyTorch 原生版本提速约 2.4 倍。

这里有一个值得注意的现象：V2 的耗时反而比 V1 高了。这不是性能倒退，而是因为 V2 的数值补偿逻辑增加了每步的计算量。如果我们只需要跑短序列，V2 的精度优势发挥不出来；但如果序列拉长到 4096，V2 的优势就开始显现。这也引出了一个问题：性能优化和数值精度需要分开评估，不能只看单一指标。

4. 真正改变性能的六个 Hackings

4.1 用 cp.async 的“搬迁”思路，而不是“预取”思路

大部分 CUDA 教程讲 cp.async 的时候，都强调它是用来“预取”——提前把下一块数据搬到共享内存，隐藏全局内存延迟。但在 Delta Attention 的 V3 里，我们的用法反过来了：把它当搬迁工具用，专门腾出寄存器空间。

具体场景是：递推循环内部，隐状态矩阵在寄存器组和共享内存之间来回切换。当一块数据计算完毕，不再需要占用寄存器时，直接用 cp.async 把它“扔”回共享内存，而不是让寄存器空着等编译器来释放。这个小改动让编译器的寄存器分配意外地宽松了很多——原本需要 168 个寄存器的 kernel，压到了 128 个，占用率从 56% 提升到了 75%。

Verifier 的 profile 报告里有一个非常明显的下降：local memory spill 从 34 次降到了 2 次。寄存器溢出一直是 kernel 性能杀手，而这个 hack 的本质就是“在寄存器不再需要的时候立即转移”，而不是“在需要的时候再规划”。

4.2 Striped 布局避免增量写入的 bank conflict

Delta Attention 的隐状态更新有一个特点：它不是重写整个矩阵，而是在旧值的基础上做增量修改。这意味着每次更新都要先读旧值、算新值、写回。如果不同 warp 同时访问共享内存的同一 bank，就会发生 bank conflict，性能直接翻车。

我们的第一版实现里，隐状态矩阵用的是标准行主序布局。实测发现，8 个 warp 同时更新隐状态时，bank conflict 率高达 37%。Verifier 建议把矩阵改成“条带布局”（striped layout）：每个 warp 负责的列被分散到不同的 bank 上，这样并发访问的 bank 就错开了。

具体来说，我们把原本 128 维的每个头拆成 4 个条带，每个条带 32 维，按warp_id * 32 + stripe_id的方式重新映射共享内存地址。这个改动让 bank conflicts 的占比从 37% 降到了 4%，kernel 整体耗时减少了约 12%。说实话，这类 hack 如果你不亲自跑ncu看报告，很难凭直觉想到——因为它在 CUDA C++ 代码层面只是共享内存索引方式的改变，看起来无关痛痒。

4.3 双缓冲的寄存器水位控制：给编译器一个台阶下

说到编译器，我学到的另一件事是：编译器不是万能的，它的优化策略在复杂循环中经常会退化成保守模式。我们的递推循环里有大量的条件分支（判断是否该延迟刷写），这些分支的存在会让编译器不敢大胆地做寄存器重排和指令调度，最终生成很“肉”的代码。

V3 的给编译器的“台阶”是双缓冲：把隐状态矩阵拆成两个缓冲区，一个负责“正在计算”的数据，一个负责“即将写入共享内存”的数据。这样做的最大好处是——循环体内的数据依赖被显式打破。编译器可以看到两个缓冲区互不干扰，从而放心地做指令级并行。

光是这个改动，就让循环体内的 IPC（每周期指令数）从 1.2 提升到了 1.6，内核耗时又降了一截。这个 hack 的思路是：不要试图告诉编译器“你可以优化这里”，而是通过代码结构设计，让优化变成一个显然而易见的选项。

4.4 延迟 flush 而不是立即 flush 的全局计数器

我前面提到 V3 引入了一个跨 block 的全局原子计数器。最初实现是每次递推都检查“所有 block 是否完成”，结果原子操作的开销吃掉了大部分性能收益。

后来我们改成“延迟 flush”策略：计数器每 8 步才检查一次。在这 8 步之间，各个 block 的隐状态增量被缓存在共享内存的额外区域里，等到 flush 时一次性合并。这样原子操作的频率降低了 8 倍，而递推逻辑的正确性不受影响。

这个 hack 的代价是共享内存占用提高了 12%，但收益非常可观——全局内存原子操作的次数从每步 1 次降到了每 8 步 1 次，相关开销从总耗时的 8.6% 降到了 2.1%。我在最初设计时完全没考虑到这个维度，是 Verifier 在分析 stall 原因时发现的：大量的st.global.atomic指令在等待 L2 缓存仲裁。

4.5 把错误的 kernel 变成诊断器的妙用

这个 hack 可能有点“歪门邪道”，但在实战中帮了大忙。KDA² 的 Verifier 会频繁把“失败”的 kernel 标记为错误——正确性测试没过，或者性能指标不达标。大多数情况下这些 kernel 会被直接丢弃，但我发现了一个有价值的用法：保留一个“故意写错”的 kernel 版本，专门用于诊断。

这个版本会把隐状态更新逻辑中的一处依赖故意漏掉，比如忘记乘衰减因子。在 Nsight Compute 里，它的行为会立刻暴露一个问题：哪些 warp 在极高延迟下空转，哪些计算单元利用率低下。对比“正确的慢版本”和“错误的快版本”，可以快速定位性能瓶颈到底是计算密集型的还是访存密集型的。

这个思路本质上是一种控制变量法。它不能直接提升性能，但它能帮你更快地找到下一个性能优化方向。我建议各位在优化复杂 kernel 时养成这个习惯——不要只看正确的代码，也可以故意造一个错误版本来观察性能特征的差异。

4.6 针对序列长度方向的 Swizzle：Sass 级优化

最后一个 hack 是在 SASS 层面做的。V3 的某个版本中，Verifier 的指令级回归报告发现，全局内存读取的指令序列中有一个奇怪的模式：每读取 128 字节的数据，编译器都会额外生成两个地址计算指令。这些地址计算指令是多余的——因为数据块在内存中的布局是固定的。

我们在 Ptxas 编译选项里加了-swizzle相关的调优，同时在 CUDA 代码层面改变了数组索引的顺序，把序列方向作为最内层维度。这个调整使得全局内存访问的合并模式变得更规整，实测内存吞吐从 78% 提升到了 91%。

SASS 层级的优化很容易被低估——“反正代码能跑，性能差不了太多”。但在这个项目中，SASS 级检查确实帮我们发现了两处编译器生成的低效模式，每一处都贡献了 3-5% 的性能提升。

5. 数字说话：优化前后的 Kernel Profile 与收益

5.1 三组 Benchmark 的对比

为了让你对 V3 的收益有一个直观认识，我整理了三组 benchmark 数据。测试环境是 H100 SXM 80GB，CUDA 12.4，模型配置为 8 头、每头 128 维，隐状态维度 1024，KV 序列长度覆盖 512 到 8192。

第一组是端到端 kernel 的推理耗时对比：

加速比稳定在 2.4 倍左右。你可以注意到，序列长度翻倍时，V3 的耗时增长也约为线性——这符合 Delta Attention 的算法复杂度特性，因为递推步骤数随序列长度线性增加。

第二组是 Nsight Compute 的关键性能指标对比（2048 序列长度）：

第三组是不同 batch size 下的表现。我们测了 batch size 从 1 到 16 的情况，V3 的性能优势在 batch size 增大时略有缩水——从单 batch 的 2.4 倍降到了 batch size=16 时的 1.9 倍。原因是 batch size 增大后，PyTorch 原生版本也能更好地利用 GPU 的多 SM 并行能力，相当于它的“默认并行度”被释放了出来。但 V3 的优势依然存在，且绝对耗时显著更低。

5.2 数值误差与偏差控制

性能之外，数值精度不容忽视。Delta Attention 的递推结构对误差极其敏感：每一步的微小偏差会被后续步骤放大。我们做了两组误差实验。

第一组是固定序列长度下，对比 V1、V2、V3 与 PyTorch FP32 参考实现的最大绝对误差：

V2 的误差控制优势随着序列长度增加越发明显。V3 的误差比 V2 略高，但仍然在同量级——它的主要优化方向是性能，数值上只要求“不劣于 V1 太多”即可。

第二组实验是检查误差的累积趋势。我把序列长度等差数列排列，观察最大误差的增长曲线。V1 大约呈线性增长，V2 则是对数增长——这正是归一化补偿策略的效果。如果你要部署到生产环境，我强烈建议用 V2 的数值逻辑搭配 V3 的性能框架，两者并不冲突，只是目前 V2 的补偿逻辑需要在计算量上做点取舍。

6. 经验与教训：Agent 设计内核这件事的边界在哪里

6.1 Agent 是出色的学者，但算不上优秀的园丁

经过这个项目，我对“Kernel Design Agents”有了一个更立体的判断：它们真的是出色的学者——能够在海量 CUDA 文档和编译器的行为模式中找到问题的线索，能快速生成多种实现方案，能自动做 profile 分析。但在“修剪”方面，它们的水平还不够。

什么是我说的“修剪”？一个 kernel 的优化过程很像照料一棵果树：Agent 擅长的是嫁接——它能提出一个全新的枝条（新的缓冲结构、新的并发模式），但嫁接完之后，它不擅长判断这个枝条是否与整体树形协调。我们 V3 的最终版本其实是一个混合体：主循环结构完全按照 Planner 的规格书实现，但其中细节性的调度参数（比如延迟刷写的间隔是 8 步而非 4 步或 16 步）是从人工调优实验里得到的。

换句话说，Agent 可以把你带到正确方向的 90%，但最后 10% 的收尾工作——细微的平衡和微调——仍然需要人来完成。这不是 Agent 能力不足的表现，而是这类工作本质上就是一个搜索过程，Agent 把最繁重的搜索工作做掉了，剩下的就交给人类的直觉。

6.2 人工介入的价值点：评估协议先行

在项目的最初几天，我犯了一个典型错误：试图让 Agent 在一轮循环里直接生成“性能最优”的内核。结果可想而知——它生成了几个正确性没问题的版本，但性能并没有明显超越基线。后来我意识到，问题不在于生成器不够好，而在于评估器太粗糙。

如果我给 Agent 的反馈只是“性能不好，继续优化”，它只能在黑暗中摸索。但当我建立了完整的评估协议——正确性测试、Nsight Compute 指标提取、SASS 指令回归、以及一组多维度 benchmark——之后，Agent 的表现立刻上了一个台阶。因为每次迭代都有明确的失败原因和优化方向，它不需要猜。

这个经验我想单独强调一下：KDA² 这类工具的真正价值不在于“写代码”的天赋，而在于“看得懂性能报告”的能力。如果你的评估体系还很原始，用 Agent 设计内核的体验会大打折扣。先把评估协议搭好，再引入 Agent，效果完全不同。

6.3 与手工编写 FlashAttention 式内核的心得对比

作为长期写 CUDA kernel 的人，我对手工实现和 Agent 实现的差异有切身体会。手工写 FlashAttention 风格的内核，每一行代码我都能解释为什么这么写——因为那是我反复试验和踩坑之后总结出来的。而这个项目的 Agent 版本中，部分代码我并不完全理解为什么这么写效果更好。

这是让人既兴奋又不安的地方。兴奋的是，Agent 可能发现了人类直觉之外的优化路径。不安的是，一旦出现诡异的性能波动（比如换了 GPU 驱动版本后性能下降 20%），我很难快速定位是 Agent 的哪个决策导致了这种脆弱性。

我的对策是：在最终发布的 kernel 里，每一个“非显然”的优化点都要求 Agent 附带解释。如果解释不清楚，我会让它换一种实现方式。这也算是人机协作中的一种“审计约束”。最终版本的 kernel 里大约有 70% 的优化点能通过这个审计，剩下的 30% 被我手工替换成了更保守但更可控的实现。

说到底，Agent 设计内核这件事，本质上是把“从大量可能方案中找到最优”的过程自动化了。剩下的仍然是工程问题：可靠性、可维护性和可解释性。我在这个项目中的最大收获不是那个 2.4 倍加速比的 kernel，而是学会了怎么和 Agent 协作——把它当成一个能力极强但又需要明确边界和审计机制的同事，而不是一个可以完全信任的黑盒工具。