@@ -733,6 +733,25 @@ 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
@@ -749,16 +768,33 @@ def _export_mlx(model, config, args):
 
     _strip_sampler_from_forward(model)
 
+    sample = getattr(args, "sample", False)
+
     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)
-    dynamic_shapes = ({1: seq_dim}, {0: seq_dim})
+
+    if sample:
+        # forward(tokens, input_pos, temperature, top_p, seed) -> token id.
+        # Scalars are static (None in dynamic_shapes); only the seq dim is dynamic.
+        model = _MLXSampleWrapper(model)
+        example_args = (
+            example_tokens,
+            example_input_pos,
+            torch.tensor(1.0, dtype=torch.float32),
+            torch.tensor(1.0, dtype=torch.float32),
+            torch.tensor(0, dtype=torch.int64),
+        )
+        dynamic_shapes = ({1: seq_dim}, {0: seq_dim}, None, None, None)
+    else:
+        example_args = (example_tokens, example_input_pos)
+        dynamic_shapes = ({1: seq_dim}, {0: seq_dim})
 
     print("Exporting with torch.export...")
     with torch.no_grad():
         exported = export(
             model,
-            (example_tokens, example_input_pos),
+            example_args,
             dynamic_shapes=dynamic_shapes,
             strict=True,
         )
@@ -781,6 +817,7 @@ def _export_mlx(model, config, args):
         "use_kv_cache": True,
         "use_sdpa_with_kv_cache": False,
         "enable_dynamic_shape": True,
+        "use_sampling": sample,
     }
     et_prog = to_edge_transform_and_lower(
         exported,
@@ -1182,6 +1219,13 @@ def main():  # noqa: C901
         action="store_true",
         help="Disable split-K (flash-decoding) SDPA for decode; use tiled SDPA instead.",
     )
