[MLX] Make SamplingHead directly exportable; drop sampler wrapper; wire runtime top-k (#20612)

Summary

SamplingHead.forward now takes the sampling params as trailing
positional args (temperature, top_k, top_p, seed), so torch.export
drives it directly and the per-model _MLXSampleWrapper is removed. As
part of going positional, top_k becomes a runtime input threaded through
export, run.py, and the C++ engine; qwen now supports per-request top-k,
and the op-level "not implemented" rejections are dropped.

  Changes

- backends/mlx/llm/sampling.py — SamplingHead.forward(self, *args)
unpacks *model_args, temperature, top_k, top_p, seed; samples (B) from
(B, vocab). No wrapper needed (the params are positional, so it's
directly exportable).
- examples/models/qwen3_5_moe/export.py — delete _MLXSampleWrapper;
model = SamplingHead(model); top_k added to example_args; dynamic_shapes
nested to mirror the single *args parameter.
- examples/models/qwen3_5_moe/run.py — _sampling_scalars feeds top_k
(<=0 → keep-all maxint); --top-k flag; the non-sample guard rejects
top_k/top_p/seed.
- backends/mlx/test/test_ops.py, test_sample.py — fixtures call
SamplingHead positionally, baking _KEEP_ALL_TOP_K/_TOP_P_OFF constants
for params they don't expose as runtime inputs.
- examples/models/qwen3_5_moe/main.cpp — --top_k flag →
SamplingConfig.top_k.
- examples/models/qwen3_5_moe/qwen35_moe_engine.cpp — top_k scalar wired
through prefill/decode (fed between temp and top_p, matching the new
forward order); 0→INT64_MAX keep-all mapping; the two "top_k is not
implemented" rejections removed.

  Testing

- Re-exported qwen, greedy output matches a pre-refactor baseline,
seeded runs are reproducible, and --top-k restricts sampling.
- MLX sample op tests pass; test_sample.py passes (incl. top-k
end-to-end).
- C++ MLX runner builds and runs the sample path; seeded reproducible
and --top_k works on-device.

Author: kiymetakdemir

backends/mlx/llm/sampling.py

@@ -15,30 +15,30 @@ class SamplingHead(nn.Module):
     Wraps a model that returns last-token logits ``(B, vocab)`` and samples a
     token id ``(B)`` on-device.
 
-        forward(*model_args, temperature, top_k=None, top_p=1.0, seed=None,
-                **model_kwargs) -> token_id
+        forward(*model_args, temperature, top_k, top_p, seed) -> token_id
+
+    The sampling params are trailing positional args so the head is directly
+    exportable (``torch.export`` drives positional inputs) without a per-model
+    wrapper.
 
       temperature: scalar float tensor, e.g. torch.tensor(0.8). Must be >= 0;
                    temperature=0 is greedy (returns argmax, no division).
-      top_k:       scalar int tensor or int; keeps only the k most likely tokens.
-                   None uses the max int default, which is clipped to the vocab
-                   size and keeps every token.
+      top_k:       scalar int tensor; keeps only the k most likely tokens. Use
+                   the max int (clipped to the vocab size) to keep every token.
       top_p:       scalar float tensor in (0, 1] for nucleus sampling. top_p=1.0
-                   (the default) keeps every token, i.e. no filtering. Pass it
-                   as a runtime input to tune per request.
+                   keeps every token, i.e. no filtering.
       seed:        scalar int tensor (seeded) or None (unseeded export)
     """
 
     def __init__(self, model: nn.Module):
         super().__init__()
         self.model = model
 
-    def forward(self, *args, temperature, top_k=None, top_p=1.0, seed=None, **kwargs):
-        logits = self.model(*args, **kwargs)  # [B, vocab]
+    def forward(self, *args):
+        *model_args, temperature, top_k, top_p, seed = args
+        logits = self.model(*model_args)  # [B, vocab]
+        if not isinstance(top_k, torch.Tensor):
+            top_k = torch.tensor(int(top_k), dtype=torch.int64)
         if not isinstance(top_p, torch.Tensor):
             top_p = torch.tensor(float(top_p))
-        if top_k is None:
-            top_k = torch.tensor(torch.iinfo(torch.int64).max, dtype=torch.int64)
-        elif not isinstance(top_k, torch.Tensor):
-            top_k = torch.tensor(int(top_k), dtype=torch.int64)
         return torch.ops.mlx.sample(logits, temperature, top_k, top_p, seed)

backends/mlx/test/test_ops.py

@@ -7700,6 +7700,12 @@ def forward(self, logits: torch.Tensor) -> torch.Tensor:
         return logits
 
 
+# Baked constants for sampling params a fixture does not expose as a runtime
+# input: keep-all top_k and top_p=off.
+_KEEP_ALL_TOP_K = torch.tensor(torch.iinfo(torch.int64).max, dtype=torch.int64)
+_TOP_P_OFF = torch.tensor(1.0)
+
+
 class SeededSampleModel(nn.Module):
     """SamplingHead with temperature AND seed as runtime forward inputs."""
 
@@ -7708,7 +7714,7 @@ def __init__(self):
         self.head = SamplingHead(_LogitsPassthrough())
 
     def forward(self, logits, temperature, seed):
-        return self.head(logits, temperature=temperature, seed=seed)
+        return self.head(logits, temperature, _KEEP_ALL_TOP_K, _TOP_P_OFF, seed)
 
 
 class UnseededSampleModel(nn.Module):
@@ -7719,7 +7725,7 @@ def __init__(self):
         self.head = SamplingHead(_LogitsPassthrough())
 
     def forward(self, logits, temperature):
-        return self.head(logits, temperature=temperature)
+        return self.head(logits, temperature, _KEEP_ALL_TOP_K, _TOP_P_OFF, None)
 
 
 class TopPSampleModel(nn.Module):
@@ -7730,7 +7736,7 @@ def __init__(self):
         self.head = SamplingHead(_LogitsPassthrough())
 
     def forward(self, logits, temperature, seed, top_p):
-        return self.head(logits, temperature=temperature, seed=seed, top_p=top_p)
+        return self.head(logits, temperature, _KEEP_ALL_TOP_K, top_p, seed)
 
 
 class TopKSampleModel(nn.Module):
@@ -7741,7 +7747,7 @@ def __init__(self):
         self.head = SamplingHead(_LogitsPassthrough())
 
     def forward(self, logits, temperature, seed, top_k):
-        return self.head(logits, temperature=temperature, seed=seed, top_k=top_k)
+        return self.head(logits, temperature, top_k, _TOP_P_OFF, seed)
 
 
 @register_test

backends/mlx/test/test_sample.py

@@ -41,6 +41,12 @@ def forward(self, logits: torch.Tensor) -> torch.Tensor:
         return logits
 
 
+# Baked constants for sampling params a fixture does not expose as a runtime
+# input: keep-all top_k and top_p=off.
+_KEEP_ALL_TOP_K = torch.tensor(torch.iinfo(torch.int64).max, dtype=torch.int64)
+_TOP_P_OFF = torch.tensor(1.0)
+
+
 class SeededSampleModel(nn.Module):
     """SamplingHead with temperature AND seed as runtime forward inputs."""
 
@@ -49,7 +55,7 @@ def __init__(self):
         self.head = SamplingHead(_LogitsPassthrough())
 
     def forward(self, logits, temperature, seed):
-        return self.head(logits, temperature=temperature, seed=seed)
+        return self.head(logits, temperature, _KEEP_ALL_TOP_K, _TOP_P_OFF, seed)
 
 
 class TopPSampleModel(nn.Module):
@@ -60,7 +66,7 @@ def __init__(self):
         self.head = SamplingHead(_LogitsPassthrough())
 
     def forward(self, logits, temperature, seed, top_p):
-        return self.head(logits, temperature=temperature, seed=seed, top_p=top_p)
+        return self.head(logits, temperature, _KEEP_ALL_TOP_K, top_p, seed)
 
 
 class TopKSampleModel(nn.Module):
@@ -71,7 +77,7 @@ def __init__(self):
         self.head = SamplingHead(_LogitsPassthrough())
 
     def forward(self, logits, temperature, seed, top_k):
-        return self.head(logits, temperature=temperature, seed=seed, top_k=top_k)
+        return self.head(logits, temperature, top_k, _TOP_P_OFF, seed)
 
 
 def _ref_gumbel_max(logits: torch.Tensor, temperature: float, seed: int):

examples/models/qwen3_5_moe/export.py

@@ -733,25 +733,6 @@ def _clean_forward(self, tokens, input_pos):
     model.forward = types.MethodType(_clean_forward, model)
 
 
-class _MLXSampleWrapper(nn.Module):
-    """Wrap the logits-producing model so ``forward`` returns a sampled token.
-
-    Temperature, top_p, and seed are runtime scalar inputs so the same .pte
-    serves any sampling request; the runner increments the seed per token.
-    """
-
-    def __init__(self, model: nn.Module):
-        super().__init__()
-        from executorch.backends.mlx.llm.sampling import SamplingHead
-
-        self.head = SamplingHead(model)
-
-    def forward(self, tokens, input_pos, temperature, top_p, seed):
-        return self.head(
-            tokens, input_pos, temperature=temperature, top_p=top_p, seed=seed
-        )
-
-
 def _export_mlx(model, config, args):
     """Export model to .pte via torch.export + MLX backend."""
     import gc
@@ -775,17 +756,22 @@ def _export_mlx(model, config, args):
     seq_dim = Dim("seq_len", min=1, max=config.max_seq_len - 1)
 
     if sample:
-        # forward(tokens, input_pos, temperature, top_p, seed) -> token id.
+        # forward(tokens, input_pos, temperature, top_k, top_p, seed) -> token id.
         # Scalars are static (None in dynamic_shapes); only the seq dim is dynamic.
-        model = _MLXSampleWrapper(model)
+        from executorch.backends.mlx.llm.sampling import SamplingHead
+
+        model = SamplingHead(model)
         example_args = (
             example_tokens,
             example_input_pos,
             torch.tensor(1.0, dtype=torch.float32),
+            torch.tensor(torch.iinfo(torch.int64).max, dtype=torch.int64),
             torch.tensor(1.0, dtype=torch.float32),
             torch.tensor(0, dtype=torch.int64),
         )
-        dynamic_shapes = ({1: seq_dim}, {0: seq_dim}, None, None, None)
+        # SamplingHead.forward takes ``*args``; dynamic_shapes mirrors that single
+        # variadic parameter as one nested tuple over the positional inputs.
+        dynamic_shapes = (({1: seq_dim}, {0: seq_dim}, None, None, None, None),)
     else:
         example_args = (example_tokens, example_input_pos)
         dynamic_shapes = ({1: seq_dim}, {0: seq_dim})

examples/models/qwen3_5_moe/main.cpp

@@ -38,6 +38,11 @@ DEFINE_string(
     "",
     "Path to file containing prompt text (overrides --prompt).");
 DEFINE_double(temperature, 0.8, "Sampling temperature (0 = greedy).");
+DEFINE_int32(
+    top_k,
+    0,
+    "Top-k sampling; keep the k most likely tokens. 0 (default) = off (keep "
+    "all). Requires a model exported with --sample (MLX on-device sampling).");
 DEFINE_double(
     top_p,
     1.0,
@@ -161,6 +166,7 @@ int main(int argc, char** argv) {
   // printed only on the first iteration (coherence check).
   llm::SamplingConfig sampling;
   sampling.temperature = static_cast<float>(FLAGS_temperature);
+  sampling.top_k = FLAGS_top_k;
   sampling.top_p = static_cast<float>(FLAGS_top_p);
   // Only a --sample model uses the seed; randomize an unset seed for those and
   // leave non-sample models at 0 so they don't trip the top_p/seed guard.

examples/models/qwen3_5_moe/qwen35_moe_engine.cpp

@@ -229,6 +229,8 @@ class Qwen35MoESession : public LLMSession {
     // across calls; only the backing scalar is rewritten per token.
     temp_tensor_mlx_ =
         from_blob(&temp_val_mlx_, {}, executorch::aten::ScalarType::Float);
+    top_k_tensor_ =
+        from_blob(&top_k_val_, {}, executorch::aten::ScalarType::Long);
     top_p_tensor_ =
         from_blob(&top_p_val_, {}, executorch::aten::ScalarType::Float);
     seed_tensor_ =
@@ -256,15 +258,12 @@ class Qwen35MoESession : public LLMSession {
     }
     float first_token_temp = temperature_;
     if (initial_sampling != nullptr) {
-      if (initial_sampling->top_k != 0) {
-        ET_LOG(Error, "prefill_tokens: top_k is not implemented");
-        return Error::NotSupported;
-      }
       if (!use_sampling_ &&
-          (initial_sampling->top_p != 1.0f || initial_sampling->seed != 0)) {
+          (initial_sampling->top_k != 0 || initial_sampling->top_p != 1.0f ||
+           initial_sampling->seed != 0)) {
         ET_LOG(
             Error,
-            "prefill_tokens: top_p/seed require a sampling model "
+            "prefill_tokens: top_k/top_p/seed require a sampling model "
             "(export with --sample); only temperature is supported otherwise");
         return Error::NotSupported;
       }
@@ -274,6 +273,7 @@ class Qwen35MoESession : public LLMSession {
           ET_LOG(Error, "prefill_tokens: top_p must be in (0, 1]");
           return Error::InvalidArgument;
         }
+        top_k_ = initial_sampling->top_k;
         top_p_ = initial_sampling->top_p;
         seed_ = initial_sampling->seed; // base seed for the kept (token 0) draw
       }
@@ -347,8 +347,9 @@ class Qwen35MoESession : public LLMSession {
 #endif
 #ifdef EXECUTORCH_BUILD_MLX
       if (use_sampling_) {
-        set_sampling_inputs(first_token_temp, top_p_, seed_);
+        set_sampling_inputs(first_token_temp, top_k_, top_p_, seed_);
         inputs.push_back(EValue(temp_tensor_mlx_));
+        inputs.push_back(EValue(top_k_tensor_));
         inputs.push_back(EValue(top_p_tensor_));
         inputs.push_back(EValue(seed_tensor_));
       }
@@ -370,14 +371,11 @@ class Qwen35MoESession : public LLMSession {
   }
 
   Result<DecodeResult> decode_one(const SamplingConfig& sampling) override {
-    if (sampling.top_k != 0) {
-      ET_LOG(Error, "Qwen35MoESession: top_k is not implemented");
-      return Error::NotSupported;
-    }
-    if (!use_sampling_ && (sampling.top_p != 1.0f || sampling.seed != 0)) {
+    if (!use_sampling_ &&
+        (sampling.top_k != 0 || sampling.top_p != 1.0f || sampling.seed != 0)) {
       ET_LOG(
           Error,
-          "Qwen35MoESession: top_p/seed require a sampling model "
+          "Qwen35MoESession: top_k/top_p/seed require a sampling model "
           "(export with --sample); only temperature is supported otherwise");
       return Error::NotSupported;
     }
@@ -395,6 +393,7 @@ class Qwen35MoESession : public LLMSession {
         ET_LOG(Error, "decode_one: top_p must be in (0, 1]");
         return Error::InvalidArgument;
       }
+      top_k_ = sampling.top_k;
       top_p_ = sampling.top_p;
     }
 
@@ -444,8 +443,9 @@ class Qwen35MoESession : public LLMSession {
 #endif
 #ifdef EXECUTORCH_BUILD_MLX
     if (use_sampling_) {
-      set_sampling_inputs(temperature_, top_p_, seed_);
+      set_sampling_inputs(temperature_, top_k_, top_p_, seed_);
       inputs.push_back(EValue(temp_tensor_mlx_));
+      inputs.push_back(EValue(top_k_tensor_));
       inputs.push_back(EValue(top_p_tensor_));
       inputs.push_back(EValue(seed_tensor_));
     }
@@ -499,8 +499,10 @@ class Qwen35MoESession : public LLMSession {
 #ifdef EXECUTORCH_BUILD_MLX
   // Rewrite the backing scalars for the reused 0-dim sampling input tensors.
   // The -1 temperature sentinel maps to 0 (greedy).
-  void set_sampling_inputs(float temp, float top_p, uint64_t seed) {
+  void
+  set_sampling_inputs(float temp, int64_t top_k, float top_p, uint64_t seed) {
     temp_val_mlx_ = (temp < 0.0f) ? 0.0f : temp;
+    top_k_val_ = (top_k <= 0) ? INT64_MAX : top_k; // 0/neg = keep all
     top_p_val_ = top_p;
     seed_val_ = static_cast<int64_t>(seed);
   }
@@ -553,6 +555,7 @@ class Qwen35MoESession : public LLMSession {
   // the seed increments per generated token for decorrelated, reproducible
   // draws.
   bool use_sampling_ = false;
+  int64_t top_k_ = 0; // 0 = off (keep all); mapped to INT64_MAX on-device
   float top_p_ = 1.0f;
   uint64_t seed_ = 0;
 
@@ -570,9 +573,11 @@ class Qwen35MoESession : public LLMSession {
 #endif
 #ifdef EXECUTORCH_BUILD_MLX
   float temp_val_mlx_ = 0.0f;
+  int64_t top_k_val_ = INT64_MAX;
   float top_p_val_ = 1.0f;
   int64_t seed_val_ = 0;
   TensorPtr temp_tensor_mlx_;
+  TensorPtr top_k_tensor_;
   TensorPtr top_p_tensor_;
   TensorPtr seed_tensor_;
 #endif