Skip to content

Metal backend: free only the destroyed delegate's memory - #24782

Open
abdelaziz-mahdy wants to merge 23 commits into
pytorch:mainfrom
abdelaziz-mahdy:fix/metal-per-delegate-cleanup
Open

Metal backend: free only the destroyed delegate's memory#24782
abdelaziz-mahdy wants to merge 23 commits into
pytorch:mainfrom
abdelaziz-mahdy:fix/metal-per-delegate-cleanup

Conversation

@abdelaziz-mahdy

@abdelaziz-mahdy abdelaziz-mahdy commented Nov 8, 2026 •
edited by pytorch-bot Bot
Loading

Copy link
Copy Markdown
Contributor

Summary

Stacked on #22957, which lists this as a known limit. Until #22957 lands, the diff includes its commits; this PR's own change is the last commit (Metal backend: free only the destroyed delegate's memory).

MetalBackend::destroy ran cleanup_memory(), which deletes every tensor handle and frees the constants of every delegate in the process. With two models loaded, or two Metal partitions in one .pte, destroying one left the other reading freed memory on its next run.

Stack

PR What Base
#22957 Bind offset views through their parent's buffer main
#22984 Keep non-packed views in place #22957
this PR destroy frees only its own delegate's memory #22957

Changes

  • Each delegate records what it creates. init and execute set the current delegate for the thread (MetalDelegateScope).
    • Handles made while it is current carry its tag in the existing tensors entry, so a run pays no extra allocation.
    • Buffers aoti_torch_mps_malloc makes while it loads (its constants) are recorded against it.
  • destroy waits for queued work, then frees only that delegate's handles and constants. When the last delegate goes it still runs the full cleanup_memory() and cleanup_tensor_metadata(), as before.
  • A failed init releases what it made so far.
  • A recorded buffer that is freed or recycled another way is forgotten, so a later allocation at the same address is never freed with the delegate.

Known limits

  • With other delegates still loaded, two caches are only reset when the last one goes:
    • the buffer pool, which keeps the destroyed delegate's recycled buffers (LRU-bounded);
    • the AOTI size/stride cache. Its entries are checked against the tensor on every read, so this has no effect on results.
  • The shims still assume the caller serializes init, execute and destroy (no locks). That is unchanged.
  • Two paths have no test: the release after a failed init (with a valid program nothing in init fails once it has allocated, so it would need fault injection), and the untagging of a handle whose delete fails.

Test plan

  • test_memory.cpp: 98 gtests, 4 of them new. Each new check was mutation-tested: 8 mutants, each failing at least one test.
  • New backend-level test, test_destroying_one_program_keeps_the_others_memory: a small runner (metal_multi_delegate_runner, built by run_metal_test.sh --build) loads two Metal programs in one process, destroys one, checks it left no handles behind, and runs the other against eager. It fails with the old unconditional cleanup and without the init ownership scope.
  • run_metal_test.sh --build + python -m unittest backends.apple.metal.tests.test_modules.TestMetalBackendModules: 167 OK (4 skipped).
  • lintrunner clean.

This came out of the review of #22957. I kept it separate so that PR stays focused. Happy to fold it in or change the approach if you'd prefer another way to scope the cleanup.

cc @nil-is-all @metascroy

@abdelaziz-mahdy
Metal backend: wait for the parent's GPU writes before reading an off…
0abc566 
…set view

aoti_torch__reinterpret_tensor read views that start partway into a buffer before the GPU had written them.

A non-packed view is copied on the CPU, and the wait for pending GPU work was guarded by metal_is_device_pointer(src). src is the offset pointer, and only a buffer's base address is registered, so the wait was skipped whenever the storage offset was non-zero. A packed view gets its own no-copy MTLBuffer over the parent's memory; Metal tracks hazards per buffer object, so nothing ordered work reading it after the pending work that writes the parent.

Decide on the wait from the base pointer, and wait before aliasing a packed view. Adds two linear-then-chunk modules that fed the second chunk to another linear and came back as zeros.
@abdelaziz-mahdy
Metal backend: bind offset views to their parent's buffer instead of …
9471b22 
…aliasing it

