Fix TOCTOU heap-buffer-overflow in SDPA per-thread scratch

[derekdixu]

Summary:
The SDPA flash-attention kernel allocates per-thread scratch space using the threadpool's current thread count, then dispatches parallel work that independently re-reads the thread count. On a 96-core host post-fork, the threadpool can be resized between the two reads, causing the parallel dispatcher to create more tasks than scratch slots were allocated. Worker threads then index past the end of the buffer, triggering heap-buffer-overflow (reproduced 55% of the time, 11/20 runs).

This adds an optional num_threads parameter to parallel_for and calc_num_tasks_and_chunk_size. When <= 0 (the default), they read the threadpool's current count as before, so existing callers are unchanged. The SDPA kernel now passes the same thread count it used to size the buffer, guaranteeing num_tasks <= num_thread and keeping every worker's ompIdx in bounds.

Reviewed By: GregoryComer

Differential Revision: D109464749

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• [ROCm] MI350 CI jobs will have longer queue times due to CI migration

❌ 1 Awaiting Approval, 4 New Failures, 3 Unrelated Failures

As of commit 10df7cc with merge base aada6d7 ():

AWAITING APPROVAL - The following workflow needs approval before CI can run:

• Build documentation (gh)

NEW FAILURES - The following jobs have failed:

• pull / test-arm-backend-no-driver (test_pytest_ops_tosa) / linux-job (gh)
  RuntimeError: Command docker exec -t a0159c35ded99f9435b21d743be49c5bf0c156a179c86e7781c1073c2c240b49 /exec failed with exit code 1
• pull / unittest / linux / linux-job (gh)
  RuntimeError: Command docker exec -t 0e902c1aa81dc29cdcf1373954b39e432afb7d2c1c2c6642e29068ad695ce7c5 /exec failed with exit code 1
• pull / unittest-editable / linux / linux-job (gh)
  RuntimeError: Command docker exec -t 7648fd40341016b06dc0f442eaf495a1f9f918baceeb7a5c2229460a015789f2 /exec failed with exit code 1
• pull / unittest-nxp-neutron / linux-job (gh)
  RuntimeError: Command docker exec -t 0e05b2235cc0c6775575bbaec88597c67899938ede98ce96e7db473c8ba10fc8 /exec failed with exit code 1

BROKEN TRUNK - The following jobs failed but were present on the merge base:

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

• pull / test-llama-runner-qnn-linux (fp32, qnn_16a16w, qnn) / linux-job (gh) (trunk failure)
• pull / unittest / macos / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
• pull / unittest-editable / macos / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1

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

[linux-foundation-easycla]

• ❌ The email address for the commit (10df7cc) is not linked to the GitHub account, preventing the EasyCLA check. Consult this Help Article and GitHub Help to resolve. (To view the commit's email address, add .patch at the end of this PR page's URL.) For further assistance with EasyCLA, please visit our EasyCLA portal and chat with our support bot.

[meta-codesync]

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

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