Skip to content

Fix steel GEMM safe load offset - #3560

Merged
jagrit06 merged 1 commit into
ml-explore:mainfrom
sofinvalery:mma
May 18, 2026
Merged

jagrit06 merged 1 commit into
ml-explore:mainfrom
sofinvalery:mma

Conversation

@sofinvalery

Copy link
Copy Markdown
Contributor

Fixed typo:)

@jagrit06 jagrit06 left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for catching that!

@jagrit06
jagrit06 merged commit 2414e5d into ml-explore:main May 18, 2026
16 checks passed
@sofinvalery
sofinvalery deleted the mma branch May 18, 2026 11:25
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants