SDPA: skip the non-attendable key columns

[pssrawat]

Summary:
Clamp the key-block width in the flash-attention custom SDPA
kernel (op_sdpa_impl.h) to the causally attendable extent.

Under causal attention the query block starting at m can only attend to
num_keys = min(m + start_pos + qBlockSize, kvSize) keys. The columns
beyond that were computed by the q @ k.T gemm, filled with -inf, run
through the softmax, and then multiplied by v -- all to contribute exactly
zero. The change is one line:

• int64_t kvBlockSize = std::min(kvSplitSize, kvSize - n);

• int64_t kvBlockSize = std::min(kvSplitSize, num_keys - n);

num_keys <= kvSize always, so the block can only get smaller, never go
out of bounds; and when is_causal is false num_keys == kvSize and
nothing changes at all. For a prefill split into qSlice query blocks the
key-column count goes from qSlice * kvSize down to the sum over blocks of
num_keys(block), a factor of 2 * qSlice / (qSlice + 1).

Worked example -- backbone prefill of 128 query tokens at start_pos = 11.
seq_len = 128 dispatches cpu_flash_attention<CTYPE, 32, 512>, so
qSplitSize = 32 and kvSplitSize = 512, and the runner hands in a key
extent trimmed to kvSize = 139:

query block m num_keys kvBlockSize before after
0 43 139 43
32 75 139 75
64 107 139 107
96 139 139 139

Key columns processed: 556 -> 364.

Three supporting hunks:

• The causal-mask guard is retightened from kvSplitSize to kvBlockSize.
  Masking is needed for a row iff m_start_pos + row < n + kvBlockSize - 1,
  which is now exactly the loop condition. The old form also ran the rows
  with last_col >= kvBlockSize, which have nothing to mask: the fill size
  kvBlockSize - last_col comes out zero or negative there and fill_stub
  no-ops either way, and the write extent was always derived from
  kvBlockSize rather than kvSplitSize, so nothing was written out of
  range. This hunk is a consistency fix rather than a bug fix -- but the
  clamp takes kvBlockSize < kvSplitSize from rare to routine, so it is
  worth having the loop bound agree with the width the rows are actually
  strided by.
• fill_stub(qk_data, ...) is sized qBlockSize * kvBlockSize rather than
  the scratch extent qSplitSize * kvSplitSize. Safe because qk_data is
  packed, not strided by kvSplitSize: every consumer indexes it as
  qk_data + row * kvBlockSize, and _qk_at_v_gemm is passed kvBlockSize
  as the leading dimension. So the tighter fill covers exactly the live
  region.
• fill_stub(dst_data, ...) likewise, at stride headSize.

This is not bit-exact. The skipped columns contribute exactly zero to the
output, but dropping them changes how the softmax row sum and the second
gemm's K extent partition across vector lanes, so the order of the floating
point accumulation moves and the last bit with it. Fourteen autoregressive
layers on that is enough to flip a near-tie argmax: on the ar suites this
diff changes generated text in 2 of 59 cases relative to its parent, and
neither of them changes a verdict.

Stacking note: this diff sits on top of the exp_u20 change rather than
under it. Applied to master on its own it regresses one benchmark case
(calling_101010000000018016_2000); stacked on exp_u20 it does not. See
the end-to-end table below.

Differential Revision: D117259224

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 7761a72 with merge base 4d3b46d ():

NEW FAILURES - The following jobs have failed:

• Cadence Build & Test / hifi-build / hifi4 (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.
• Cadence Build & Test / vision-build / vision (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.

[meta-codesync]

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

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