aoti_torch__reinterpret_tensor gave every packed view with a storage offset its own no-copy MTLBuffer over the parent's memory. Metal treats the two buffers as unrelated: a write through one followed by a read of the other in the same serial compute encoder sees stale data every time. Inductor does exactly that when it fills the result of a cat by writing through views of it, and the opposite when an op reads a chunk of a buffer another op just wrote, so models built on split/cat (YOLO's C2f blocks) ran but were wrong.

Register such a view against the buffer it lives in and bind that buffer at the view's offset, in compute encoders and in blits. This replaces the wait-before-aliasing from the previous commit, which only covered reads. MPSGraphTensorData cannot address into a buffer, so MPSGraph ops still get an alias for an offset view, now with the stream synchronized on both sides of the graph.

Adds cat variants of the chunk tests. The last-dim one is skipped: its slices are non-packed views, which reinterpret_tensor materializes into a copy, so writes through them are still lost.
@abdelaziz-mahdy
Cover a chunk of a chunk in the Metal offset-view tests
0a4a1f5 
Inductor flattens nested views into a single reinterpret of the base buffer with the combined offset, so this exercises a larger packed offset rather than a view of a view, but it pins that behavior down.
@abdelaziz-mahdy
Merge remote-tracking branch 'upstream/main' into local/offset-views-…
4ad8cea 
…retain
@abdelaziz-mahdy
Keep a Metal view registered while any handle to it is alive
a3acce5 
A view registration is counted, but only aoti_torch__reinterpret_tensor with a
non-zero offset took a count. aoti_torch_new_tensor_handle on a view, and a
reinterpret of a view at offset 0, both make another handle at the same
address without one, while aoti_torch_delete_tensor_object gives one back for
every handle. Deleting either handle therefore unregistered the view for the
other: a kernel writing through the surviving handle fell back to a temporary
and the write never reached the buffer, and MPSGraph ops failed to find the
tensor's Metal buffer.

metal_retain_view takes a count for an address that is a registered view and
does nothing otherwise; both paths call it.

Tests cover both ways of making the second handle, in both deletion orders.
They need the test_metal_memory target, hence the merge of main.
@abdelaziz-mahdy
Arm the post-graph sync only once the alias buffer exists
89dbe34 
get_mtl_buffer asked the stream to wait after the next graph before it knew
the alias MTLBuffer could be created. On failure it throws, and the flag would
have made the next, unrelated graph wait.

Also reword a comment in materialize_packed that claimed only base addresses
are known as device pointers; constant sub-buffers are mapped at interior
addresses too.
@abdelaziz-mahdy
Tie the alias wait to its own graph, and keep CPU-backed and nested v…
67d8d5d 
…iews working

Review of 89dbe34 found two problems, and a third came up checking the
memory accounting around them.

- The wait that an aliasing buffer needs was a flag on the stream, armed by
  get_mtl_buffer and consumed by whichever graph ran next. Another thread's
  graph could consume it, and an op failing between get_mtl_buffer and its
  graph left it armed for an unrelated one. get_mtl_buffer now only reports
  that it made an alias, the op passes that to its own executeMPSGraph call,
  and the stream waits before and after encoding that graph inside the same
  block on its serial queue, so nothing else can come in between. The stream
  no longer holds any such state.

- A packed offset view of CPU memory lost the no-copy buffer it used to get,
  so MPSGraph ops could not find it and kernels wrote to a temporary. It gets
  one again. Unlike before, the buffer is counted with the view's handles and
  released with the last of them, so that it does not outlive the CPU memory
  under it; a longer view at the same address gets a buffer that covers it.

- A view at an offset took a count on its parent allocation that deleting it
  never gave back, and a view of a view, or a handle copied from a view, took
  none. The parent was then either never freed, or freed while such a handle
  still pointed into it. Every handle into memory the runtime owns now holds a
  count on the allocation it lives in, found through the handle it was made
  from, and gives it back when deleted.

