RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered

快速结论：该报错出现在 xFormers 使用 Triton 后端执行 _xformers_tiled_matmul_kernel 处理大矩阵（尤其是 TP+SP 长序列 + 大词表场景）时，属于 CUDA 非法内存访问。Issue 中已将 PR #1346 认定为解决方案。

适用环境：xFormers 0.0.30+56be3b5.d20250810；PyTorch 2.7.0a0+79aa17489c.nv25.04；Python 3.12.3；nvcc 12.9.41；pytorch-triton 3.2.0+git4b3bb1f8b.nvinternal；使用 NCCL 后端、2 卡 torchrun 复现。

最快修复方案：升级/应用包含 PR #1346 的 xFormers 提交或版本，这是 Issue 中已明确指向的修复。

注意事项：Issue 中多位用户仅回复 “Same error”，未逐一给出各自环境的修复验证结果；合入 PR 后仍需用原始大矩阵用例做前向 + 反向回归确认。

问题场景

在 xFormers 中调用 sequence_parallel_leading_matmul（TP+SP，长序列、大 Vocab）时，xFormers 在 Triton 路径下启动 _xformers_tiled_matmul_kernel 处理大矩阵，触发 CUDA illegal memory access。Issue 复现代码使用 seq_len_global=131072、hidden_in=2048、hidden_out=64512、bf16，并执行 y.mean().backward()。命令为 CUDA_LAUNCH_BLOCKING=1 CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node=2 reproduce_error.py。

报错原文

RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered

报错定位描述：在 _xformers_tiled_matmul_kernel 中发生，触发路径为 _launch_triton_matmul。

原因分析

从 Issue 描述看，问题集中在 xFormers Triton 矩阵乘法内核 _xformers_tiled_matmul_kernel 处理大矩阵时的索引/边界处理，导致非法内存访问。Issue 作者指出“TP+SP 长序列 + 大 Vocab”训练时很容易触发。解决方案被指向 PR #1346。不过，Issue 中未给出该 PR 具体改动的逐行解释，故内核层面的根因细节尚未在讨论中完整展开。

环境排查

确认 xFormers 版本是否为 0.0.30+56be3b5.d20250810 或包含 PR #1346 的提交/版本。
确认 PyTorch 版本：2.7.0a0+79aa17489c.nv25.04。
确认 Python 版本：3.12.3。
确认 pytorch-triton 版本：3.2.0+git4b3bb1f8b.nvinternal。
确认 nvcc 版本：12.9.41。
确认 GPU 数量与分布式设置：Issue 中使用 CUDA_VISIBLE_DEVICES=0,1、NCCL 后端、torchrun --nproc_per_node=2。
确认是否处于 TP+SP 长序列、大 Vocab 的实际训练负载。

解决步骤

查看 PR #1346，确认其合入状态与对应提交。
将 xFormers 升级到包含该 PR 的版本或提交（可优先尝试）。Issue 中已将 #1346 认定为解决方案，但未给出具体版本号或安装命令。
用 Issue 提供的复现脚本，在相同环境（bvfloat16 输入、同样的矩阵尺寸、同样的 torchrun 命令）下复跑前向与反向。
若仍有报错，按评论建议构造最小 (M, N, K) 边界用例，找出 _xformers_tiled_matmul_kernel 从正确执行切换到 illegal access 的最小形状，以帮助隔离残余索引问题。
分别验证三组场景：最小边界形状、原始大 Vocab 用例、前向 + 反向。
如三组场景均通过，可在验证版本/提交下关闭该 Issue；否则以最小边界用例继续定位。

验证方法

使用 Issue 中的复现脚本，确认在相同矩阵尺寸、TP+SP 设置与 Triton 后端下不再出现 RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered，且前向与反向均可正常完成。至少覆盖评论提出的三类验证：最小边界形状、原始大 Vocab 用例、前向 + 反向。