[ExecuTorch][WebGPU] SDPA: skip QK contraction for fully-masked causal tiles

[JulianCloudNTH]

Stack from ghstack (oldest at bottom):

• [ExecuTorch][WebGPU] SDPA: branchless aligned/tail loads in the QK/AV kernels #20493
• -> [ExecuTorch][WebGPU] SDPA: skip QK contraction for fully-masked causal tiles #20492
• [ExecuTorch][WebGPU] Coalesce SDPA AV V-cache reads along contiguous head-dim #20459
• [ExecuTorch][WebGPU] Register-tile the SDPA QK/AV kernels #20405

Skip the QK contraction for fully-masked causal tiles — at S=128 prefill ~48% of the (query, key) tiles are entirely above the diagonal and contribute nothing; this elides their dot products (prefill-only; bit-identical output).

Problem: For causal prefill, ~half the (query S-tile, key context-tile) pairs are entirely above the diagonal, yet the kernel still computes their full d4 dot product before masking the result to NEG_INF.

Solution: Skip the contraction for fully-masked tiles; the existing per-element mask still writes the sentinel:

• Before: every (s0, c0) tile runs the full d4 dot-product loop, then store_qk masks above-diagonal elements to NEG_INF.
• After: a fully-masked tile (c0 > s0 + TM-1 + input_pos) breaks the d4 loop immediately (acc stays 0); store_qk masks every element to NEG_INF exactly as before.

Implementation:

• Add skip_tile = c0 > s0 + (TM - 1) + params.input_pos, folded into the d4 loop break condition.
• Store loop unchanged — runs unconditionally, so no scratch entry is left stale.
• Mirrors Vulkan sdpa_compute_attn_weights_tiled.glsl (tile_in_mask_region).

Constraints:

• No KV-cache, host, dispatch, or uniform change (all tiles still launch; the skip is in-shader).
• Prefill-only: decode S=1 never triggers it (c0 <= input_pos < input_pos + TM - 1).
• NEG_INF stays the WGSL-safe -1.0e30 (WGSL forbids a literal -inf); does not copy Vulkan's -1.0/0.0.

Co-authored with Claude Code.
@exported-using-ghexport

Differential Revision: D109517773

Differential Revision: D109517773

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 2 Active SEVs

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

• [ROCm] MI350 CI runner label rename: rebase PRs using old linux.rocm.gpu.gfx950.* labels
• [Housekeeping] Rename uw1 cluster to meta-prod-aws-uw1

❌ 2 New Failures, 1 Unrelated Failure

As of commit 173487b with merge base e03f777 ():

NEW FAILURES - The following jobs have failed:

• pull / test-llama-runner-qnn-linux (fp32, qnn_16a16w, qnn) / linux-job (gh)
  RuntimeError: Command docker exec -t 7071cd2c3ce3f3047964eccadfece7c7607b320effc835909074470b1b678fe6 /exec failed with exit code 137
• pull / unittest-buck / linux / linux-job (gh)
  RuntimeError: Command docker exec -t 6e24395ccd72e7ad6cd56fae0ec41acbc82501c76e7eb56617de5c9478458df1 /exec failed with exit code 3

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

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

• pull / unittest-buck / macos / macos-job (gh) (trunk failure)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 3

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.

[SS-JIA]

Review automatically exported from Phabricator review in Meta.