Merged upstreamFeatureMerged Jun 29, 2026
Add top-k support to MLX sample
Added top-k filtering support to PyTorch ExecuTorch's MLX sampling path so on-device LLM sampling can restrict candidates to the k most likely tokens.
pytorch/executorch · #20564
Problem
The MLX backend already supported temperature and top-p nucleus sampling, but SamplingHead reserved top_k without implementing it. Passing a top_k value raised NotImplementedError, leaving ExecuTorch users without a standard decoding control for MLX-backed on-device generation.
Approach
Threaded top_k through SamplingHead and the mlx::sample custom op, clipped runtime k to vocabulary size, applied top-k filtering before top-p so probabilities are renormalized over the top-k subset, and lowered the behavior through MLX graph primitives without C++ or schema changes.
Impact and scope
- Completes a common LLM decoding control in the ExecuTorch MLX backend for mobile, embedded, and edge inference workflows.
- Makes top-k compose with top-p filtering, giving users predictable candidate restriction and nucleus sampling together.
- Preserves no-filter behavior through the default max-int sentinel while supporting runtime top_k inputs.
Validation
- Added eager sampling tests for top-k restriction, default keep-all behavior, vocab-size clipping, and top-k plus top-p composition.
- Added export and end-to-end coverage proving top-k survives the ExecuTorch export and MLX runtime path.
- PR validation covered focused unittest suites for sample ops/export behavior plus compileall checks on the changed MLX backend modules.