test_metal_memory: 9 new tests. All but AliasedGraphSettlesItsOwnWork fail on
89dbe34; that one covers what the change must keep working.
@abdelaziz-mahdy
Wait for the GPU before a host/device copy touches the memory
090bf1a 
metal_copy_memory copied with memcpy and only then waited for the GPU, and
only when the source was on the device. A copy to the host could read what a
pending command buffer had not written yet, and a copy to the device could
overwrite memory a pending command buffer was still to read. It now waits
before the copy whenever either side is on the device; with nothing pending
the wait does nothing.

The backend itself waits before copying a model's outputs, so exported models
were not affected; direct aoti_torch_copy_ calls were. Two tests, one per
direction, fail without the change.
@abdelaziz-mahdy
Merge remote-tracking branch 'upstream/main' into fix/metal-offset-views
d108d69 
# Conflicts:
#	backends/apple/metal/runtime/shims/memory.cpp
#	backends/apple/metal/runtime/test/test_memory.cpp
@abdelaziz-mahdy
Keep CPU memory until queued GPU work using it is done
110723a 
Deleting the last handle on CPU memory freed it right away, even when a
command still queued on the stream reads or writes it through the no-copy
buffer of a view of it. That buffer does not own the memory, so keeping the
buffer alive does not keep the memory. The stream is now settled before CPU
memory is freed; with nothing queued the wait does nothing. Before the
ownership change the memory under such a view was never freed at all, which
is why this did not show before.

materialize_packed likewise waited only when the parent was in a Metal
buffer. Graphs and kernels can write CPU memory through those no-copy buffers
too, so a non-packed view of it copied what was there before the write. It
now always waits before copying.

QueuedCpuViewKeepsBackingStorageAlive is the reported case; it and
MaterializingCpuMemoryWaitsForGpuWrites fail without the change.
@abdelaziz-mahdy
Bind views of CPU memory into one buffer per region, tied to the memory
ba664e1 
A view of CPU memory got a no-copy buffer of its own. Overlapping views of
the same memory were then separate MTLBuffers, which Metal does not order,
the problem this PR fixes for views of Metal buffers; and a longer view at
the same address replaced the buffer while work was still queued on it.

Views of CPU memory are now bound, at an offset, into one no-copy buffer over
the region they belong to, as views of a Metal allocation are bound into its
buffer. For memory the runtime allocated the region is the whole allocation
(its size is now recorded), and its buffer lives until the memory is freed:
views are deleted and made again all the time, and a new buffer for the same
memory would not be ordered with work queued on the old one. The allocation's
own address resolves to that buffer too. Memory the runtime did not allocate
can be freed and reused behind its back, so its region is the tensor the view
is made from and goes, after a wait, with its last view. A region is only
replaced by a longer buffer when a view reaches past an extent that was not
known, after a wait. The buffers use the default CPU cache mode.

Along the way, from an adversarial review of the whole change:

- An empty view at the end of a buffer was registered at the next
  allocation's address. It now points at its parent's start.
- Freeing CPU memory waited for the GPU every time; it now waits only if the
  memory has a region, the only way the GPU reaches it.
- aoti_torch_copy_ into or out of CPU memory with a buffer was queued as a
  blit and returned before it ran. Copies involving CPU memory now wait and
  are made on the CPU.
- A handle whose view registration failed was left behind; handles are now
  published only once registration succeeded.
- The mark at a view's address is removed with its last handle, so a later
  allocation there is not taken for a view.

A graph fed a CPU view now gets an alias and settles the stream around
itself, so the tests for queued work use a kernel, which binds the buffer at
an offset without waiting. Each new test fails without its change.
@abdelaziz-mahdy
Wait before a packed copy only for GPU work that can touch it
1d6b4c4 
materialize_packed waited for the whole stream before every copy. It now
waits only when queued work may write the source, because the source
overlaps a Metal buffer or a CPU region with a buffer, or may still use
the destination, because it came from the pool and the stream has not
waited since its buffer was freed. Pool entries record the stream's
completed wait count for that.

