Repository navigation
fix(moe): keep right_sorted on the SwitchGLU sorted prefill path - #2257
Merged
Merged
Conversation
PR #2137 shared one explicit identity `lhs_indices` across gate, up and down on the sorted prefill path of `SwitchGLU::forward`. Upstream `gather_qmm` and `gather_mm` build their primitive with `right_sorted = sorted_indices && !lhs_indices`, and every backend keys its sorted MoE prefill kernel on that flag (the ROCm expert-batched kernel and sorted-rhs schedule in `qmm.hip`, Metal's `gather_qmm_rhs` and sorted `gather_mm`, CUDA's grouped GEMM), so the explicit lhs silently moved affine MoE prefill back to the per-row kernels: 1.7x to 4.8x slower on gfx1151 (issue #2192). The sorted path now shares only the `uint32` expert ids and leaves `lhs_indices` to MLX, whose default is the same arange; the unsorted (decode) sharing of #2137 is unchanged, so its decode-graph node cut stays. `SwitchLinear::forward_indexed` documents the constraint and debug-asserts it. `mlxcel_core::gather_sorted_flags` walks an unevaluated graph and reports the `left_sorted` / `right_sorted` state of every `GatherQMM` / `GatherMM` node, so `sorted_prefill_keeps_right_sorted_on_every_projection` pins the flag on all three projections for quantized and regular experts on every backend without running a kernel; it fails on the previous `main`. Measured on gfx1151 at pp512 against the pre-fix build, three guarded rounds with arms alternated (docs/benchmark_results/data/rocm-moe-prefill-right-sorted-gfx1151-2026-10-09/): granite-4.0-h-tiny-4bit prefill 606.91 to 918.41 tok/s at the median (2026-10-05 page: 918.59), Mixtral-8x7B 25.33 to 122.21 (page: 125.63), Qwen3-30B-A3B 291.26 to 318.36, gpt-oss-20b 465.74 to 500.10; decode unchanged. The rocprofv3 trace shows `gather_qmv_expert_batched_kernel` back at 240 launches (42.7 percent of kernel time) where the pre-fix build had none. Refs #2241
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.
Summary
SwitchGLU::forwardlostright_sortedwhen perf(moe): share gather indices across SwitchGLU projections #2137 shared one explicit identitylhs_indicesacross gate, up and down: upstreamgather_qmm/gather_mmbuild the primitive withright_sorted = sorted_indices && !lhs_indices, and every backend's sorted MoE prefill kernel keys on that flag, so affine MoE prefill fell back to per-row kernels, 1.7x to 4.8x slower on gfx1151.What changed
src/models/switch_layers.rs: the sorted path now shares only theuint32expert ids and leaveslhs_indicesto MLX, whose default is the same arange; the unsorted (decode) sharing of perf(moe): share gather indices across SwitchGLU projections #2137 is unchanged, so its decode-graph node cut stays.SwitchLinear::forward_indexeddocuments the constraint and debug-asserts it.src/lib/mlxcel-core: newgather_sorted_flagsbridge accessor walks an unevaluated graph and reports theleft_sorted/right_sortedstate of everyGatherQMM/GatherMMnode.docs/environment-variables.md: theMLX_ROCM_GATHER_QMV_EXPERT_BATCHEDrow records that the op setsright_sortedonly when the caller passessorted_indicesand nolhs_indices.docs/benchmark_results/data/rocm-moe-prefill-right-sorted-gfx1151-2026-10-09/.Benchmarks (gfx1151, pp512, medians of 3 guarded runs, arms alternated)
gather_qmv_expert_batched_kernelback at 240 launches (42.7 percent of kernel time) where the pre-fix build had noneReproduction
models::switch_layers::gather_index_tests::sorted_prefill_keeps_right_sorted_on_every_projectionpinsright_sortedon all three projections for quantized and regular experts without running a kernel. On the previousmainit fails with "projection 0 lost right_sorted: the sorted path passed an explicit lhs_indices"; with the fix it passes (1 passed, 0 failed, 9071 filtered out, both runs on gfx1151).Metal and CUDA cannot be measured on this Linux host; the flags are computed in the op constructor, backend-independent, and the test pins them for every backend.
Verified locally with
make verify-rocmon gfx1151 at d3faa75: 12,199 passed, 0 failed, 403 ignored.Closes #2241