[MLX] Gemma4-31B ondevice sampling

[kiymetakdemir]

Summary

Lets the MLX-exported Gemma 4 31B model sample the next token on-device instead of returning logits for host-side sampling. Sampling is opt-in at export (--sample); temperature, top_k, top_p, and seed are runtime inputs, and the runner increments the seed per token.

Changes

• export.py — --sample wraps the model in SamplingHead so forward(tokens, input_pos, temperature, top_k, top_p, seed) → int64 token; records a use_sampling constant-method flag. Non-sample export unchanged.
• gemma4_31b_engine.cpp — reads use_sampling from metadata; when set, consumes the int64 token directly instead of logits_to_token, feeds the scalar inputs (temperature, top_k, top_p, seed) in prefill/decode (across the min/max prefill chunking), and manages the per-token seed schedule. top_k (0 = keep all → INT64_MAX on-device) and top_p are now supported; top_p is range-checked to (0, 1]; top_k/top_p/seed are rejected on non-sample models.
• main.cpp — --top_k / --top_p / --seed flags wired into SamplingConfig; an unset seed is randomized only for sampling models.
• tests/test_mlx_pipeline.py — test_export_to_pte_with_sampling: tiny-model MLX export with --sample, asserting the use_sampling flag, int64 token output, and same-seed reproducibility.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20561

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[metascroy]

Ideally SamplingHead from backends/mlx/llm would be directly usable. The fact that we keep defining wrappers suggest that maybe it should be refactored to work?

In terms of shape, I expect sampling to work on (B, vocab) and return (B) (tokens).

I don't see a case for (B, S, vocab).

[kiymetakdemir]

I've made SamplingHead operate on (B, vocab) → (B) so it's model-agnostic, but it still needs a positional adapter for torch.export since its sampling args are keyword-only.

SamplingHead wraps the whole model (self.model(*args)) behind a generic, keyword-only forward(*args, temperature, seed, …), so I think each model needs a small adapter that maps its fixed forward(tokens, input_pos, temperature, top_p, seed) onto that signature.

I checked the existing heads and they avoid this by being tensor layers, lm_head (nn.Linear), FlowMatchingHead take a tensor in tensor out and are called inside the model's forward, so they ride on the model's already-positional export entry point. Would we want to change SamplingHead to a tensor head like this? The CUDA path avoids it too with a plain free function sampler.py:sample(logits, temperature) is called inside model.forward (return sample(logits, temperature)).

[metascroy]

What if we did:

def forward(self, *args, **kwargs):
    *model_args, temperature, top_k, top_p, seed = args
    logits = self.model(*model_args, **kwargs) # [B, V]
    return torch.ops.mlx.sample(last, temperature, top_k, top_p, seed)

Then in Qwen code (and here), we could just do:

model = SamplingHead(model)
example_args = (example_tokens, example_input_pos, temp, top_k, top_p, seed)

With no wrapper?

[kiymetakdemir]

This should work, thank you!

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 2m 18s —— View job

---

Code review

• Read changed files (export.py, engine, main.cpp, tests)
• Analyze sampling logic, seed handling, guards
• Address reviewer's question re: reusing SamplingHead
• Post review feedback

Overall this is a clean, well-guarded port of the Qwen3.5 MoE on-device sampling pattern (#20545). The runtime guards, range-checks, and the per-token seed schedule are consistent with that PR. A few notes below.

---

On @metascroy's question: why not reuse SamplingHead?

You're right, and the shape mismatch is the reason a wrapper is needed today:

• backends/mlx/llm/sampling.py:36-40 — SamplingHead.forward expects the wrapped model to return (B, S, vocab) and does last = logits[:, -1, :] to get (B, vocab) before calling torch.ops.mlx.sample.
• The Gemma MLX source transform already returns the last-token logits: mlx_source_transformations.py:150 (last = self.lm_head(x[:, -1, :])) produces (B, vocab) — there is no S axis. Feeding that into SamplingHead would make [:, -1, :] index the wrong axis (or fail), so _MLXSampleWrapper (export.py:312) skips the slice and calls mlx.sample(logits, ...) directly.

So the two runners disagree on the head's input contract: Qwen feeds (B, S, vocab) and relies on SamplingHead to slice; Gemma pre-slices to (B, vocab). That's exactly the (B,S,vocab) vs (B,vocab) split you flagged.

The clean fix is the one you suggested: make SamplingHead operate on (B, vocab) → (B) and drop the internal [:, -1, :] slice (which is really a model-output concern, not a sampling concern). Then:

• Gemma can use SamplingHead directly (its MLX forward already yields (B, vocab)).
• Qwen keeps its last-token slice in the model forward (or a tiny adapter), matching how mlx.sample itself is documented (custom_ops.py:405 — logits: [B, vocab]).

That removes _MLXSampleWrapper here and the divergent contract. Worth doing as a small follow-up so the next runner doesn't add a third wrapper. If you'd like, I can take that refactor on (touching SamplingHead, the Gemma export, and the Qwen export).

Other observations

1. Sampling runs on every prefill chunk, output discarded. In run_prefill_chunk (gemma4_31b_engine.cpp:527+), each chunk feeds the scalar inputs and the baked-in sampler runs, but only the final chunk's token is used. The seed is held constant across chunks and incremented once after prefill (:355), so this is correct, just mildly wasteful. Sampling cost is negligible vs. the model, so fine to leave — flagging only for awareness.

2. No dtype/shape check on the int64 token read. read_sampled_token (:104-105) does output.const_data_ptr<int64_t>()[0] when use_sampling. This trusts the metadata flag matches the actual output tensor. Given use_sampling_ is read from the model's own constant method that's a safe coupling, but a defensive ET_CHECK on scalar_type() == Long would fail loudly instead of reinterpreting bytes if the two ever drift.

3. Test executes forward on the lowered program directly. test_export_to_pte_with_sampling guards on MLXPartitioner import but then runs forward.execute(...) via Runtime. The reproducibility assertion relies on the host reference mlx.sample (custom_ops.py:416-419 explicitly notes it is not bit-identical to the on-device graph). On non-Mac CI this only exercises the CPU reference — worth a comment in the test that same-seed reproducibility is being checked against the host reference, not the delegate, so the assertion's scope is clear.

4. Seed semantics look correct. -1 → random only for sample models, non-sample stays at 0 to avoid the guard (main.cpp:178-188), explicit --seed 0 on a sample model is reproducible (not re-randomized), and top_p/seed are correctly rejected on non-sample models in both prefill_tokens and decode_one. valid_top_p of (0, 1] matches the op's documented domain. 👍

None of these are blockers. The main actionable item is the SamplingHead refactor you raised — I'd recommend landing this as-is if it unblocks you and doing that refactor as a fast follow across both runners.
· branch gemma-ondevice-sampling

[metascroy]

I think topk is supported now?

[kiymetakdemir]

Yes, I will update this.