viable/strict/1791014076: [MPS] Reduce strided inputs in place for full reductions (#198646)
PyTorch 合并 PR #198646，为 MPS 后端新增 FlatStrided 归约计划。此前 #198494 移除 Strided pass-1 kernel 后，dim=None 的全归约对非连续输入会先物化连续副本再跑 Flat，副本搬运字节多于归约读取，导致 fp32 下 x[:, ::2].sum() 比重构前慢 2 倍，max/min/all/any 自迁出 MPSGraph 后也一直承担该拷贝。新方案用 at::collapse_dims 将输入压成 [rows, run] 并原地遍历，覆盖 sum、mean、nansum、count_nonzero、max、min、all、any、var/std。
Development
- First Reportviable/strict/1791014076: [MPS] Reduce strided inputs in place for full reductions (#198646)PyTorch Core
- Industry Responsetrunk/9d75d51a07ee04d24dde44688281b35ead889dce: [MPS] Reduce strided inputs in place for full reductions (#198646)PyTorch Core
- Current Assessment该改动属于框架后端性能回归修复，不涉及模型能力或商业格局。它反映 PyTorch 在 MPS 路径上持续补齐与 CUDA 侧对等的归约实现，并愿意为特定访存模式单独设定线程组预算。可验证信号：后续 release note 或 PR 是否把同类 strided 原地策略推广到其他归约维度或后端。Agent Pulse · analysis
该 PR 针对 MPS 全归约的非连续输入路径做原地化改造：at::collapse_dims 把输入折叠为 [rows, run] 视图，reduction_flat_strided 直接在该视图上遍历，线程组按行 grid-stride 且每行一次 divmod，lane 以固定步长推进 run；短行多个打包进同一线程组，长 run 均分为等长块以便一维 strided 视图仍能铺满整个 grid。实现基于共享 OP 模板，因此 sum、mean、nansum、count_nonzero、max、min、all、any 与 var/std 均可复用。折叠后为 3 个以上 block 的布局，以及无法均分的 run（如质数长度）仍走 .contiguous() + Flat 路径。strided kernel 使用独立线程组预算：每组 2048 元素而非 Flat 的 8192，理由是短行 strided 遍历受延迟约束。
这是一次针对内存带宽与延迟权衡的 kernel 调度调整，而非算子语义变化。判断依据是证据明确给出「副本搬运字节多于归约读取」这一根因，以及 2048 与 8192 两档线程组预算的取舍。可验证的下一信号：在 MPS 上对 x[:, ::2].sum() 等非连续全归约做基准，确认相对 #198494 之后版本恢复或超过重构前水平，并观察 3+ block 与质数长度 run 的退化路径是否仍被触发。
该改动属于框架后端性能回归修复，不涉及模型能力或商业格局。它反映 PyTorch 在 MPS 路径上持续补齐与 CUDA 侧对等的归约实现，并愿意为特定访存模式单独设定线程组预算。可验证信号：后续 release note 或 PR 是否把同类 strided 原地策略推广到其他归约维度或后端。
对在 Apple 芯片上跑 PyTorch 的团队，非连续张量全归约的延迟与带宽开销可能下降，尤其影响逐列切片求和、掩码统计等常见预处理。价值大小取决于工作负载中非连续归约的占比，需以本地基准而非 PR 描述为准。
若该 FlatStrided 计划在真实负载中稳定，PyTorch 可能把「原地遍历非连续视图」作为 MPS 归约的默认优先路径，并逐步缩小 .contiguous() + Flat 的适用范围。观察点是非均分 run 与多 block 布局这两类回退是否在后续 PR 中被继续收敛。