[Operator Mechanism]Add CUDA small-stride index gradient kernel

[feixi139]

PR Category

Operator Mechanism

PR Types

Bug fixes

Description

本 PR 修复 GPU 高级索引（advanced indexing）反向的三个问题：

1. 低精度 dtype 下重复下标的累加逐步舍入，x[:, idx] 这类表达式的梯度精度明显低于 x[idx]；
2. None（newaxis）出现在高级索引之前时，索引被绑定到错误的轴，前向 shape 就是错的；
3. index_elementwise_get_grad 没有校验 slice_offset 和被索引视图的跨度，非法值会变成越界写。

1. 高级索引反向的低精度累加

torch 的 index_put_(accumulate=True) CUDA 实现完全没有 atomicAdd，只有排序路径，所以它的结果确定、可复现，且重复下标只舍入一次。本节的三处改动合起来把 Paddle 的所有 accumulate 反向也导向这条路径：1.1 放开路由，1.2 补上排序路径缺失的 kernel，1.3 让排序路径能处理放开后新进来的跨步视图。

1.1 路由判据用错了维度，x[:, idx] 落到逐项舍入的 atomicAdd

问题

index_elementwise_get_grad 原来的路由条件是：

if (accumulate && index.size() == 1 && !is_combined) {
  IndexPutWithSortKernel<T, int64_t>(...);
  return;
}

is_combined 来自 pybind：eager_method.cc 里 index_size = PyTuple_GET_SIZE(index_ptr) 是 __getitem__ 索引元组的元素个数（slice / None / Ellipsis 都算），slice_utils.h 里 is_combined = (index_size == 1) ? false : true。因此真实判据是「索引元组只有一个元素」：

• x[idx] → 走排序 kernel（IndexingBackwardKernelStride1 / IndexingBackwardKernel），重复下标在 opmath_t（fp32）寄存器里归约，只在写回时舍入一次；
• x[:, idx]、x[..., idx]、x[idx, None] → is_combined=true，落到 GPUIndexElementwiseGetGrad，IndexEleGetGradAccKernel 内是 CudaAtomicAdd(reinterpret_cast<T*>(...), T value)，即 fp16/bf16 原生 atomicAdd，每加一次就舍回低精度，误差随重复次数增长，且累加顺序不确定、run-to-run 不可复现。

触发条件是同一下标重复 ≥3 次：N=2 时低精度逐步累加恒等于正确舍入（0+a 精确，a+b 只舍一次），所以必须重复三次以上才会暴露。[16, 8] 配 x[:, idx]（索引长 5）200 组随机数据统计：

最大重复次数	fp16 与 torch 不一致率	bf16
1	0/46	0/46
2	0/127	0/127
3	26/27	27/27

改动

路由门从 accumulate && index.size() == 1 && !is_combined 改为 if (accumulate)（index_elementwise_get_grad_kernel.cu:724）。放开之后有两类输入是排序路径原来处理不了的，分别由 1.2、1.3 补上。

1.2 排序路径缺 1 < sliceSize <= 32 的 kernel

问题

sliceSize（被索引轴之后的元素个数）落在 1 < sliceSize <= 32 时，即使走到排序路径，也只有 IndexingBackwardKernelStride1（仅 sliceSize == 1）和按 feature 展开的通用 IndexingBackwardKernel 两个选择，后者对每个重复项都做一次 grad_weight 的 read-modify-write，同样是每步舍入。torch 在这里有专门的 indexing_backward_kernel_small_stride，且它真正的分派维度就是 sliceSize 与 warp size 的比较，而不是索引表达式的写法。

改动

新增 IndexingBackwardKernelSmallStride（:344），覆盖 1 < sliceSize <= WARP_SIZE：一个线程独占一个 feature 列，该下标的所有重复项在 opmath_t（fp32）寄存器内归约，写回时才舍入一次。分派处（:638）由此按 sliceSize 三分：== 1 走 IndexingBackwardKernelStride1，<= WARP_SIZE 走新 kernel，其余走通用 IndexingBackwardKernel。

1.3 排序 kernel 无法寻址跨步视图

问题

排序 kernel 把 grad_weight 当作一块平坦 buffer 寻址；原 IndexPutWithSortKernel 自己做 transpose / inverse-perm 并用 funcs::makeLinearIndex 构造线性索引，只在「被索引对象就是整块连续 x_grad」时成立。1.1 放开路由后，x[1::2, idx] 这类先切片再索引也会进来，此时被索引对象是 x_grad 内的一块跨步子区域，直接跑排序 kernel 会写到错误位置。

