在 FlashQLA 升级 TileLang 的过程中，一个状态修正 kernel 出现了明显的数值错误。关闭 warp specialization 可以恢复正确性，却损失了原本希望保留的性能。更棘手的是：同一内核在完整调用链中出错，抽出来单独重放有时又完全正确。

最后，问题收敛到共享内存双缓冲的槽位复用：消费者尚未安全取得本轮数据，生产者已经开始覆盖同一槽位。最终修复在保留自动 warp specialization 的前提下，用显式槽位读取依赖把 membar.cta 放在自动释放之前。32K、32 heads 的重点前向路径收回约 7 微秒，相比关闭 WS 降低约 20% 的修正 kernel 耗时。

从版本升级后的数值偏差开始先恢复正确性：只关闭问题 kernel 的 WS 理解三个缓冲区，才能判断哪里需要同步第一个弯路：重放通过，不代表原位竞争消失生成代码中的变化，也是线索而非结论决定性证据：错误值来自同一槽位的下一代修复尝试：哪些动作有效，哪些只是改变时序第二个转折：fence 写在源码后面，为什么仍然失败？把保护操作绑定到 h 槽位三个可用源码方案的具体差别方案一：绑定 h 槽位的 async proxy fence 方案二：绑定 h 槽位的 membar 方案三：延长 h 槽位生命周期最终修复：保留 WS，补齐 h 槽位保护验证分别回答了什么问题性能结果：收回 kernel 损失，但不夸大完整算子收益 32K、32 heads 的修正 kernel 全配置与完整算子从这次修复中留下的方法参考资料

从版本升级后的数值偏差开始

对比环境分别为 TileLang 0.1.9 / TVM FFI 0.1.9，以及 TileLang 0.1.12 / TVM FFI 0.1.11。实验使用 Hopper H100，输入 Q/K/V 为 BF16，状态与累加沿用算子的 FP32 路径。

错误首先暴露在上下文并行（CP）的状态修正阶段。长序列被划分为多个片段后，各片段的局部结果还需要结合跨片段传播的状态。FlashQLA 的 tilelang_correct_h0 同时承担前向初始状态修正和反向终止状态修正，因此一个底层问题会同时影响前向与反向。

逐层比较中间结果后，warmup、fallback 判定、局部状态 ht 和传播矩阵 mt 都与旧版本一致，第一次差异出现在修正后的状态。一次代表性观测中，前向 vk 的相对 L2 误差约为 41.47%，反向 kv 约为 9.12%。误差幅度会随调用改变，远超正常舍入差异。

这里的 kv/vk 表示状态矩阵的两种存储布局：kv 为 K×V，vk 为 V×K。布局会改变矩阵乘法的方向和转置方式，因此定位时必须同时看不同布局，而不能只验证一条路径。

先恢复正确性：只关闭问题 kernel 的 WS

最早有效的修改非常小：

它只关闭该 kernel 的自动 warp specialization，保留 CP 选择、两级流水线和原有 GEMM。中间状态与最终输出恢复了参考结果，说明问题与这条 WS 执行路径紧密相关。

但这一修改也改变了内核结构：观察到的 WS 版本采用 256 线程、TMA 搬运和 producer/consumer 分工；关闭 WS 后采用 128 线程及另一条搬运路径。重点前向修正 kernel 出现约 7 微秒的额外耗时。接下来的目标是保留原有并行结构，把缺失的数据交接约束补回来。

理解三个缓冲区，才能判断哪里需要同步

把状态递推简化后，可以写成：

其中 H_i 是片段局部状态，M_i 是传播矩阵，propagate 根据布局和方向选择矩阵乘法。与这一递推对应的共享内存并非只有一块。

缓冲区	写入与读取	必须保护的关系
`h_shared`	TMA producer 写入局部状态；consumer 复制到寄存器 fragment	搬运完成后才能读；消费者读完之前不能复用槽位
`m_shared`	TMA producer 写入传播矩阵；consumer 的 GEMM 读取	搬运完成后计算；GEMM 使用结束后才能复用槽位
`hd_shared`	consumer 写入旧递推状态；consumer 的 GEMM 读取	消费者内部的跨 warp 写后读与下一轮覆盖顺序

h_shared 和 m_shared 使用两个物理槽位，hd_shared 则是单份缓冲。自动 WS 把搬运和计算分给不同线程组，使它们可以重叠执行。TileLang 的 WS 示例介绍了这种 producer/consumer 结构。

这张图对应两种不同的正确性要求：新数据什么时候可以读，以及旧数据什么时候允许被覆盖。等待 TMA 完成，只回答了前一个问题。

第一个弯路：重放通过，不代表原位竞争消失

最初，一个现象很容易把调查引向输入封装或上游同步：完整 CP 链路出错，而保存输入、重新上传后，独立修正 kernel 可以得到精确结果。

