MLX-LM custom sampling and logits processors
Parent: MLX-LM Python generation sampling and evaluation · Published reference · snapshot 2026-10-23 · viability 4.40 / 5 (research)
↓ Facts as markdownall context files
Source-anchored research for MLX-LM custom sampling and logits processors. Run status: COMPLETED.
These notes link each claim to its source. Some citations identify archived research reports that are no longer hosted on this site. A published reference means the content is available; it does not certify independent review or accuracy.Read the editorial policy and follow the sources before relying on a claim.
Definitions
- apply_top_k rejects top_k unless it is an integer strictly greater than zero and strictly less than logprobs.shape[-1] (the vocabulary size), raising ValueError otherwise. make_sampler separately treats top_k=0 as disabling its top-k filter. (low confidence) (verified as of 2026-10-23) [source]
How it works
- make_sampler returns greedy argmax immediately when temp == 0, so the filtering chain is bypassed at zero temperature. At nonzero temperature, the factory applies enabled top-p, min-p, XTC, then top-k filters to log probabilities before categorical sampling. (high confidence) (verified as of 2026-10-23) [source]
- A report against mlx-lm 0.31.3 described TypeError: 'NoneType' object is not iterable when continuous batches mixed requests with and without logits processors. Upstream PR #1826 merged on 2026-09-04, normalizing per-sequence samplers and processors and closing the report. (high confidence) (verified as of 2026-10-23) [source]
Facts and statements
- MLX-LM's Python generate and stream_generate functions accept separate sampler and logits_processors callables. Each processor receives token history and current logits; processors run in list order before normalization, and current generate.py passes normalized log-probabilities to the sampler. (high confidence) (verified as of 2026-10-23) [source]
- sample_utils.make_sampler returns greedy argmax for temp=0; otherwise it applies enabled top-p, min-p, XTC, and top-k filters before categorical sampling at the requested temperature. make_logits_processors builds optional additive logit bias followed by enabled repetition, presence, and frequency penalties. The independent csm-mlx README demonstrates make_sampler(temp=0.8, top_k=50) and logit-bias use. (high confidence) (verified as of 2026-10-23) [source]
- The independent Exo MLX wrapper composes an application-defined ban_token_ids(eos_ids) processor ahead of built-in penalties and passes the processor list per sequence to MLX-LM's BatchGenerator.insert, providing a concrete pattern for combining custom filters with make_logits_processors in a batched path. (high confidence) (verified as of 2026-10-23) [source]
- April 2026 reports for mlx-lm 0.31.2 describe stateful processors seeing stale token history in stream or MTP paths: one report attributes the observation to lazy token arrays, and another says MTP draft calls omitted the backbone's just-sampled token. The MTP report author later described a vendored fix that snapshots token context and restores it after a rejected draft; the earlier report says stateless repetition penalties were unaffected. Verify this version-specific behavior against the release in use. (medium confidence) (verified as of 2026-04-28) [source]
- The current generate_step API accepts optional sampler and logits_processors arguments, and stream_generate forwards its extra keyword arguments to generate_step. Each processor receives (tokens, logits) and returns revised logits; generation applies processors before normalizing logits and sampling. This lets callers pass custom callables as well as the supplied factories. (high confidence) (verified as of 2026-10-23) [source]
- make_logits_processors can build processors for additive logit_bias, sign-aware repetition penalty, presence penalty, and frequency penalty. Its three penalty context-size defaults are 20 tokens; presence subtracts once for tokens in the context, while frequency subtracts per occurrence. (high confidence) (verified as of 2026-10-23) [source]
- Independent consumers demonstrate the intended integration pattern: csm-mlx passes MLX-LM's sampler to its generator and documents logit bias through logits_processors; Exo's MLX backend imports both factories and builds sampler and processor callables for batched generation. (high confidence) (verified as of 2026-10-23) [source]
- A Xinference issue opened on 2025-04-09 records that importing make_logits_processors and make_sampler from mlx_lm.utils failed; a follow-up shows the integration importing them from mlx_lm.sample_utils. Current MLX-LM source defines the factories in sample_utils.py, so integrations should check the installed package's module path. (medium confidence) (verified as of 2026-10-23) [source]
- An open MLX-LM issue reports that mtp_generate_step in mlx-lm 0.31.2 passed stale prev_tokens to stateful processors on speculative draft calls. The author says the current main MTP path had since been restructured, so this version-specific report does not establish a defect in ordinary serial generate_step. (low confidence) (verified as of 2026-10-23) [source]
- A Sampler callable receives log-probabilities and returns token indices as an MLX array. In generate_step, logits processors first modify raw logits in list order; the code then subtracts mx.logsumexp(logits, keepdims=True) and passes the resulting log-probabilities to sampler. (high confidence) (verified as of 2026-10-23) [source]
- min_p removes candidates whose probability is below min_p times the most probable token probability. top_p instead filters by cumulative probability mass. min_tokens_to_keep protects a minimum number of the highest-probability candidates from min_p filtering and must be a positive integer. (low confidence) (verified as of 2026-10-23) [source]
- make_logits_processors takes logit_bias as a dictionary mapping integer token IDs to additive bias values. It constructs index and value arrays from the dictionary keys and values, then adds those values to the corresponding logits. (low confidence) (verified as of 2026-10-23) [source]
- For tokens in its selected context, repetition_penalty multiplies negative logits by the penalty and divides nonnegative logits by the penalty. presence_penalty subtracts its penalty once for a token present in context, while frequency_penalty subtracts a penalty scaled by that token's occurrence count. (low confidence) (verified as of 2026-10-23) [source]
Children
- Stateful logits processor state rollback for speculative decoding (frontier)
- Random sampling behavior across MLX worker threads (frontier)
Frontier under this node: Random sampling behavior across MLX worker threads, Stateful logits processor state rollback for speculative decoding