Skip to content

[cuda backend] skip fully-masked KV blocks calculation in SDPA#20198

Open
Gasoonjia wants to merge 1 commit into
g4-opt-int4-vecloadfrom
g4-opt-prefill-window-sdpa
Open

[cuda backend] skip fully-masked KV blocks calculation in SDPA#20198
Gasoonjia wants to merge 1 commit into
g4-opt-int4-vecloadfrom
g4-opt-prefill-window-sdpa

Conversation

@Gasoonjia

@Gasoonjia Gasoonjia commented Jun 10, 2026

Copy link
Copy Markdown
Contributor

This PR skip attention calculations entirely masked by

  1. sliding-window
  2. causal via start_n > max_seq_pos (lower tranigle of attention calculation)

result:

prompt len before after gain (vs before) llama.cpp (cuBLAS prefill t/s) gain (vs llama.cpp)
256 729.9 1065.7 +46% 1150.5 −7.4%
512 935.2 1582.7 +69% 1609.9 −1.7%
1024 1116.9 2024.5 +81% 1601.2 +26.4%
2048 1238.0 2325.5 +88% 1573.7 +47.8%

Prefill +46-88% all lengths; decode remain the same;
Surpass llama.cpp by 48% on 2k input while on par (-7.4%) on short prompt
Profling on gemma4-31b: SDPA takes 58.1% ->18.5% of e2e prefilling time.
Numerically bf16-exact vs dense+mask (unit test).

Block-sparse early-exit in _sdpa_fwd_kernel_body: skip KV blocks that are
entirely masked (sliding-window via HAS_MASK sum==0, causal via start_n>max_seq_pos).
Exact (skipped blocks are x1,+0 no-ops). Prefill +46-88% all lengths; decode safe;
SDPA nsys 58.1%->18.5%. Numerically bf16-exact vs dense+mask (unit test).
@pytorch-bot

pytorch-bot Bot commented Jun 10, 2026

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❌ 3 New Failures, 2 Unrelated Failures

As of commit ac2968b with merge base a79f3e4 (image):

NEW FAILURES - The following jobs have failed:

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

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

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

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 10, 2026
@linux-foundation-easycla

Copy link
Copy Markdown

CLA Missing ID

  • ❌ The email address for the commit (ac2968b) 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.

@Gasoonjia Gasoonjia changed the title [cuda][prefill] window-aware SDPA: skip fully-masked KV blocks (idea #1) [cuda][prefill] window-aware SDPA: skip fully-masked KV blocks Jun 10, 2026
@Gasoonjia Gasoonjia changed the title [cuda][prefill] window-aware SDPA: skip fully-masked KV blocks [cuda backend]window-aware SDPA: skip fully-masked KV blocks Jun 10, 2026
@Gasoonjia Gasoonjia marked this pull request as ready for review June 10, 2026 21:35
@Gasoonjia Gasoonjia changed the title [cuda backend]window-aware SDPA: skip fully-masked KV blocks [cuda backend] skip fully-masked KV blocks calculation in SDPA Jun 10, 2026
@Gasoonjia Gasoonjia requested a review from mergennachin June 10, 2026 21:40
@mergennachin mergennachin requested a review from digantdesai June 10, 2026 21:46
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant