Fuse the causal mask into the SDPA fallback softmax - #3919
Open
spokvulcan wants to merge 1 commit into
Open
Conversation
When scaled_dot_product_attention takes the ops fallback with a causal mask, the chain greater_equal -> where(finfo.min) -> softmax(precise) materializes a full masked-scores intermediate (512 MB per layer chunk at 32K, bf16) plus a 32 MB bool mask, all to feed a softmax that immediately reduces over the same data. This folds the causal select into a looped softmax kernel: the same precise reduction tree with N_READS=4, the select applied at both load sites, so scores are read once and probabilities written once. Bit-identical to the chain it replaces: the masked value is exactly bf16 finfo.min as f32 and exp(min - max) underflows to exact zero, so per-element arithmetic is unchanged. Verified bitwise on 14/14 configurations (3 seeds x sequence 8192/32768 x offsets) and token-identical over 48/48 end-to-end A/B pairs, and a fallback-forcing equality test in test_fast_sdpa.py pins the bit-identity in CI. Gated conservatively: causal only, no array mask, no sinks, bf16, axis above the looped-softmax limit, GPU stream with Metal available, never inside a function transform. Everything else keeps the stock chain. Kernel time -45% at 32K and -61% at 8K on the softmax stage; end to end +2.8% prefill at 32K on a 35B MoE model (M3 Max), and the giant intermediate is gone.
spokvulcan
force-pushed
the
perf/fused-causal-softmax-sdpa-fallback
branch
from
July 31, 2026 18:41
c08dbda to
94e1c6a
Compare
spokvulcan
changed the title
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
1 participant
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.
Long-context prefill takes the SDPA ops fallback for shapes the fast kernels don't cover. With a causal mask, that chain is greater_equal -> where(finfo.min) -> softmax(precise): at 32K it materializes 512 MB of masked scores plus a 32 MB bool mask per layer chunk, all to feed a softmax that immediately reduces over the same data. The chain costs about 6.5 ms per chunk, of which the mask build and the where are about 3 ms.
This folds the causal condition into the looped softmax kernel: the same precise reduction tree with N_READS=4, with the causal select applied at the load sites. Scores get read once and probabilities written once, and the two mask passes plus the giant intermediate disappear.
The result is bit-identical to the chain it replaces, not approximately. The masked value is exactly bf16 finfo.min upcast to f32, and exp(min - max) underflows to exact zero, so per-element arithmetic is unchanged. Verified bitwise on 14 of 14 configurations (3 seeds, sequence lengths 8192 and 32768, several mask offsets), and token-identical over 48 of 48 interleaved end-to-end A/B pairs on a 35B model. A test in test_fast_sdpa.py forces the fallback and asserts exact equality against the reference chain, so the bit-identity claim is enforced in CI rather than taken on faith.
The gate is deliberately conservative: causal only, no array mask, no sinks, bf16, softmax axis > 4096, GPU stream with Metal available, and never inside a function transform, so grad, vmap, and compile keep the differentiable stock chain. Anything outside the gate takes the stock chain exactly as before.
Measured: the mask+softmax stage drops 45% at 32K and 61% at 8K, which lands on the softmax's memory-traffic floor. End to end that is +2.8% prefill at 32K on Qwen3.6-35B-A3B (M3 Max), plus the memory-pressure relief from never allocating the masked-scores tensor.
The kernel body is deliberately a frozen copy of softmax_looped's body with the select applied at the load sites, since that is what the bitwise verification covered. If you would rather see this as a causal template parameter on softmax_looped itself (or a proper primitive with per-backend gating), happy to rework it that way as a follow-up and re-verify.