+    parser.add_argument(
+        "--sample",
+        action="store_true",
+        help="MLX only: sample the next token on-device (Gumbel-max with "
+        "temperature/top_p/seed runtime inputs) instead of returning logits "
+        "for host-side sampling.",
+    )
     parser.add_argument(
         "--moe-activation-dtype",
         choices=["bf16", "int8"],
@@ -1230,6 +1274,9 @@ def main():  # noqa: C901
         if args.turboquant:
             parser.error("--turboquant is not supported with --backend mlx")
 
+    if args.sample and args.backend != "mlx":
+        parser.error("--sample is only supported with --backend mlx")
+
     if args.backend == "metal":
         if args.turboquant:
             parser.error("--turboquant is not supported with --backend metal")
 
@@ -21,6 +21,7 @@
 #include <algorithm>
 #include <cinttypes>
 #include <fstream>
+#include <random>
 #include <string>
 #include <vector>
 
@@ -37,6 +38,17 @@ DEFINE_string(
     "",
     "Path to file containing prompt text (overrides --prompt).");
 DEFINE_double(temperature, 0.8, "Sampling temperature (0 = greedy).");
+DEFINE_double(
+    top_p,
+    1.0,
+    "Nucleus sampling top_p in (0, 1]; 1.0 = off. Requires a model exported "
+    "with --sample (MLX on-device sampling).");
+DEFINE_int64(
+    seed,
+    -1,
+    "Base RNG seed for on-device sampling; the runner increments it per token. "
+    "-1 (default) draws a random seed each run; set a value for reproducible "
+    "output. Requires a model exported with --sample.");
 DEFINE_int32(max_new_tokens, 128, "Maximum tokens to generate.");
 DEFINE_int32(
     warmup,
@@ -149,6 +161,20 @@ 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_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.
+  const auto& md = engine->metadata();
+  const auto us_it = md.find("use_sampling");
+  const bool model_samples = us_it != md.end() && us_it->second != 0;
+  uint64_t base_seed = FLAGS_seed < 0 ? 0 : static_cast<uint64_t>(FLAGS_seed);
+  if (model_samples && FLAGS_seed < 0) {
+    base_seed = static_cast<uint64_t>(std::random_device{}());
+  }
+  sampling.seed = base_seed;
+  if (model_samples) {
+    ET_LOG(Info, "Sampling base seed: %" PRIu64, base_seed);
+  }
   const int total_iters = FLAGS_warmup + std::max(1, FLAGS_num_iters);
   std::vector<double> prefill_tps_samples;
   std::vector<double> decode_tps_samples;
 
@@ -57,11 +57,18 @@ constexpr const char* kDecodeMethod = "decode";
 // `forward` method was compiled for. Read into the metadata map in create().
 constexpr const char* kMaxPrefillChunk = "get_max_prefill_chunk";
 
+// Constant method (MLX --sample exports) signalling that `forward` samples a
+// token id on-device and returns it directly. Read into the metadata map in
+// create().
+constexpr const char* kUseSampling = "use_sampling";
+
 Result<uint64_t> read_sampled_token(
     const executorch::aten::Tensor& output,
-    float temperature) {
+    float temperature,
+    bool use_sampling) {
 #ifdef EXECUTORCH_BUILD_CUDA
   (void)temperature;
+  (void)use_sampling;
   const void* ptr = output.const_data_ptr();
   cudaPointerAttributes attrs;
   const bool on_device = cudaPointerGetAttributes(&attrs, ptr) == cudaSuccess &&
@@ -78,6 +85,11 @@ Result<uint64_t> read_sampled_token(
   }
   return static_cast<uint64_t>(llrintf(val));
 #else
+  // On-device sampling (MLX --sample): `forward` already returns the int64
+  // token id, so read it directly instead of sampling host-side from logits.
+  if (use_sampling) {
+    return static_cast<uint64_t>(output.const_data_ptr<int64_t>()[0]);
+  }
   return static_cast<uint64_t>(
       logits_to_token(output, temperature < 0.0f ? 0.0f : temperature));
 #endif
@@ -208,6 +220,19 @@ class Qwen35MoESession : public LLMSession {
 #ifdef EXECUTORCH_BUILD_CUDA
     temp_tensor_ =
         from_blob(&temp_val_, {1}, executorch::aten::ScalarType::Float);
+#endif
+#ifdef EXECUTORCH_BUILD_MLX
+    if (auto it = metadata_.find(kUseSampling); it != metadata_.end()) {
+      use_sampling_ = it->second != 0;
+    }
+    // 0-dim scalar inputs appended to `forward` when sampling on-device. Reused
+    // across calls; only the backing scalar is rewritten per token.
+    temp_tensor_mlx_ =
+        from_blob(&temp_val_mlx_, {}, executorch::aten::ScalarType::Float);
+    top_p_tensor_ =
+        from_blob(&top_p_val_, {}, executorch::aten::ScalarType::Float);
+    seed_tensor_ =
+        from_blob(&seed_val_, {}, executorch::aten::ScalarType::Long);
 #endif
   }
 
@@ -231,15 +256,27 @@ class Qwen35MoESession : public LLMSession {
     }
     float first_token_temp = temperature_;
     if (initial_sampling != nullptr) {
-      if (initial_sampling->top_p != 1.0f || initial_sampling->top_k != 0 ||
-          initial_sampling->seed != 0) {
+      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)) {
         ET_LOG(
             Error,
-            "prefill_tokens: only temperature is supported; top_p/top_k/seed "
-            "are not yet implemented");
+            "prefill_tokens: top_p/seed require a sampling model "
+            "(export with --sample); only temperature is supported otherwise");
         return Error::NotSupported;
       }
       first_token_temp = initial_sampling->temperature;
+      if (use_sampling_) {
+        if (!valid_top_p(initial_sampling->top_p)) {
+          ET_LOG(Error, "prefill_tokens: top_p must be in (0, 1]");
+          return Error::InvalidArgument;
+        }
+        top_p_ = initial_sampling->top_p;
+        seed_ = initial_sampling->seed; // base seed for the kept (token 0) draw
+      }
     }
     if (!valid_temperature(first_token_temp)) {
       ET_LOG(Error, "prefill_tokens: temperature must be -1 or in [0, 1]");
@@ -307,6 +344,14 @@ class Qwen35MoESession : public LLMSession {
 #ifdef EXECUTORCH_BUILD_CUDA
       set_temp(first_token_temp);
       inputs.push_back(EValue(temp_tensor_));
+#endif
+#ifdef EXECUTORCH_BUILD_MLX
+      if (use_sampling_) {
+        set_sampling_inputs(first_token_temp, top_p_, seed_);
+        inputs.push_back(EValue(temp_tensor_mlx_));
+        inputs.push_back(EValue(top_p_tensor_));
+        inputs.push_back(EValue(seed_tensor_));
+      }
 #endif
       auto sampled =
           run_locked(method, inputs, first_token_temp, /*sync_after=*/true);
@@ -316,15 +361,24 @@ class Qwen35MoESession : public LLMSession {
     }
     pending_ = sampled_token;
     prev_decode_token_.reset();
+#ifdef EXECUTORCH_BUILD_MLX
+    if (use_sampling_) {
+      seed_ += 1; // first decode token draws with base+1
+    }
+#endif
     return Error::Ok;
   }
 
   Result<DecodeResult> decode_one(const SamplingConfig& sampling) override {
-    if (sampling.top_p != 1.0f || sampling.top_k != 0 || sampling.seed != 0) {
+    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)) {
       ET_LOG(
           Error,
-          "Qwen35MoESession: only temperature is supported; top_p/top_k/seed "
-          "are not implemented");
+          "Qwen35MoESession: top_p/seed require a sampling model "
+          "(export with --sample); only temperature is supported otherwise");
       return Error::NotSupported;
     }
     if (!valid_temperature(sampling.temperature)) {
@@ -336,6 +390,13 @@ class Qwen35MoESession : public LLMSession {
         InvalidState,
         "decode_one requires a pending token; call prefill_tokens() first");
     temperature_ = sampling.temperature;
+    if (use_sampling_) {
+      if (!valid_top_p(sampling.top_p)) {
+        ET_LOG(Error, "decode_one: top_p must be in (0, 1]");
+        return Error::InvalidArgument;
+      }
+      top_p_ = sampling.top_p;
+    }
 
     if (stop_.load(std::memory_order_relaxed)) {
       return DecodeResult{0, "", /*is_eos=*/false, /*is_terminal=*/true};
@@ -380,13 +441,26 @@ class Qwen35MoESession : public LLMSession {
 #ifdef EXECUTORCH_BUILD_CUDA
     set_temp(temperature_);
     inputs.push_back(EValue(temp_tensor_));
+#endif
+#ifdef EXECUTORCH_BUILD_MLX
+    if (use_sampling_) {
+      set_sampling_inputs(temperature_, top_p_, seed_);
+      inputs.push_back(EValue(temp_tensor_mlx_));
+      inputs.push_back(EValue(top_p_tensor_));
+      inputs.push_back(EValue(seed_tensor_));
+    }
 #endif
     auto sampled =
         run_locked(kDecodeMethod, inputs, temperature_, /*sync_after=*/false);
     ET_CHECK_OK_OR_RETURN_ERROR(sampled.error());
     pending_ = sampled.get();
     prev_decode_token_ = token;
     pos_ += 1;
+#ifdef EXECUTORCH_BUILD_MLX
+    if (use_sampling_) {
+      seed_ += 1; // next token draws with the next seed in the schedule
+    }
+#endif
     return DecodeResult{
         token, std::move(text_piece), /*is_eos=*/false, /*is_terminal=*/false};
   }
@@ -412,12 +486,26 @@ class Qwen35MoESession : public LLMSession {
     return temperature == -1.0f || (temperature >= 0.0f && temperature <= 1.0f);
   }
 
+  static bool valid_top_p(float top_p) {
+    return top_p > 0.0f && top_p <= 1.0f;
+  }
+
 #ifdef EXECUTORCH_BUILD_CUDA
   void set_temp(float t) {
     temp_val_ = (t <= 0.0f) ? 1e-6f : t;
   }
 #endif
 
+#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) {
+    temp_val_mlx_ = (temp < 0.0f) ? 0.0f : temp;
+    top_p_val_ = top_p;
+    seed_val_ = static_cast<int64_t>(seed);
+  }
+#endif
+
   Result<uint64_t> run_locked(
       const char* method,
       std::vector<EValue>& inputs,
@@ -434,7 +522,8 @@ class Qwen35MoESession : public LLMSession {
     auto res = module_->execute(method, inputs);
 #endif
     ET_CHECK_OK_OR_RETURN_ERROR(res.error());
-    auto sampled = read_sampled_token(res.get()[0].toTensor(), temperature);
+    auto sampled =
+        read_sampled_token(res.get()[0].toTensor(), temperature, use_sampling_);
     ET_CHECK_OK_OR_RETURN_ERROR(sampled.error());
 #ifdef EXECUTORCH_BUILD_CUDA
     if (sync_after && cudaDeviceSynchronize() != cudaSuccess) {
@@ -460,6 +549,13 @@ class Qwen35MoESession : public LLMSession {
   float temperature_ = -1.0f;
   std::atomic<bool> stop_{false};
 
+  // On-device sampling state (MLX --sample). top_p_/seed_ are engine-managed:
+  // the seed increments per generated token for decorrelated, reproducible
+  // draws.
+  bool use_sampling_ = false;
+  float top_p_ = 1.0f;
+  uint64_t seed_ = 0;
+
   int64_t decode_token_data_[1] = {0};
   int64_t decode_pos_data_[1] = {0};
   TensorPtr decode_tokens_;
@@ -472,6 +568,14 @@ class Qwen35MoESession : public LLMSession {
   float temp_val_ = 1e-6f;
   TensorPtr temp_tensor_;
 #endif
+#ifdef EXECUTORCH_BUILD_MLX
+  float temp_val_mlx_ = 0.0f;
+  float top_p_val_ = 1.0f;
+  int64_t seed_val_ = 0;
+  TensorPtr temp_tensor_mlx_;
+  TensorPtr top_p_tensor_;
+  TensorPtr seed_tensor_;
+#endif
 };
 
 } // namespace
@@ -513,6 +617,11 @@ Result<std::unique_ptr<Qwen35MoEEngine>> Qwen35MoEEngine::create(
   if (auto mpc = meta_module->get(kMaxPrefillChunk); mpc.ok()) {
     metadata_result.get()[kMaxPrefillChunk] = mpc->toScalar().to<int64_t>();
   }
+  // get_llm_metadata doesn't harvest use_sampling; surface it so the session
+  // knows `forward` returns a token id (on-device) rather than logits.
+  if (auto us = meta_module->get(kUseSampling); us.ok()) {
+    metadata_result.get()[kUseSampling] = us->toScalar().to<int64_t>();
+  }
 #endif
   auto eos_ids = get_eos_ids(tokenizer.get(), meta_module.get());
   // This export's metadata doesn't carry the chat-turn EOS (config.json has no