一个 GPU kernel 里有两个 block。各自写完数据,再等对方,最后都应该算出 1 + 2 = 3。但 TileLang 的 CuTeDSL 后端在二维线程布局下,曾经会算出 [1, 3]:一个 block 提前越过了同步点,读到另一个 block 还没更新的数据。

问题出在一个很短的判断:tid.x == 0。我提交的 TileLang PR #3400 修正了它,已于 2026 年 10 月 10 日合并。

谁来代表整个 block?

这段 sync_grid() 用一个全局计数器记录已经到达同步点的 block 数量。每个 block 派出一个线程,把计数器加一;最后到达的 block 重置计数器,并推进 generation,通知其他 block 可以继续。

它依赖一个前提:每个 block 只能上报一次。

一维 block 里,tid.x == 0 确实只会选中一个线程。换成 (32, 2),坐标 (0, 0, 0) 和 (0, 1, 0) 都符合条件。两个线程属于同一个 block,却都去增加了全局计数。

如果整个 grid 恰好只有两个 block,第一个 block 就可能把两次“到达”全报了。计数器到了二,另一个 block 却还没写数据。

这里即使使用原子加法也不够。原子操作保证计数更新不会互相覆盖,却不负责检查这两次更新是否来自不同的 block。计数器记录得完全正确,上报次数本身仍然可能是错的。需要收紧的是谁能更新它。

让第二个 block 晚一点

为了让错误显现,回归用例在 block 1 写数据前,加入了一段有限的 device-clock 循环,让 block 0 先走到 barrier。两个输入初始为零,block 0 写入 1,block 1 写入 2,同步后各自读取两者的和。

用例的流程可以简化成这样:

text
block 0:A[0] = 1 → sync_grid() → B[0] = A[0] + A[1]
block 1:延迟 → A[1] = 2 → sync_grid() → B[1] = A[0] + A[1]

预期结果:B = [3, 3]

这里故意安排了到达顺序。否则两个 block 都写完了才触发同步,即使计数有错,也可能碰巧得到正确答案。只检查输出“通常对不对”,很容易漏掉这种问题。

用例中的读写也限定在 (0, 0, 0)。这样,同一个 block 内不会再有多个线程抢着写结果,失败就能集中到 barrier 的行为上。输入只有两个整数,既不用比较浮点误差,也不需要从复杂算子的中间结果里猜是哪一步先出了问题。

当时在 RTX 5090 上运行原生 CuTeDSL 测试,基线先通过了一维配置,随后在 (32, 2) 下得到 [1, 3]。测试使用 -x,遇到这次失败就停止,没有继续跑完基线的另外两种配置。

把三个坐标一起算进去

修复后的 leader 是物理线程 (0, 0, 0)。实际 PTX 将 %tid.x、%tid.y、%tid.z 按位或,再判断结果是否为零。三个坐标都为零,才进入上报路径。

用更容易读的伪代码表示,就是:

python
# 旧判断:二维、三维 block 可能选中多个线程
is_leader = tid_x == 0

# 修复后的语义:每个 block 只选中一个线程
is_leader = (tid_x == 0 and tid_y == 0 and tid_z == 0)

改动只涉及 leader 选择。原有的计数器与 generation 协议、内存 fence,以及 cooperative launch 要求都保留。完整代码和新增测试可以直接看 PR diff。

四种布局,重复调用八次

测试覆盖 32、(32, 2)、(32, 1, 2) 和 (16, 2, 2),分别检查一维、y 方向展开、z 方向展开,以及 y/z 同时展开的情况。

每种配置都调用 kernel 八次,每次先清零输入输出,再检查结果是否为 [3, 3]。这也会检查 barrier 状态能否在多次调用间正常复用。修复后四个用例全部通过,并在第二个新进程里再次通过。

每轮清零是有用的。如果沿用上一轮已经写好的输入,提前读也可能读到正确的值,重复调用就失去了检查意义。一维配置同样保留,用来确认修复没有把原本正常的路径改坏;另外三种布局则分别让遗漏 y、遗漏 z 的判断有机会暴露。

当时的环境是 RTX 5090、Python 3.12、Torch 2.11.0+cu130 和 CuTeDSL 4.7.0。两组对照都使用 TileLang 0.1.15 的原生编译器与运行时;基线 helper 与当时 main 的 994b44e 逐字节一致,候选组只替换修改过的 Python DSL helper。验证包含真实编译和 cooperative GPU launch,但没有完整重编当前 main 的 C++ 部分。

最终 A/B 还关闭了 kernel 缓存,避免换了 helper 却继续加载旧 cubin。环境与运行结果保留在 原 PR 的验证记录 中,本文依据当时的测试整理。

修复范围是多维 block 的同步正确性,并发 stream 共用 barrier 状态仍是另一个问题。

另一篇 ptxas 报告解析器的修复记录 讲的是工具链里的数据混读:两份各自正确的 kernel 报告,被解析器拼成了一份错误结果。

本文使用 AI 工具辅助整理。实验设置和测试结果见正文。

相关原始记录 ↗讨论这个问题 ↗