改动

IndexPutWithSortKernel 改为接收一个描述被索引视图的 SortedPathLayout（:401），并直接用 funcs::computeLinearIndex 在该视图上构造线性索引。被索引轴的位置由 layout 给出，且必然是连续的一段，因此 transposeToFrontAndInvPerm 那一步在此路径上不再需要，也省掉了 permute 后的 Contiguous 与写回时的反向 transpose。

视图信息的恢复方式（DeriveSortedPathLayout，:437）：kernel 参数描述的是 restride 之后的视图 —— restride_src()（paddle/fluid/pybind/slice_utils.h）把被索引轴换成广播后的 index shape 并给这些轴零 stride，所以 input_strides 里连续的一段 0 就标出了被索引块，起点即 dims_before；而 index_dims（即 AdvancedIndex::indexed_sizes）开头正是这些轴的 extent，于是视图在索引之前的 shape 及其在 x_grad 内的 stride 都可以还原。据此分三种情况：

• 视图恰好是整块连续 x_grad 时，直接在 x_grad 上跑排序 kernel；
• 视图是跨步子区域时（x[1::2, idx]），先分配一个与该视图同形的连续 scratch，在其上归约，再用 StridedCopyKernel 按 view_strides / view_offset 搬回 x_grad —— 这与 torch 用 slice_backward 组合 index_backward 的做法一致。x_grad 已经清零，所以搬回时无需累加。注意 phi::StridedCopyKernel 会覆写 out 的 meta，调用方需自行 save / restore；
• IsNonOverlapping（:411）拒绝自重叠的 as_strided 式视图，这类布局回落到原来的 elementwise 路径（精度更低、run-to-run 不确定），与上游行为保持一致。

slice_offset / index_strides 是字节偏移（slice_utils.h 里 stride * element_size_bytes），派生 layout 时统一按 elesize 换算回元素单位，不能整除即视为非法并回落。

2. None 出现在高级索引之前导致轴绑定错位

问题

None 位于高级索引之前时，前向 shape 就与 numpy / torch 不一致：

表达式（x 为 [16, 8]，i 长 3）	numpy / torch	修改前 Paddle
x[None, :, i]	(1, 16, 3)	(1, 3, 8)
x[:, None, i]	(16, 1, 3)	(16, 3, 8)
x[None, i]	(1, 3, 8)	(3, 16, 8)
x[i, None]	(3, 1, 8)	(3, 1, 8)（正确）

ParseIndex 里的 estimated_dim 是基础索引输出张量的轴号（slice 保留 → ++，整数 → 不 ++（会 decrease），Ellipsis → += rank - specified，bool tensor → += rank），advanced_index_dim[estimated_dim] = estimated_dim 用它记录高级索引落在哪个轴。但 Py_None 分支只 push none_axes 而不推进 estimated_dim，而 getTensorWithBasicIndexing 结尾的 unsqueeze_ad_func(out, *none_axes) 确实在输出里插入了这个轴（none_axes 已按 decrease_axis 修正，就是基础输出的坐标系）。两个坐标系因此错开 none_count 个位置，索引被绑到错误的轴上。

同分支的单 bool（True / 0-D bool tensor）处理是对的：它 push none_axes 的同时也 estimated_dim++，正好印证 None 分支缺这一步。

静态图 parse_index 有同构的 bug，且更早暴露：replace_none 先把 None 从 indices 里剥掉，循环根本看不到 None，修改前静态图 x[None, i] 直接在 gather_nd 报 index out of range。

改动

• slice_utils.h 的 ParseIndex Py_None 分支补 estimated_dim++；
• eager_method.cc 两处 advanced_index_dim 容量从 rank * 2 改为 rank * 2 + index_size（setitem 侧为 + size）。否则 rank 2 张量上 x[None, None, None, None, None, i] 会让 estimated_dim 到 5，发生越界写；
• variable_index.py 的 parse_index 改成 _, none_axes = replace_none(indices) 保留原 indices，循环里遇到 None 时 estimated_dim += 1; continue，并把 advanced_index 扩容 len(indices)。

3. slice_offset 与视图跨度的越界校验

问题

