diff --git a/backends/mlx/builder/op_helpers.py b/backends/mlx/builder/op_helpers.py
index debe266d8ad..a147bbf94ae 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 002cda892f3..13790870b13 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,326 @@ 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
+
+    # 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
+
+
+@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 3c3c2c323a8..a285d321277 100644
--- a/backends/mlx/runtime/MLXInterpreter.h
+++ b/backends/mlx/runtime/MLXInterpreter.h
@@ -872,7 +872,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));
 }
 
@@ -895,6 +898,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,
@@ -906,7 +912,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 281199a8002..3420a6e28dc 100644
--- a/backends/mlx/serialization/schema.fbs
+++ b/backends/mlx/serialization/schema.fbs
@@ -949,6 +949,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 {
@@ -964,6 +969,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 a6cfc98d2f3..1aaa31ce906 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)
 
 
@@ -7927,3 +7988,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")