From b1f941fb589d83209aa0dc493e2da01739e9d1ce Mon Sep 17 00:00:00 2001 From: "Axel.Cffrd.Dnty" Date: Sun, 5 Jul 2026 17:04:43 +0200 Subject: [PATCH] backends/mlx: runtime MoE expert-sort for decode (issue #20554) Replace the compile-time sort_experts: bool flag in SwitchMLP with a runtime decision made inside two new custom ops (moe_gather_inputs, moe_scatter_outputs). A single exported .pte now handles both prefill (sorted, coalesced gather_mm) and decode (unsorted, no argsort overhead) without requiring separate exports. Schema uses wire-compatible sorted_indices_flag: IntOrVid (appended to GatherMmNode/GatherQmmNode) with static sorted_indices: bool fallback. Fixes TakeNode.index typing in MoE handlers (IntOrVidOrTid.from_tid, not raw Tid) so export serialization succeeds. Test plan: - Windows: python backends/mlx/test/validate_moe_20554.py (all passed) - CI: test-mlx job on macos-14-xlarge (run_all_tests) Fixes #20554 PR authored with Claude. Co-authored-by: Cursor --- backends/mlx/builder/op_helpers.py | 21 ++ backends/mlx/custom_ops.py | 87 +++++- backends/mlx/llm/switch.py | 41 ++- backends/mlx/ops.py | 156 +++++++++- backends/mlx/runtime/MLXInterpreter.h | 16 +- backends/mlx/serialization/generate.py | 15 +- backends/mlx/serialization/schema.fbs | 7 + backends/mlx/test/test_ops.py | 251 +++++++++++++++- backends/mlx/test/validate_moe_20554.py | 282 ++++++++++++++++++ examples/models/qwen3_5_moe/export.py | 2 +- .../qwen3_5_moe/mlx_source_transformations.py | 16 +- 11 files changed, 846 insertions(+), 48 deletions(-) create mode 100644 backends/mlx/test/validate_moe_20554.py diff --git a/backends/mlx/builder/op_helpers.py b/backends/mlx/builder/op_helpers.py index 8c278cdd0ff..bdc4eb05570 100644 --- a/backends/mlx/builder/op_helpers.py +++ b/backends/mlx/builder/op_helpers.py @@ -353,6 +353,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 5605b59c543..8a033841cf8 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: @@ -381,7 +387,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] @@ -454,3 +460,74 @@ 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) + inv_order = torch.empty(0, dtype=torch.int32) # sentinel: never read + 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 the unread sentinel) (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 28d408cbd71..2bf5d8b50ac 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 @@ -131,12 +132,17 @@ 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.") @@ -193,6 +199,7 @@ def __init__( activation=None, bias: bool = False, fuse_gate_up: bool = False, + sort_cutoff: int = 1, ): super().__init__() if activation is None: @@ -201,6 +208,9 @@ 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. + self.sort_cutoff = sort_cutoff if fuse_gate_up: self.gate_up_proj = SwitchLinear( @@ -223,7 +233,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. @@ -232,25 +241,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) @@ -262,11 +263,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 39167445c37..d9f4ea6cffb 100644 --- a/backends/mlx/ops.py +++ b/backends/mlx/ops.py @@ -20,10 +20,12 @@ import torch from executorch.backends.mlx.builder.op_helpers import ( + emit_floordiv, emit_if_else, emit_lifted_constant, emit_quantized_biases, emit_shape, + emit_sub_int, parse_dequant_node, to_mlx_qparams, torch_dtype_to_scalar_type, @@ -1789,6 +1791,22 @@ 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.""" + from executorch.backends.mlx.serialization.mlx_graph_schema import ( + IntOrVid, + ItemIntNode, + ) + + if isinstance(sorted_indices, Slot): + _, item_slot = P.make_tmp_value_slot() + P.emit(ItemIntNode(x=P.slot_to_tid(sorted_indices), out=P.slot_to_vid(item_slot))) + return P.to_int_or_vid(item_slot) + return IntOrVid.from_literal(1 if sorted_indices else 0) + + @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.""" @@ -1801,7 +1819,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( @@ -1811,7 +1830,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 @@ -1839,7 +1859,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) @@ -1882,12 +1903,139 @@ 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. Both branches pre-allocate + and write every output slot, since emit_if_else requires the two branches + to write the same fixed output slots. + """ + from executorch.backends.mlx.serialization.mlx_graph_schema import ( + ArgsortNode, TakeNode, FloorDivideNode, ExpandDimsNode, + RepeatNode, IdCopyNode, IntOrVid, IntOrVidOrTid, + ) + + 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 + + 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) + + def emit_sorted(): + _, order_slot = P.make_tmp_slot() + P.emit(ArgsortNode(x=P.slot_to_tid(expert_indices), out=P.slot_to_tid(order_slot), axis=0)) + + _, inv_order_slot = P.make_tmp_slot() + P.emit(ArgsortNode(x=P.slot_to_tid(order_slot), out=P.slot_to_tid(inv_order_slot), axis=0)) + + _, idx_slot = P.make_tmp_slot() + P.emit(TakeNode(x=P.slot_to_tid(expert_indices), out=P.slot_to_tid(idx_slot), + index=IntOrVidOrTid.from_tid(P.slot_to_tid(order_slot)), axis=0)) + + top_k_const = emit_lifted_constant(P, top_k, torch.int32) + _, row_idx_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_slot))) + + _, 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)) + + _, x_input_slot = P.make_tmp_slot() + P.emit(ExpandDimsNode(x=P.slot_to_tid(x_gathered_slot), out=P.slot_to_tid(x_input_slot), axis=-2)) + + one_const = emit_lifted_constant(P, 1, torch.int32) + + P.emit(IdCopyNode(x=P.slot_to_tid(x_input_slot), out=P.slot_to_tid(out_slots[0]))) + P.emit(IdCopyNode(x=P.slot_to_tid(idx_slot), out=P.slot_to_tid(out_slots[1]))) + P.emit(IdCopyNode(x=P.slot_to_tid(one_const), out=P.slot_to_tid(out_slots[2]))) + P.emit(IdCopyNode(x=P.slot_to_tid(inv_order_slot), out=P.slot_to_tid(out_slots[3]))) + + def emit_unsorted(): + top_k_iov = IntOrVid.from_literal(top_k) + _, 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=top_k_iov, axis=0)) + + _, x_input_slot = P.make_tmp_slot() + P.emit(ExpandDimsNode(x=P.slot_to_tid(x_rep_slot), out=P.slot_to_tid(x_input_slot), axis=-2)) + + zero_const = emit_lifted_constant(P, 0, torch.int32) + + # 0-element sentinel: never read on the unsorted path (see + # moe_scatter_outputs), so a plain static constant suffices. + empty_inv_order = P.make_or_get_constant( + "_moe_inv_order_sentinel", torch.empty(0, dtype=torch.int32) + ) + + P.emit(IdCopyNode(x=P.slot_to_tid(x_input_slot), out=P.slot_to_tid(out_slots[0]))) + P.emit(IdCopyNode(x=P.slot_to_tid(expert_indices), out=P.slot_to_tid(out_slots[1]))) + P.emit(IdCopyNode(x=P.slot_to_tid(zero_const), out=P.slot_to_tid(out_slots[2]))) + P.emit(IdCopyNode(x=P.slot_to_tid(empty_inv_order), out=P.slot_to_tid(out_slots[3]))) + + # 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)) + 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]. Reads sort_experts (0-d + tensor) to a Vid via ItemIntNode, then emit_if_else on it: the sorted + branch gathers with inv_order before reshaping, the unsorted branch + reshapes directly. + """ + from executorch.backends.mlx.serialization.mlx_graph_schema import ( + ItemIntNode, SqueezeNode, TakeNode, ReshapeNode, IntOrVid, IntOrVidOrTid, + ) + + args = P.args(n) + down, sort_experts, inv_order = args[0], args[1], args[2] + top_k = args[3] # static int + + 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_val_slot = P.make_tmp_value_slot() + P.emit(ItemIntNode(x=P.slot_to_tid(sort_experts), out=P.slot_to_vid(cond_val_slot))) + + n_top_k_iov = emit_shape(P, n.args[0], down_sq_slot, end_dim=1)[0] + 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, P.to_int_or_vid(cond_val_slot), 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..cb1e07bd978 100644 --- a/backends/mlx/runtime/MLXInterpreter.h +++ b/backends/mlx/runtime/MLXInterpreter.h @@ -872,7 +872,16 @@ 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 +904,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 +918,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" if fbs_type == "Vid": return "std::optional" + 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 66db9ead8b1..6d7383251a9 100644 --- a/backends/mlx/test/test_ops.py +++ b/backends/mlx/test/test_ops.py @@ -6617,12 +6617,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 @@ -6631,13 +6647,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 @@ -6650,16 +6673,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 @@ -6669,10 +6696,18 @@ 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, ...]: @@ -6694,10 +6729,16 @@ def __init__( in_features: int, out_features: int, group_size: int = 32, + 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_ @@ -6738,12 +6779,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 @@ -6757,6 +6803,7 @@ def __init__( batch_size: int = 2, group_size: int = 32, dtype: torch.dtype = torch.float32, + sorted_indices: bool = False, ): self.num_experts = num_experts self.in_features = in_features @@ -6764,6 +6811,7 @@ def __init__( self.batch_size = batch_size self.group_size = group_size self.dtype = dtype + self.sorted_indices = sorted_indices parts = [ "gather_qmm", @@ -6774,6 +6822,8 @@ def __init__( ] if dtype != torch.float32: parts.append(str(dtype).split(".")[-1]) + if sorted_indices: + parts.append("sorted") self.name = "_".join(parts) @classmethod @@ -6783,6 +6833,9 @@ def get_test_configs(cls) -> List["GatherQmmTest"]: 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 get_edge_compile_config(self): @@ -6796,6 +6849,7 @@ def create_model(self) -> nn.Module: self.in_features, self.out_features, self.group_size, + sorted_indices=self.sorted_indices, ) return model.to(self.dtype) @@ -7927,3 +7981,190 @@ 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 + ) + return x_input + + +@register_test +class MoeGatherInputsTest(OpTestCase): + """Test case for mlx::moe_gather_inputs. + + Covers both N=1 (decode, unsorted) and N>sort_cutoff (prefill, sorted) + paths via separate configs, matching get_test_configs() convention. + """ + + 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, + 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 + + 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) + + @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"), + ] + + def get_expected_node_counts(self) -> Optional[Dict[str, int]]: + if self.batch_size > self.sort_cutoff: + # sorted path: 2x ArgsortNode, no RepeatNode, no IfNode + return { + "ArgsortNode": 2, + "TakeNode": 2, + "FloorDivideNode": 1, + "ExpandDimsNode": 1, + "IdCopyNode": 5, + "RepeatNode": 0, + "IfNode": 0, + } + else: + # unsorted path: RepeatNode, no ArgsortNode, no IfNode + return { + "ArgsortNode": 0, + "RepeatNode": 1, + "ExpandDimsNode": 1, + "IdCopyNode": 5, + "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) + + +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] + NK = x_input.shape[0] + down = x_input[..., : self.hidden_out].contiguous() # cheap slice as proxy + if down.shape[-1] < self.hidden_out: + down = down.repeat(1, 1, self.hidden_out // down.shape[-1] + 1)[ + ..., : self.hidden_out + ] + down = down.view(NK, 1, self.hidden_out) + 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. + """ + + 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, + 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 + + 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) + + @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"), + ] + + 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) + diff --git a/backends/mlx/test/validate_moe_20554.py b/backends/mlx/test/validate_moe_20554.py new file mode 100644 index 00000000000..ab485089e33 --- /dev/null +++ b/backends/mlx/test/validate_moe_20554.py @@ -0,0 +1,282 @@ +#!/usr/bin/env python3 +"""Local validation for issue #20554 (no MLX runner / no full pip install).""" + +from __future__ import annotations + +import os +import sys +import types +from pathlib import Path + +REPO = Path(__file__).resolve().parents[3] +SRC = REPO / "src" + + +def _bootstrap_executorch() -> None: + """Wire executorch.* to repo paths (Windows git symlinks are often broken).""" + import importlib.util + + sys.path.insert(0, str(SRC)) + sys.path.insert(0, str(REPO)) + + import executorch + + def _load_pkg(name: str) -> None: + entry = SRC / "executorch" / name + if entry.is_file(): + target = entry.read_text(encoding="utf-8").strip().replace("/", os.sep) + real_root = (entry.parent / target).resolve() + else: + real_root = (REPO / name).resolve() + init_py = real_root / "__init__.py" + if not init_py.exists(): + return + full_name = f"executorch.{name}" + spec = importlib.util.spec_from_file_location( + full_name, + init_py, + submodule_search_locations=[str(real_root)], + ) + if spec is None or spec.loader is None: + return + mod = importlib.util.module_from_spec(spec) + sys.modules[full_name] = mod + setattr(executorch, name, mod) + spec.loader.exec_module(mod) + + def _load_subpkg(parent: str, name: str, root: Path) -> None: + init_py = root / "__init__.py" + if not init_py.exists(): + return + full_name = f"{parent}.{name}" + spec = importlib.util.spec_from_file_location( + full_name, + init_py, + submodule_search_locations=[str(root)], + ) + if spec is None or spec.loader is None: + return + mod = importlib.util.module_from_spec(spec) + sys.modules[full_name] = mod + parent_mod = sys.modules[parent] + setattr(parent_mod, name, mod) + spec.loader.exec_module(mod) + + # exir.tracer needs executorch.extension.pytree (pybindings optional). + ext = types.ModuleType("executorch.extension") + sys.modules["executorch.extension"] = ext + executorch.extension = ext + _load_subpkg("executorch.extension", "pytree", REPO / "extension" / "pytree") + + _load_pkg("exir") + _load_pkg("backends") + + +def _load_custom_ops(): + import importlib.util + + spec = importlib.util.spec_from_file_location( + "mlx_custom_ops", + REPO / "backends" / "mlx" / "custom_ops.py", + ) + assert spec and spec.loader + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def test_eager_moe_ops() -> None: + import torch + + _load_custom_ops() + + torch.manual_seed(0) + top_k = 2 + D, H = 32, 16 + + # --- decode path (N=1, sort_cutoff=1) --- + x = torch.randn(1, D) + expert_indices = torch.tensor([[1, 3]], dtype=torch.int64) + x_in, idx, sort_flag, inv = torch.ops.mlx.moe_gather_inputs( + x, expert_indices, top_k, 1 + ) + assert x_in.shape == (2, 1, D) + assert idx.shape == (2,) + assert sort_flag.item() == 0 + assert inv.numel() == 0 + + down = torch.randn(2, 1, H) + out = torch.ops.mlx.moe_scatter_outputs(down, sort_flag, inv, top_k) + assert out.shape == (1, top_k, H) + assert torch.allclose(out, down.squeeze(-2).reshape(1, top_k, H).clone()) + + # --- prefill path (N=4, sort_cutoff=1) --- + x = torch.randn(4, D) + expert_indices = torch.tensor([[2, 0], [1, 3], [0, 2], [3, 1]], dtype=torch.int64) + x_in, idx, sort_flag, inv = torch.ops.mlx.moe_gather_inputs( + x, expert_indices, top_k, 1 + ) + assert x_in.shape == (8, 1, D) + assert sort_flag.item() == 1 + assert inv.shape == (8,) + + # Reference sorted path (mirror eager in custom_ops) + flat = expert_indices.flatten() + order = flat.argsort().to(torch.int32) + inv_ref = order.argsort().to(torch.int32) + idx_ref = flat[order].to(torch.int32) + x_ref = x[(order // top_k).to(torch.int64)].unsqueeze(-2) + assert torch.equal(idx, idx_ref) + assert torch.equal(inv, inv_ref) + assert torch.allclose(x_in, x_ref) + + down = torch.randn(8, 1, H) + out = torch.ops.mlx.moe_scatter_outputs(down, sort_flag, inv, top_k) + down_sq = down.squeeze(-2) + ref = down_sq[inv_ref].reshape(4, top_k, H) + assert torch.allclose(out, ref) + + print("PASS: eager moe_gather_inputs / moe_scatter_outputs") + + +def test_opcheck() -> None: + import torch + from torch.library import opcheck + + _load_custom_ops() + + x = torch.randn(4, 32) + expert_indices = torch.randint(0, 8, (4, 2)) + opcheck(torch.ops.mlx.moe_gather_inputs, (x, expert_indices, 2, 1)) + down = torch.randn(8, 1, 16) + sort_experts = torch.tensor(1, dtype=torch.int32) + inv_order = torch.arange(8, dtype=torch.int32) + opcheck( + torch.ops.mlx.moe_scatter_outputs, + (down, sort_experts, inv_order, 2), + ) + print("PASS: torch.library.opcheck") + + +def test_export_traces_moe_ops() -> None: + import torch + import torch.nn as nn + from torch.export import export + + _load_custom_ops() + + class MoeModel(nn.Module): + def forward(self, x, expert_indices): + gathered = torch.ops.mlx.moe_gather_inputs(x, expert_indices, 2, 1) + return torch.ops.mlx.moe_scatter_outputs( + torch.randn(gathered[0].shape[0], 1, 16), + gathered[2], + gathered[3], + 2, + ) + + ep = export( + MoeModel(), + (torch.randn(4, 32), torch.randint(0, 4, (4, 2))), + ) + targets = { + str(n.target) + for n in ep.graph.nodes + if n.op == "call_function" + } + assert any("moe_gather_inputs" in t for t in targets), targets + assert any("moe_scatter_outputs" in t for t in targets), targets + print("PASS: torch.export traces moe ops as leaf nodes") + + +def _count_mlx_nodes(mlx_graph) -> dict[str, int]: + from collections import Counter + + return dict( + Counter( + type(instr.op).__name__ + for chain in mlx_graph.instruction_chains + for instr in chain.instructions + ) + ) + + +def test_export_lowering_node_counts() -> None: + import torch + import torch.nn as nn + from torch.export import export + + from executorch.backends.mlx import custom_ops # noqa: F401 + from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder + from executorch.exir import EdgeCompileConfig, to_edge + + class MoeGather(nn.Module): + def forward(self, x, expert_indices): + return torch.ops.mlx.moe_gather_inputs(x, expert_indices, 2, 1)[0] + + cfg = EdgeCompileConfig(_check_ir_validity=False) + + ep = export(MoeGather(), (torch.randn(4, 32), torch.randint(0, 4, (4, 2)))) + prefill = _count_mlx_nodes( + MLXProgramBuilder(to_edge(ep, compile_config=cfg).exported_program()).build() + ) + assert prefill.get("IfNode", 0) == 0, prefill + assert prefill.get("ArgsortNode", 0) == 2, prefill + assert prefill.get("RepeatNode", 0) == 0, prefill + print(f"PASS: MLX lowering (prefill): {prefill}") + + ep1 = export(MoeGather(), (torch.randn(1, 32), torch.randint(0, 4, (1, 2)))) + decode = _count_mlx_nodes( + MLXProgramBuilder(to_edge(ep1, compile_config=cfg).exported_program()).build() + ) + assert decode.get("IfNode", 0) == 0, decode + assert decode.get("ArgsortNode", 0) == 0, decode + assert decode.get("RepeatNode", 0) == 1, decode + print(f"PASS: MLX lowering (decode): {decode}") + + +def test_switch_mlp_forward() -> None: + import torch + import torch.nn as nn + from executorch.backends.mlx import custom_ops # noqa: F401 + from executorch.backends.mlx.llm.switch import SwitchMLP, pack_all_switch_linears + + mlp = SwitchMLP(32, 64, num_experts=4, sort_cutoff=1) + for mod in mlp.modules(): + if hasattr(mod, "experts"): + for e in mod.experts: + nn.init.uniform_(e.weight, -0.1, 0.1) + pack_all_switch_linears(mlp) + + x = torch.randn(4, 32) + weights = torch.softmax(torch.randn(4, 2), dim=-1) + indices = torch.randint(0, 4, (4, 2)) + + out_prefill = mlp(x, weights, indices, top_k=2) + assert out_prefill.shape == (4, 32) + + x1 = torch.randn(1, 32) + w1 = torch.softmax(torch.randn(1, 2), dim=-1) + i1 = torch.randint(0, 4, (1, 2)) + out_decode = mlp(x1, w1, i1, top_k=2) + assert out_decode.shape == (1, 32) + print("PASS: SwitchMLP forward (prefill + decode)") + + +def main() -> int: + test_eager_moe_ops() + test_opcheck() + test_export_traces_moe_ops() + _bootstrap_executorch() + test_switch_mlp_forward() + try: + test_export_lowering_node_counts() + except Exception as e: + print(f"FAIL: MLX lowering: {e}") + return 1 + print("\nAll validations passed.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/models/qwen3_5_moe/export.py b/examples/models/qwen3_5_moe/export.py index d2c3914c6a2..c317462ed6b 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: 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")