perf: bound DRY and penalty per-step cost on long histories #2259

Problem / Background

With a full-history DRY window, the host-side DRY pass costs more per step as the output grows, and on a repeating stream that per-step cost is quadratic in the window. Both DRY and the windowed frequency/presence path also allocate and upload a vocab-sized f32 buffer on every step.

This matters for server requests with long outputs:

• The request field dry_penalty_last_n is a usize clamped only to i32::MAX (src/server/request_options.rs:312-317), so a client can ask for a window that covers the whole history.
• --dry-penalty-last-n -1 maps to DRY_FULL_HISTORY (src/server/startup.rs:1112-1121).
• SamplingConfig::default() uses DRY_FULL_HISTORY (src/lib/mlxcel-core/src/generate.rs:995).
• The history includes the prompt (DecodeState::new, engine/direct_decode.rs).

The security review of #2256 (Part 1 of #2229) flagged this as a CPU-time amplifier. #2256 does not add any of this cost. Its per-row pipeline now overlaps the cost with the next forward, but the synchronous step and the scheduler still pay it serially. Paths below are relative to src/lib/mlxcel-core/src/, with line numbers from main 25baebcc.

Current Behavior

• apply_dry_penalty (sampling.rs:1879-1998) builds a HashMap<i32, Vec<usize>> over the whole window on every step (:1913-1916).
• For each earlier occurrence of the last token, it then walks backward to measure the match (:1921-1946).
• On a stream with period p and window N, that is about N/p occurrences, each matching back almost to the window start. One step therefore costs about N^2 / (2p) comparisons, and a whole generation is cubic.
• It then fills vec![0.0f32; batch * vocab] and uploads it to add the penalties (:1985-1997).
• apply_frequency_presence_penalty (sampling.rs:1808-1834) is the rebuild path, taken when penalty_last_n > 0 or there is no SamplerState. It also allocates vec![0.0f32; vocab] (:1826) per step.
• The full-history incremental path SamplerState::apply_frequency_presence (sampling.rs:2196-2240) already avoids that buffer by using sparse take/put.
• apply_dry_penalty has one caller, preprocess_penalty_stages (sampling.rs:763). Every decode path reaches DRY through it: the engine's synchronous step, both B=1 pipelines, the scheduler, the CLI and speculative decoding.

Proposed Solution

1. DRY in linear time per step, with identical results. Rewrite the matching in apply_dry_penalty as a stateless pass that costs O(window) per step, keeping the current semantics exactly. For window w, last index n, and an occurrence pos < n with w[pos] == w[n], the current loop computes match_len = min(L(pos), pos - b(pos)), where:

   • L(pos) is the longest common suffix of w[..=pos] and w[..=n].
   • b(pos) is the largest index < pos where dry_breaks_at(config, w, b) holds, or -1 if there is none.

   Compute L(pos) as Z[n - pos], using the Z-function of the reversed window, which is the technique upstream llama.cpp's DRY sampler uses. Compute b with one forward pass that records the most recent breaker index. The occurrence filter, the >= dry_allowed_length test, the exponent cap and the max-per-next-token merge stay as they are. Move the host computation into a pure fn dry_penalty_map(window: &[i32], config: &SamplingConfig) -> Vec<(i32, f32)> with no MLX calls, so it can be tested without a device.

2. No vocab-sized host buffers.

   • Apply DRY penalties sparsely: promote the logits to f32, then take_along_axis / add / put_along_axis over the touched ids. Build the indices and values with one row per batch row, without relying on broadcasting.
   • Move the sparse body of SamplerState::apply_frequency_presence into a free apply_frequency_presence_sparse(logits, idx, val), and call it from both that method and apply_frequency_presence_penalty.
   • Both changes produce the same values as today. The current add(logits, f32_buffer) already promotes to f32, the touched ids get the same f32 operation, and untouched ids keep their promoted value. The only possible difference is -0.0 versus +0.0, which no sampling operation can tell apart.

Rejected alternatives.

