Qwen 3.5 MoE: bypass on-device sampler conditional in non-CUDA export

Qwen35MoE.forward currently routes through an Optional[Tensor] temperature
parameter and an if/else that picks between the on-device fused Gumbel-max
sampler (CUDA) and raw logits (non-CUDA). The sampling branch is dead code
for MLX and Metal exports, since those backends sample on the host.

Even though torch.export statically eliminates the branch when temperature
defaults to None, the parameter, default value, and unused else-branch leak
into the exported program: extra placeholder nodes, different graph hashes,
and shifted kernel selection in the lowered MLX/Metal graph. On the tiny
test model this slows MLX prefill ~9-37% and decode ~5-19%, and shows up as
~10-25% noise on Metal.

Bind model.forward to a minimal (tokens, input_pos) -> logits variant inside
_export_mlx and _export_metal before torch.export, so the captured program
matches what the backend kernels are tuned for. Eager-mode callers and the
CUDA export path are unaffected.

Author: Gasoonjia

examples/models/qwen3_5_moe/export.py

@@ -554,6 +554,44 @@ def export_and_lower(model, config, args):
         _export_cuda(model, config, args)
 
 
+def _strip_sampler_from_forward(model):
+    """Replace ``Qwen35MoE.forward`` with a minimal variant that only returns
+    raw logits, for backends that do not use the on-device fused sampler.
+
+    The default ``Qwen35MoE.forward`` accepts an optional sampling temperature
+    and dispatches between two paths: when a temperature is provided it slices
+    to the last position, casts to float32, and runs an on-device Gumbel-max
+    sampler that returns a sampled token id; when no temperature is provided
+    it returns the full ``[B, T, V]`` logits tensor. The on-device sampler is
+    only used by the CUDA path (the runner passes a temperature input);
+    non-CUDA backends sample on the host side after reading logits, so the
+    sampling branch is dead code at trace time.
+
+    The presence of the ``Optional[Tensor]`` parameter, the matching default
+    value in the exported signature, and the unused else-branch all leak into
+    the program ``torch.export`` produces -- adding placeholder nodes,
+    perturbing graph hashes used by downstream pattern matching, and shifting
+    kernel selection in the lowered graph. Empirically this slows MLX prefill
+    and decode by 10-30% even though the conditional is statically dead.
+
+    This helper rebinds ``model.forward`` to a clean two-argument variant
+    (tokens, input_pos) -> [B, T, V] logits before tracing, so the program
+    captured by ``torch.export`` for non-CUDA backends matches the shape the
+    backend kernels are tuned for. The original ``forward`` is unaffected for
+    other callers (e.g. eager inference, the CUDA export path).
+    """
+    import types
+
+    def _clean_forward(self, tokens, input_pos):
+        x = self.embed_tokens(tokens)
+        for layer in self.layers:
+            x = layer(x, input_pos)
+        x = self.norm(x)
+        return self.lm_head(x)
+
+    model.forward = types.MethodType(_clean_forward, model)
+
+
 def _export_mlx(model, config, args):
     """Export model to .pte via torch.export + MLX backend."""
     import gc
@@ -568,6 +606,10 @@ def _export_mlx(model, config, args):
     from executorch.exir.passes import MemoryPlanningPass
     from torch.export import Dim, export
 
+    # Use the minimal forward variant for non-CUDA export. See
+    # _strip_sampler_from_forward docstring for why this matters.
+    _strip_sampler_from_forward(model)
+
     example_tokens = torch.tensor([[0, 1]], dtype=torch.long)
     example_input_pos = torch.tensor([0, 1], dtype=torch.long)
     seq_dim = Dim("seq_len", min=1, max=config.max_seq_len - 1)
@@ -651,6 +693,10 @@ def _export_metal(model, config, args):
     inductor_config.coordinate_descent_tuning = False
     inductor_config.aot_inductor.compile_wrapper_opt_level = "O0"
 
+    # Use the minimal forward variant for non-CUDA export. See
+    # _strip_sampler_from_forward docstring for why this matters.
+    _strip_sampler_from_forward(model)
+
     # --- Decode method (T=1, static shape) ---
     print("Exporting decode method...")
     decode_tokens = torch.tensor([[0]], dtype=torch.long)