Issue Station
Sign In
Tools hereToolbox

pytorch - ✅(Solved) Fix [MPS] aten::copy_ into strided view silently wraps writes at element offset > 2^32 [1 pull requests, 1 participants]

Official PRs (…)
ON THIS PAGE

Recommended Tools

×6

Utilities matched from this issue’s tags and category — try them while you read without losing context.

GitHub issue graph ai analysis

Paste a GitHub issue URL. We fetch that issue, discover linked issues from bodies/comments/timeline, collect linked pull requests, and produce a structured English report.

The report is written in English Markdown for sharing and archival.

Helpful · Quick feedback

Loading…
GitHub stats
pytorch/pytorch#182052•Fetched 2026-05-01 05:32:40
View on GitHub
Comments
0
Participants
1
Timeline
30
Reactions
0
Author
jhavukainenjhavukainen
Participants
jhavukainenjhavukainen
Timeline (top)
mentioned ×12subscribed ×12labeled ×5cross-referenced ×1

Fix Action

Fixed

PR fix notes

PR #182054: [MPS] Fix uint32 offset overflow in scatter/gather kernels for strided views crossing 2^32 elements

Description (problem / solution / changelog)

Fixes issue: #182052

Addresses the wrapping around in scatter/gather kernels when relative linear offset in storage greater than 2^32. In aten/src/ATen/mps/IndexKernels.h scatter and gather variants use uint32 stride arithmetic, so stride[d] * idx[d] truncates to 32 bits before summing to the accumulator. CPU works correctly.

Adds regression tests against the issue. Note that the tests require a relatively hefty amount of memory due to the error only reproducing at the stride values crossing the uint32 boundary. Especially if the second test fails (like the current nightly) since the assert framework computes the statistics to report. Let me know if we should work around this somehow to avoid issues with the CI, but the current CPU vs. MPS formulation of the test seemed to be the cleanest way to demonstrate the issue here.

Changed files

  • aten/src/ATen/mps/IndexKernels.h (modified, +216/-22)
  • aten/src/ATen/native/mps/operations/View.mm (modified, +45/-13)
  • test/test_mps.py (modified, +38/-0)

Code Example

import torch

# [U=32, L=48, Hkv=8, C=4096, Eh=128] fp16 = 12 GiB.
SHAPE = (32, 48, 8, 4096, 128)

cpu_x = torch.zeros(SHAPE, dtype=torch.float16)
mps_x = cpu_x.to("mps")
update = torch.ones((SHAPE[0], 1, SHAPE[2], 1, SHAPE[4]), dtype=torch.float16)

cpu_x[:, 0:1, :, 0:1, :] = update
mps_x[:, 0:1, :, 0:1, :] = update.to("mps")

mismatched = (cpu_x != mps_x.cpu()).sum().item()
print(f"mismatched elements: {mismatched}  (expected 0)")

# On broken torch: 20480
#   - 10,240 at layer  0 (users 22..31 never written)
#   - 10,240 at layer 32 (users 0..9 clobbered by the wrap)
RAW_BUFFERClick to expand / collapse

🐛 Describe the bug

On MPS, aten::copy_ into a strided destination view silently routes writes incorrectly when any element of the view has a relative linear offset in storage greater than 2^32. The MPS scatter/gather kernels in aten/src/ATen/mps/IndexKernels.h use uint32 stride arithmetic, so stride[d] * idx[d] truncates to 32 bits before summing to the accumulator. CPU gives the correct result.

Plain-slice __setitem__ lowers to aten::slice + aten::copy_ so any x[slices] = update into a large enough destination is affected. Contiguous full-tensor copies are unaffected as they use MTLBlit path which tracks 64-bit offsets apparently.

Minimal repro that fails with the current nightly

import torch

# [U=32, L=48, Hkv=8, C=4096, Eh=128] fp16 = 12 GiB.
SHAPE = (32, 48, 8, 4096, 128)

cpu_x = torch.zeros(SHAPE, dtype=torch.float16)
mps_x = cpu_x.to("mps")
update = torch.ones((SHAPE[0], 1, SHAPE[2], 1, SHAPE[4]), dtype=torch.float16)

cpu_x[:, 0:1, :, 0:1, :] = update
mps_x[:, 0:1, :, 0:1, :] = update.to("mps")

mismatched = (cpu_x != mps_x.cpu()).sum().item()
print(f"mismatched elements: {mismatched}  (expected 0)")

# On broken torch: 20480
#   - 10,240 at layer  0 (users 22..31 never written)
#   - 10,240 at layer 32 (users 0..9 clobbered by the wrap)

I'll file a PR for utilizing 64-bit dtypes in the stride + offset computation to fix the issue.

Versions

torch==2.13.0.dev20260430 OS: macOS 26.4 (arm64) CPU: Apple M2 Max

cc @ezyang @gchanan @kadeng @msaroufim @kulinseth @malfet @DenisVieriu97 @aditvenk

extent analysis

TL;DR

The issue can be fixed by utilizing 64-bit dtypes in the stride and offset computation for MPS scatter/gather kernels.

Guidance

  • The root cause of the issue is the use of uint32 stride arithmetic in the MPS scatter/gather kernels, which truncates to 32 bits before summing to the accumulator.
  • To verify the issue, run the provided minimal repro code and check if the number of mismatched elements is greater than 0.
  • The issue can be mitigated by avoiding large enough destinations in x[slices] = update operations or by using contiguous full-tensor copies which are unaffected.
  • A fix can be implemented by modifying the aten/src/ATen/mps/IndexKernels.h file to use 64-bit dtypes for stride and offset computation.

Example

No code snippet is provided as the fix requires modifying the PyTorch source code.

Notes

This issue is specific to the MPS backend and does not affect CPU or other backends. The fix will require modifying the PyTorch source code and may not be applicable to all versions of PyTorch.

Recommendation

Apply workaround by avoiding large enough destinations in x[slices] = update operations until a fixed version of PyTorch is available.

Vote matrix · Quick signals

Works
Did the solution work? Tap to confirm.
Easy Fix
Was it a quick fix?
Time Saver
Did it save you time?
Blocking
Was it severely blocking?
Common Issue
Are others likely hitting this too?
Flaky / Intermittent
Is it intermittent?
Verified / Reproducible
Can you reproduce it reliably?
Loading…
#dependency conflict#environment setup#docker error#permission error#memory optimization

Still need to ship something?

×6

Another batch ranked right after the header list — different links, same matching logic.

Back to top recommendations

TRENDING