Skip to content

[Executorch][LLM] Use caching allocator for runner #15730

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
kimishpatel wants to merge 77 commits into
base: main
Choose a base branch
from
Open

[Executorch][LLM] Use caching allocator for runner #15730

kimishpatel wants to merge 77 commits into from
+20 −3

Conversation

@kimishpatel
Copy link
Contributor

@kimishpatel kimishpatel commented Nov 11, 2025 •
edited
Loading

Stack from ghstack (oldest at bottom):

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: D86120038

@kimishpatel
[Executorch] parallelize op_choose_qparams
3f789b8 
When doing prefill for quantized kv cache, with large prefill length, parallelizing this op helps.

Differential Revision: [D84962234](https://our.internmc.facebook.com/intern/diff/D84962234/)

**NOTE FOR REVIEWERS**: This PR has internal Meta-specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D84962234/)!

[ghstack-poisoned]
@kimishpatel
[Executorch] Add simd path for op quantize
08dd980 
Reason this doesnt directly use Vectorize class is because the equivalent APIs dont exist in Vectorize class

Differential Revision: [D84962236](https://our.internmc.facebook.com/intern/diff/D84962236/)

**NOTE FOR REVIEWERS**: This PR has internal Meta-specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D84962236/)!

[ghstack-poisoned]
@kimishpatel
[Executorch] Add multithreading for op_quantize
27fc8b1 
As the title

Differential Revision: [D84962233](https://our.internmc.facebook.com/intern/diff/D84962233/)

**NOTE FOR REVIEWERS**: This PR has internal Meta-specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D84962233/)!

[ghstack-poisoned]
@kimishpatel
Reduce allocation overhead in quantized sdpa
ae61ab4 
For small models dequantizing portions of v cache causes extra alloc overhead.

Probably a better way to handle this is to dequantize entire v cache outside the model

There isnt significant perf advantage from this yet but subsequent diffs will use caching allocator where this refactor help.

Differential Revision: [D85532077](https://our.internmc.facebook.com/intern/diff/D85532077/)

[ghstack-poisoned]
@kimishpatel
[Executorch] Introduce caching cpu memory allocator
ea16e15 
Meant to use this for temp allocator for kernels. Specifically for sdpa, it seems that on iOS there is a significant overhead coming from allocations

Differential Revision: [D85532079](https://our.internmc.facebook.com/intern/diff/D85532079/)

[ghstack-poisoned]
@kimishpatel
Update base for Update on "[Executorch] Introduce caching cpu memory …
c3ed4b2 
…allocator"

Meant to use this for temp allocator for kernels. Specifically for sdpa, it seems that on iOS there is a significant overhead coming from allocations

Differential Revision: [D85532079](https://our.internmc.facebook.com/intern/diff/D85532079/)

[ghstack-poisoned]
@kimishpatel
Update on "[Executorch] Introduce caching cpu memory allocator"
08ab552 
Meant to use this for temp allocator for kernels. Specifically for sdpa, it seems that on iOS there is a significant overhead coming from allocations

Differential Revision: [D85532079](https://our.internmc.facebook.com/intern/diff/D85532079/)

[ghstack-poisoned]
@kimishpatel
Update base for Update on "[Executorch] Introduce caching cpu memory …
dbf63cc 
…allocator"

Meant to use this for temp allocator for kernels. Specifically for sdpa, it seems that on iOS there is a significant overhead coming from allocations

Differential Revision: [D85532079](https://our.internmc.facebook.com/intern/diff/D85532079/)

[ghstack-poisoned]
@kimishpatel
Update on "[Executorch] Introduce caching cpu memory allocator"
f9ce984 
Meant to use this for temp allocator for kernels. Specifically for sdpa, it seems that on iOS there is a significant overhead coming from allocations

Differential Revision: [D85532079](https://our.internmc.facebook.com/intern/diff/D85532079/)

[ghstack-poisoned]
@kimishpatel
Update base for Update on "[Executorch] Introduce caching cpu memory …
86c7c4b 
…allocator"

Meant to use this for temp allocator for kernels. Specifically for sdpa, it seems that on iOS there is a significant overhead coming from allocations

Differential Revision: [D85532079](https://our.internmc.facebook.com/intern/diff/D85532079/)

[ghstack-poisoned]
@kimishpatel
Update on "[Executorch] Introduce caching cpu memory allocator"
0c23c32 
Meant to use this for temp allocator for kernels. Specifically for sdpa, it seems that on iOS there is a significant overhead coming from allocations

Differential Revision: [D85532079](https://our.internmc.facebook.com/intern/diff/D85532079/)

[ghstack-poisoned]
@kimishpatel
Update base for Update on "[Executorch] Introduce caching cpu memory …
68d76d3 
…allocator"

Meant to use this for temp allocator for kernels. Specifically for sdpa, it seems that on iOS there is a significant overhead coming from allocations

Differential Revision: [D85532079](https://our.internmc.facebook.com/intern/diff/D85532079/)

[ghstack-poisoned]
@kimishpatel
Update on "[Executorch] Introduce caching cpu memory allocator"
79bb135 
Meant to use this for temp allocator for kernels. Specifically for sdpa, it seems that on iOS there is a significant overhead coming from allocations

Differential Revision: [D85532079](https://our.internmc.facebook.com/intern/diff/D85532079/)

[ghstack-poisoned]
@kimishpatel
[Executorch] Use temp allocator for allocating scratch memory
351a400 
This allows us to leverage temp memory allocator and if that allocator is caching allocator it reduces the allocaiton overhead.

Differential Revision: [D85532076](https://our.internmc.facebook.com/intern/diff/D85532076/)

[ghstack-poisoned]
@kimishpatel
[Executorch] Make module constructors uniform across
b4fdc22 
Existing constructors dont compose well such that if you want data loader or data files constructor then you cannot get to override memory allocator.
Fix that.

Differential Revision: [D86120037](https://our.internmc.facebook.com/intern/diff/D86120037/)

[ghstack-poisoned]
@kimishpatel
[Executorch][LLM] Use caching allocator for runner
00fffa1 
We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

[ghstack-poisoned]
This was referenced Nov 10, 2025
Merged
Merged
@pytorch-bot
Copy link

pytorch-bot bot commented Nov 11, 2025 •
edited
Loading

🔗 Helpful Links

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

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

❌ 1 New Failure, 1 Unrelated Failure

As of commit 3cd0176 with merge base 78a73d4 (image):

NEW FAILURE - The following job has failed:

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

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

This was referenced Nov 10, 2025
Merged
Merged
Merged
Merged
Merged
Merged
kimishpatel added a commit that referenced this pull request Nov 11, 2025
@kimishpatel
[Executorch][LLM] Use caching allocator for runner
48ff224 
We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

ghstack-source-id: 322321964
Pull Request resolved: #15730
@kimishpatel
Update on "[Executorch][LLM] Use caching allocator for runner"
90d3d57 
We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

[ghstack-poisoned]
kimishpatel added a commit that referenced this pull request Dec 4, 2025
@kimishpatel
[Executorch][LLM] Use caching allocator for runner
e89edd7 
Pull Request resolved: #15730

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.
ghstack-source-id: 327191616
@exported-using-ghexport

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)
@kimishpatel
Update base for Update on "[Executorch][LLM] Use caching allocator fo…
834171f 
…r runner"

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

[ghstack-poisoned]
@kimishpatel
Update on "[Executorch][LLM] Use caching allocator for runner"
be88d80 
We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

[ghstack-poisoned]
kimishpatel added a commit that referenced this pull request Dec 4, 2025
@kimishpatel
[Executorch][LLM] Use caching allocator for runner
10f7191 
Pull Request resolved: #15730

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.
ghstack-source-id: 327215751
@exported-using-ghexport

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)
@kimishpatel
Update base for Update on "[Executorch][LLM] Use caching allocator fo…
bae4829 
…r runner"

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

[ghstack-poisoned]
@kimishpatel
Update on "[Executorch][LLM] Use caching allocator for runner"
4082b28 
We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

[ghstack-poisoned]
kimishpatel added a commit that referenced this pull request Dec 5, 2025
@kimishpatel
[Executorch][LLM] Use caching allocator for runner
1563045 
Pull Request resolved: #15730

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.
ghstack-source-id: 327388019
@exported-using-ghexport

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)
@kimishpatel
Update base for Update on "[Executorch][LLM] Use caching allocator fo…
71cc532 
…r runner"

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

[ghstack-poisoned]
@kimishpatel
Update on "[Executorch][LLM] Use caching allocator for runner"
54f9381 
We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

[ghstack-poisoned]
kimishpatel added a commit that referenced this pull request Dec 5, 2025
@kimishpatel
[Executorch][LLM] Use caching allocator for runner
d73e781 
Pull Request resolved: #15730

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.
ghstack-source-id: 327424796
@exported-using-ghexport

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)
@kimishpatel
Update base for Update on "[Executorch][LLM] Use caching allocator fo…
230cd24 
…r runner"

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

[ghstack-poisoned]
@kimishpatel
Update on "[Executorch][LLM] Use caching allocator for runner"
494bbd5 
We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

[ghstack-poisoned]
kimishpatel added a commit that referenced this pull request Dec 5, 2025
@kimishpatel
[Executorch][LLM] Use caching allocator for runner
1f55d4d 
Pull Request resolved: #15730

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.
ghstack-source-id: 327477384
@exported-using-ghexport

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)
@kimishpatel
Update base for Update on "[Executorch][LLM] Use caching allocator fo…
997b5e2 
…r runner"

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

[ghstack-poisoned]
@kimishpatel
Update on "[Executorch][LLM] Use caching allocator for runner"
4092750 
We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

[ghstack-poisoned]
kimishpatel added a commit that referenced this pull request Dec 5, 2025
@kimishpatel
[Executorch][LLM] Use caching allocator for runner
1b891c2 
Pull Request resolved: #15730

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.
ghstack-source-id: 327485432
@exported-using-ghexport

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)
@kimishpatel
Update base for Update on "[Executorch][LLM] Use caching allocator fo…
7590e9c 
…r runner"

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

[ghstack-poisoned]
@kimishpatel
Update on "[Executorch][LLM] Use caching allocator for runner"
4e0b339 
We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

[ghstack-poisoned]
kimishpatel added a commit that referenced this pull request Dec 5, 2025
@kimishpatel
[Executorch][LLM] Use caching allocator for runner
915636b 
Pull Request resolved: #15730

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.
ghstack-source-id: 327518045
@exported-using-ghexport

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)
@kimishpatel
Update base for Update on "[Executorch][LLM] Use caching allocator fo…
f06f5ba 
…r runner"

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

[ghstack-poisoned]
@kimishpatel
Update on "[Executorch][LLM] Use caching allocator for runner"
d63ffbd 
We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

[ghstack-poisoned]
@kimishpatel
Update base for Update on "[Executorch][LLM] Use caching allocator fo…
e22cb35 
…r runner"

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

[ghstack-poisoned]
@kimishpatel
Update on "[Executorch][LLM] Use caching allocator for runner"
7608f53 
We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

[ghstack-poisoned]
kimishpatel added a commit that referenced this pull request Dec 7, 2025
@kimishpatel
[Executorch][LLM] Use caching allocator for runner
52132b3 
Pull Request resolved: #15730

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.
ghstack-source-id: 327688519
@exported-using-ghexport

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)
@kimishpatel kimishpatel changed the base branch from gh/kimishpatel/213/base to main December 8, 2025 00:03
@kimishpatel
Merge branch 'main' into gh/kimishpatel/213/head
fed9aea
@kimishpatel
Update base for Update on "[Executorch][LLM] Use caching allocator fo…
251b270 
…r runner"

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

[ghstack-poisoned]
@kimishpatel
Update on "[Executorch][LLM] Use caching allocator for runner"
3cd0176 
We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)

[ghstack-poisoned]
kimishpatel added a commit that referenced this pull request Dec 9, 2025
@kimishpatel
[Executorch][LLM] Use caching allocator for runner
06ad343 
Pull Request resolved: #15730

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.
ghstack-source-id: 328001114
@exported-using-ghexport

Differential Revision: [D86120038](https://our.internmc.facebook.com/intern/diff/D86120038/)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

Copilot code review Copilot Copilot left review comments

@larryliu0820 larryliu0820 Awaiting requested review from larryliu0820 larryliu0820 is a code owner

@kirklandsign kirklandsign Awaiting requested review from kirklandsign kirklandsign is a code owner

@jackzhxng jackzhxng Awaiting requested review from jackzhxng jackzhxng is a code owner

@mergennachin mergennachin Awaiting requested review from mergennachin mergennachin is a code owner

+1 more reviewer

@derekxu derekxu derekxu approved these changes

Reviewers whose approvals may not affect merge requirements

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. fb-exported meta-exported

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@kimishpatel @derekxu