docs: document softmax's precise argument - #4178
Merged
Merged
Conversation
`precise` is bound but absent from both the type signature and the docstring, so it is not discoverable from Python. It selects the accumulation type: without it, `AccT` defaults to the input type (softmax.h), so a bf16 softmax accumulates in bf16. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
apocryphx
added a commit
to apocryphx/Apertura
that referenced
this pull request
Aug 11, 2026
…ually runs
The attn_oproj probe showed fused at 38-81% within 1 ULP against unfused
at 98.8-100%, and I attributed that to mx::fast::scaled_dot_product_attention.
That attribution was wrong: at the fixture's seq_len no fused kernel runs
at all. sdpa_full supports head_dim {64,80,96,128} (Apertura is 256/512)
and sdpa_vector requires seq <= 8, so use_fallback returns true and
"fused" is MLX's fallback COMPOSITION against Apertura's manual one — two
composed graphs differing mainly in GQA handling (MLX broadcasts 5-D,
Apertura materialises repeatKV 3-D), hence different GEMM tiling.
The regime that matters was never measured. At decode, seq=1 with gqa 2
satisfies (seq * gqa) <= 32 and head_dim 256 is supported, so LOCAL layers
hit the real vector kernel while global layers (512) still fall back. That
mix is what production generation runs.
--fused-lockstep drives both configurations along one forced token stream
with a populated cache and reports |dlogit| per step. gemma-4-31b-it,
prefill 512, 64 steps:
prefill-shape (both fall back) max |dlogit| 1.030e+01
decode-shape max |dlogit| 6.516e+00
mean 6.772e-01
argmax flips 0/64
So the real fused kernel is no worse than the fallback composition, and
the earlier framing overstated its role. What survives is that the fused
path diverges from Apertura's reference-faithful path by roughly 10x more
than that path diverges from PyTorch (mean 0.677 vs 0.063) — while
changing no tokens across 64 decode steps.
Also checked and ruled out: MLX's fallback already uses precise softmax
(fast.cpp, softmax(scores, {-1}, true, s)), so the accumulator trap
documented in ml-explore/mlx#4178 is not the cause here.
aptransformerTests 16/16.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
zcbenz
approved these changes
Aug 11, 2026
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.
preciseis bound inpython/src/ops.cppbut appears in neither thenb::sigsignature nor the docstring, so it isn't discoverable from Python.It selects the accumulation type —
AccTdefaults to the input type insoftmax.h, so a bf16 softmax accumulates in bf16, which can lose precision over long reduction axes.Docs only.