于是，把观察点放回原始调用位置，保留同一组实参，逐一改变执行条件：

实验	为什么做	观察结果与作用
相同实参立即重放	区分固定数值误差与依赖执行时序的问题	相同输入的输出仍会变化，竞争问题成为主要方向
设备同步后重放	检查是否只是上游计算尚未完成	同步后仍然出错，不能用调用前同步解决
只换输出缓冲区、逐个或全部克隆输入	观察地址、对象生命周期与布局的影响	误差有时减轻，但没有稳定消失；张量连续性和数值相同并不足以解释现象
输入经过 CPU 往返后重放	检查重新保存和分配是否可靠消除问题	仍能复现错误，保存输入本身不是修复
交替输入并扰动缓存	扩大执行时序变化，观察候选屏障的效果	连未修改版本也变得正确，因此这组条件下的候选通过不能证明修复有效

这一步改变了后续实验的判断方式：修复候选必须放回能够复现原版错误的调用链，在同场原版仍然失败时比较。否则，一个让竞争暂时不容易发生的执行条件，很容易被误认为真正的同步修复。

生成代码中的变化，也是线索而非结论

两版编译结果的 producer/consumer 线程组位置发生了交换。检查地址表达式后，可以看到新版通过局部线程编号或常量偏移做了补偿；线程组换位本身不足以解释错误。

同样，某个位置没有旧版的 partial barrier，也不能直接推导出“新版完全没有保护递推状态”。继续检查会发现其他位置仍有消费者屏障。针对存档 CUDA 重新生成的 PTX，也没有提供关键 shared 读写被简单提升到屏障之前的直接证据。

因此，调查没有停在“新版本少了一条同步”这个描述上，而是继续寻找究竟哪一块数据被谁覆盖。

决定性证据：错误值来自同一槽位的下一代

接下来不再同时修改多个同步点，而是隔离输入路径。

把 fallback 全部关闭后，不执行状态传播 GEMM，错误依然存在。把 mt 或原始状态清零也不能消除错误；只有把 ht 清零，输出才恢复精确。这把问题集中到了：

然后在不执行 GEMM 的条件下，将错误输出与不同片段、相同空间坐标的 ht 逐元素比较。结果非常直接：错误值完全等于同一物理槽位下一次复用时写入的数据。

双缓冲用 i & 1 选择槽位，因此执行顺序相差两轮的数据共用同一槽位。前向遍历时，错误元素匹配应读片段之后的 +2 片段；反向遍历时，匹配 −2 片段。这里是逐元素完全相等，不是误差分布相似。

示意图表达的是由数值匹配得到的竞争关系，不是硬件指令级时间线。能够确定的是：本轮读取与下一代 TMA 覆盖之间缺少有效保护。

进一步的单变量修改也与这个判断一致：保留 producer，把 consumer 对 h_shared 的读取改为对应位置的全局 ht 读取，错误消失；只把 h 槽位释放延后，错误也消失。直接读全局内存的版本用于定位，最终没有采用这种绕开共享内存的实现。

修复尝试：哪些动作有效，哪些只是改变时序

在明确 h 槽位是问题所在后，先直接调整生成 CUDA 中的交接位置，区分几类机制。

尝试	具体改动	结果与解释
Compiler fence	在 h release 前加入空内联汇编和 `memory` clobber	原位仍错；仅增加编译器内存约束不足以解决该竞争
Consumer 线程汇合	在 release 前加入 `bar.sync 4,128`	原位仍错；消费者汇合不能替代这里所需的完整交接约束
延迟 producer	producer 获得空槽后执行 `__nanosleep(100)`	部分反向调用改善，前向仍错；改变时间窗口不是稳定修复
延后 h release	分别移到 m wait 后、GEMM 后或 m release 后	各位置均恢复精确结果，支持槽位复用顺序是关键
Async proxy fence	在 consumer 释放 h 前，或 producer 获得空槽后加入 `fence.proxy.async.shared::cta`	恢复精确结果；去掉该实验中汇编的 clobber 后，硬件 fence 仍有效
CTA memory fence	在 h release 前加入 `membar.cta` 并保留 clobber	恢复精确结果，为后续源码实现提供另一条路径

这些实验不能简单归结为“只有 async proxy fence 才能修复”。延后释放和 CTA memory fence 都有效，说明需要分别观察缓冲区生命周期、生成顺序和硬件内存顺序。

membar.cta 是 CTA 范围的内存顺序指令；fence.proxy.async 处理 generic 与 async proxy 之间的顺序。二者不是同义指令，也不能从本实验推导出所有 TMA 场景都可以互相替代。PTX ISA 的 membar/fence 与 async proxy 说明给出了各自语义。

