[Executorch] Use temp allocator for allocating scratch memory

[kimishpatel]

Stack from ghstack (oldest at bottom):

• [Executorch][LLM] Use caching allocator for runner #15730
• [Executorch] Make module constructors uniform across #15729
• -> [Executorch] Use temp allocator for allocating scratch memory #15728
• [Executorch] Introduce caching cpu memory allocator #15611
• Reduce allocation overhead in quantized sdpa #15610

This allows us to leverage temp memory allocator and if that allocator is caching allocator it reduces the allocaiton overhead.

Differential Revision: D85532076

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures, 2 Unrelated Failures

As of commit af048bd with merge base a09a4b7 ():

NEW FAILURES - The following jobs have failed:

• pull / test-static-llama-qnn-linux (stories_260k_bc) / linux-job (gh)
  RuntimeError: Command docker exec -t bec67b6cc57065a05e31127c12446355e22c85aa8bb286051abb2a5ac029593d /exec failed with exit code 1
• pull / unittest-arm-backend-with-no-fvp (test_pytest_models) / linux-job (gh)
  RuntimeError: Command docker exec -t d42e42c94b4a15ef37ea4ca6cd1d312906b499d1e6458dd5d245612ccf22806a /exec failed with exit code 1

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

• periodic / test-models-linux (cmake, vit, xnnpack-quantization-delegation, linux.2xlarge, 90) / linux-job (gh) (matched linux rule in flaky-rules.json)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.

BROKEN TRUNK - The following job failed but was present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / android / run-emulator (gh) (trunk failure)
  Timeout waiting for emulator to boot.

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

[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.

[Copilot]

Pull Request Overview

This PR refactors memory allocation in the Flash Attention implementation to use the temporary memory allocator from RuntimeContext instead of stack-allocated std::vector objects. This enables the use of caching allocators when available, reducing allocation overhead.

• Adds RuntimeContext& ctx parameter to cpu_flash_attention function
• Replaces stack-allocated vectors with ctx.allocate_temp() calls with fallback to heap allocation
• Removes unnecessary buf_reduced allocation (dead code for unsupported reduced types)
• Updates all call sites to pass the RuntimeContext

Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
extension/llm/custom_ops/op_sdpa_impl.h	Modified cpu_flash_attention signature to accept RuntimeContext, replaced vector allocations with temp allocator calls
extension/llm/custom_ops/op_sdpa.cpp	Updated all call sites (6 locations) to pass ctx parameter to cpu_flash_attention

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

This static vector scratch_for_quant_dequant_vec is declared but never used in the code. It appears to be a leftover from the refactoring where the local vector was replaced with the temp allocator approach. This should be removed.

Suggested change

static std::vector<char> scratch_for_quant_dequant_vec;