slice_offset 是前向记录的「被索引视图在 x_grad buffer 内的字节偏移」。排序路径和 elementwise 回落路径都会把它加到 x_grad 的基址上，因此一个非法值（负数，或超出 buffer）会直接变成越界写，而不是报错。

同样地，跨步视图是稀疏的，它触达的最远位置由跨度决定而不是元素个数决定，只检查 numel <= grad_numel 不足以保证不越界。

改动

• 在 IndexElementwiseGetGradKernel 入口校验 0 <= slice_offset < x_grad 字节数；
• 在 DeriveSortedPathLayout 里校验视图触达的最后一个位置 slice_offset / elesize + Σ (view_dims[k] - 1) * view_strides[k] < x_grad->numel()，报错信息带上 offset、view shape 和 strides，便于定位。

修改后的行为

类别	之前	之后
fp16/bf16 x[:, idx] 等重复下标反向	每个重复项舍入一次，误差随重复次数增长，顺序不确定	fp32 寄存器内归约，写回舍入一次，与 torch 逐位一致
1 < sliceSize <= 32 的排序路径	走按 feature 展开的通用 kernel，逐项 read-modify-write	新增 small-stride kernel，单线程独占一列
先切片再高级索引（x[1::2, idx]）	走 elementwise atomicAdd 回落	连续 scratch 归约后按视图 stride 搬回
x[None, :, i] 等 None 在前	前向 shape 错、__setitem__ 多数直接报 broadcast 失败、静态图 gather_nd 越界	与 torch 一致
多个连续 None	advanced_index_dim 可能越界写	容量按 index_size 扩展
非法 slice_offset / 越界视图	静默越界写	明确报错

自重叠的 as_strided 式视图仍回落到原 elementwise 路径，行为与之前相同。

是否引起精度变化

是

[risemeup1111]

Failed CI看板

流水线名称	问题标签	修复建议	日志片段
33257579405	未解析	### 日志分析报告 当前 CI 快照抓取超时，未能取得失败 job 列表和日志；PR diff 也不可用。因此无法确认失败流水线、失败测试或根本原因，也不能据此判断 CI 已通过。 失败的测试 case: text 未提供可验证的失败 job、测试 case 或错误日志。 根本原因分析: 现有数据不足，不能区分 PR 代码问题、CI 基础设施问题或上游依赖问题，也不能建立 job 之间的因果链。 修复建议: 1. 重新获取 CI Run 33257579405 Attempt 2 的失败 job 列表和完整失败日志。 2. 同时获取 PR head 1d23f7520eb35ceff46d168b69b84946659b8037 的实际 diff。 3. 在获得可验证日志前，不建议修改代码或重跑特定 job。	未提供

日志分析报告

当前 CI 快照抓取超时，未能取得失败 job 列表和日志；PR diff 也不可用。因此无法确认失败流水线、失败测试或根本原因，也不能据此判断 CI 已通过。

失败的测试 case:

未提供可验证的失败 job、测试 case 或错误日志。

根本原因分析:

现有数据不足，不能区分 PR 代码问题、CI 基础设施问题或上游依赖问题，也不能建立 job 之间的因果链。

修复建议:

1. 重新获取 CI Run 33257579405 Attempt 2 的失败 job 列表和完整失败日志。
2. 同时获取 PR head 1d23f7520eb35ceff46d168b69b84946659b8037 的实际 diff。
3. 在获得可验证日志前，不建议修改代码或重跑特定 job。

Powered by Nyanpasu with gpt-5.6-luna 默认推理级别, please check the suggestions carefully.

[risemeup1111]

Paddle-Bot Review Board （review完成)

序号	位置	优先级	规则来源	状态
1	步长切片梯度写回		仓库规则：算子与内存正确性	🚧

^{Powered by Nyanpasu with gpt-5.6-sol 默认推理级别, please check the suggestions carefully.}

[risemeup1111]

这里把 layout.view_offset（来自 DealWithIndex 对 sub_tensor.data() - tensor.data() 的结果）传给 StridedCopyKernel。对于 x[::2, idx] 这类步长切片，strided_slice kernel 会为 sub_tensor 单独分配连续 buffer，并不与 tensor 共享 storage；此时指针差不是 x_grad 内的合法字节偏移，可能为负并在 DenseTensorMeta::offset 中转成超大无符号值。写回会导致梯度写错或越界。请仅在确认共享 storage 且偏移有效时走该路径，步长切片保留 elementwise/slice-backward 回退，并补充回归测试。