第二个转折：fence 写在源码后面，为什么仍然失败？

直接改生成 CUDA 可以证明某个位置有效，却不能作为可维护的源码修复。因此，下一步把操作写回 TileLang。

两个最直观的写法都失败了：

检查实际生成代码后发现，自动 WS 仍在 h 的最后一次可见读取之后释放槽位，而上面两个调用没有显式表明它们依赖 h_shared。生成顺序相当于：

fence 确实存在，却没有处在需要保护的交接之前。问题不只是“用哪条指令”，还有“如何让编译器把这条指令关联到正确的 buffer 生命周期”。

把保护操作绑定到 h 槽位

有效写法用 T.access_ptr(h_shared, "r") 显式描述读取依赖，再把它传给外部 helper：

这里的指针是编译依赖的一部分，helper 不需要真的再加载一次数据。在这次生成结果中，自动释放被安排在 helper 调用之后：

这还解释了为什么 asm volatile、memory clobber 和 TileLang 的 buffer 依赖需要分别理解。volatile 保留汇编的执行，clobber 声明编译器可见的内存副作用；它们并不自动为 TileLang 的 WS 生命周期分析提供一个具体的 h 槽位读取关系。NVIDIA Inline PTX 文档也区分了 volatile 与隐式内存副作用的声明。

三个可用源码方案的具体差别

方案一：绑定 h 槽位的 async proxy fence

在 h 复制到 fragment 后立即调用 helper，helper 内部执行：

读取依赖把调用留在 h 槽位释放前，proxy fence 提供对应硬件顺序操作。该方案恢复了正确性，也恢复了重点前向路径的大部分 WS 性能。它最先被采用，随后作为完整性能比较的候选保留。

方案二：绑定 h 槽位的 membar

调用位置与上一个方案相同，只把 helper 中的指令换成：

两者具有相同的显式 h 读取依赖与 clobber，硬件指令不同。在这次状态修正流水线中，membar 同样恢复精确结果，且重点 kv 反向路径耗时明显低于 proxy 方案。

方案三：延长 h 槽位生命周期

第三个方案不在 helper 中增加硬件 fence，而是使用空汇编：

与前两个方案不同，这次把带 h 读取依赖的 helper 放到本轮循环末尾，即条件 GEMM 之后。这样自动释放跟着推迟，producer 必须更晚才能复用 h 槽位。

它与先前失败的“release 前空 compiler fence”并不矛盾：关键改变在于显式依赖的最后使用位置，以及由此改变的槽位生命周期。该方案恢复了正确性，但可能减少搬运与计算的重叠机会；实际是否更慢仍由计时决定。

源码方案	helper 位置	核心约束	本次结果
直接 proxy fence	源码中的 h copy 后	未声明 h 读取依赖	生成释放仍在 fence 前，失败
直接 threadfence_block	源码中的 h copy 后	未声明 h 读取依赖	生成释放仍在 fence 前，失败
proxy	h copy 后	h 读取依赖 + async proxy fence	正确，进入性能比较
membar	h copy 后	h 读取依赖 + CTA memory fence	正确，最终采用
delay-release	本轮条件 GEMM 后	把 h 的显式最后使用延长到循环末尾	正确，进入性能比较

最终修复：保留 WS，补齐 h 槽位保护

最终落地的是 membar 方案。关键结构如下，省略了形状、布局分支和具体递推计算：

代码中的省略号表示原有算子逻辑，以上是修复结构展示。原本用于 hd_shared 路径的 T.fence_proxy_async() 仍然保留；本次 membar 保护的是 h 槽位，不能把其他缓冲区的 fence 一并替换或删除。

生成 CUDA 中仍然保留 256 线程、TMA 和 producer/consumer 分工，helper 位于 h 槽位释放之前。最终源码与完成数值和性能验证的 membar 候选在 Python AST 层面一致，没有在选型后引入另一种 kernel 实现。

验证分别回答了什么问题

修复验证没有通过放宽误差阈值来掩盖问题。原有数值判定保持不变，并额外观察与旧版本参考结果的逐元素关系。比“通过了多少项”更重要的是，每类验证分别排除了什么解释。

验证方式	作用与动机	得到的结论
原位前向 vk 与反向 kv 对照	保留真实调用条件，避免独立重放暂时掩盖竞争	原版仍错时，三个带槽位依赖的方案恢复精确结果
真实失败输入固定重放	固定数值输入，检查修复能否处理已捕获的错误样本	前后向修正结果与参考逐元素一致
完整算子与旧版本输入重放	检查修正是否能传递到最终输出，覆盖不同布局、初始状态、变长和 cache 路径	有效输出与旧版本参考逐元素一致，原有数值判定保持不变
重复执行与独立 CP 开关	检查连续调用及前后向组合，避免只修好一种配置	恢复相应数值行为，没有依赖每次调用前的全局同步或输入复制
实际生成 CUDA 核查	确认优化没有被意外关闭，并检查保护位置	WS 分工仍在，helper 调用先于 h 槽位释放
计时前后输出比较	确认测到的快实现仍然正确，包含 Graph/Event 实际执行后的输出	参与选型的三个候选均满足正确性要求，独立修正结果保持精确一致

