[Bug已解决] slice_scatter 在 CUDA 上重新原地-reinplace-失效 DISABLED test 绕过方案解决方案

[Bug已解决] slice_scatter 在 CUDA 上重新原地reinplace失效DISABLED test绕过方案解决方案一、现象长什么样你用了 PyTorch 的torch.slice_scatter把一段切片写回原张量的指定位置并且场景涉及reinplace重新原地化 / functionalization 下的 in-place 重放——常见于torch.compile的 functional 化路径、或torch.func/ functorch 的变换里。结果遇到一个被官方DISABLED的测试DISABLED test_slice_scatter_reinplace_cuda (__main__.GPUTests)即 pytorch/pytorch#145189。含义slice_scatter在 CUDA 上当走「reinplace把 functional 视图重放成真正的 in-place 写回」路径时行为不正确 / 崩溃测试被禁用。 对你来说实际表现可能是torch.compile后的模型里用到slice_scatter或tensor[i:j] x这种切片区赋值被编译成 slice_scatter在 CUDA 上结果错或崩而 CPU 上正常。 本文聚焦slice_scatter 是什么、reinplace 是什么、为什么 CUDA 上这组合会坏、怎么绕过用 clone 替代 in-place、关 functionalization 的 reinplace 等。二、背景slice_scatter 与 reinplacetorch.slice_scatter(input, src, dim, start, end)返回一个新张量是在input的dim维度[start:end]位置放入src。它是input[..., start:end, ...] src这种「切片区赋值」的函数式版本不修改 input返回新张量。import torch x torch.zeros(4, 4) src torch.ones(2, 4) out torch.slice_scatter(x, src, dim0, start1, end3) # out[1:3] 1其余为 0reinplace重新原地化PyTorch 2.x 的functionalization函数式化机制把「in-place 操作」先记录成「functional 视图」等需要时再「reinplace重放成真正的原地写回」。这在torch.compile/torch.func里很常见——编译器为了能 safely 做变换先把 in-place 变 functional最后再决定要不要真的 in-place。 问题slice_scatter的 reinplace 路径在 CUDA 上有缺陷——functional 记录与真正重放时的设备/步长/别名处理不一致导致重放写错位置或崩溃。三、为什么 CUDA 上 slice_scatter reinplace 会坏functionalization 的 reinplace 要处理「源张量 src 与目标切片是同一内存别名」的情况。在 CUDA 上步长/偏移计算在 reinplace 时错位slice_scatter的 reinplace 需要按dim/start/end把 src 拷回 input 的对应位置。CUDA 内核在算目标偏移时若 input 是非连续或 src 与 input 别名偏移算错 → 写错位置 / 越界。CUDA 异步 别名检测失效functionalization 靠「别名分析」判断 src 是否和 input 共享存储CUDA 上某些视图的别名信息在 reinplace 重放时没正确传递导致「该原地写却新建」或「该新建却原地覆盖」→ 结果错。与编译优化的交互torch.compile可能把slice_scatter融合进其它算子reinplace 边界被打乱CUDA 内核拿到错误的源/目标指针。 本质slice_scatter的 reinplacefunctional→in-place 重放在 CUDA 下的别名/偏移处理有 bugCPU 路径没事CUDA 路径出错。四、最小可运行复现带守卫下面演示slice_scatter的正常用法CPU 稳以及「被 compile 函数式化后 CUDA 可能出问题」的示意import torch def demo_slice_scatter_cpu(): x torch.zeros(4, 4) src torch.ones(2, 4) out torch.slice_scatter(x, src, dim0, start1, end3) print(slice_scatter 结果:\n, out) # out[1:3] 应为 1 def demo_with_compile(): if not torch.cuda.is_available(): print([skip] 无 GPU仅说明 compile slice_scatter) return def f(x, src): return torch.slice_scatter(x, src, dim0, start1, end3) x torch.zeros(4, 4, devicecuda) src torch.ones(2, 4, devicecuda) try: g torch.compile(f) out g(x, src) print(compile slice_scatter 形状:, out.shape) except Exception as e: if slice_scatter in str(e).lower() or reinplace in str(e).lower(): print([确认] 命中 slice_scatter reinplace CUDA 问题 (#145189):, e) else: print(其它错误:, e) if __name__ __main__: demo_slice_scatter_cpu() demo_with_compile()要点CPU 上 slice_scatter 稳定bug 在「CUDA functionalization reinplace」组合。下面给绕过。五、解决方案一避免 in-place 语义用 clone 替代最根本的绕过不要用「切片区原地赋值」的语义改用「先 clone 再 functional 写」避免 reinplaceimport torch def safe_slice_write(x, src, dim0, start1, end3): # 先 clone再 slice_scatter返回新张量绝不原地改 x out torch.slice_scatter(x.clone(), src, dimdim, startstart, endend) return out # 用法需要原地语义时手动赋值回原变量 x torch.zeros(4, 4, devicecuda if torch.cuda.is_available() else cpu) src torch.ones(2, 4, devicex.device) x safe_slice_write(x, src) # 用返回值覆盖而非原地原理绕开 functionalization 的 reinplace它只在「真 in-place 记录」时触发用slice_scatter的纯 functional 返回值不触发 reinplace 路径。六、解决方案二关掉 torch.compile 的 reinplace / functionalization 相关优化若 bug 在 compile 路径可关掉触发 reinplace 的优化import torch # 用 fullgraphFalse 容错或关 in-place 重放相关 compiled torch.compile(f, fullgraphFalse) # 更细某些版本可用选项关闭 functionalization 的 reinplace # torch._functorch.config.* 视版本常见做法是避免在该算子前后做 in-place另一个思路把用到slice_scatter的函数从编译图里摘出来用torch.compiler.disable包住import torch torch.compiler.disable def slice_op(x, src): return torch.slice_scatter(x, src, dim0, start1, end3) def model(x, src): y other_compute(x) y slice_op(y, src) # 这一段不进编译图避开 reinplace 缺陷 return y七、解决方案三用等价的 index_put / 索引赋值替代slice_scatter可用index_put_/ 高级索引替代避开该算子的 reinplace 路径import torch def slice_scatter_via_index(x, src, dim0, start1, end3): out x.clone() # 用切片索引赋值等价语义 if dim 0: out[start:end] src elif dim 1: out[:, start:end] src return out x torch.zeros(4, 4) src torch.ones(2, 4) out slice_scatter_via_index(x, src) print(out[1:3]) # 全 1out[start:end] src这种 Python 切片赋值在 functionalization 下走的是不同的更成熟的reinplace 路径不一定触发 slice_scatter 那个缺陷。八、解决方案四升级并关注该测试启用#145189 是内部测试被禁用官方会修CUDA 下 slice_scatter 的 reinplace 别名/偏移处理。升级pip install --upgrade torch --index-url https://download.pytorch.org/whl/cu124判断修复该 DISABLED 测试被启用且通过即代表 CUDA 下 slice_scatter reinplace 修好。修复前用「clone 替代 in-place / disable 编译段 / 等价 index 赋值」绕过。九、排查清单slice_scatter在 CUDA 上结果错/崩且涉及 compile / torch.func → 确认 #145189CUDA reinplace 缺陷。CPU 正常、CUDA 异常 → 基本可定位到 CUDA reinplace 路径。绕开用x.clone()slice_scatter返回值纯 functional不触发 reinplace用torch.compiler.disable把该算子移出编译图等价替代用out[start:end] src切片赋值或index_put升级关注该 DISABLED 测试启用即代表修复。十、小结DISABLED test_slice_scatter_reinplace_cuda#145189的本质是torch.slice_scatter切片区写回在 CUDA 上当走 functionalization 的「reinplacefunctional 视图重放成真正 in-place 写回」路径时别名/偏移处理有 bug——src 与 input 别名、或非连续布局下重放写错位置或崩溃。CPU 路径没事CUDA 路径出错。 应对纯 functional 绕过先x.clone()再slice_scatter取返回值绝不原地改不触发 reinplace移出编译图torch.compiler.disable包住该算子避开 compile 的 reinplace 缺陷等价替代用out[start:end] src切片赋值或index_put走更成熟的 reinplace 路径等升级关注该 DISABLED 测试启用即代表修复。 记住slice_scatter 的 reinplace 在 CUDA 下有别名/偏移 bug。想要原地语义就先 clone 拿返回值再覆盖原变量——既安全又等价。