Reduce allocation overhead in quantized sdpa

[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

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

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 1 Unrelated Failure

As of commit 4e529d1 with merge base 8af8252 ():

NEW FAILURE - The following job has failed:

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

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

• Test Metal Backend / export-model-metal-artifact (openai, whisper-large-v3-turbo, non-quantized) / macos-job (gh) (matched macos rule in flaky-rules.json)
  File doesn't exist

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

[Copilot]

Pull Request Overview

This PR refactors the quantized scaled dot-product attention (SDPA) implementation to reduce allocation overhead by moving the dequantization buffer allocation from inside the dequant_and_gemm function to the outer cpu_flash_attention scope. Instead of allocating a new std::vector for each dequantization operation, a pre-allocated per-thread scratch buffer is now shared across iterations.

Key changes:

• Added buf_qdq_ptr parameter to dequant_and_gemm and _qk_at_v_gemm functions to accept externally allocated dequantization buffers
• Allocated a shared scratch buffer (scratch_for_quant_dequant) in cpu_flash_attention with per-thread partitioning
• Removed the local std::vector<float> dequantized_v_data allocation from dequant_and_gemm

---

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

[Copilot]

The comment says "align to 64 bytes" but kAlignment = 32 aligns to 32 elements. Since size_per_thread_qdq_vec is an element count (not byte count), and assuming accum_t is float (4 bytes), this aligns to 128 bytes (32 * 4), not 64 bytes.

Either:

1. Change kAlignment to 16 if 64-byte alignment is desired, or
2. Update the comment to say "align to 32 elements" or "align to 128 bytes (for float)"

Suggested change

	// Lets align size_per_thread_qdq_vec to 64 bytes, for coalesced cache reads,
	// Lets align size_per_thread_qdq_vec to 32 elements (128 bytes for float), for coalesced cache reads,

[Copilot]

The alignment calculation is incorrect. The formula (x + kAlignment - 1) & (-(kAlignment - 1)) uses the wrong mask.

For aligning to a power-of-2 boundary, the correct formula is:

(size_per_thread_qdq_vec + kAlignment - 1) & (-kAlignment)

or equivalently:

(size_per_thread_qdq_vec + kAlignment - 1) & ~(kAlignment - 1)

The current code uses -(kAlignment - 1) which equals -31 = 0xFFFFFFE1, but the correct mask should be -32 = 0xFFFFFFE0 to properly zero out the bottom 5 bits.

Suggested change

	(size_per_thread_qdq_vec + kAlignment - 1) & (-(kAlignment - 1));
	(size_per_thread_qdq_vec + kAlignment - 1) & -kAlignment;

[mergennachin]

@copilot are you sure? Please double check again

[Copilot]

The buffer size calculation appears to be larger than necessary. The dequantize operation needs kvBlockSize * headSize elements (at most kvSplitSize * headSize), but this allocates qSplitSize * kvSplitSize * headSize. The extra qSplitSize factor seems unnecessary and wastes memory per thread.

Consider changing to:

int64_t size_per_thread_qdq_vec = kvSplitSize * headSize;

Suggested change

	int64_t size_per_thread_qdq_vec = qSplitSize * kvSplitSize * headSize;
	int64_t size_per_thread_qdq_vec = kvSplitSize * headSize;

[Copilot]

@mergennachin I've opened a new pull request, #15852, to work on those changes. Once the pull request is ready, I'll request review from you.