[alirezamonfared]

Summary:

TL;DR

Give independently loaded ExecuTorch Modules independent QNN manager lifetimes while preserving context sharing between methods within one module. Key the delegate cache by (context signature, runtime allocator), count method references, and release complete manager state on final destruction or initialization failure.

Context

QnnExecuTorchBackend is process-global and caches expensive QnnManager instances across methods. A manager is placement-allocated from a module's runtime allocator; methods in that module share its allocator, but independent modules can load identical .pte bytes and be destroyed separately. D122851825 hardens this backend contract without changing FCB export contents or artifact selection.

Problem

A context signature alone also matched independent modules, allowing one module to destroy manager state still used by another. Calling Destroy() alone also omitted non-trivial manager member destruction, and competing initializers could publish or abandon duplicate candidates. A process-global HTP spill-fill handle did not represent the lifetime of its owning context.

Before: process cache[signature] -> Manager M (allocator A)

Module A methods ----> M <---- Module B, same program bytes
release Module A ----> M destroyed
Module B executes ---> stale context

Solution

Use a cache key containing the signature and allocator, with mutex-protected lookup, membership, publication and release. Each loaded method holds a reference; try_emplace preserves the first published manager and competing initializers destroy their unused candidate. A destructor-only ownership guard covers backend, context, DLC compilation and tensor allocation failures. Full QnnManager destruction also releases its input/output tensor wrappers and other non-trivial members; their storage was not released by calling Destroy() alone. On final release, erase cache state under the mutex and destroy the manager after unlocking. A missing forward entry is reported and its stale reverse entry removed safely.

After: cache[(signature, allocator A)] -> Manager A -> A's methods
       cache[(signature, allocator B)] -> Manager B -> B's methods

release A's last method -> unlock cache -> ~Manager A -> Destroy()
B's methods            -> Manager B remains alive

Scope HTP spill-fill groups by allocator and configured buffer size. Thread the group handle explicitly through context initialization instead of using a process-global handle. Serialize initialization within each group, retain its first manager until all members and pending initializations are gone, and destroy member contexts before the owner. The registry holds weak group references, preventing an ownership cycle and allowing cleanup followed by a fresh group.

weak registry[(allocator, spill-fill bytes)] -> Group
method entry -> member manager + Group
Group        -> first manager (HTP owner)

last member/pending initializer releases Group
  -> member context destroyed
  -> first context destroyed
  -> registry's weak reference expires

Keep the registered backend alive for the process lifetime so loaded modules cannot destroy QNN contexts through global cache teardown after allocators or the QNN library have finalized. Normal method destruction still releases manager state; the explicit defaulted backend destructor avoids conflicting empty-body formatting between local and OSS ClangFormat versions. Python lifecycle imports follow the ordering reported by OSS UFMT.

Compatibility and evidence limits

Separate live modules allocate separate QNN contexts, which can increase HTP memory use for duplicate blobs. Methods within one module still share a manager. FCB weight sharing inside an artifact does not supply shared runtime ownership across independent modules. Concurrent execution through one module and execution racing destruction require client coordination.

QAI Hub shard spill-fill cost

gen_pte_from_ctx_bin() calls update_spill_fill_size() over the shard programs before writing separate .pte files. The helper sets each participating context's max_sf_buf_size to the largest requirement in the export group. The QAI Hub Llama runner constructs a separate Module for each shard and loads every shard method; default Module construction gives each one a separate runtime allocator.

On device, HtpContextCustomConfig::CreateContextCustomConfig() registers an HTP group using firstGroupHandle and maxSpillFillBuffer. The former process-global handle allowed shard contexts to join one group. With this diff, the (allocator, max_sf_buf_size) key isolates those independently allocated Modules, so each shard starts its own spill-fill group even when the exported size matches.

N resident shard Modules, common configured S bytes (S > 0)

Before: shard 1 ... shard N -> one HTP spill-fill group -> S bytes
After:  shard i / allocator i -> group i              -> S bytes each

Expected configured reservation: S -> N * S bytes
Expected additional reservation: (N - 1) * S bytes

This is a projection of the spill-fill reservation for the described shard layout, not measured RSS, PSS, DSP heap usage or peak device memory. SDK allocation behavior, alignment and other context/model memory affect the observed totals. The existing host simulator tests do not issue target spill-fill group registration and cannot measure this device cost. No connected physical device was available for this follow-up, so the magnitude on a representative QAI Hub shard model remains unmeasured.

