From e42552c68028a602f11b5d4f9341fa8e4d473559 Mon Sep 17 00:00:00 2001
From: "Axel.Cffrd.Dnty" <AxelNoun@users.noreply.github.com>
Date: Thu, 23 Jul 2026 01:34:59 +0200
Subject: [PATCH 1/4] backends/mlx: runtime MoE expert-sort for decode

Replace the compile-time sort_experts flag with a runtime decision
(sort when M > sort_cutoff) via mlx::moe_gather_inputs /
mlx::moe_scatter_outputs. The FlatBuffer schema keeps the existing
sorted_indices bool and appends an optional sorted_indices_flag
(IntOrVid) last on GatherMmNode/GatherQmmNode, so existing .pte
files stay wire-compatible; the interpreter falls back to the bool
when the flag is absent. Handlers write directly into their output
slots (no IdCopyNode), the flag's ItemIntNode is memoized on the
builder, and static-shape exports fold both branches (no IfNode).

Fixes #20554
---
 backends/mlx/builder/op_helpers.py            |  21 ++
 backends/mlx/custom_ops.py                    |  89 ++++-
 backends/mlx/llm/switch.py                    |  45 ++-
 backends/mlx/ops.py                           | 346 +++++++++++++++++-
 backends/mlx/runtime/MLXInterpreter.h         |  10 +-
 backends/mlx/serialization/generate.py        |  15 +-
 backends/mlx/serialization/schema.fbs         |   7 +
 backends/mlx/test/test_ops.py                 | 318 +++++++++++++++-
 examples/models/qwen3_5_moe/export.py         |   4 +-
 .../qwen3_5_moe/mlx_source_transformations.py |  16 +-
 10 files changed, 821 insertions(+), 50 deletions(-)

