[AlexCheema]

This PR fixes which tokens the repetition, presence and frequency penalties count. Logits are unchanged when no penalty is set, and so is greedy decoding without penalties.

The bug

The batch engine never penalised the token just generated.

• Its step (opt_batch_gen._patched_step) handed logits processors self.tokens, the tokens already in the KV cache.
• The token being fed in a step joins that list only at the end of the step.
• So penalties never saw the most recent token, and an immediate repeat ("the the", "\n\n\n") went unpenalised.
• mlx-lm's own batch step and generate_step both include the fed token.

What the penalties counted also depended on the engine.

• The batch engine's history started with one stray prompt token, because exo prefills the prompt itself and inserts only the last two prompt tokens.
• The sequential engine's history started with the last prompt token (or the last two, depending on mlx-lm's version).
• Neither had the rest of the prompt.

Recorded on Qwen3.5-2B (greedy, repetition penalty 1.1, presence penalty 1.5), for the first generated token, the processors saw:

	presence/frequency count	repetition counts
batch engine, before	the second-to-last prompt token	the second-to-last prompt token
sequential engine, before	the last prompt token	the last prompt token
both engines, after	nothing yet (no token generated)	the last 20 prompt tokens

At every later step, the batch engine's history also lacked the token being fed.

What each penalty should count

How common implementations define them:

• OpenAI's API: presence and frequency penalties are counts of the tokens sampled so far, i.e. the generated text.
• Hugging Face transformers: repetition_penalty is applied over input_ids, which for decoder-only models includes the prompt (an option excludes it). It has no presence or frequency penalty.
• Some open-source servers: feed the prompt into the sampler's history, so repeat, presence and frequency penalties all cover the last N tokens of prompt and generated text.
• Other widely used open-source inference servers: presence and frequency penalties over the generated text only; repetition penalty over the prompt and the generated text.
• mlx-lm: all three over the history its generation loop builds, which includes the prompt.

This PR follows the OpenAI definition for presence and frequency, which is what exo's OpenAI-compatible API exposes and what card values such as Qwen's presence_penalty = 1.5 are tuned for. It follows the transformers definition for the repetition penalty:

• presence and frequency penalties: the generated tokens, the latest included;
• repetition penalty: the prompt and the generated tokens, the latest included.

The windows are unchanged: mlx-lm's last 20 tokens for presence and frequency, and repetition_context_size (default 20) for repetition.

Counting the prompt in the presence penalty would penalise every number and word of the question at 1.5, which matters for tasks that copy from the prompt. The quality check below includes that variant.

The change

• Batch step: processors get mlx-lm's per-sequence token context (GenerationBatch._token_context): the tokens in the cache plus the one being fed, as mlx-lm's own step does. It is evaluated with the step, and it costs one append per step instead of converting the whole history to an array every step.
• penalty_processors(...) in generator/generate.py: used by both engines. It builds mlx-lm's processors with the same values, the same skipping of unset or zero penalties and the same order, each wrapped in a PenaltyProcessor that hands its penalty the tokens it counts.
  • How it finds the prompt boundary. Engines call processors with the end of the prompt followed by the generated tokens. A processor learns on its first call how many prompt tokens its engine's history starts with: the first call is for the first generated token, when the history holds only prompt tokens.
  • What each penalty then gets. Presence and frequency get the rest of the history. The repetition penalty also gets the prompt tokens before the engine's history, so it sees prompt plus generated.
  • Result. Both engines, and mlx-lm versions that start the history differently, give the same result.

Evidence

Unit tests (src/exo/worker/engines/mlx/tests/test_penalty_history.py):

• Recording processors see exactly the intended tokens, for engines whose histories start with 1, 2 or all the prompt tokens, and for repetition windows of 3, 20 and 0 (0 keeps the whole history).
• Penalty values, skipping and order match mlx-lm's make_logits_processors.
• The batch step hands processors the tokens in the cache plus the one being fed.

Both engines, recorded, Qwen3.5-2B (greedy, repetition 1.1, presence 1.5, 12 tokens):

• Before: the four checks (batch/sequential × repetition/presence) all differed from the intended semantics, from the first token.
• After: all four matched at every step, and the two engines generated the same tokens.

Quality: GSM8K at temperature 1. Qwen3.5-2B (8-bit) with its card's sampling defaults: temperature 1, top-p 0.95, top-k 20, presence penalty 1.5, repetition penalty 1.0. 150 problems × 3 seeds, 8 requests decoding at once through the batch engine, thinking off (to keep answers short), at most 2048 tokens. Variants:

• before: this PR's base;
• after: this PR;
• prompt counted: after, but with presence and frequency also counting the prompt (the mlx-lm definition).

variant	accuracy (450 answers, 95% CI)	per seed	answers cut off at 2048 tokens
before	73.3% ± 4.1	72.7 / 72.7 / 74.7	69
after	78.7% ± 3.8	78.7 / 80.0 / 77.3	58
prompt counted	77.6% ± 3.9	76.7 / 76.7 / 79.3	60

Paired by problem and seed, with an exact two-sided sign test on the answers only one variant got right:

• after vs before: 53 vs 29, p = 0.011.
• prompt counted vs before: 51 vs 32, p = 0.048.
• after vs prompt counted: 34 vs 29, p = 0.62.

So counting the latest token (and dropping the stray prompt token) is a significant improvement here, with fewer answers running into the token limit. Whether presence also counts the prompt made no measurable difference on this test, so the choice between the two rests on the definitions above.

Interplay with #2524

Both PRs touch the batch step's processor loop. Once this lands, #2524's windowed history is redundant (the token context here is already one append per step), and its repetition-penalty-1.0 skip moves into penalty_processors. Whichever lands second gets rebased.

Checks

• uv run basedpyright: 0 errors.
• uv run ruff check: clean.
• nix fmt: no changes.
• uv run pytest: 721 passed, 3 skipped. The image-generation tests need torch and weren't run.

🤖 Generated with Claude Code