If device evidence shows a material cost, an opt-in cross-Module group ID is a possible follow-up. Such a group must keep its first context and allocator-owned manager storage alive until all shard Modules and pending initializations release it; changing the key to an ID alone would not satisfy that lifetime requirement. This diff does not introduce that opt-in.

Exported lifecycle tests cover both module destruction orders, multiple methods, cleanup/reload and sibling unload/reload on the HTP simulator, including multi-context mode with a nonzero spill-fill size. C++ tests check handle identity, allocator isolation, reference accounting, concurrent publication, stale-map cleanup, spill-fill ownership and process exit. Target SDK spill-fill registration still requires device verification; the recorded FCB AOT failure occurred before runtime module loading and does not establish FCB runtime compatibility.

Authored with AI assistance from Codex.

Differential Revision: D122851825

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: alirezamonfared / name: Ali Monfared (8de26ed)

[meta-codesync]

@alirezamonfared has exported this pull request. If you are a Meta employee, you can view the originating Diff in D122851825.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[psiddh]

@alirezamonfared another suggestion while we're here: if InitBackend/InitContext/AllocateTensor fails below, we return without calling ~QnnManager(). Since init failed, ExecuTorch never calls destroy() for this handle, so the QNN backend/context handles leak. Matters for apps that retry or fall back to another backend. A scope guard right after the placement new, dismissed just before add_cached_delegate, would cover all three error paths. wdyt ?

[qti-horodnic]

Thanks @alirezamonfared, left a few comments. As I mentioned in one of them, for modules that load the same blob, memory usage might increase since there will multiple contexts. Can you confirm that no FCB flow loads the same blob bytes into several Modules (e.g. weight-sharing bins split across .pte files)? I think it should be fine but I want to make sure there are no FCB regressions

[qti-horodnic]

Two init calls with the same key can both miss the cache. The second call overwrites delegate_map_[key], but delegate_map_rev_ keeps both handles, so a later erase removes the wrong entry. Can we use try_emplace instead? And if the key exists, destroy the new manager and return the cached handle.

[qti-horodnic]

This is a pre-existing issue, but if InitBackend, InitContext, CompileDlc or AllocateTensor fails, ~QnnManager() doesn't run. Can we call qnn_manager->~QnnManager() before each error return or use a scope guard?

[qti-horodnic]

Modules that load the same blob no longer share one QNN context, so HTP memory can increase. Also, destroy() now runs the full destructor, so the tensor wrappers are freed. Please mention this in the PR description.

[qti-horodnic]

Can we add a test for this branch? A unit test that loads the same .pte in 2 programs, deletes the first, and executes the second. Also can we check that 2 methods of one Module still get the same handle?

[qti-horodnic]

Also, looking a bit more at this, the methods of one Module share a handle but nothing counts the users. So there can be case where unload_method("prefill") destroys the manager while decode is loaded, then decode fails with "DelegateHandle has been deleted". This is another existing issue which not caused by this pr, do you want to handle this here too? Maybe keeping a user count per cache entry and run ~QnnManager() only when the count is 0? This will make the solution more robust, but I'm fine if you want to keep the scope more contained.

[qti-horodnic]

Minor: pre-existing typo, I guess should be temporary, not temporal.

[psiddh]

@alirezamonfared both CLA & linters need a fix

[psiddh]

@claude can you review mine and Arik's comments on this PR and check if they all have been addressd

[claude]

Claude finished @psiddh's task in 2m 15s —— View job

---

