fix(generate): penalties in CxxGenerator see a history one token stale #2090

Problem / Background

Every pipelined decode loop in CxxGenerator (src/lib/mlxcel-core/src/generate.rs) builds the next step's sample before it pushes the token it has just read into token_history. With a history-reading sampler (repetition, frequency, or presence penalty, DRY: anything where SamplingConfig::needs_token_history() is true), the token at step t+1 is sampled against a history that is missing token t. The first decode token after the prefill sample is sampled against the prompt alone.

Expected behavior is that the penalty for the token after t sees t, as in mlx-lm's generate_step (its _step concatenates the consumed input token onto tokens before calling the logits processors) and llama.cpp (common_sampler_accept runs after each sample, before the next).

Impact: mlxcel generate, mlxcel run, and interactive chat output under --repetition-penalty, --dry-multiplier (and frequency/presence penalties where configurable) is computed against the wrong window. A token can repeat immediately after itself without being penalized, and with a small --repeat-last-n the window is shifted by one. Greedy output without penalties is unaffected.

Current Behavior

Verified at main 78af3e2. In each loop the step t+1 sample (sample_token_optimized_with_state(&next_logits, sampling, &token_history, &mut sampler_state)) is built at the top of the iteration, and token_val = ffi::item_i32(&y) followed by token_history.push(token_val) comes later in the same iteration:

MLXCEL_FORCE_SYNC=1 (generate.rs:1752) does not change it: it only swaps async_eval_pair for eval, and the sample is still built before the push.

Evidence: found while hardening PR #2074 (prompt-lookup decoding). Its test drives a cache-dependent toy target (InductionModel, a one-token induction head over the KV cache, vocab 6, prompt from the seeded lcg_tokens) with repetition_penalty: 1.3, penalty_last_n: 3 at greedy. PromptLookupGenerator, which samples each verify position against the up-to-date history, emits 5, 4, 0, 3, 5, 4, 0, 3, ...; CxxGenerator emits 5, 5, 5, 0, 3, 5, 1, 5, 2, 4, .... A host-side Python simulation of the same model reproduces the first stream exactly with the current history and the second exactly with a history lagging by one token, so the divergence is the lag, not prompt lookup.

Proposed Solution

In all four loops, on the needs_history branch only, read token t (item_i32(&y)), apply the existing EOS check, and push it into token_history before calling sample_token_optimized_with_state for step t+1. reshape_token_for_forward(&y) and model.forward for step t+1 may stay ahead of the read so the forward still overlaps the host sync. Read token_val once per iteration and reuse it for generated_tokens, loop detection, and the callback. The !needs_history branch keeps its current order and overlap unchanged.

Rejected: a device-side lazy history (mlx-lm's approach). SamplerState (sampling.rs:2080) and the DRY path consume a host &[i32], so that would be a sampler rewrite for a fix that only needs the ordering corrected.

Scope

In scope: the four loops above, a regression test, and the server batch scheduler check below.

Out of scope: SpeculativeGenerator (speculative/mod.rs) and the XLA batch path (mlxcel-xla/src/batch.rs); file separately if a similar lag is found there.

Implementation Notes

• Server batch scheduler (preliminary reading, confirm in the PR): penalty rows fail config_supports_fused_batch_except_bias (sampling.rs:1047-1048), so they never enter the lookahead pipeline (server/batch/scheduler/decode_tick.rs:367-405). The per-row paths sample, try_eval, read, and push within one tick (decode_tick.rs:923 / 957 / 1049 in execute_batched_decode, 1455 / 1487 / 1563 in decode_single_step), and prefill pushes the first token at prefill.rs:1504 before the first decode tick. If that holds, state it in the PR; if not, fix it in the same PR.
• Edge cases: EOS at step t still breaks before t is pushed (the already-built forward for t+1 is discarded, as today); n + 1 >= max_tokens builds no next step, so nothing changes there; drop the later push on the needs_history path so token t is not pushed twice.
• Reuse: PR feat(speculative): add prompt-lookup decoding to mlxcel generate #2074's sequential_reference, InductionModel, INDUCTION_VOCAB, and lcg_tokens in src/lib/mlxcel-core/src/speculative/prompt_lookup_tests.rs. If feat(speculative): add prompt-lookup decoding to mlxcel generate #2074 has merged, move them into a shared #[cfg(test)] module that both test files import instead of copying them; if not, add them with this test and let feat(speculative): add prompt-lookup decoding to mlxcel generate #2074 reuse them on rebase.

Acceptance Criteria

• In each of the four loops the token read at step t is in the history used to sample step t+1, without giving up pipelining for the no-history path (the no-penalty path stays byte-identical and keeps its overlap).
• A failing-first test: for seeds 1..=6, InductionModel with repetition_penalty: 1.3, penalty_last_n: 3 at greedy, each of the four entry points (the embeddings variants with None, None) returns exactly sequential_reference(&InductionModel, prompt, 96, &sampling). It fails on current code.
• The existing greedy no-penalty parity (CxxGenerator equals sequential_reference) still passes for all four entry points.
• Server batch scheduler checked; fixed in the same PR if affected, or the PR states why not.
• The PR notes that CLI outputs under penalties change (intended).

Verification

cargo fmt --all -- --check
cargo clippy --workspace --all-targets --features metal,accelerate -- -D warnings   # Linux CUDA: --features cuda
cargo test --workspace --profile test-fast --features metal,accelerate -- --test-threads=1   # Linux CUDA: --features cuda

On a real checkpoint, mlxcel generate -m <model> -p "<prompt>" --max-tokens 128 --temp 0 must produce byte-identical output before and after the change, and the same command with --repetition-penalty 1.3 --repeat-last-n 3 must no longer emit a token immediately after itself where the penalty should have suppressed it. Measure no-penalty decode throughput once at the end, on a quiet box, and confirm it is unchanged.

Refs #2074