性能结果：收回 kernel 损失，但不夸大完整算子收益

比较包含关闭 WS、proxy、membar、delay-release 四个实现。它们使用相同输入与实参，在同一张 H100 上计时。覆盖的配置包含 4K、8K、32K 序列，4、16、32 heads，两种状态布局，以及初始状态、变长、fallback 和 cache 复用组合。

每个阶段测五轮，每轮取 15 个样本的中位数，再对五轮中位数取中位数。Graph 捕获 16 次调用并折算单次耗时；完整算子还使用 Event 方法测量。编译、预热与正确性检查不计入延迟，候选沿用一致的缓存扰动方式。

32K、32 heads 的修正 kernel

下表单位为微秒，数值越低越好，均为 Graph 结果。

路径	关闭 WS	proxy	membar	delay-release
kv 前向	34.604	27.376	27.282	28.030
kv 反向	59.134	59.968	53.982	62.944
vk 前向	37.270	29.982	29.812	30.052
vk 反向	57.602	51.532	51.474	51.880

membar 相比关闭 WS，kv/vk 前向分别减少 7.322/7.458 微秒，对应耗时下降 21.2%/20.0%。相比 proxy，最大的差异出现在 kv 反向：减少 5.986 微秒，耗时下降约 10.0%。

这处反向差异也明显超过轮间变化：proxy 的五轮范围为 59.566～60.144 微秒，membar 为 53.672～54.812 微秒。其余重点路径中，两者范围存在重叠，因此更适合描述为接近，而不是给每个微小差异都贴上稳定提升的标签。

作为历史版本参考，0.1.9 WS 在另一张同型号卡上的 kv/vk 前向分别为 27.220/29.706 微秒，反向分别为 53.964/51.090 微秒。最终候选已经回到相近延迟水平；候选间选型依据则来自上表的同卡比较。

全配置与完整算子

在所有适用配置中，membar 的独立前向、反向修正 kernel 中位数均低于 proxy 和 delay-release。相对 proxy，前向耗时变化为 −3.36%～−0.34%，反向为 −9.98%～−0.11%。

但修正 kernel 只是完整算子的一部分。下面记录完整算子相对 proxy 的变化范围，负值表示 membar 更快；范围来自不同配置，不是置信区间。

完整算子阶段	Graph 耗时变化	Event 耗时变化
前向	−1.78%～+0.09%	−3.30%～+1.78%
反向	−2.56%～+0.73%	−2.84%～+0.14%
反向 cache 复用	−1.93%～+0.07%	−1.54%～+0.30%
前向加反向	−3.27%～+0.82%	−1.15%～+2.75%
前向加反向 cache 复用	−2.01%～+0.13%	−1.12%～+0.08%

完整算子的变化较小，也不是所有配置都更快。选择 membar 的理由是保持正确 WS 结构、重点路径收益明确，以及修正 kernel 的整体中位数表现更好；不是声称完整算子都能获得 20% 加速。

从这次修复中留下的方法

这次调查最有价值的证据，不是一开始看到的版本代码差异，而是错误值精确对应到了双缓冲槽位的下一代内容。它把一个容易泛化为“同步有问题”的现象，变成了可操作的读写关系：谁还在读、谁开始覆盖、哪次槽位交接需要被保护。

另一个关键是把源码顺序和生成顺序分开。直接插入 fence 并不保证它出现在自动释放之前；显式声明具体 buffer 的读取依赖，才让保护操作参与到 WS 的生命周期安排中。随后再比较不同硬件 fence 与延后释放的性能，才能得到既正确又保留优化收益的实现。

最终采用的补丁很小，但决定它是否正确的证据来自完整链路：原位复现、缓冲区隔离、下一代数值匹配、生成代码检查，以及正确输出下的性能比较。

参考资料

FlashQLA：本文讨论的算子项目；本文给出的补丁结构与实验数值来自本次实测，不代表该修复已合入上游。

TileLang 的 warp specialization 示例：producer/consumer 与 TMA 搬运结构。

NVIDIA PTX ISA 9.0：mbarrier、membar、fence 与 async proxy 的指令语义。

NVIDIA Inline PTX：Incorrect Optimization：volatile 与 memory clobber 的作用。

Antarx Memos

一次 TileLang Warp Specialization 数值错误的定位与修复：从双缓冲竞争到 membar