Triton地址索引跨过int32边界导致CUDA error

前言

一次 Qwen3.5-4B 训练在约 23.4 万 packed tokens 时触发：

CUDA error: an illegal memory access was encountered

错误最初出现在把 cu_seq_lens 从 GPU 拷回 CPU 的地方，看起来像 packed sequence 边界坏了。打开同步执行后，真正的调用链收敛到了 BatchLoRA 的 Triton 融合线性层。

根因很纯粹：kernel 使用 int32 计算扁平元素地址，token行数 × 输出宽度 刚好越过 $2^{31}-1$，索引发生回绕，最终访问错误 GPU 地址。

生产故障的形状

失败请求为：

packed token rows M = 234624
MLP投影宽度 N = 9216

扁平元素数：

234624 × 9216 = 2,162,294,784

有符号 int32 上限：

2^31 - 1 = 2,147,483,647

刚好多出：

14,811,137 elements

对于宽度 9216，安全 token 行数边界是：

floor((2^31 - 1) / 9216) = 233016

生产输入的 234624 只比边界多约 0.7%，所以较小测试一直正常，大输入却稳定崩溃。

为什么Python堆栈指错了地方

最初报错位置是：

cu_seq_lens.detach().to(device="cpu", dtype=torch.int64).tolist()

这段代码只是在做一次 GPU 到 CPU 的同步。CUDA kernel 默认异步执行，Triton kernel 非法访存后，错误状态可能等到下一次同步、拷贝或 collective 才被观察到。

设置：

CUDA_LAUNCH_BLOCKING=1

后，调用链变成：

Qwen3.5 decoder layer
→ MLP gate_proj
→ BatchLoRA Triton path
→ FusedLoRAFunction.forward
→ fused_multi_lora_xw_sb
→ illegal memory access

这和另一个越界 token 导致 embedding gather assert 的问题一样：GPU 报错堆栈经常只是观察点，不是犯错点。

最小实验

直接调用 fused_multi_lora_xw_sb，只保留 token 行数和输出宽度：

token行数M	M×9216	是否超过int32	修复前
232960	2,146,959,360	否	通过
234624	2,162,294,784	是	稳定非法访存

Custom Loss为什么更容易触发

forward_backward_custom 是两阶段协议：

1. 完整forward取得current logprobs
2. 客户端计算自定义loss和dL/dlogprob
3. 再提交forward-backward

第二阶段 forward-backward 已按 DDP 分片，每卡约 6 万 token；第一阶段纯 forward 当时没有分片，每个 rank 都处理完整 234624 tokens。

所以现象看起来像“使用 DAPO 自定义 loss 就崩”，实际在第一阶段还没执行用户 loss 时就已经失败。自定义 loss 只是额外触发了一次未分片的大 forward。

这也解释了为什么分片修复能降低触发概率，但不是 kernel 根治。单卡、单条超长样本或分片后仍越界的输入依然可能跨过 int32 地址边界。

修复方案

应急方案可以：

• 把单 rank token 限制在 233016 以下
• 设置 BATCHLORA_FORCE_PYTHON=1 绕开 Triton

但模型投影宽度会变化，换一个更宽的模型，安全 token 数又不同。关闭融合 kernel 也会损失性能。

正式修复是在计算大张量行地址前显式提升到 tl.int64：

row_index(int64) × row_stride

forward 的输入、输出地址和 backward 的 dy、dx 都要检查，不能只修当前最先崩的写入。

修复验证

token行数M	输出宽度N	元素数	结果
232960	9216	2,146,959,360	forward通过
234624	9216	2,162,294,784	forward通过
333312	9216	3,071,803,392	forward通过
234624	9216	2,162,294,784	backward通过，梯度finite

随后按原故障条件重新运行 DDP4、7 个 LoRA、234624 packed tokens 的真实请求，完整完成，所有 LoRA 也能正常卸载。

最后

int32 能表示 21 亿，看起来对 tensor shape 绰绰有余；但扁平地址计算乘的是所有维度。23 万行乘 9216 宽度，就已经跨过边界。

写 GPU kernel 时，shape 参数本身没有越界不代表地址表达式安全。任何 row * stride + col 都要按最大物理 tensor 规模检查中间乘积，而不是只看最终 pointer 类型。

相关记录：

• Trio-Loop Issue #105