Gemma MLX: reuse SamplingHead for on-device sampling; wire runtime top-k

Author: kiymetakdemir

examples/models/gemma4_31b/export.py

@@ -309,24 +309,6 @@ def _export_cuda(
     print("Done.")
 
 
-class _MLXSampleWrapper(nn.Module):
-    """Wrap the model so ``forward`` returns a sampled token id.
-
-    The MLX source transforms make ``forward`` return last-token logits
-    ``(B, vocab)``, so sample directly. 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__()
-        self.model = model
-
-    def forward(self, tokens, input_pos, temperature, top_p, seed):
-        logits = self.model(tokens, input_pos)
-        return torch.ops.mlx.sample(logits, temperature, top_p, seed)
-
-
 def _export_mlx(
     model: Gemma4_31B,
     config: Gemma4_31BConfig,
@@ -383,15 +365,23 @@ def _export_mlx(
     example_tokens = torch.tensor([[0, 1]], dtype=torch.long)
     example_input_pos = torch.tensor([0, 1], dtype=torch.long)
     if sample:
-        model = _MLXSampleWrapper(model)
+        # forward(tokens, input_pos, temperature, top_k, top_p, seed) -> token id.
+        # gemma's MLX forward already returns last-token logits (B, vocab), so
+        # SamplingHead is used directly with no per-model wrapper.
+        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/gemma4_31b/gemma4_31b_engine.cpp

@@ -304,6 +304,8 @@ class Gemma4_31BSession : public LLMSession {
     }
     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_ =
@@ -329,15 +331,12 @@ class Gemma4_31BSession : 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;
       }
@@ -347,6 +346,7 @@ class Gemma4_31BSession : 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;
       }
@@ -398,14 +398,11 @@ class Gemma4_31BSession : public LLMSession {
   }
 
   Result<DecodeResult> decode_one(const SamplingConfig& sampling) override {
-    if (sampling.top_k != 0) {
-      ET_LOG(Error, "Gemma4_31BSession: 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,
-          "Gemma4_31BSession: top_p/seed require a sampling model "
+          "Gemma4_31BSession: top_k/top_p/seed require a sampling model "
           "(export with --sample); only temperature is supported otherwise");
       return Error::NotSupported;
     }
@@ -423,6 +420,7 @@ class Gemma4_31BSession : 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;
     }
 
@@ -474,8 +472,9 @@ class Gemma4_31BSession : public LLMSession {
     inputs.push_back(EValue(decode_pos_));
 #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_));
     }
@@ -522,8 +521,10 @@ class Gemma4_31BSession : public LLMSession {
   }
 
 #ifdef EXECUTORCH_BUILD_MLX
-  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);
   }
@@ -564,8 +565,9 @@ class Gemma4_31BSession : 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_));
     }
@@ -705,6 +707,7 @@ class Gemma4_31BSession : public LLMSession {
   std::atomic<bool> stop_{false};
 
   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;
 
@@ -727,9 +730,11 @@ class Gemma4_31BSession : 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

examples/models/gemma4_31b/main.cpp

@@ -50,6 +50,11 @@ DEFINE_string(
     "",
     "Path to file containing prompt text (overrides --prompt).");
 DEFINE_double(temperature, 0.8, "Sampling temperature (0 = near-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,
@@ -174,6 +179,7 @@ int main(int argc, char** argv) {
 
   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/gemma4_31b/tests/test_mlx_pipeline.py

@@ -283,6 +283,7 @@ def sample(seed):
                         tokens,
                         input_pos,
                         torch.tensor(0.8, dtype=torch.float32),
+                        torch.tensor(torch.iinfo(torch.int64).max, dtype=torch.int64),
                         torch.tensor(0.9, dtype=torch.float32),
                         torch.tensor(seed, dtype=torch.int64),
                     ]