commit() now keeps the buffer it commits, so that the next commitAndWait
waits for it. A wait after a plain commit used to return before the
committed work was done.

The extent and is_packed_strides now multiply in int64.
@abdelaziz-mahdy
Keep the buffer flush() commits for the next wait
ddfadcd 
With commit-and-continue enabled, flush() released the buffer it
committed, so the next commitAndWait did not wait for that work, while
the buffer pool counts that wait as settling everything queued before
it. flush() now keeps the buffer, as commit() does. With
commit-and-continue disabled, it also releases the one it kept before
instead of leaking it.
@abdelaziz-mahdy
Register blobs that lie inside memory the runtime binds
903f9f7 
A tensor made from a blob at an offset into a Metal buffer, or into CPU
memory the runtime allocated or wraps in a region, was not known to be
there: kernels bound it by copy, graphs refused it, and a view of it got
a no-copy buffer of its own over bytes another buffer covers. Such a blob
is now registered as a view of that memory when it is made, and holds the
allocation it lies in, as a view does. A blob that runs past the end of
its buffer is an error, and so is a CPU view that would make two regions
overlap.

Several handles can share an address, e.g. a blob and a view. The address
now stays tracked until the last of them is deleted, and each handle gives
back only the view registration it took.

The buffer pool only hands a freed buffer back to the stream it was freed
on: a wait on one stream does not cover work queued on another.
@abdelaziz-mahdy
Drop a destroyed stream's buffers from the pool
f76c5e4 
Buffers freed on a stream are only reused on that stream, so once it is
gone they only took up the pool until evicted, and a later stream at the
same address would have been handed them. The stream now takes them out
of the pool when it is destroyed.

The test for materializing memory inside a region again covers a blob
made before a region covers it, which blob registration no longer covers
on its own.
@abdelaziz-mahdy
Refuse what a tensor cannot hold, and fix memory handling main got wrong
f31c7ec 
Every shim now checks its arguments at entry and returns an error instead
of aborting in make_tensor_ptr or truncating: sizes and strides that do
not fit int32 (a dim of size 1 never steps, so its stride may be
anything), strides no dense tensor has, counts and offsets that
overflow, and views outside the memory they are made from. Views with
overlapping dims are materialized like views with holes. Empty tensors
can be allocated and copied, and a constant of size 0 (a null data
pointer) loads.

Blobs are registered against the constant they lie in, never reaching
into the next one, and a CPU region is refused where it would take in
another tensor's memory or a runtime allocation. CPU regions record when
queued work last bound them, and are released or grown with a wait only
if work bound them since the stream last waited.

Already broken on main, fixed here too: aoti_torch_mps_memcpy and
aoti_torch_mps_free wait for queued readers, refuse buffers that are not
constants buffers, and clean up the constants' own buffers;
cleanup_memory waits before freeing; aoti_torch_copy_ no longer reads
past the dims of tensors of different ranks.
@abdelaziz-mahdy
Test that freeing the first constant does not wait for its readers
2240830 
Queued work reads the constant at the start of a constants buffer through
that buffer itself, which owns its memory and which the command buffer
retains until the work is done. Only the constants' own no-copy buffers,
which do not own their memory, need aoti_torch_mps_free to wait first.
@abdelaziz-mahdy
Retire handles in freed memory, and fix op error paths that freed own…
34bf7ee 
…ed memory

aoti_torch_mps_free now forgets the views in any buffer it frees, not only
in constants buffers with nested ones, and retires the handles still in it:
deleting one later only drops the handle, so it cannot unregister or uncount
a later allocation at its address, and shims refuse it. It refuses memory a
tensor owns and a constant's own buffer, which go with the tensor and with
the buffer of all constants.

The CPU region stamp no longer records a stream, which the one-stream model
does not have.

