Repository navigation
[Metal] Add sdpa_split kernels for mid rows for speculative decoding - #4723
Merged
Merged
Conversation
zcbenz
added
the
await verification
dhiltgen
added a commit
to dhiltgen/ollama
that referenced
this pull request
Oct 20, 2026
SDPA fast paths for speculative verify: up to 8 query tokens whose query heads per KV head stack to 33-128 rows. Applied to the pinned MLX until the pin includes it; drop this commit then.
dhiltgen
added a commit
to dhiltgen/ollama
that referenced
this pull request
Oct 22, 2026
SDPA fast paths for speculative verify: up to 8 query tokens whose query heads per KV head stack to 33-128 rows. Applied to the pinned MLX until the pin includes it; drop this commit then.
zcbenz
reviewed
Oct 22, 2026
| q, k, v, has_mask, has_arr_mask, do_causal, output_logsumexp, s); | ||
| if (force_fused) { | ||
| if (!has_fused) { | ||
| if (!has_fused && !split) { |
Member
There was a problem hiding this comment.
If the supported dims of sdpa_split is a subset of other fused sdpa kernel, we shouldn't need to check it here, if not then we should have has_fused_kernel also check use_split.
Contributor
Author
There was a problem hiding this comment.
Moved into has_fused_kernel
zcbenz
removed
the
await verification
dhiltgen
force-pushed
the
sdpa-nax-split
branch
from
October 23, 2026 21:03
ca54e5f to
eb5207a
Compare
zcbenz
force-pushed
the
sdpa-nax-split
branch
3 times, most recently
from
October 25, 2026 10:40
5260f32 to
7a00a76
Compare
zcbenz
approved these changes
Oct 25, 2026
zcbenz
changed the title
zcbenz
force-pushed
the
sdpa-nax-split
branch
from
October 25, 2026 11:36
7a00a76 to
b56c627
Compare
zcbenz
force-pushed
the
sdpa-nax-split
branch
from
October 25, 2026 11:50
b56c627 to
b3fbbe3
Compare
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
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.
This PR aims to speed up speculative decoding. The existing SDPA has 2 fast paths: The vector kernel handles up to 8 query tokens provided tokens × query heads per KV head is 32 or less; the full-attention kernel covers over 8 tokens. This leaves a gap which makes checking more drafted tokens fall into the slow path (e.g. 4+ drafts on 7-8 heads.) This PR implements fast paths for NAX and pre-NAX for up to 8 tokens with 33-128 rows.
Performance:
Using a 32k code prompt on M5 Max:
Best result examples on a few test systems: