Vulkan: fuse Q@K^T and softmax on the SDPA decode path - #22941
Open
msluszniak wants to merge 2 commits into
Open
msluszniak wants to merge 2 commits into
msluszniak wants to merge 2 commits into
Conversation
pytorch-bot
Bot
added
the
module: vulkan
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22941
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ❌ 1 Awaiting Approval, 2 New FailuresAs of commit 8c348f9 with merge base 3887eea ( NEW FAILURES - The following jobs have failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
At sequence length 1 the attn-weights matmul and the softmax that follows it are two dispatches with a round trip through attn_weights in between. Fuse them: one work group per (s, q_h) walks the whole context row, keeps it in shared memory, and writes softmax-normalized weights directly. The sequence length is only known once sizes are current, because a dynamic-shape export reports the max bound at graph build time. Adding or dropping a node cannot be deferred to resize time the way a shader pick can, so all three nodes are built and each global work group picker collapses to zero for the path it is not serving; DispatchNode::encode skips a zero dispatch. The fused shader holds the attn_weights row in shared memory, so it is only selected while that row fits (see kFusedSoftmaxMaxContextTexels). Qwen3-0.6B on Adreno 840, same binary, arms interleaved and selected by an env toggle: 1668 -> 1640 dispatches per token, one per layer for 28 layers, and decode 90.9 -> 97.1 tok/s at a 16-token prompt. At a 400-token prompt the two are within noise of each other (69.4 vs 67.5), which is what the loss of context-axis parallelism predicts.
The fused QK+softmax dispatch reduces a whole attn_weights row inside one work group, so its parallelism is capped at num_q_heads work groups while the separate pair gets one work group per context texel. Measured on Qwen3-0.6B/Adreno, it wins 3.6-6.2% below ~192 tokens and costs 9-11% beyond ~550, crossing over around 220. The gate now reads the live context length instead of the K cache bound, so it re-decides every token: the pickers are re-run from DynamicDispatchNode::trigger_resize.
msluszniak
force-pushed
the
ms/vulkan-sdpa-fused-decode
branch
from
September 21, 2026 09:31
fddd6ff to
8c348f9
Compare
executorch-triage
Bot
added
the
community: contribution
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
At sequence length 1,
sdpa_implruns the attn-weights matmul and the softmax that follows it as two dispatches, with a round trip throughattn_weightsin between. This fuses them for the decode path: one work group per(s, q_h)walks the whole context row, keeps it in shared memory, reduces to row max and exp sum, and writes softmax-normalized weights straight toattn_weights_softmax.Selecting it
The sequence length is only known once sizes are current: a dynamic-shape export reports the max bound at graph build time, never 1. A shader pick can be deferred to
trigger_resize, but adding or dropping a node cannot, so all three nodes are built and each global work group picker collapses to zero for the path it is not serving.DispatchNode::encodealready skips a zero dispatch.use_fused_qk_softmax()is the single predicate all three pickers call, so the fused node and the separate pair can never both run.Bound
use_fused_qk_softmax()also gates on the live context length, read from theinput_possymint at resize, so the choice is remade every token. The crossover measured at roughly 220 tokens on Adreno 840, so the cutoff is 256 and past it the separate pair is selected. Astatic_assertkeeps that cutoff inside the shader's sharedattn_weightsrow (kFusedSoftmaxMaxContextTexels, matchingMAX_CONTEXT_TEXEL_LEN).Measured
Qwen3-0.6B, Adreno 840 (Galaxy S26 Ultra), same binary and same
.pte, arms interleaved and selected by an env toggle so nothing else differs. Medians.28 fewer dispatches is exactly one per layer for 28 layers.
The gain is at short context: trading context-axis parallelism for a single work group only pays while the row is short, which is what the cutoff encodes. With the gate in place decode runs the fused path below 256 tokens and the separate pair above it, taking the better of the two columns in either regime.
cc @SS-JIA @manuelcandales @digantdesai @cbilgin