[Cherry-Pick] Support iteration space beyond INT32_MAX in GPU index_elementwise_put kernels(#79810) - #79811
Open
omoYang wants to merge 3 commits into
Open
[Cherry-Pick] Support iteration space beyond INT32_MAX in GPU index_elementwise_put kernels(#79810)#79811omoYang wants to merge 3 commits into
omoYang wants to merge 3 commits into
Conversation
omoYang
changed the title
omoYang
changed the title
risemeup1111
approved these changes
Sep 22, 2026
risemeup1111
approved these changes
Sep 22, 2026
The forward kernel rejected numel > INT32_MAX unconditionally, so a large indexed region failed even on the 64-bit offset path. The grad kernel had no check at all and silently wrapped its 32-bit linear index instead. The offset type is chosen from the operands' byte spans, which does not bound the iteration space: a small index tensor broadcast over a large input still produces numel > INT32_MAX. Re-dispatch the 32-bit instantiation to the 64-bit one in that case, and replace the numel guard with a grid-size check, which is the real limit once the linear index is 64-bit wide. Also restrict the value_grad numel guards in the grad kernel to the 32-bit instantiation.
…rload GPUIndexElementwisePutWithTensorKernel was left out of the previous commit, but it has the same exposure: the offset type is picked from the operands' byte spans, and a broadcast `value` keeps its span small while the iteration space can still exceed INT32_MAX, so the 32-bit instantiation would truncate the linear index passed to OffsetCalculator::get(). Re-dispatch to the 64-bit instantiation and add the grid-size check.
OffsetCalculator's constructor already validates every operand's byte reach against the offset type it was instantiated with (CheckOffsetRange, called whenever signed_strides is set), and that bound is both tighter and more precise than comparing value_grad's element count to INT32_MAX: value_grad here is always contiguous and input-shaped, so its byte reach is (numel - 1) * sizeof(T).
omoYang
force-pushed
the
cherry-pick-79810-release3.4
branch
from
September 24, 2026 08:30
d4f3376 to
a725a12
Compare
risemeup1111
approved these changes
Sep 24, 2026
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
PR Category
Operator Mechanism
PR Types
Bug fixes
Description
dev PR: #79810
x[index] = value这类 advanced indexing 的 setitem 在 GPU 上走index_elementwise_put/index_elementwise_put_with_tensor及其反向。当被索引区域的迭代空间numel超过INT32_MAX时：N = numel做了无条件的PADDLE_ENFORCE，与实际使用的 offset 位宽无关，即使调度端已选 64 位 offset 也会直接报the value of N should be in [0, std::numeric_limits<int32_t>::max()]。实际触发场景是大 padding 率样本下input_embeds[indices] = value，len(indices) * hidden_size越界导致训练中断。offset_calc.get()的形参就是uint32_t，线性下标被静默截断，写入落到错误地址。改动内容：
numel > INT32_MAX时重新派发到 64 位实例化，位置在构造 offset calculator 之前（32 位的构造函数本身会对越界偏移抛错）。numel守卫，改为检查 grid 维度——线性下标放宽到 64 位后真正的硬上限是gridDim.x，且dim3构造存在int64_t到unsigned int的隐式窄化。value_grad->numel()守卫：OffsetCalculator构造函数（signed_strides = true下）的CheckOffsetRange已经逐操作数校验偏移是否装得进 offset 类型，且更精确。本 PR 是 #79810 到 release/3.4 的等价 cherry-pick。#79731 合入后 release/3.4 的
index_elementwise代码已与 develop 对齐（MakeOffsetCalculatorPut<3, true>、CheckOffsetRange、IsInInt32Range(..., IndexOperandByteSpan, StridedOperandByteSpan)判据均在），因此三个 commit 无冲突原样应用，与 #79810 逐字节一致（仅phi::SizeOf命名空间限定符按各分支既有风格保留）。复现
buf取[2, 1000000]float32（8 MiB，字节跨度远在 uint32 内）、索引取长度rep的 int64 张量、行 1 只出现在索引尾部时：rep = 4075（迭代空间 4.075e9）正常；rep = 4100（4.1e9）在未修复版本上直接CUDA error(700) illegal memory access。阈值低于
2**32：32 位实例化的线性下标走IntDivider<unsigned int>::div（t = __umulhi(n, m1); return (t + n) >> shift;），其中t + n是 32 位加法，对最大 extentd可用的下标上界是2**32 * d / 2**ceil(log2 d)，最坏情况接近2**31。这也是重派发阈值取INT32_MAX的依据。是否引起精度变化
否。改动只影响下标与偏移的计算位宽，以及守卫条件的作用范围，不改变任何计算逻辑。