[Big Tensor] Support iteration space beyond INT32_MAX in GPU index_elementwise_put kernels

[omoYang]

PR Category

Operator Mechanism

PR Types

Bug fixes

Description

x[index] = value 这类 advanced indexing 的 setitem 在 GPU 上走 index_elementwise_put / index_elementwise_put_with_tensor 及其反向。当被索引区域的迭代空间 numel 超过 INT32_MAX 时，当前实现有两个问题：

前向：无条件拒绝

index_elementwise_put_kernel.cu 里对 N = numel 做了无条件的 PADDLE_ENFORCE，要求 N <= INT32_MAX，与实际使用的 offset 位宽无关。因此即使调度端已经选择了 64 位 offset 的实例化，调用一样会失败：

PreconditionNotMetError: the value of N should be in [0, std::numeric_limits<int32_t>::max()]

实际触发场景：大 padding 率样本下 input_embeds[indices] = value，len(indices) * hidden_size 越过 INT32_MAX 直接报错中断训练。

反向：静默截断

index_elementwise_put_grad_kernel.cu 连这个检查都没有。32 位实例化下 OffsetCalculator::get() 的形参类型就是 uint32_t，线性下标会被截断且不报错，写入落到错误地址。

本 PR 的改动

offset 位宽是调度端按各操作数的字节跨度选的，而字节跨度并不约束迭代空间大小：一个很小的 index 张量广播到大 input 上，跨度全在 int32 以内，但 numel 可以远超 INT32_MAX。所以：

• 在 32 位实例化里发现 numel > INT32_MAX 时，重新派发到 64 位实例化，而不是报错或截断。放在构造 OffsetCalculator 之前，因为 32 位的构造函数本身会对越界偏移抛错。
• 删掉前向的 numel 守卫，改为对 grid 维度做检查：线性下标放宽到 64 位后，真正的硬上限是 gridDim.x，而且 dim3 的构造存在从 int64_t 到 unsigned int 的隐式窄化，需要显式挡一层。
• 删除反向里两处 IsInInt32Range(value_grad->numel()) enforce守卫：MakeOffsetCalculatorPut里已经对其逐维检查过

CPU / XPU 后端的同名守卫本 PR 未改动。

精度测试

paddle.Tensor.__setitem__(Tensor([100000, 1024],"float32"), Tensor([100000],"int64"), Tensor([1024],"float32"), )
paddle.Tensor.__setitem__(Tensor([2, 1000000],"float32"), Tensor([4100],"int64"), Tensor([1000000],"float32"), )
paddle.Tensor.__setitem__(Tensor([600000, 4096],"float32"), Tensor([600000],"int64"), Tensor([4096],"float32"), )
paddle.Tensor.__setitem__(Tensor([100000, 1024],"float32"), Tensor([100000],"int64"), 1, )
paddle.Tensor.__setitem__(Tensor([2, 1000000],"float32"), Tensor([5000],"int64"), 1, )
paddle.Tensor.__setitem__(Tensor([600000, 4096],"float32"), Tensor([600000],"int64"), 1, )

可以直接复制到cases.txt中，在PaddleAPITest中通过--api_conig_file="cases.txt" + --accuract_stable=true测试精度
前三条覆盖index_elementwise_put_with_tensor，，后三条覆盖index_elementwise_put，内部依次覆盖纯int32索引/跨度小→launch GPUxxxKernel时offsetT=uint32，但 N>INT32_MAX / 跨度>INT32_MAX

是否引起精度变化

否
改动只影响下标与偏移的计算位宽，以及守卫条件的作用范围，不改变任何计算逻辑。

[risemeup1111]

Paddle-Bot Review Board （review完成)

序号	位置	优先级	规则来源	状态
1	张量赋值前向 GPUIndexElementwisePutWithTensorKernel 大迭代空间处理（已在 29e30c2d 重派发到 64 位并补 grid 守卫）		仓库规则：算子 Kernel 前向反向与整数宽度边界	✅

_{Powered by Nyanpasu claude with Opus 4.8 默认推理级别, please check the suggestions carefully.}

[risemeup1111]

本 PR 修复了标量前向 GPUIndexElementwisePutKernel 与反向 GPUIndexElementwisePutGradKernel 的大迭代空间问题，但同文件内的张量赋值前向 GPUIndexElementwisePutWithTensorKernel（约 161-257 行）未做同样处理。

• 调度端 IndexElementwisePutWithTensorKernel（约 356 行）仍按各操作数的字节跨度选择 32 位实例化；而字节跨度并不约束迭代空间大小（正是本 PR 描述里的核心论点）。因此 x[indices] = value_tensor 在小 x + 大 indices 广播场景下会选到 32 位实例化。
• 32 位实例化下 OffsetCalculator<3, uint32_t, true>::get() 的形参是 uint32_t（见 index_elementwise.cu.h:216）。当迭代空间 numel 越过 32 位索引范围时，第 238 行传入的 int64_t idx 会被静默截断，写入错误地址且不报错——与本 PR 为反向修复的失效模式完全一致。构造函数里的 CheckOffsetRange 只校验偏移字节跨度，拦不住线性下标的截断。
• 另外第 230 行 const dim3 grid((N + block.x * vt - 1) / (block.x * vt)); 同样缺少 grid 维度上限检查和 int64_t → unsigned int 的显式窄化保护，与本 PR 在两个已修复 kernel 中新增的 grid 守卫不一致。

建议对 GPUIndexElementwisePutWithTensorKernel 施加与标量前向相同的「numel 超限时重派发到 64 位实例化 + grid 守卫」处理，否则张量赋值前向仍会在大迭代空间下静默产生错误结果。

[risemeup1111]

已在 29e30c2 验证修复：GPUIndexElementwisePutWithTensorKernel 现已在 32 位实例化下对 numel > INT32_MAX 重派发到 64 位实例化（置于 MakeOffsetCalculatorPut 之前），并补充了 grid 维度上限检查与 int64_t → unsigned int 的显式窄化保护，与标量前向/反向一致。问题已解决。

[wanghuancoder]

LGTM