Top-K Sampling Support

[mikemykhaylov]

Hello,

I find that for some models, like Mistral Nemo, it is very beneficial to restrict the number of considered completions to improve the coherence of the output. However, compared to llama.cpp, MLX does not seem to support top K sampling in LM Studio. It does look like the underlying library supports that, so implementing that would be much appreciated.

llama.cpp models support Top-K

MLX models do not support Top-K

[neilmehta24]

Thanks for bringing this to our attention. The MLX core library does indeed support a top-k matrix operation, but the MLX LLM library does not support top-k sampling. Here are the supported generation/sampling options as of today https://github.com/ml-explore/mlx-examples/blob/dfa4dd6/llms/mlx_lm/utils.py#L200-L215 . Please track this issue ml-explore/mlx-examples#1167 for adding support in mlx_lm