diff --git a/backends/mlx/builder/op_helpers.py b/backends/mlx/builder/op_helpers.py
index a902e987ab8..c846513f708 100644
--- a/backends/mlx/builder/op_helpers.py
+++ b/backends/mlx/builder/op_helpers.py
@@ -359,6 +359,27 @@ def emit_ceil_div(
     return P.to_int_or_vid(out_slot)
 
 
+def emit_floordiv(
+    P: "MLXProgramBuilder",
+    a: "IntOrVid",
+    b: "IntOrVid",
+) -> "IntOrVid":
+    """Emit ``a // b`` (floor division), folding when both operands are
+    static.
+    """
+    from executorch.backends.mlx.serialization.mlx_graph_schema import (
+        FloorDivideIntNode,
+        IntOrVid,
+    )
+
+    if not a.is_vid and not b.is_vid:
+        return IntOrVid.from_literal(a.literal // b.literal)
+
+    _, out_slot = P.make_tmp_value_slot()
+    P.emit(FloorDivideIntNode(a=a, b=b, out=P.slot_to_vid(out_slot)))
+    return P.to_int_or_vid(out_slot)
+
+
 def emit_if_else(
     P: "MLXProgramBuilder",
     cond: "IntOrVid",
diff --git a/backends/mlx/custom_ops.py b/backends/mlx/custom_ops.py
index 3f268369d5e..ecf829902f1 100644
--- a/backends/mlx/custom_ops.py
+++ b/backends/mlx/custom_ops.py
@@ -14,7 +14,7 @@
 can execute efficiently but may not have direct PyTorch equivalents.
 """
 
-from typing import Optional
+from typing import Optional, Tuple
 
 import torch
 from torch import Tensor
@@ -285,7 +285,7 @@ def gather_mm(
     b: Tensor,  # [E, K, N] or [..., K, N]
     rhs_indices: Optional[Tensor] = None,  # Expert selection indices
     lhs_indices: Optional[Tensor] = None,  # Optional LHS gather indices
-    sorted_indices: bool = False,
+    sorted_indices: Optional[Tensor] = None,  # 0-d int; None/0 = unsorted
 ) -> Tensor:
     """
     Gather matrix multiply — matches mlx::core::gather_mm semantics exactly.
@@ -295,6 +295,10 @@ def gather_mm(
 
     For MoE: a=[N_tokens, 1, K], b=[E, K, out], rhs_indices=[N_tokens]
     → output=[N_tokens, 1, out]. Caller squeezes dim -2.
+
+    sorted_indices is layout-only (a correctness contract for the MLX kernel
+    at runtime); numerics are identical either way, so the eager reference
+    ignores it.
     """
     if rhs_indices is not None:
         b_sel = b[rhs_indices]
@@ -309,7 +313,7 @@ def gather_mm_fake(
     b: Tensor,
     rhs_indices: Optional[Tensor] = None,
     lhs_indices: Optional[Tensor] = None,
-    sorted_indices: bool = False,
+    sorted_indices: Optional[Tensor] = None,
 ) -> Tensor:
     # Matches MLX: output = indices.shape + [M, N]
     # For simplicity, use matmul shape rules after gather
@@ -334,7 +338,7 @@ def gather_qmm(
     group_size: int = 32,
     bits: int = 4,
     mode: str = "affine",
-    sorted_indices: bool = False,
+    sorted_indices: Optional[Tensor] = None,  # 0-d int; None/0 = unsorted
 ) -> Tensor:
     """
     Gather quantized matrix multiply — matches mlx::core::gather_qmm semantics.
@@ -343,6 +347,8 @@ def gather_qmm(
 
     For MoE: x=[N_tokens, 1, K], w=[E, out, K_packed], rhs_indices=[N_tokens]
     → output=[N_tokens, 1, out]. Caller squeezes dim -2.
+
+    sorted_indices is layout-only; ignored here (see gather_mm docstring).
     """
     # Eager fallback: gather, dequantize, matmul
     if rhs_indices is not None:
@@ -392,7 +398,7 @@ def gather_qmm_fake(
     group_size: int = 32,
     bits: int = 4,
     mode: str = "affine",
-    sorted_indices: bool = False,
+    sorted_indices: Optional[Tensor] = None,
 ) -> Tensor:
     # Matches MLX: output = indices.shape + [M, N]
     M = x.shape[-2]
@@ -465,3 +471,76 @@ def sample(
 @torch.library.register_fake("mlx::sample")
 def sample_fake(logits, temperature, top_k, top_p, seed=None):
     return logits.new_empty(logits.shape[:-1], dtype=torch.long)
+
+
+# ---------------------------------------------------------------------
+# Runtime MoE expert-sort for decode (MLX backend)
+# ---------------------------------------------------------------------
+
+
+@torch.library.custom_op("mlx::moe_gather_inputs", mutates_args=())
+def moe_gather_inputs(
+    x: Tensor, expert_indices: Tensor, top_k: int, sort_cutoff: int
+) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+    """Branch on M on purpose — this is the executable spec the lowering
+    handler (ops.py) mirrors branch-for-branch. Sorting is an invertible
+    permutation (identical numerics either way); the two paths exist for
+    the lowering's sake, not the math's."""
+    N = x.shape[0]
+    if N > sort_cutoff:  # SORTED path  (handler: emit_sorted)
+        flat = expert_indices.flatten()
+        order = flat.argsort().to(torch.int32)
+        inv_order = order.argsort().to(torch.int32)
+        idx = flat[order].to(torch.int32)  # [N*top_k]
+        x_input = x[(order // top_k).to(torch.int64)].unsqueeze(-2)  # [N*top_k, 1, D]
+        sort_experts = torch.ones((), dtype=torch.int32)
+    else:  # UNSORTED path (handler: emit_unsorted)
+        x_input = x.repeat_interleave(top_k, dim=0).unsqueeze(-2)  # [N*top_k, 1, D]
+        idx = expert_indices.flatten().to(torch.int32)  # [N*top_k]
+        sort_experts = torch.zeros((), dtype=torch.int32)
+        # Identity permutation: inverse of "no reorder". Safe if scatter
+        # accidentally takes the sorted branch (Take becomes a no-op).
+        inv_order = torch.arange(N * top_k, dtype=torch.int32)
+    return x_input, idx, sort_experts, inv_order
+
+
+@torch.library.register_fake("mlx::moe_gather_inputs")
+def moe_gather_inputs_fake(
+    x: Tensor, expert_indices: Tensor, top_k: int, sort_cutoff: int
+) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+    """Must NOT branch on M (symbolic SymInt under export — data-dependent
+    control flow on it is illegal). One shape for all M: the sorted-path
+    shape for x_input/idx/inv_order."""
+    N = x.shape[0]
+    D = x.shape[-1]
+    NK = N * top_k
+    x_input = x.new_empty((NK, 1, D))
+    idx = expert_indices.new_empty((NK,), dtype=torch.int32)
+    sort_experts = x.new_empty((), dtype=torch.int32)
+    inv_order = x.new_empty((NK,), dtype=torch.int32)
+    return x_input, idx, sort_experts, inv_order
+
+
+@torch.library.custom_op("mlx::moe_scatter_outputs", mutates_args=())
+def moe_scatter_outputs(
+    down: Tensor, sort_experts: Tensor, inv_order: Tensor, top_k: int
+) -> Tensor:
+    down = down.squeeze(-2)  # [N*top_k, H]
+    if sort_experts.item():  # prefill: scatter back   (handler: emit_then)
+        down = down[inv_order]
+    # decode: no scatter; inv_order is identity (handler: emit_else).
+    # .clone(): output must not alias the input under mutates_args=() —
+    # required by torch.library.opcheck's aliasing check on the no-op
+    # (unsorted) reshape path.
+    return down.reshape(down.shape[0] // top_k, top_k, -1).clone()  # [N, top_k, H]
+
+
+@torch.library.register_fake("mlx::moe_scatter_outputs")
+def moe_scatter_outputs_fake(
+    down: Tensor, sort_experts: Tensor, inv_order: Tensor, top_k: int
+) -> Tensor:
+    """Shape derived only from down + top_k — no branching needed, no
+    dependency on inv_order's shape."""
+    NK = down.shape[0]
+    H = down.shape[-1]
+    return down.new_empty((NK // top_k, top_k, H))
diff --git a/backends/mlx/llm/switch.py b/backends/mlx/llm/switch.py
index 69fcb6eeee7..eba48f283aa 100644
--- a/backends/mlx/llm/switch.py
+++ b/backends/mlx/llm/switch.py
@@ -41,6 +41,7 @@
 """
 
 import logging
+from typing import Optional
 
 import torch
 import torch.nn as nn
@@ -171,15 +172,20 @@ def forward(
         self,
         x: torch.Tensor,
         indices: torch.Tensor,
-        sorted_indices: bool = False,
+        sorted_indices: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         """Forward without unsqueeze/squeeze — caller manages dimensions.
 
         Used by UnfusedMoEExperts which passes x as [N, 1, 1, D]
         and indices as [N, top_k] to handle all experts at once.
+
+        sorted_indices: None, or a 0-d int tensor where 0 means the expert
+        gather is unsorted and any nonzero value means it is sorted (see
+        gather_mm/gather_qmm docstrings in custom_ops.py). Passed straight
+        through to those ops.
         """
         if not self._packed:
-            raise RuntimeError("SwitchLinear.pack() must be called before forward_raw.")
+            raise RuntimeError("SwitchLinear.pack() must be called before forward.")
 
         if self._is_quantized:
             return torch.ops.mlx.gather_qmm(
@@ -233,6 +239,7 @@ def __init__(
         activation=None,
         bias: bool = False,
         fuse_gate_up: bool = False,
+        sort_cutoff: int = 1,
     ):
         super().__init__()
         if activation is None:
@@ -241,6 +248,11 @@ def __init__(
         self.num_experts = num_experts
         self.intermediate_size = intermediate_size
         self.fuse_gate_up = fuse_gate_up
+        # Static export-time threshold, compared against M=N inside
+        # moe_gather_inputs to decide sort/no-sort at runtime.
+        if sort_cutoff < 1:
+            raise ValueError(f"sort_cutoff must be >= 1, got {sort_cutoff}")
+        self.sort_cutoff = sort_cutoff
 
         if fuse_gate_up:
             self.gate_up_proj = SwitchLinear(
@@ -263,7 +275,6 @@ def forward(
         expert_weights: torch.Tensor,
         expert_indices: torch.Tensor,
         top_k: int,
-        sort_experts: bool = False,
     ) -> torch.Tensor:
         """Forward pass through the gated MoE MLP.
 
@@ -272,25 +283,17 @@ def forward(
             expert_weights: Routing weights [N, top_k] (already softmaxed)
             expert_indices: Expert assignments [N, top_k]
             top_k: Number of experts per token
-            sort_experts: Sort tokens by expert index for coalesced memory
-                access during prefill. No effect on decode (single token).
 
         Returns:
             Output tensor [N, D]
+
+        Sort/no-sort is a runtime decision (M vs self.sort_cutoff) made
+        inside moe_gather_inputs, rather than a compile-time flag. Configure
+        the threshold once via SwitchMLP(..., sort_cutoff=...).
         """
-        N = x.shape[0]
-
-        if sort_experts:
-            flat_indices = expert_indices.flatten()
-            order = flat_indices.argsort().to(torch.int32)
-            inv_order = order.argsort().to(torch.int32)
-            sorted_idx = flat_indices[order].to(torch.int32)
-            x_sorted = x[(order // top_k).to(torch.int64)]
-            x_input = x_sorted.unsqueeze(-2)
-            idx = sorted_idx
-        else:
-            x_input = x.unsqueeze(-2).unsqueeze(-2)
-            idx = expert_indices
+        x_input, idx, sort_experts, inv_order = torch.ops.mlx.moe_gather_inputs(
+            x, expert_indices, top_k, self.sort_cutoff
+        )
 
         if self.fuse_gate_up:
             gate_up = self.gate_up_proj(x_input, idx, sorted_indices=sort_experts)
@@ -302,11 +305,7 @@ def forward(
         h = self.activation(gate) * up
         down = self.down_proj(h, idx, sorted_indices=sort_experts)
 
-        if sort_experts:
-            down = down.squeeze(-2)
-            down = down[inv_order].reshape(N, top_k, -1)
-        else:
-            down = down.squeeze(-2)
+        down = torch.ops.mlx.moe_scatter_outputs(down, sort_experts, inv_order, top_k)
 
         return (down * expert_weights.unsqueeze(-1)).sum(dim=-2)
 
diff --git a/backends/mlx/ops.py b/backends/mlx/ops.py
index 4053688b121..479987bb78c 100644
--- a/backends/mlx/ops.py
+++ b/backends/mlx/ops.py
@@ -20,10 +20,13 @@
 
 import torch
 from executorch.backends.mlx.builder.op_helpers import (
+    emit_floordiv,
     emit_if_else,
     emit_lifted_constant,
+    emit_product,
     emit_quantized_biases,
     emit_shape,
+    emit_sub_int,
     parse_dequant_node,
     regroup_affine_scales,
     to_mlx_qparams,
@@ -1790,6 +1793,39 @@ def _split_with_sizes_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     return output_slots
 
 
+def _resolve_sorted_indices_flag(P: MLXProgramBuilder, sorted_indices):
+    """Convert a gather_mm/gather_qmm ``sorted_indices`` arg to an IntOrVid for
+    the ``sorted_indices_flag`` schema field. Accepts either a static bool
+    (legacy call sites) or a 0-d runtime tensor (Slot) carrying 0/1.
+
+    Slot results are memoized on the builder (keyed by slot identity):
+    _moe_gather_inputs_handler pre-seeds the memo with the shape-derived
+    condition Vid, so the common MoE case reads the flag without emitting an
+    ItemIntNode (and without its device sync); otherwise a single ItemIntNode
+    is emitted per distinct flag tensor and reused by every consumer.
+    """
+    from executorch.backends.mlx.serialization.mlx_graph_schema import (
+        IntOrVid,
+        ItemIntNode,
+    )
+
+    if not isinstance(sorted_indices, Slot):
+        return IntOrVid.from_literal(1 if sorted_indices else 0)
+
+    memo = getattr(P, "_sorted_indices_flag_memo", None)
+    if memo is None:
+        memo = P._sorted_indices_flag_memo = {}
+    cached = memo.get(id(sorted_indices))
+    if cached is not None:
+        return cached
+
+    _, item_slot = P.make_tmp_value_slot()
+    P.emit(ItemIntNode(x=P.slot_to_tid(sorted_indices), out=P.slot_to_vid(item_slot)))
+    flag = P.to_int_or_vid(item_slot)
+    memo[id(sorted_indices)] = flag
+    return flag
+
+
 @REGISTRY.register(target=[torch.ops.mlx.gather_mm.default])
 def _gather_mm_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     """Handle mlx::gather_mm — fused gather + matmul for MoE experts."""
@@ -1802,7 +1838,8 @@ def _gather_mm_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     b = args[1]
     rhs_indices = args[2] if len(args) > 2 else kwargs.get("rhs_indices")
     lhs_indices = args[3] if len(args) > 3 else kwargs.get("lhs_indices")
-    sorted_indices = args[4] if len(args) > 4 else kwargs.get("sorted_indices", False)
+    sorted_indices = args[4] if len(args) > 4 else kwargs.get("sorted_indices")
+    sorted_indices_iov = _resolve_sorted_indices_flag(P, sorted_indices)
 
     out = P.make_or_get_slot(n)
     P.emit(
@@ -1812,7 +1849,8 @@ def _gather_mm_handler(P: MLXProgramBuilder, n: Node) -> Slot:
             out=P.slot_to_tid(out),
             lhs_indices=P.slot_to_tid(lhs_indices) if lhs_indices is not None else None,
             rhs_indices=P.slot_to_tid(rhs_indices) if rhs_indices is not None else None,
-            sorted_indices=sorted_indices,
+            sorted_indices=False,
+            sorted_indices_flag=sorted_indices_iov,
         )
     )
     return out
@@ -1840,7 +1878,8 @@ def _gather_qmm_handler(P: MLXProgramBuilder, n: Node) -> Slot:
     group_size = args[7] if len(args) > 7 else kwargs.get("group_size", 32)
     bits = args[8] if len(args) > 8 else kwargs.get("bits", 4)
     mode = args[9] if len(args) > 9 else kwargs.get("mode", "affine")
-    sorted_indices = args[10] if len(args) > 10 else kwargs.get("sorted_indices", False)
+    sorted_indices = args[10] if len(args) > 10 else kwargs.get("sorted_indices")
+    sorted_indices_iov = _resolve_sorted_indices_flag(P, sorted_indices)
 
     # Convert quantized weights to MLX format
     w_target, w_data = P.get_placeholder_target_and_tensor(w_node)
@@ -1888,12 +1927,311 @@ def _gather_qmm_handler(P: MLXProgramBuilder, n: Node) -> Slot:
             group_size=group_size,
             bits=bits,
             mode=mode,
-            sorted_indices=sorted_indices,
+            sorted_indices=False,
+            sorted_indices_flag=sorted_indices_iov,
         )
     )
     return out
 
 
+@REGISTRY.register(target=[torch.ops.mlx.moe_gather_inputs.default])
+def _moe_gather_inputs_handler(P: MLXProgramBuilder, n: Node):
+    """Lowering for mlx::moe_gather_inputs. Mirrors the moe_gather_inputs
+    eager reference (custom_ops.py) branch-for-branch: emit_sorted for the
+    N > sort_cutoff case, emit_unsorted otherwise. emit_if_else requires both
+    branches to write the same fixed output slots, so each branch's final
+    producer targets out_slots directly (no IdCopy indirection).
+    """
+    from executorch.backends.mlx.serialization.mlx_graph_schema import (
+        ARangeNode,
+        ArgsortNode,
+        AsTypeNode,
+        FloatOrVid,
+        FloorDivideNode,
+        FullNode,
+        IntOrVid,
+        IntOrVidOrTid,
+        RepeatNode,
+        ReshapeNode,
+        TakeNode,
+    )
+
+    args = P.args(n)
+    x, expert_indices = args[0], args[1]
+    top_k = args[2]  # static int
+    sort_cutoff = args[3]  # static int, consulted here at emission
+    if not isinstance(top_k, int) or top_k < 1:
+        raise ValueError(f"moe_gather_inputs: top_k must be a static int >= 1, got {top_k!r}")
+    if not isinstance(sort_cutoff, int) or sort_cutoff < 1:
+        raise ValueError(
+            f"moe_gather_inputs: sort_cutoff must be a static int >= 1, got {sort_cutoff!r}"
+        )
+
+    out_slots = P.make_or_get_slots(n)  # (x_input, idx, sort_experts, inv_order)
+
+    m_iov = emit_shape(P, n.args[0], x, end_dim=1)[0]  # M = N (token count)
+    int32_st = torch_dtype_to_scalar_type(torch.int32)
+
+    def emit_flag(value: int) -> None:
+        # 0-d int32 flag written straight into out_slots[2].
+        P.emit(
+            FullNode(
+                shape=[],
+                v=FloatOrVid.from_literal(float(value)),
+                scalar_type=int32_st,
+                out=P.slot_to_tid(out_slots[2]),
+            )
+        )
+
+    def emit_sorted():
+        # Mirror eager: flat = expert_indices.flatten() before argsort.
+        # Argsorting a [N, top_k] tensor on axis=0 sorts columns independently
+        # and yields the wrong permutation / output layout.
+        _, flat_slot = P.make_tmp_slot()
+        P.emit(
+            ReshapeNode(
+                x=P.slot_to_tid(expert_indices),
+                out=P.slot_to_tid(flat_slot),
+                shape=[IntOrVid.from_literal(-1)],
+            )
+        )
+
+        # MLX argsort returns uint32; cast to int32 like eager (.to(torch.int32)).
+        _, order_u_slot = P.make_tmp_slot()
+        P.emit(
+            ArgsortNode(
+                x=P.slot_to_tid(flat_slot), out=P.slot_to_tid(order_u_slot), axis=0
+            )
+        )
+        _, order_slot = P.make_tmp_slot()
+        P.emit(
+            AsTypeNode(
+                x=P.slot_to_tid(order_u_slot),
+                out=P.slot_to_tid(order_slot),
+                scalar_type=int32_st,
+            )
+        )
+
+        _, inv_order_u_slot = P.make_tmp_slot()
+        P.emit(
+            ArgsortNode(
+                x=P.slot_to_tid(order_slot),
+                out=P.slot_to_tid(inv_order_u_slot),
+                axis=0,
+            )
+        )
+        P.emit(
+            AsTypeNode(
+                x=P.slot_to_tid(inv_order_u_slot),
+                out=P.slot_to_tid(out_slots[3]),
+                scalar_type=int32_st,
+            )
+        )
+
+        # idx = flat[order].to(int32) — the cast matches the eager reference
+        # and the int32 dtype advertised by moe_gather_inputs_fake.
+        _, idx_raw_slot = P.make_tmp_slot()
+        P.emit(
+            TakeNode(
+                x=P.slot_to_tid(flat_slot),
+                out=P.slot_to_tid(idx_raw_slot),
+                index=IntOrVidOrTid.from_tid(P.slot_to_tid(order_slot)),
+                axis=0,
+            )
+        )
+        P.emit(
+            AsTypeNode(
+                x=P.slot_to_tid(idx_raw_slot),
+                out=P.slot_to_tid(out_slots[1]),
+                scalar_type=int32_st,
+            )
+        )
+
+        top_k_const = emit_lifted_constant(P, top_k, torch.int32)
+        # FloorDivideNode routes integer tensors through mlx::divide, which
+        # promotes to float (at_least_float). Cast back before Take — otherwise
+        # gather rejects the indices ("must be integral").
+        _, row_idx_f_slot = P.make_tmp_slot()
+        P.emit(
+            FloorDivideNode(
+                a=P.slot_to_tid(order_slot),
+                b=P.slot_to_tid(top_k_const),
+                out=P.slot_to_tid(row_idx_f_slot),
+            )
+        )
+        _, row_idx_slot = P.make_tmp_slot()
+        P.emit(
+            AsTypeNode(
+                x=P.slot_to_tid(row_idx_f_slot),
+                out=P.slot_to_tid(row_idx_slot),
+                scalar_type=int32_st,
+            )
+        )
+
+        _, x_gathered_slot = P.make_tmp_slot()
+        P.emit(
+            TakeNode(
+                x=P.slot_to_tid(x),
+                out=P.slot_to_tid(x_gathered_slot),
+                index=IntOrVidOrTid.from_tid(P.slot_to_tid(row_idx_slot)),
+                axis=0,
+            )
+        )
+        P.emit(
+            ExpandDimsNode(
+                x=P.slot_to_tid(x_gathered_slot),
+                out=P.slot_to_tid(out_slots[0]),
+                axis=-2,
+            )
+        )
+
+        emit_flag(1)
+
+    def emit_unsorted():
+        _, x_rep_slot = P.make_tmp_slot()
+        P.emit(
+            RepeatNode(
+                x=P.slot_to_tid(x),
+                out=P.slot_to_tid(x_rep_slot),
+                repeats=IntOrVid.from_literal(top_k),
+                axis=0,
+            )
+        )
+        P.emit(
+            ExpandDimsNode(
+                x=P.slot_to_tid(x_rep_slot),
+                out=P.slot_to_tid(out_slots[0]),
+                axis=-2,
+            )
+        )
+
+        # idx = expert_indices.flatten().to(int32) — mirror the eager
+        # reference: [N, top_k] -> [N*top_k], int32. Passing the unflattened
+        # 2-D tensor through would break gather_mm's index broadcasting for
+        # 1 < N <= sort_cutoff and diverge from the fake's advertised meta.
+        _, idx_flat_slot = P.make_tmp_slot()
+        P.emit(
+            ReshapeNode(
+                x=P.slot_to_tid(expert_indices),
+                out=P.slot_to_tid(idx_flat_slot),
+                shape=[IntOrVid.from_literal(-1)],
+            )
+        )
+        P.emit(
+            AsTypeNode(
+                x=P.slot_to_tid(idx_flat_slot),
+                out=P.slot_to_tid(out_slots[1]),
+                scalar_type=int32_st,
+            )
+        )
+
+        emit_flag(0)
+
+        # Identity permutation (inverse of no-reorder). Semantically correct
+        # for unsorted, and safe if scatter ever takes the sorted branch:
+        # Take(down, arange) is a no-op instead of crashing on empty(0).
+        nk_iov = emit_product(P, [m_iov, IntOrVid.from_literal(top_k)])
+        P.emit(
+            ARangeNode(
+                out=P.slot_to_tid(out_slots[3]),
+                start=IntOrVid.from_literal(0),
+                stop=nk_iov,
+                step=IntOrVid.from_literal(1),
+                scalar_type=int32_st,
+            )
+        )
+
+    # cond = (M - 1) // sort_cutoff: 0 (-> else/unsorted) for M <= sort_cutoff,
+    # >= 1 (-> then/sorted) for M > sort_cutoff. The IfNode rule is nonzero ->
+    # then. If M is a compile-time literal, both fold and emit_if_else picks
+    # one branch — no IfNode emitted.
+    cond = emit_floordiv(
+        P,
+        emit_sub_int(P, m_iov, IntOrVid.from_literal(1)),
+        IntOrVid.from_literal(sort_cutoff),
+    )
+
+    # Seed the sorted_indices_flag memo: cond has the same truthiness as the
+    # 0-d flag tensor in out_slots[2], so downstream gather_mm/gather_qmm and
+    # moe_scatter_outputs can consume it as an IntOrVid directly instead of
+    # each emitting an ItemIntNode (which forces a device sync per read).
+    memo = getattr(P, "_sorted_indices_flag_memo", None)
+    if memo is None:
+        memo = P._sorted_indices_flag_memo = {}
+    memo[id(out_slots[2])] = cond
+
+    emit_if_else(P, cond, emit_sorted, emit_unsorted)
+    return out_slots
+
+
+@REGISTRY.register(target=[torch.ops.mlx.moe_scatter_outputs.default])
+def _moe_scatter_outputs_handler(P: MLXProgramBuilder, n: Node) -> Slot:
+    """Lowering for mlx::moe_scatter_outputs. down [N*top_k, 1, H] -> squeeze
+    -> (gather if sorted) -> reshape [N, top_k, H]. The sorted/unsorted
+    condition comes from _resolve_sorted_indices_flag: a memoized shape-derived
+    Vid when sort_experts was produced by moe_gather_inputs (folding both ops
+    onto the same branch under static shapes), an ItemIntNode otherwise.
+    """
+    from executorch.backends.mlx.serialization.mlx_graph_schema import (
+        IntOrVid,
+        IntOrVidOrTid,
+        ReshapeNode,
+        SqueezeNode,
+        TakeNode,
+    )
+
+    args = P.args(n)
+    down, sort_experts, inv_order = args[0], args[1], args[2]
+    top_k = args[3]  # static int
+    if not isinstance(top_k, int) or top_k < 1:
+        raise ValueError(
+            f"moe_scatter_outputs: top_k must be a static int >= 1, got {top_k!r}"
+        )
+
+    out_slot = P.make_or_get_slot(n)
+
+    _, down_sq_slot = P.make_tmp_slot()
+    P.emit(
+        SqueezeNode(x=P.slot_to_tid(down), out=P.slot_to_tid(down_sq_slot), dims=[-2])
+    )
+
+    cond = _resolve_sorted_indices_flag(P, sort_experts)
+
+    n_top_k_iov = emit_shape(P, n.args[0], down, end_dim=1)[0]  # N * top_k
+    n_iov = emit_floordiv(P, n_top_k_iov, IntOrVid.from_literal(top_k))
+
+    def emit_then():  # prefill: scatter back
+        _, gathered_slot = P.make_tmp_slot()
+        P.emit(
+            TakeNode(
+                x=P.slot_to_tid(down_sq_slot),
+                out=P.slot_to_tid(gathered_slot),
+                index=IntOrVidOrTid.from_tid(P.slot_to_tid(inv_order)),
+                axis=0,
+            )
+        )
+        P.emit(
+            ReshapeNode(
+                x=P.slot_to_tid(gathered_slot),
+                out=P.slot_to_tid(out_slot),
+                shape=[n_iov, IntOrVid.from_literal(top_k), IntOrVid.from_literal(-1)],
+            )
+        )
+
+    def emit_else():  # decode: no scatter
+        P.emit(
+            ReshapeNode(
+                x=P.slot_to_tid(down_sq_slot),
+                out=P.slot_to_tid(out_slot),
+                shape=[n_iov, IntOrVid.from_literal(top_k), IntOrVid.from_literal(-1)],
+            )
+        )
+
+    emit_if_else(P, cond, emit_then, emit_else)
+    return out_slot
+
+
+
 @REGISTRY.register(
     target=[torch.ops.aten.split.Tensor, torch.ops.aten.split_copy.Tensor]
 )
diff --git a/backends/mlx/runtime/MLXInterpreter.h b/backends/mlx/runtime/MLXInterpreter.h
index 51946e966a6..ad15d3748dc 100644
--- a/backends/mlx/runtime/MLXInterpreter.h
+++ b/backends/mlx/runtime/MLXInterpreter.h
@@ -941,7 +941,10 @@ exec_gather_mm(const GatherMmNode& n, ExecutionState& st, StreamOrDevice s) {
     rhs_idx = st.const_tensor_ref(*n.rhs_indices);
   }
 
-  array Y = gather_mm(A, B, lhs_idx, rhs_idx, n.sorted_indices, s);
+  bool sorted = n.sorted_indices_flag.has_value()
+      ? (resolve_int(*n.sorted_indices_flag, st) != 0)
+      : n.sorted_indices;
+  array Y = gather_mm(A, B, lhs_idx, rhs_idx, sorted, s);
   st.set_tensor(n.out, std::move(Y));
 }
 
@@ -964,6 +967,9 @@ exec_gather_qmm(const GatherQmmNode& n, ExecutionState& st, StreamOrDevice s) {
     rhs_idx = st.const_tensor_ref(*n.rhs_indices);
   }
 
+  bool sorted = n.sorted_indices_flag.has_value()
+      ? (resolve_int(*n.sorted_indices_flag, st) != 0)
+      : n.sorted_indices;
   array Y = gather_qmm(
       X,
       Wq,
@@ -975,7 +981,7 @@ exec_gather_qmm(const GatherQmmNode& n, ExecutionState& st, StreamOrDevice s) {
       n.group_size,
       n.bits,
       n.mode,
-      n.sorted_indices,
+      sorted,
       s);
   st.set_tensor(n.out, std::move(Y));
 }
diff --git a/backends/mlx/serialization/generate.py b/backends/mlx/serialization/generate.py
index 46e89526bbc..778b8f466b2 100755
--- a/backends/mlx/serialization/generate.py
+++ b/backends/mlx/serialization/generate.py
@@ -828,6 +828,10 @@ def _emit_py_prebuild(kind: str, fld: FBSField) -> List[str]:
             return [f"        {n}_vec = {expr}"]
         else:
             return [f"        {n}_vec = {expr} if op.{n} is not None else None"]
+    if kind == "optional_int_or_vid":
+        return [
+            f"        {n}_off = self._build_int_or_vid(builder, op.{n}) if op.{n} is not None else None"
+        ]
     if kind in _PY_PREBUILD_OFFSET:
         suffix = "_off"
         expr = _PY_PREBUILD_OFFSET[kind].format(name=n)
@@ -890,6 +894,12 @@ def _emit_py_add(
             f"        if {n}_off is not None:",
             f"            {add}(builder, {n}_off)",
         ]
+    # Optional compound offset (e.g. sorted_indices_flag: IntOrVid, no default)
+    if kind == "optional_int_or_vid":
+        return [
+            f"        if {n}_off is not None:",
+            f"            {add}(builder, {n}_off)",
+        ]
     return None
 
 
@@ -927,7 +937,7 @@ def _get_field_kind(fld: FBSField, table: FBSTable) -> str:  # noqa: C901
     if t == "Vid":
         return "optional_vid" if not fld.required else "vid"
     if t == "IntOrVid":
-        return "int_or_vid"
+        return "optional_int_or_vid" if not fld.required else "int_or_vid"
     if t == "FloatOrVid":
         return "float_or_vid"
     if t == "VidOrTid":
@@ -1099,6 +1109,8 @@ def _fbs_type_to_cpp(
             return "std::optional<Tid>"
         if fbs_type == "Vid":
             return "std::optional<Vid>"
+        if fbs_type in FBS_COMPOUND_TYPES:
+            return f"std::optional<{cpp_type}>"
         if fld is not None and fld.default == "null" and fbs_type in FBS_TO_CPP:
             return f"std::optional<{cpp_type}>"
 
@@ -1378,6 +1390,7 @@ def _fixup_flatc_imports() -> None:
     "vid": "vid",
     "optional_vid": "vid",
     "int_or_vid": "int_or_vid",
+    "optional_int_or_vid": "int_or_vid",
     "float_or_vid": "float_or_vid",
     "vid_or_tid": "vid_or_tid",
     "int_or_vid_or_tid": "int_or_vid_or_tid",
diff --git a/backends/mlx/serialization/schema.fbs b/backends/mlx/serialization/schema.fbs
index c69f5f48dda..87ab9aec10a 100644
--- a/backends/mlx/serialization/schema.fbs
+++ b/backends/mlx/serialization/schema.fbs
@@ -962,6 +962,11 @@ table GatherMmNode {
     lhs_indices: Tid;           // optional - LHS gather indices
     rhs_indices: Tid;           // optional - RHS gather indices (expert selection)
     sorted_indices: bool = false;
+    // Runtime (possibly data-dependent) override of `sorted_indices` above.
+    // Appended last to preserve wire compatibility with existing .pte files:
+    // when absent, deserializers fall back to the static `sorted_indices`
+    // bool. When present, it takes precedence.
+    sorted_indices_flag: IntOrVid;
 }
 
 table GatherQmmNode {
@@ -977,6 +982,8 @@ table GatherQmmNode {
     group_size: int32;
     bits: int32;
     sorted_indices: bool = false;
+    // See GatherMmNode.sorted_indices_flag.
+    sorted_indices_flag: IntOrVid;
 }
 
 table ScanNode {
diff --git a/backends/mlx/test/test_ops.py b/backends/mlx/test/test_ops.py
index 888c96a0731..8e3e7a7e1e8 100644
--- a/backends/mlx/test/test_ops.py
+++ b/backends/mlx/test/test_ops.py
@@ -6603,12 +6603,28 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]:
 class GatherMmModel(nn.Module):
     """Model using mlx::gather_mm for expert selection + matmul."""
 
-    def __init__(self, num_experts: int, in_features: int, out_features: int):
+    def __init__(
+        self,
+        num_experts: int,
+        in_features: int,
+        out_features: int,
+        sorted_indices: bool = False,
+    ):
         super().__init__()
         self.register_buffer(
             "weight",
             torch.randn(num_experts, out_features, in_features),
         )
+        # sorted_indices is Optional[Tensor] (0-d int32) rather than a bool.
+        # Store as buffer so it is part of the exported graph and exercises
+        # the IntOrVid runtime path in the handler.
+        if sorted_indices:
+            self.register_buffer(
+                "sorted_flag",
+                torch.ones((), dtype=torch.int32),
+            )
+        else:
+            self.sorted_flag = None
 
     def forward(self, x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
         import executorch.backends.mlx.custom_ops as _  # noqa
@@ -6617,13 +6633,20 @@ def forward(self, x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
         # Transpose weight from [E, out, in] to [E, in, out]
         # gather_mm returns [N, 1, out], squeeze dim -2
         return torch.ops.mlx.gather_mm(
-            x.unsqueeze(-2), self.weight.transpose(-1, -2), rhs_indices=indices
+            x.unsqueeze(-2),
+            self.weight.transpose(-1, -2),
+            rhs_indices=indices,
+            sorted_indices=self.sorted_flag,
         ).squeeze(-2)
 
 
 @register_test
 class GatherMmTest(OpTestCase):
-    """Test case for mlx::gather_mm."""
+    """Test case for mlx::gather_mm.
+
+    Includes a sorted=True config to exercise the Optional[Tensor] ->
+    IntOrVid runtime path in _gather_mm_handler.
+    """
 
     name = "gather_mm"
     rtol = 1e-4
@@ -6636,16 +6659,20 @@ def __init__(
         out_features: int = 128,
         batch_size: int = 2,
         dtype: torch.dtype = torch.float32,
+        sorted_indices: bool = False,
     ):
         self.num_experts = num_experts
         self.in_features = in_features
         self.out_features = out_features
         self.batch_size = batch_size
         self.dtype = dtype
+        self.sorted_indices = sorted_indices
 
         parts = ["gather_mm", f"e{num_experts}", f"i{in_features}", f"o{out_features}"]
         if dtype != torch.float32:
             parts.append(str(dtype).split(".")[-1])
+        if sorted_indices:
+            parts.append("sorted")
         self.name = "_".join(parts)
 
     @classmethod
@@ -6655,15 +6682,28 @@ def get_test_configs(cls) -> List["GatherMmTest"]:
             cls(num_experts=8, in_features=128, out_features=256),
             cls(dtype=torch.bfloat16),
             cls(batch_size=1),
+            # Exercise sorted_indices=Tensor (IntOrVid runtime path)
+            cls(sorted_indices=True),
+            cls(sorted_indices=True, dtype=torch.bfloat16),
         ]
 
     def create_model(self) -> nn.Module:
-        model = GatherMmModel(self.num_experts, self.in_features, self.out_features)
+        model = GatherMmModel(
+            self.num_experts,
+            self.in_features,
+            self.out_features,
+            sorted_indices=self.sorted_indices,
+        )
         return model.to(self.dtype)
 
     def create_inputs(self) -> Tuple[torch.Tensor, ...]:
         x = torch.randn(self.batch_size, self.in_features, dtype=self.dtype)
         indices = torch.randint(0, self.num_experts, (self.batch_size,))
+        # sorted_indices=True is a contract with the MLX kernel: indices must
+        # actually be sorted. Eager ignores the flag; unsorted + sorted=True
+        # yields large localized numeric errors (max_diff ~ tens).
+        if self.sorted_indices:
+            indices, _ = torch.sort(indices)
         return (x, indices)
 
 
@@ -6681,10 +6721,16 @@ def __init__(
         out_features: int,
         group_size: int = 32,
         packed: bool = False,
+        sorted_indices: bool = False,
     ):
         super().__init__()
         self.out_features = out_features
         self.group_size = group_size
+        # Same pattern as GatherMmModel
+        if sorted_indices:
+            self.register_buffer("sorted_flag", torch.ones((), dtype=torch.int32))
+        else:
+            self.sorted_flag = None
 
         # Create per-expert nn.Linear, quantize, extract inner tensors
         from executorch.backends.mlx.llm.quantization import quantize_model_
@@ -6729,12 +6775,17 @@ def forward(self, x: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
             biases=self.zero_point,
             rhs_indices=indices,
             group_size=self.group_size,
+            sorted_indices=self.sorted_flag,
         ).squeeze(-2)
 
 
 @register_test
 class GatherQmmTest(OpTestCase):
-    """Test case for mlx::gather_qmm."""
+    """Test case for mlx::gather_qmm.
+
+    Includes a sorted=True config to exercise the Optional[Tensor] ->
+    IntOrVid runtime path in _gather_qmm_handler.
+    """
 
     name = "gather_qmm"
     rtol = 0.1
@@ -6749,6 +6800,7 @@ def __init__(
         group_size: int = 32,
         dtype: torch.dtype = torch.float32,
         packed: bool = False,
+        sorted_indices: bool = False,
     ):
         self.num_experts = num_experts
         self.in_features = in_features
@@ -6757,6 +6809,7 @@ def __init__(
         self.group_size = group_size
         self.dtype = dtype
         self.packed = packed
+        self.sorted_indices = sorted_indices
 
         parts = [
             "gather_qmm",
@@ -6769,6 +6822,8 @@ def __init__(
             parts.append("packed")
         if dtype != torch.float32:
             parts.append(str(dtype).split(".")[-1])
+        if sorted_indices:
+            parts.append("sorted")
         self.name = "_".join(parts)
 
     @classmethod
@@ -6782,6 +6837,9 @@ def get_test_configs(cls) -> List["GatherQmmTest"]:
             # to_mlx_qparams prepacked (view -> uint32) lowering path.
             cls(packed=True),
             cls(packed=True, dtype=torch.bfloat16),
+            # Exercise sorted_indices=Tensor (IntOrVid runtime path)
+            cls(sorted_indices=True),
+            cls(sorted_indices=True, dtype=torch.bfloat16),
         ]
 
     def get_edge_compile_config(self):
@@ -6796,12 +6854,15 @@ def create_model(self) -> nn.Module:
             self.out_features,
             self.group_size,
             packed=self.packed,
+            sorted_indices=self.sorted_indices,
         )
         return model.to(self.dtype)
 
     def create_inputs(self) -> Tuple[torch.Tensor, ...]:
         x = torch.randn(self.batch_size, self.in_features, dtype=self.dtype)
         indices = torch.randint(0, self.num_experts, (self.batch_size,))
+        if self.sorted_indices:
+            indices, _ = torch.sort(indices)
         return (x, indices)
 
 
@@ -8005,3 +8066,250 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]:
             torch.tensor(self.temperature),
             torch.tensor(0, dtype=torch.int64),
         )
+
+
+# ---------------------------------------------------------------------------
+# MoE runtime expert-sort for decode: moe_gather_inputs / moe_scatter_outputs
+# ---------------------------------------------------------------------------
+
+
+class MoeGatherInputsModel(nn.Module):
+    """Wraps moe_gather_inputs to make it exportable as a single-output model.
+    Returns only x_input (the first of the four outputs) so OpTestCase can
+    compare it against the eager reference using its standard allclose check.
+    The remaining outputs (idx, sort_experts, inv_order) are validated in
+    MoeScatterOutputsModel below via the round-trip prefill test.
+    """
+
+    def __init__(self, top_k: int = 2, sort_cutoff: int = 1):
+        super().__init__()
+        self.top_k = top_k
+        self.sort_cutoff = sort_cutoff
+
+    def forward(self, x: torch.Tensor, expert_indices: torch.Tensor) -> torch.Tensor:
+        import executorch.backends.mlx.custom_ops as _  # noqa: F401
+
+        x_input = torch.ops.mlx.moe_gather_inputs(
+            x, expert_indices, self.top_k, self.sort_cutoff
+        )[0]
+        return x_input
+
+
+@register_test
+class MoeGatherInputsTest(OpTestCase):
+    """Test case for mlx::moe_gather_inputs.
+
+    Static configs cover the N=1 (decode, unsorted) and N>sort_cutoff
+    (prefill, sorted) folded paths; the dynamic config keeps N symbolic so
+    both branches and the runtime IfNode selection are exercised from a
+    single exported artifact (the actual point of the runtime sort).
+    """
+
+    name = "moe_gather_inputs"
+    rtol = 1e-5
+    atol = 1e-5
+
+    def __init__(
+        self,
+        batch_size: int = 4,
+        hidden_size: int = 32,
+        num_experts: int = 8,
+        top_k: int = 2,
+        sort_cutoff: int = 1,
+        dynamic_batch: bool = False,
+        tag: str = "",
+    ):
+        self.batch_size = batch_size
+        self.hidden_size = hidden_size
+        self.num_experts = num_experts
+        self.top_k = top_k
+        self.sort_cutoff = sort_cutoff
+        self.dynamic_batch = dynamic_batch
+
+        parts = ["moe_gather_inputs", f"N{batch_size}", f"E{num_experts}", f"k{top_k}"]
+        if tag:
+            parts.append(tag)
+        self.name = "_".join(parts)
+        self.expected_node_counts = self.get_expected_node_counts()
+
+    @classmethod
+    def get_test_configs(cls) -> List["MoeGatherInputsTest"]:
+        return [
+            cls(batch_size=4, tag="prefill"),  # N > sort_cutoff -> sorted path
+            cls(batch_size=1, tag="decode"),  # N <= sort_cutoff -> unsorted path
+            cls(batch_size=4, num_experts=16, top_k=4, tag="top4"),
+            # Symbolic N: emits both branches behind an IfNode; exported at
+            # N=4 (sorted at runtime), re-run at N=1 (unsorted at runtime).
+            cls(batch_size=4, dynamic_batch=True, tag="dyn"),
+        ]
+
+    def get_expected_node_counts(self) -> Optional[Dict[str, int]]:
+        if self.dynamic_batch:
+            # Both branches are present behind an IfNode; per-branch node
+            # counts depend on how the counter treats branch chains, so only
+            # the runtime-selection structure is asserted elsewhere.
+            return None
+        if self.batch_size > self.sort_cutoff:
+            # sorted path only (condition folds at export time)
+            return {
+                "ArgsortNode": 2,
+                "TakeNode": 2,
+                "FloorDivideNode": 1,
+                "RepeatNode": 0,
+                "ARangeNode": 0,
+                "IfNode": 0,
+            }
+        # unsorted path only
+        return {
+            "ArgsortNode": 0,
+            "TakeNode": 0,
+            "RepeatNode": 1,
+            "ARangeNode": 1,
+            "IfNode": 0,
+        }
+
+    def create_model(self) -> nn.Module:
+        return MoeGatherInputsModel(top_k=self.top_k, sort_cutoff=self.sort_cutoff)
+
+    def create_inputs(self) -> Tuple[torch.Tensor, ...]:
+        x = torch.randn(self.batch_size, self.hidden_size)
+        expert_indices = torch.randint(
+            0, self.num_experts, (self.batch_size, self.top_k)
+        )
+        return (x, expert_indices)
+
+    def create_test_inputs(self) -> Tuple[torch.Tensor, ...]:
+        if not self.dynamic_batch:
+            return self.create_inputs()
+        # Decode-shaped inputs: forces the runtime IfNode onto the unsorted
+        # branch of the artifact exported at batch_size.
+        x = torch.randn(1, self.hidden_size)
+        expert_indices = torch.randint(0, self.num_experts, (1, self.top_k))
+        return (x, expert_indices)
+
+    def get_dynamic_shapes(self) -> Optional[Dict[str, any]]:
+        if not self.dynamic_batch:
+            return None
+        batch_dim = Dim("moe_batch", min=1, max=16)
+        return {
+            "x": {0: batch_dim},
+            "expert_indices": {0: batch_dim},
+        }
+
+
+class MoeScatterOutputsModel(nn.Module):
+    """Round-trip: moe_gather_inputs -> identity down_proj -> moe_scatter_outputs.
+    Returns the final [N, top_k, hidden] tensor so OpTestCase can verify
+    the full gather/sort/scatter pipeline end-to-end.
+    """
+
+    def __init__(self, top_k: int = 2, sort_cutoff: int = 1, hidden_out: int = 16):
+        super().__init__()
+        self.top_k = top_k
+        self.sort_cutoff = sort_cutoff
+        self.hidden_out = hidden_out
+
+    def forward(self, x: torch.Tensor, expert_indices: torch.Tensor) -> torch.Tensor:
+        import executorch.backends.mlx.custom_ops as _  # noqa: F401
+
+        x_input, idx, sort_experts, inv_order = torch.ops.mlx.moe_gather_inputs(
+            x, expert_indices, self.top_k, self.sort_cutoff
+        )
+        # Simulate a down_proj output: [N*top_k, 1, hidden_out]. A plain
+        # slice keeps the per-row values distinct, so a wrong or missing
+        # inverse permutation in the lowering shows up as a mismatch.
+        down = x_input[..., : self.hidden_out].contiguous()
+        return torch.ops.mlx.moe_scatter_outputs(
+            down, sort_experts, inv_order, self.top_k
+        )
+
+
+@register_test
+class MoeScatterOutputsTest(OpTestCase):
+    """Test case for mlx::moe_scatter_outputs.
+
+    Validates the round-trip shape and that the unsorted (decode) path
+    produces a result consistent with the sorted (prefill) path. The dynamic
+    config runs both runtime branches from one artifact.
+    """
+
+    name = "moe_scatter_outputs"
+    rtol = 1e-4
+    atol = 1e-4
+
+    def __init__(
+        self,
+        batch_size: int = 4,
+        hidden_size: int = 32,
+        hidden_out: int = 16,
+        num_experts: int = 8,
+        top_k: int = 2,
+        sort_cutoff: int = 1,
+        dynamic_batch: bool = False,
+        tag: str = "",
+    ):
+        self.batch_size = batch_size
+        self.hidden_size = hidden_size
+        self.hidden_out = hidden_out
+        self.num_experts = num_experts
+        self.top_k = top_k
+        self.sort_cutoff = sort_cutoff
+        self.dynamic_batch = dynamic_batch
+
+        parts = [
+            "moe_scatter_outputs",
+            f"N{batch_size}",
+            f"E{num_experts}",
+            f"k{top_k}",
+        ]
+        if tag:
+            parts.append(tag)
+        self.name = "_".join(parts)
+        self.expected_node_counts = self.get_expected_node_counts()
+
+    @classmethod
+    def get_test_configs(cls) -> List["MoeScatterOutputsTest"]:
+        return [
+            cls(batch_size=4, tag="prefill"),
+            cls(batch_size=1, tag="decode"),
+            cls(batch_size=4, num_experts=16, top_k=4, hidden_out=32, tag="top4"),
+            cls(batch_size=4, dynamic_batch=True, tag="dyn"),
+        ]
+
+    def get_expected_node_counts(self) -> Optional[Dict[str, int]]:
+        if self.dynamic_batch:
+            return None
+        # With a static batch the sort condition folds at export time, and
+        # the scatter reuses the same folded condition (via the builder's
+        # sorted_indices_flag memo) — so no IfNode survives in either op.
+        return {"IfNode": 0}
+
+    def create_model(self) -> nn.Module:
+        return MoeScatterOutputsModel(
+            top_k=self.top_k,
+            sort_cutoff=self.sort_cutoff,
+            hidden_out=self.hidden_out,
+        )
+
+    def create_inputs(self) -> Tuple[torch.Tensor, ...]:
+        x = torch.randn(self.batch_size, self.hidden_size)
+        expert_indices = torch.randint(
+            0, self.num_experts, (self.batch_size, self.top_k)
+        )
+        return (x, expert_indices)
+
+    def create_test_inputs(self) -> Tuple[torch.Tensor, ...]:
+        if not self.dynamic_batch:
+            return self.create_inputs()
+        x = torch.randn(1, self.hidden_size)
+        expert_indices = torch.randint(0, self.num_experts, (1, self.top_k))
+        return (x, expert_indices)
+
+    def get_dynamic_shapes(self) -> Optional[Dict[str, any]]:
+        if not self.dynamic_batch:
+            return None
+        batch_dim = Dim("moe_batch", min=1, max=16)
+        return {
+            "x": {0: batch_dim},
+            "expert_indices": {0: batch_dim},
+        }
diff --git a/examples/models/qwen3_5_moe/export.py b/examples/models/qwen3_5_moe/export.py
index 81e091ce8b3..a8dd37d945c 100644
--- a/examples/models/qwen3_5_moe/export.py
+++ b/examples/models/qwen3_5_moe/export.py
@@ -53,7 +53,7 @@ def _prepare_and_quantize_mlx(model, config, args):
         model,
         model_dtype=torch.bfloat16,
         config=config,
-        sort_experts=True,
+        sort_cutoff=1,
         fuse_gate_up=False,
     )
     if args.qlinear or args.qembedding:
@@ -345,7 +345,7 @@ def load_prequantized_model_mlx(
             model,
             model_dtype=torch.bfloat16,
             config=config,
-            sort_experts=True,
+            sort_cutoff=1,
             fuse_gate_up=False,
         )
 
diff --git a/examples/models/qwen3_5_moe/mlx_source_transformations.py b/examples/models/qwen3_5_moe/mlx_source_transformations.py
index 3c460fc9c54..0921d042d17 100644
--- a/examples/models/qwen3_5_moe/mlx_source_transformations.py
+++ b/examples/models/qwen3_5_moe/mlx_source_transformations.py
@@ -70,7 +70,6 @@ def _sparse_moe_forward(self, x):
         expert_weights,
         expert_indices,
         self.top_k,
-        sort_experts=getattr(self, "_sort_experts", False),
     )
 
     shared_out = self.shared_expert(x_flat)
@@ -215,7 +214,7 @@ def _exportable_gated_delta_net_forward(self, x, input_pos):
     return self.out_proj(output)
 
 
-def _swap_moe_experts(model, fuse_gate_up):
+def _swap_moe_experts(model, fuse_gate_up, sort_cutoff=1):
     """FusedMoEExperts → SwitchMLP."""
     from executorch.backends.mlx.llm.switch import SwitchMLP
 
@@ -229,6 +228,7 @@ def _swap_moe_experts(model, fuse_gate_up):
             module.intermediate_size,
             module.num_experts,
             fuse_gate_up=fuse_gate_up,
+            sort_cutoff=sort_cutoff,
         )
         switch_mlp.to(dtype=module.w1_weight.dtype)
 
@@ -321,12 +321,11 @@ def _swap_rms_norm(model):
     return count
 
 
-def _swap_sparse_moe(model, sort_experts):
+def _swap_sparse_moe(model):
     """SparseMoE → no .float() on expert_weights."""
     count = 0
     for _name, module in model.named_modules():
         if isinstance(module, SparseMoE):
-            module._sort_experts = sort_experts
             module.forward = types.MethodType(_sparse_moe_forward, module)
             count += 1
     return count
@@ -336,7 +335,7 @@ def mlx_source_transformations(
     model,
     model_dtype=torch.bfloat16,
     config=None,
-    sort_experts=False,
+    sort_cutoff=1,
     fuse_gate_up=False,
 ):
     """Replace all Triton-dependent modules with MLX-compatible equivalents.
@@ -353,15 +352,16 @@ def mlx_source_transformations(
         model: The Qwen 3.5 MoE model to transform.
         model_dtype: Target dtype for the model (default: bf16).
         config: Model config (Qwen35MoEConfig).
-        sort_experts: Sort tokens by expert index for coalesced memory access.
+        sort_cutoff: Token-count threshold for runtime MoE expert sorting.
+            Sort when M > sort_cutoff (default 1 = sort on prefill only).
         fuse_gate_up: Fuse gate+up into single SwitchLinear.
     """
-    count_moe = _swap_moe_experts(model, fuse_gate_up)
+    count_moe = _swap_moe_experts(model, fuse_gate_up, sort_cutoff)
     count_gdn = _swap_gated_delta_net(model, model_dtype)
     count_attn = _swap_full_attention(model, config)
     count_kv = _swap_kv_cache(model, model_dtype)
     count_norm = _swap_rms_norm(model)
-    count_moe_fwd = _swap_sparse_moe(model, sort_experts)
+    count_moe_fwd = _swap_sparse_moe(model)
 
     logger.info(f"Replaced {count_moe} FusedMoEExperts → SwitchMLP")
     logger.info(f"Replaced {count_gdn} GatedDeltaNet → exportable PyTorch forward")

From dcb095c0d15e270785180bd6018e3d861b75ff88 Mon Sep 17 00:00:00 2001
From: "Axel.Cffrd.Dnty" <AxelNoun@users.noreply.github.com>
Date: Sat, 25 Jul 2026 20:32:58 +0200
Subject: [PATCH 2/4] backends/mlx: seed sorted-flag memo through getitem slots

The getitem handler copies tuple outputs into fresh slots (IdCopyNode),
so a memo keyed on the producer's out_slots[2] never matched the slot
consumers actually receive. Seed it under the getitem users' slots as
well, so gather_mm/gather_qmm/moe_scatter_outputs resolve the flag
without emitting an ItemIntNode, and static-shape exports fold the
IfNode away as the test configs assert.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 backends/mlx/ops.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/backends/mlx/ops.py b/backends/mlx/ops.py
index 479987bb78c..a44d1befb4f 100644
--- a/backends/mlx/ops.py
+++ b/backends/mlx/ops.py
@@ -2160,6 +2160,21 @@ def emit_unsorted():
         memo = P._sorted_indices_flag_memo = {}
     memo[id(out_slots[2])] = cond
 
+    # The flag reaches consumers through operator.getitem, whose handler copies
+    # the tuple element into a fresh slot (IdCopyNode) rather than aliasing it.
+    # Seed the memo under those slots too, so gather_mm/gather_qmm and
+    # moe_scatter_outputs resolve the flag without emitting an ItemIntNode --
+    # and, under static shapes, fold their IfNode away entirely (the property
+    # the static test configs assert).
+    for user in n.users:
+        if (
+            user.target is operator.getitem
+            and len(user.args) > 1
+            and user.args[1] == 2
+            and len(user.users) > 0
+        ):
+            memo[id(P.make_or_get_slot(user))] = cond
+
     emit_if_else(P, cond, emit_sorted, emit_unsorted)
     return out_slots
 

From b003ef2423d8c4100c6791a606a42dd48e9a3e77 Mon Sep 17 00:00:00 2001
From: "Axel.Cffrd.Dnty" <AxelNoun@users.noreply.github.com>
Date: Tue, 28 Jul 2026 22:42:11 +0200
Subject: [PATCH 3/4] backends/mlx: satisfy lintrunner (UFMT) on ops.py

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 backends/mlx/ops.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/backends/mlx/ops.py b/backends/mlx/ops.py
index a44d1befb4f..e0912d8c26d 100644
--- a/backends/mlx/ops.py
+++ b/backends/mlx/ops.py
@@ -1961,7 +1961,9 @@ def _moe_gather_inputs_handler(P: MLXProgramBuilder, n: Node):
     top_k = args[2]  # static int
     sort_cutoff = args[3]  # static int, consulted here at emission
     if not isinstance(top_k, int) or top_k < 1:
-        raise ValueError(f"moe_gather_inputs: top_k must be a static int >= 1, got {top_k!r}")
+        raise ValueError(
+            f"moe_gather_inputs: top_k must be a static int >= 1, got {top_k!r}"
+        )
     if not isinstance(sort_cutoff, int) or sort_cutoff < 1:
         raise ValueError(
             f"moe_gather_inputs: sort_cutoff must be a static int >= 1, got {sort_cutoff!r}"
@@ -2246,7 +2248,6 @@ def emit_else():  # decode: no scatter
     return out_slot
 
 
-
 @REGISTRY.register(
     target=[torch.ops.aten.split.Tensor, torch.ops.aten.split_copy.Tensor]
 )

From cb3f835800a1156921d40d40daed40776f899697 Mon Sep 17 00:00:00 2001
From: Scott Roy <scroy@meta.com>
Date: Fri, 31 Jul 2026 17:18:10 -0700
Subject: [PATCH 4/4] up

---
 backends/mlx/CMakeLists.txt                   |  1 +
 .../mlx/patches/mlx_gather_mm_rhs_lda.patch   | 34 +++++++++++++++++++
 extension/llm/export/load.py                  | 22 +++++++++---
 3 files changed, 53 insertions(+), 4 deletions(-)
 create mode 100644 backends/mlx/patches/mlx_gather_mm_rhs_lda.patch

diff --git a/backends/mlx/CMakeLists.txt b/backends/mlx/CMakeLists.txt
index c6eb2ac3156..15683790f83 100644
--- a/backends/mlx/CMakeLists.txt
+++ b/backends/mlx/CMakeLists.txt
@@ -187,6 +187,7 @@ message(
 set(_mlx_patches
     ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_nax_has_include.patch
     ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_qmm_splitk_bk_align.patch
+    ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_gather_mm_rhs_lda.patch
 )
 ExternalProject_Add(
   mlx_external
diff --git a/backends/mlx/patches/mlx_gather_mm_rhs_lda.patch b/backends/mlx/patches/mlx_gather_mm_rhs_lda.patch
new file mode 100644
index 00000000000..8430d83e9cd
--- /dev/null
+++ b/backends/mlx/patches/mlx_gather_mm_rhs_lda.patch
@@ -0,0 +1,34 @@
+Fix the activation row stride in sorted RHS gather_mm.
+
+The specialized gather_mm_rhs paths flatten all leading activation dimensions
+into M, but derive lda from the original second-to-last dimension. For a valid
+[N, 1, K] view produced by expand_dims, that singleton dimension can have
+stride 1. The kernel then reads row n from A + n instead of A + n * K.
+
+The activation is made row-contiguous before this calculation, so flattened
+rows are K elements apart. Use K as lda in both the Steel and NAX paths.
+
+Upstream candidate.
+
+diff --git a/mlx/backend/metal/matmul.cpp b/mlx/backend/metal/matmul.cpp
+index 87d2bf52..4fc2c1de 100644
+--- a/mlx/backend/metal/matmul.cpp
++++ b/mlx/backend/metal/matmul.cpp
+@@ -1897,7 +1897,7 @@ void gather_mm_rhs(
+   int K = a.shape(-1);
+   int M = a.size() / K;
+   int N = b.shape(-1);
+-  int lda = a.strides()[a.ndim() - 2]; // should be K
++  int lda = K;
+ 
+   // Define the dispatch blocks
+   int bm = 16, bn = 64, bk = 16;
+@@ -2030,7 +2030,7 @@ void gather_mm_rhs_nax(
+   int K = a.shape(-1);
+   int M = a.size() / K;
+   int N = b.shape(-1);
+-  int lda = a.strides()[a.ndim() - 2]; // should be K
++  int lda = K;
+   int E = b.shape(0);
+ 
+   // Define the dispatch blocks
diff --git a/extension/llm/export/load.py b/extension/llm/export/load.py
index 77c8f8ee74b..8175ee43e47 100644
--- a/extension/llm/export/load.py
+++ b/extension/llm/export/load.py
@@ -164,19 +164,33 @@ def _iter_safetensors(path: str) -> Iterable[Pair]:
             unflatten_tensor_state_dict,
         )
 
+        # Use exact data keys and per-tensor metadata to avoid O(N²) checkpoint loading.
         for name in tensor_names:
+            tensor_metadata = json.loads(metadata[name])
+            tensor_type = tensor_metadata.get("_type")
             parts = name.rsplit(".", 1)
             module_fqn = parts[0] if len(parts) > 1 else ""
             weight_name = parts[-1]
             prefix = (
                 f"{module_fqn}._{weight_name}_" if module_fqn else f"_{weight_name}_"
             )
+            if tensor_type in ("Tensor", "torch.Tensor"):
+                data_keys = [name]
+            else:
+                data_keys = [
+                    prefix + suffix
+                    for suffix in tensor_metadata.get("_tensor_data_names", [])
+                ]
             partial = {
-                key: handle.get_tensor(key)
-                for key, handle in key_to_handle.items()
-                if key == name or key.startswith(prefix)
+                key: key_to_handle[key].get_tensor(key)
+                for key in data_keys
+                if key in key_to_handle
             }
-            result, _ = unflatten_tensor_state_dict(partial, metadata)
+            single_metadata = {
+                "tensor_names": json.dumps([name]),
+                name: metadata[name],
+            }
+            result, _ = unflatten_tensor_state_dict(partial, single_metadata)
             for fqn, value in result.items():
                 yield fqn, value