• A server-side cap on the effective window. It changes the output of any request that asks for a longer window, and is unnecessary once each step is linear.
• An incremental DRY index stored in SamplerState. It would need to rebuild when history diverges (SamplerState::sync), rebase when the window slides, and recheck breakers with tails whose result changes as the window grows. All of that only improves O(window) to O(occurrences).

Scope

In scope: sampling.rs (apply_dry_penalty, the new dry_penalty_map, apply_frequency_presence_penalty, the new apply_frequency_presence_sparse, SamplerState::apply_frequency_presence) and the tests in the same file.

Out of scope:

• The windowed repetition penalty's per-step to_vec + sort (sampling.rs:1764). It is bounded by the caller's penalty_last_n and costs O(W log W), not quadratic.
• The XLA sampler (src/lib/mlxcel-xla/src/sampler.rs).
• Any change to request validation or defaults.

Implementation Notes

• Reuse: dry_breaks_at (sampling.rs:1851) unchanged, the sparse take/put body from sampling.rs:2226-2234, and the LCG history generator style of lcg_logits (sampling.rs:4261).
• Edge cases:
  • History shorter than 2, dry_base < 1.0, a window of 0, and a window shorter than 2 keep their early returns.
  • A breaker at pos itself does not stop the match, because only positions before pos are checked. The new code must keep this.
  • A breaker head whose tail runs past the window end does not break.
  • An empty penalty map returns ffi::copy(logits), as today.
• Allocation: O(window) scratch per step (the Z array and the breaker array) is acceptable. Do not add any O(vocab) host allocation.

Acceptance Criteria

• Identical logits. Keep the current DRY code verbatim as #[cfg(test)] fn apply_dry_penalty_reference. Over at least 500 LCG-generated cases, the new apply_dry_penalty matches it element-wise (comparing f32 values, so that -0.0 == +0.0). The cases must cover:
  • vocab sizes 8 and 64
  • periodic and random streams, with history lengths 0 to 2048
  • windows DRY_FULL_HISTORY, 0, 1, 2, 17 and 64
  • dry_allowed_length 1 to 4, and dry_base 1.0 and 1.75
  • exact-id breakers, and breaker heads with empty and non-empty tails
  • batch 1 and 2, with f32 and bf16 logits
• Step cost. dry_penalty_map counts its element comparisons and breaker checks under cfg(test). On a period-3 stream with a full-history window of 4096, one step counts at most 8 * 4096. Doubling the window to 8192 at most doubles the count, plus a constant. The current code does about 4096^2 / 6 on the same input.
• Frequency/presence parity. The rebuild path matches a #[cfg(test)] copy of the current vocab-buffer implementation for penalty_last_n 1 and 64, at batch 1 and 2, with f32, f16 and bf16 logits. The existing parity tests (sampling.rs:3780-3830) still pass.
• No vocab-sized buffers. Neither DRY nor apply_frequency_presence_penalty allocates a vocab-sized host buffer any more.
• Existing behavior intact. The existing DRY tests (sampling.rs:2880, :5530-5568, :5911-5952) and the engine stream-equality tests pass unchanged. Identical logits give identical greedy and seeded tokens.
• Timing evidence. An #[ignore] test, dry_step_cost_report, prints microseconds per step for the reference and the new implementation at windows of 1k, 4k and 16k on a period-3 stream. Run it once on a quiet machine and paste the numbers into the PR.

Verification

cargo fmt --check
cargo clippy --profile test-fast -p mlxcel-core --lib --tests -- -D warnings
cargo test --profile test-fast -p mlxcel-core --lib sampling -- --test-threads=1
cargo test --profile test-fast -p mlxcel-core --lib engine:: -- --test-threads=1
cargo test --profile test-fast -p mlxcel-core --lib sampling::tests::dry_step_cost_report -- --ignored --nocapture --test-threads=1

On Linux CUDA hosts, add --features cuda. On GB10, run the GPU tests under gpu-lock run --tag <purpose> --.

Refs #2229, #2256. Related to #822 only through the same review.