topk freed the memory of its values handle on later errors, and SDPA leaked
its output handle and buffers; both now delete the handle once it owns the
memory. topk refuses a negative k and returns empty tensors for an empty
result, which copy_ now copies out whatever strides either side has. topk
returns its scratch to the pool instead of going through mps_free.
@abdelaziz-mahdy
Merge remote-tracking branch 'upstream/main' into fix/metal-offset-views
67e0c98
@abdelaziz-mahdy
Merge remote-tracking branch 'upstream/main' into fix/metal-offset-views
64917ca
@abdelaziz-mahdy
Keep the offset of an empty view for views made from it
093b315 
An empty view is bound at its parent's start, so that one at the end of its
buffer is not taken for a view of the next allocation. That also dropped its
storage offset: reinterpreting an empty view at offset 8 of an 8-element
allocation as 4 elements at offset 0 addressed elements 0-3 instead of being
refused past the end. The address an empty view stands for is now recorded
(empty_view_address) and views of it, and of its copies, are offset from
there and bounds-checked against it.
@abdelaziz-mahdy
Metal backend: free only the destroyed delegate's memory
6815df1 
MetalBackend::destroy ran cleanup_memory(), which deletes every tensor
handle and frees the constants of every delegate in the process, so
destroying one model left any other loaded model reading freed memory.

Each delegate now records the handles made while it loads or runs, and
the buffers aoti_torch_mps_malloc makes while it loads (its constants).
destroy waits for queued work and frees only those; the full
cleanup_memory() still runs when the last delegate is destroyed. A
failed init releases what it made so far.
Copilot AI balanced review requested due to automatic review settings November 8, 2026 21:54
@pytorch-bot

pytorch-bot Bot commented Nov 8, 2026 •
edited
Loading

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24782

Note: Links to docs will display an error until the docs builds have been completed.

❌ 3 Awaiting Approval, 88 New Failures, 1 Pending, 1 Unrelated Failure, 9 Unclassified Failures

As of commit 36a6283 with merge base 5c788db (image):

AWAITING APPROVAL - The following workflows need approval before CI can run:

NEW FAILURES - The following jobs have failed:

UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:

FLAKY - The following job failed but was likely due to flakiness present on trunk:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Nov 8, 2026
This was referenced Nov 8, 2026
Open
Open
Copilot AI reviewed Nov 8, 2026
View reviewed changes

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🟡 Changes recommended

Regression tests do not exercise backend-level multi-delegate lifecycle or initialization-failure cleanup, nor verify that queued work completes before memory is freed, leaving key use-after-free safeguards unprotected.

3 open issues
What changed in this PR

Scopes Metal memory cleanup to the delegate being destroyed, preserving other live models or partitions. The stacked diff also includes #22957’s offset-view binding and memory-safety changes.

Changes:

  • Tracks delegate-owned handles and constants, with cleanup on destruction and failed initialization.
  • Binds offset views through parent buffers and strengthens synchronization, validation, and error cleanup.
  • Adds memory-lifetime regressions and chunk/view model tests.
File Description
backends/​apple/​metal/​tests/​test_modules.py Adds chunk, concatenation, and nested-view models.
backends/​apple/​metal/​runtime/​test/​test_memory.cpp Adds ownership, view-lifetime, synchronization, and validation tests.
backends/​apple/​metal/​runtime/​shims/​shim_mps.mm Tracks delegate constants and safely releases their buffers.
backends/​apple/​metal/​runtime/​shims/​memory.h Declares delegate ownership and scoped tracking interfaces.
backends/​apple/​metal/​runtime/​shims/​memory.cpp Implements scoped cleanup, view accounting, and argument validation.
backends/​apple/​metal/​runtime/​shims/​et_metal.mm Implements shared view bindings and stream-aware buffer reuse.
backends/​apple/​metal/​runtime/​shims/​et_metal.h Exposes view, buffer, and synchronization helpers.
backends/​apple/​metal/​runtime/​ops/​op_topk.mm Handles empty results, alias synchronization, and failure cleanup.
backends/​apple/​metal/​runtime/​ops/​op_sdpa.mm Releases output handles and buffers on failure.
backends/​apple/​metal/​runtime/​ops/​op_mm.mm Synchronizes graph execution around offset aliases.
backends/​apple/​metal/​runtime/​ops/​op_convolution.mm Adds alias synchronization and failed-output handle cleanup.
backends/​apple/​metal/​runtime/​ops/​op_bmm.mm Passes alias synchronization requirements to graph execution.
backends/​apple/​metal/​runtime/​ops/​op_addmm.mm Passes alias synchronization requirements to graph execution.
backends/​apple/​metal/​runtime/​ops/​common.mm Creates temporary graph buffers for offset views.
backends/​apple/​metal/​runtime/​ops/​common.h Updates the graph-buffer helper contract.
backends/​apple/​metal/​runtime/​metal_backend.cpp Wires ownership scopes and per-delegate lifecycle cleanup.

