Skip to content

Fix fence tracking for donated dynamic slice offsets - #4099

Merged
zcbenz merged 4 commits into
ml-explore:mainfrom
michaelellis003:fix/3880-dynamic-offset-fence-tracking
Aug 11, 2026
Merged

Fix fence tracking for donated dynamic slice offsets#4099
zcbenz merged 4 commits into
ml-explore:mainfrom
michaelellis003:fix/3880-dynamic-offset-fence-tracking

Conversation

@michaelellis003

@michaelellis003 michaelellis003 commented Aug 9, 2026

Copy link
Copy Markdown
Contributor

Proposed changes

Fixes #3880.

Metal's compute_dynamic_offset can donate the dynamic-start input buffer to
its one-element offset output. The donated alias is currently registered as
an encoder temporary. During CommandEncoder::end_encoding(), temporary
cleanup removes that shared pointer from the input and output dependency sets
before producer fences are resolved, so a dynamic slice can use a stale start
after an encoder boundary.

Register the offset as a temporary only when its storage was newly allocated.
The donated path remains enabled, and the primitive input retains the shared
storage through command-buffer completion. This matches the CPU path and the
temporary-lifetime pattern fixed for vector SDPA in #3121.

Per review, the committed regression is now the public slice_update
reproducer from #3880 (no internal APIs).

Testing, on an M3 Pro (36 GB, macOS 26.2), base
8c28c38:

  • Committed regression: on pristine main it fails 10/10 fresh processes
    with default buffer limits (the 512 MB reduction crosses the per-buffer
    memory limit, so the command-buffer boundary lands between the start's
    producer and the slice) and 10/10 with MLX_MAX_OPS_PER_BUFFER=0; with
    the fix, 0/20 in both modes.
  • Public slice reproduction from [BUG] Metal slice_update can use a stale array-valued start across command buffers #3880, fresh processes with
    MLX_MAX_OPS_PER_BUFFER=0: pristine 20/20 wrong; fixed 0/50 wrong.
  • Same reproductions under METAL_DEVICE_WRAPPER_TYPE=1
    METAL_DEBUG_ERROR_MODE=0: pristine 20/20 wrong; fixed 0/20 wrong.
  • Full C++ suite on Metal: 264 test cases, all passing, plain and under the
    Metal validation wrapper.
  • Full Python suite: 810 tests, OK (46 skipped).
  • pre-commit run --all-files: passing.

Checklist

  • I have read the
    CONTRIBUTING
    document
  • I have run pre-commit run --all-files to format my code / installed
    pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature
    works
  • I have updated the necessary documentation (if needed — no public API
    change)

@zcbenz zcbenz left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Adding a test relying on internal platform-dependent API is quite tricky, I think using your test in the original issue should be enough.

@zcbenz
zcbenz force-pushed the fix/3880-dynamic-offset-fence-tracking branch from b448941 to d0b6177 Compare August 11, 2026 03:46

@zcbenz zcbenz left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks!

@zcbenz
zcbenz merged commit 7e8b4cc into ml-explore:main Aug 11, 2026
28 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] Metal slice_update can use a stale array-valued start across command buffers

2 participants