SIGN IN SIGN UP

backends/mlx: runtime MoE expert-sort for decode (issue #20554) (#20685)

## Summary

Replace the compile-time `sort_experts: bool` flag in `SwitchMLP` with a
runtime decision inside two new custom ops (`moe_gather_inputs`,
`moe_scatter_outputs`). A single exported `.pte` now handles both
prefill (sorted, coalesced `gather_mm`) and decode (unsorted, no argsort
overhead) without separate exports.

**Wire-compatible schema:** `sorted_indices: bool` is retained on
`GatherMmNode`/`GatherQmmNode`; a new optional `sorted_indices_flag:
IntOrVid` field (appended last) carries runtime 0/1 values.
`MLXInterpreter.h` prefers the flag when present, otherwise falls back
to the static bool.

**Serialization fix:** `TakeNode.index` expects `IntOrVidOrTid`; MoE
handlers now pass `IntOrVidOrTid.from_tid(...)` instead of a raw `Tid`
(fixes export-time FlatBuffer serialization failure).

`MLXLoader.{h,cpp}` and FlatBuffer bindings are regenerated
automatically by `generate.py` + `flatc` during the CMake build on Mac
CI — not included in this commit, per repo convention.

## Test plan

- [x] Windows: `python backends/mlx/test/validate_moe_20554.py` (all
passed)
- [x] Windows: export → lowering → FlatBuffer serialization validated
for MoE + GatherMm/GatherQmm (Python path, no Metal)
- [ ] CI: `test-mlx` job on `macos-14-xlarge` (`run_all_tests`)

Fixes #20554

PR authored with Claude.


cc @metascroy

---------

Co-authored-by: Axel.Cffrd.Dnty <AxelNoun@users.noreply.github.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Scott Roy <scroy@meta.com>
A
Axel.Cffrd.Dnty committed
fdb88c3d332dac95d9703e4dfb138282c4e92dea
Parent: 097b9aa
Committed by GitHub <noreply@github.com> on 8/1/2026, 1:09:13 AM