Review effort: Balanced

Comment thread backends/apple/metal/runtime/metal_backend.cpp
Comment on lines +372 to +375
executorch::backends::aoti::ScopeGuard release_on_failure([&]() noexcept {
if (!initialized) {
metal_release_delegate(handle);
}

@abdelaziz-mahdy abdelaziz-mahdy Nov 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

I looked for a way to make init fail after it allocates. With a valid program, nothing in init fails once it has allocated; a bad weights blob makes AOTI read out of bounds instead of returning an error. So this would need a fault-injection hook in init. The release the guard runs is the one the shim tests cover. Happy to add a hook if you'd like one.

Comment thread backends/apple/metal/runtime/metal_backend.cpp
Comment on lines +701 to +703
if (metal_release_delegate(handle)) {
cleanup_memory();
executorch::backends::aoti::cleanup_tensor_metadata();

@abdelaziz-mahdy abdelaziz-mahdy Nov 8, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Good point, the module tests run one program per process. I added test_destroying_one_program_keeps_the_others_memory: a small runner loads two Metal programs in one process, destroys one, checks it left no handles behind, and runs the other against eager. It fails with the old unconditional cleanup and without the init ownership scope. Dropping the execute scope isn't caught there, because no handle outlives a run unless runtime constant folding is on, which the Metal export doesn't enable; the shim-level test covers that case. (36a6283)

Comment thread backends/apple/metal/runtime/shims/memory.cpp Resolved
@abdelaziz-mahdy
Metal backend: test destroying one of two programs in one process
36a6283 
The module tests run one program per process, so nothing exercised a
delegate being destroyed while another stays loaded. A small runner now
loads two Metal programs, destroys one, checks it left no handles behind
and runs the other, whose output is compared against eager.

Also corrects the comment on the wait in metal_release_delegate: the
frees after it do not depend on it.
Copilot AI balanced review requested due to automatic review settings November 8, 2026 22:22
Copilot AI reviewed Nov 8, 2026
View reviewed changes

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🔵 Needs a closer look

The cross-cutting Metal buffer, no-copy view, and delegate lifetime changes require hardware-backed validation by a Metal maintainer on macOS, which this review environment cannot perform.

2 open issues
1 resolved since last review

Review effort: Balanced

@executorch-triage executorch-triage Bot added the community: contribution PRs coming from community (excluding hardware partners) label Nov 9, 2026
@nil-is-all nil-is-all added partner: apple For backend delegation, kernels, demo, etc. from the 3rd-party partner, Apple release notes: apple Changes to the Apple backend delegate module: metal Issues related to the AOTI Metal backend labels Nov 10, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

Copilot code review Copilot
Copilot review effort
Applies to this pull request for everyone.Learn more about Copilot code review.
Copilot left review comments
@shoumikhin shoumikhin Awaiting requested review from shoumikhin shoumikhin is a code owner
@metascroy metascroy Awaiting requested review from metascroy metascroy is a code owner
@larryliu0820 larryliu0820 Awaiting requested review from larryliu0820 larryliu0820 is a code owner

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. community: contribution PRs coming from community (excluding hardware partners) module: metal Issues related to the AOTI Metal backend partner: apple For backend delegation, kernels, demo, etc. from the 3rd-party partner, Apple release notes: apple Changes to the Apple backend delegate

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@abdelaziz-mahdy @nil-is-all