[feixi139]

risemeup1111 感谢 review。这里的前提需要澄清一下：strided_slice 在 FLAGS_use_stride_kernel 打开时返回的是共享 storage 的视图，不会单独分配连续 buffer。

paddle/phi/kernels/stride/strided_slice_kernel.cc:110 就是 out->ResetHolder(input.Holder())，实测各种跨步切片都共享 holder、offset 合法非负：

表达式 | offset | _is_shared_buffer_with(x) -- | -- | -- x[::2] | 0 B | True x[1:5] | 24 B | True x[1::2] | 24 B | True x[2:7:2] | 48 B | True x[::3] | 0 B | True x[1::2, ::2] | 24 B | True

flag 关闭时切片确实是拷贝，但那种状态下这段代码根本不执行：index_elementwise_get 全仓库只有 4 个真实调用点，全部位于 slice_utils.h:810 / :1334 的 FLAGS_use_stride_kernel 分支内，flag 为 0 时 getitem 会退回 gather / gather_nd，index_elementwise_get_grad 不入图，layout.view_offset 也就不会被构造。也就是说"strided_slice 单独分配 buffer"与"执行这段 grad kernel"这两个条件互斥。

关于建议的"步长切片保留 elementwise/slice-backward 回退"：这条不能采纳。slice_utils.h:725 在 FLAGS_use_stride_kernel && pos_of_new_dim != 0 时不做 transpose（transed_tensor = tensor），所以在 stride 路径上退回 gather_nd(transed_tensor, ...) 会 gather 到前导轴，得到形状正确但数值错误的结果 —— 例如 x[1:5, idx] 会变成按行 gather。同构的问题在 bool 分支上已经存在（x[::2][:, mask] 实测得到 [2,4]，正确应为 [3,2]）。要加这种回退必须先真正 transpose，否则是引入静默错误。

边界检查这一条很有价值，已按此补上（commit 见上）：

1. kernel 入口校验 slice_offset，同时覆盖排序路径和 elementwise 回落两条路：

const int64_t grad_bytes = x_grad->numel() *
    static_cast<int64_t>(sizeof(T));
PADDLE_ENFORCE_GE(slice_offset, 0, ...);
PADDLE_ENFORCE_LT(slice_offset, grad_bytes, ...);

2. DeriveSortedPathLayout 里补视图跨度检查。原来只有 numel > grad_numel → return false，对跨步视图不够 —— 稀疏视图的元素数小但跨度可能超出缓冲区：

int64_t last = slice_offset / elesize;
for (size_t i = 0; i < view.size(); ++i) last += (view[i] - 1) * vstride[i];
PADDLE_ENFORCE_LT(last, grad_numel, ...);  // 报错信息带 view / strides / slice_offset

回归测试也补了 TestIndexElementwiseGetGradSlicedView：5 个表达式（含非零 offset 的 x[1::2, idx]、x[2:7:2, idx]）× fp32/fp16/bf16 × FLAGS_use_stride_kernel 开关两态，断言覆盖整个 x.grad，所以"梯度块被搬到错误位置"会直接暴露为未触及行非零。

验证情况：

• test_index_elementwise_grad 15 / test_index_elementwise 8 全过，新增 ENFORCE 无误报

• 回归全绿：test_getitem 435 / test_setitem 323 / getitem_appendix 18 / setitem_appendix 12 / test_index_put_op 86 / test_set_value_op 163

• paddle-vs-torch 反向逐位 sweep（--atol=0 --rtol=0）792/792 完全相等，含 x[1::2, idx]、x[2:7:2, idx]、x[::2, idx, ::2] 等跨步视图，4 种 shape × 3 dtype × 重复度 {2,8,40} × 正/负下标

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (develop@0489926). Learn more about missing BASE report.

Additional details and impacted files

@@             Coverage Diff             @@
##             develop    #79716   +/-   ##
===========================================
  Coverage           ?   100.00%           
===========================================
  Files              ?         2           
  Lines              ?         8           
  Branches           ?         0           
===========================================
  Hits               ?         8           
  Misses             ?         0           
  Partials           ?         0           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[liuhao2638]

@Paddle-Bot review此pr

[liuhao2638]

@Paddle-Bot review此pr

[wanghuancoder]

LGTM