Repository navigation
[Metal] gather_qmm improvement - #4572
Merged
Merged
Conversation
davidtai
pushed a commit
to Layr-Labs/mlx
that referenced
this pull request
Sep 28, 2026
(cherry picked from commit 9008944)
This was referenced Sep 28, 2026
davidtai
added a commit
to Layr-Labs/mlx
that referenced
this pull request
Oct 6, 2026
… call The gather_qmm row-tile backport (ml-explore/mlx ml-explore#4572) makes the sorted gather_qmm_rhs kernels read per-expert offsets. Those kernels give a correct result only for sorted indices. Before the backport, the kernel read the index of each row, so the expert-tile route could send a call with unsorted indices (a sortedness retract) to it. gather_qmm_rhs now returns false for a retracted call, and GatherQMM::eval_gpu then runs the unsorted gather route. All other calls are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Gajesh2007
pushed a commit
to Layr-Labs/mlx
that referenced
this pull request
Oct 6, 2026
…explore/mlx ml-explore#4567, ml-explore#4572) (#27) * [Metal] Gather mm improvement (ml-explore#4567) (cherry picked from commit a2a09fd) * docs(forkdiff): describe the upstream gather_mm row-tile backport Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> * [Metal] gather_qmm improvement (ml-explore#4572) (cherry picked from commit 9008944) * fix(metal): run the unsorted gather route for a retracted expert-tile call The gather_qmm row-tile backport (ml-explore/mlx ml-explore#4572) makes the sorted gather_qmm_rhs kernels read per-expert offsets. Those kernels give a correct result only for sorted indices. Before the backport, the kernel read the index of each row, so the expert-tile route could send a call with unsorted indices (a sortedness retract) to it. gather_qmm_rhs now returns false for a retracted call, and GatherQMM::eval_gpu then runs the unsorted gather route. All other calls are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Anastasiia Filippova <a_filippova@apple.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
davidtai
added a commit
to Layr-Labs/mlx
that referenced
this pull request
Oct 7, 2026
Bring in #27, the gather_mm and gather_qmm row-tile scheduling backport (ml-explore/mlx ml-explore#4567, ml-explore#4572). - quantized_nax.h, affine_gather_qmm_rhs_nax: take main's kernel. Each threadgroup now holds rows of one expert (schedule_row_tile), so the per-tile expert-segment loop is gone. This branch's segment elision ran inside that loop. Without the loop, seg_lo is 0 and seg_hi is sgp_sm, so seg_partial is always false and seg_empty equals main's !sg_active. The sorted-endpoint probe has no probe to skip. Remove kGatherRhsSegmentElide, kGatherRhsSortedEndpointElide and the gather_rhs_load_frag_row / gather_rhs_mma_frag_row helpers, which have no caller. Keep this branch's row-strip simdgroup layout (SGM/SGN, DARKBLOOM_GEMMA4_NAX_GATHER_TILING) on main's kernel, and rewrite the note about how it composes with the row-tile schedule. - fork.yaml: auto-merged; replace the segment-elision bullet with the row-strip layout of the NAX qmm_t and gather-rhs kernels. - quantized.h: auto-merged; main changed the gather_qmm_rhs kernels and this branch changed the QMV kernels, which do not overlap. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
jonathan308
added a commit
to Layr-Labs/mlx
that referenced
this pull request
Oct 10, 2026
Upstream ml-explore/mlx c7ff35d ("Skip unnecessary simdgroup computations for quantised MOE matmuls on NAX", ml-explore#4352) reached this fork through the v0.32.2 squash (734241b), not by ancestry, in the form that ml-explore#4572 later gave it: `const bool sg_active = sgp_sm > 0;` and two `if (sg_active)` blocks around the MMA loop in each of fp_gather_qmm_rhs_nax (fp_quantized_nax.h) and affine_gather_qmm_rhs_nax (quantized_nax.h). This removes the flag and runs the MMA in every simdgroup again, as before ml-explore#4352. It is the same change as the earlier revert of ml-explore#4352 on 0.32.2 (e9debbe) and its re-application on the 0.32.3 kernels, hand-ported because the surrounding lines differ here. Why: these kernels run only on GPUs with NAX. On a two-Mac JACCL pair with one NAX GPU, mixture-of-experts serving with the skip present lost all-reduce completions, and 92 to 96 GiB stayed wired on the host until it was rebooted. The builds that served that pair afterwards carried the revert. It is therefore kept out of every build that runs distributed on mixed hardware. Output is unchanged: stores stay bounded by sgp_sm, and the rows past the tile that the MMA now reads are either inside M (rows_in_bounds) or loaded as zeros by load_safe with sgp_sm rows. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
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.
Same as in #4567
Given #4567 I asked Claude code to do the same for
gather_qmm.qwen3.5-397b-a17b: