Repository navigation
Fix steel GEMM safe load offset - #3560
Merged
Merged
Conversation
magicnight
added a commit
to magicnight/mlx
that referenced
this pull request
Aug 8, 2026
sinistersoundworks
pushed a commit
to sinistersoundworks/Mac-MLX
that referenced
this pull request
Aug 25, 2026
Second triage round over the release scan's candidate list. Three were worth weighing; this carries one and records why the other two are not. ml-explore/mlx#3560 is described upstream as "Fixed typo :)" and is one character, but the line is the bounds-checked fragment load in steel/gemm/mma.h: it indexed the source with off_x where off_y belongs. That is the load steel GEMM uses for a tile hanging off the edge of a matrix, so a wrong column can be read into an edge fragment with no bounds check firing — a silent wrong-value path in any dense matmul with a partial tile. Carried WITHOUT a local repro, unlike the fixes in magicnight#110. A probe across five ragged shapes on an M5 Max passed, and I could not establish that it exercised the off_x != off_y case, so absence of a reproduction here is not evidence the path is unreachable. The asymmetry decided it: the source is plainly wrong, the path is one every dense matmul uses, and the fix is one character, while the failure it guards against is wrong numbers rather than a crash. SteelGemmSafeLoadParityTests keeps the probe as a regression guard. Not carried, with reasons recorded rather than left implicit: - mlx#3609 (int32 overflow in matvec row offset) needs vocab * hidden above 2^31. No model we support crosses it, but Cohere2 sits at 0.98x, so this becomes reachable the moment a larger one lands. Worth re-checking whenever a model with a bigger embedding table is added. - mlx#3720 (compiled-kernel correctness for negative strides) needs a negative-strided input. mlx-swift-lm does use compile() in the MoE path, but nothing in inference produces reversed slices. Regenerating the JIT copies needed a different technique than magicnight#110 did. mma.h is embedded in conv.cpp and gemm.cpp, and its content is split by nested #line directives rather than sitting in one contiguous block, so replacing the section between separators left the original text in place — visible as a 2272-line insertion with the old expression still present. The fix applies the same one-character edit to the generated copies directly.
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.
Fixed typo:)