[llm][custom_ops] Loosen sdpa_with_kv_cache test tolerance to fix host-dependent flakiness

[psiddh]

Summary

24 tests in SDPATestForLargeSeqLength/SDPATestForSpeculativeDecode (test_sdpa_with_kv_cache.py) fail intermittently (~60% failure rate observed Sep 2–16) on allclose(ref, op, atol=1e-6).

The custom blocked CPU kernel legitimately diverges from the ATen reference SDPA at the ~1e-6 level: accumulation order differs by tiling, host BLAS path (gemm_uses_blas() varies by CI host), and thread count. With zero headroom at the assertion threshold, the identical test binary passes or fails depending on which CI host runs it. For the seq-130 variants with +/-15 scaled tensors, the divergence is amplified to ~1e-3.

Ruled out as triggers (failures span all kernel versions in the window):

• ring-attention optimization + revert (Sep 9)
• SDPA: skip the non-attendable key columns (#22112) #22112 column-skip (Sep 10, postdates earliest failures)
• Aug 28/31 BF16 optimizations (bf16-guarded or purely additive; float32 path untouched)

The sibling test_quantized_sdpa.py in the same directory already documents this phenomenon ("when scaling tensors, the test fails with smaller atol") and uses atol=1e-5 by default / atol=1e-3 for scaled tensors — and those tests are green on the same hosts. This change matches that precedent in _test_sdpa_common: atol = 1e-3 if scale_tensors else 1e-5.

Classification: TEST_BUG (no product change).

Test Plan

• pytest --collect-only: 34 tests collected, syntax OK.
• Full numerical verification via CI (the relevant OSS CI jobs exercising these tests).

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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

As of commit 682fafb with merge base b03dec2 ():

NEW FAILURES - The following jobs have failed:

• Build documentation / build (buck2) / Build doc (gh)
  Could not load credentials from any providers
• Cadence Build & Test / Resolve CI docker image / resolve (gh)
  ##[error]Refusing to check out fork pull request code from a 'pull_request_target' workflow. This workflow runs with the base repository's GITHUB_TOKEN, secrets, default-branch cache scope, and runner access. Fetching and executing a fork's code in that trusted context commonly leads to "pwn request" vulnerabilities. To opt in, review the risks at https://gh.io/securely-using-pull_request_target and set 'allow-unsafe-pr-checkout: true' on the actions/checkout step.

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]

🟢 Approval recommended

The test-only tolerance update addresses documented numerical variation without product changes.

Pull request overview

This PR adjusts SDPA KV-cache test tolerances to prevent host-dependent numerical flakiness.

Changes:

• Uses 1e-5 tolerance normally and 1e-3 for scaled tensors.
• Applies the tolerance to both SDPA assertions.

File summaries

File	Summary
extension/llm/custom_ops/test_sdpa_with_kv_cache.py	Adds scale-aware numerical tolerances.

Review details

• Files reviewed: 1/1 changed files
• Comments generated: 0
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.