All of @qti-horodnic's and @psiddh's code comments are addressed in the current revision, but lintrunner is still failing (psiddh's last comment), and the CLA check still needs the author to link their commit email.

Comment-by-comment status

@qti-horodnic (review, CHANGES_REQUESTED)

#	Comment	Status	Where
1	Use try_emplace so two racing inits don't overwrite delegate_map_ and leave delegate_map_rev_ out of sync; destroy the losing manager and return the cached handle	✅ Addressed	add_cached_delegate uses try_emplace, adds to delegate_map_rev_ only on insert, and returns the winner's handle. init returns that handle. The loser's manager_guard goes out of scope and runs ~QnnManager(). ConcurrentPublicationKeepsOneManager tests this.
2	~QnnManager() doesn't run if InitBackend/InitContext/CompileDlc/AllocateTensor fails	✅ Addressed	std::shared_ptr<QnnManager> manager_guard(qnn_manager, destroy_manager) is created right after the placement new, so every early ET_CHECK_OR_RETURN_ERROR return destroys the manager. ~QnnManager() calls Destroy() (QnnManager.cpp:66).
3	Mention the possible HTP memory increase and the full-destructor change in the PR description	✅ Addressed	The description now covers both ("Separate live modules now allocate separate QNN contexts…" / "Full QnnManager destruction also releases its tensor wrappers…").
4	Add a test that loads the same .pte twice, deletes the first, and runs the second; also check that two methods of one Module share a handle	✅ Addressed, with a caveat	test_qnn_shared_delegate.py covers both destruction orders, same-module sharing (it checks the "Use cached delegate handle … shift" log), cleanup/reload, and sibling unload/reload. It also runs with spill-fill enabled. The new QnnExecuTorchBackendTest.cpp covers the cache logic. Caveat: both targets exist only in backends/qualcomm/tests/BUCK, so the OSS "Test QNN Backend" job never runs them. Its green result says nothing about these tests.
5	Typo: temporal should be temporary	✅ Fixed	QnnExecuTorchBackend.cpp:71
6	Count each method's use of a shared handle so unload_method("prefill") doesn't destroy the manager while decode is still loaded	✅ Addressed	DelegateEntry::users goes up on every cache hit in get_cached_delegate and on a lost race in add_cached_delegate. erase_cached_delegate only drops the entry when the count reaches 0. SharesUntilLastMethodIsReleased and test_method_survives_sibling_unload_and_reload test this.
7	Confirm no FCB flow loads the same blob bytes into several Modules	⚠️ Answered in the description, not confirmed in a reply	The description says "FCB export contents and selection are unchanged; FCB weight sharing within an artifact does not provide shared runtime ownership across independent modules." I spot-checked the in-tree runners. qnn_llama_runner.cpp:301-311 and qnn_multimodal_runner.cpp:271-288 each create one Module per different .pte, so their signatures already differ. extension::Module shares one memory_allocator_ across its methods (module.cpp:630), so prefill/decode in one Module still share a context. I found no in-tree regression. Out-of-tree apps that load the same .pte into two Modules will now use two contexts. @qti-horodnic may want the author to confirm this explicitly.

@psiddh

#	Comment	Status	Where
1	Use a scope guard right after the placement new, dismissed before add_cached_delegate, to cover the init error paths	✅ Addressed	Same fix as horodnic #2. It's a shared_ptr guard rather than a dismissable one: on success, ownership is shared with the cache entry (and with group->owner for the first spill-fill context), so dropping the local copy does not destroy the manager.
2	CLA and linters need a fix	❌ Not addressed	Lint (run 41780739632) fails in two places: (a) CLANGFORMAT on QnnExecuTorchBackend.h:29, because the PR changed ~QnnExecuTorchBackend(){}; to ~QnnExecuTorchBackend() {};. Reverting it also removes an unrelated diff line. (b) UFMT on test_qnn_shared_delegate.py: the torch.testing.assert_close(...) calls at lines 136-139 and 143-146 should be on one line, and class TestQnnSpillFillSharedDelegate at line 155 needs two blank lines before it. lintrunner -a fixes both. CLA: EasyCLA says the commit email for 2f58db1 isn't linked to the author's GitHub account. Only the author can fix this.

Remaining items

• Run lintrunner -a and push the result. Fix this →
• Fix the CLA email link (author).
• Optional: say in the PR how the new BUCK-only tests were run, since OSS CI doesn't run them.

[psiddh]

Can you fix Lint / lintrunner issues and should be good to land (@qti-horodnic let me know if that is not the case)

[qti-horodnic]

Thanks @alirezamonfared, I think this pr can proceed. The comment below is non-blocking but something we should consider if possible

[qti-horodnic]

The qaihub shard Modules are exported to share one spill-fill buffer (update_spill_fill_size). The allocator in the key now gives each shard its own buffer, so the extra memory is (N−1) × max_sf_buf_size. Can we measure this on device and mention it in the description? If the cost is large, we can add an opt-in cross-Module sharing in a follow-up, for example with a group ID.

[psiddh]

Need the linter fix, accepting to unblock