[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 与 reinplace

torch.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, dim=0, start=1, end=3) # out[1:3] == 1，其余为 0

reinplace（重新原地化）：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 上：

1. 步长/偏移计算在 reinplace 时错位：slice_scatter的 reinplace 需要按dim/start/end把 src 拷回 input 的对应位置。CUDA 内核在算目标偏移时，若 input 是非连续或 src 与 input 别名，偏移算错 → 写错位置 / 越界。
2. CUDA 异步 + 别名检测失效：functionalization 靠「别名分析」判断 src 是否和 input 共享存储；CUDA 上某些视图的别名信息在 reinplace 重放时没正确传递，导致「该原地写却新建」或「该新建却原地覆盖」→ 结果错。
3. 与编译优化的交互：torch.compile可能把slice_scatter融合进其它算子，reinplace 边界被打乱，CUDA 内核拿到错误的源/目标指针。 本质：slice_scatter的 reinplace（functional→in-place 重放）在 CUDA 下的别名/偏移处理有 bug，CPU 路径没事，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, dim=0, start=1, end=3) 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, dim=0, start=1, end=3) x = torch.zeros(4, 4, device="cuda") src = torch.ones(2, 4, device="cuda") 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 写」避免 reinplace：

import torch def safe_slice_write(x, src, dim=0, start=1, end=3): # 先 clone，再 slice_scatter（返回新张量），绝不原地改 x out = torch.slice_scatter(x.clone(), src, dim=dim, start=start, end=end) return out # 用法：需要原地语义时，手动赋值回原变量 x = torch.zeros(4, 4, device="cuda" if torch.cuda.is_available() else "cpu") src = torch.ones(2, 4, device=x.device) x = safe_slice_write(x, src) # 用返回值覆盖，而非原地

原理：绕开 functionalization 的 reinplace（它只在「真 in-place 记录」时触发），用slice_scatter的纯 functional 返回值，不触发 reinplace 路径。

六、解决方案二：关掉 torch.compile 的 reinplace / functionalization 相关优化

若 bug 在 compile 路径，可关掉触发 reinplace 的优化：

import torch # 用 fullgraph=False 容错，或关 in-place 重放相关 compiled = torch.compile(f, fullgraph=False) # 更细：某些版本可用选项关闭 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, dim=0, start=1, end=3) 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, dim=0, start=1, end=3): 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]) # 全 1

out[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 赋值」绕过。

九、排查清单

1. slice_scatter在 CUDA 上结果错/崩且涉及 compile / torch.func → 确认 #145189（CUDA reinplace 缺陷）。
2. CPU 正常、CUDA 异常 → 基本可定位到 CUDA reinplace 路径。
3. 绕开：用x.clone()+slice_scatter返回值（纯 functional，不触发 reinplace）；
4. 用@torch.compiler.disable把该算子移出编译图；
5. 等价替代：用out[start:end] = src切片赋值或index_put；
6. 升级：关注该 DISABLED 测试启用即代表修复。

十、小结

DISABLED test_slice_scatter_reinplace_cuda（#145189）的本质是：torch.slice_scatter（切片区写回）在 CUDA 上，当走 functionalization 的「reinplace（functional 视图重放成真正 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 拿返回值再覆盖原变量——既安全又等价。