sampling: index penalties by token id instead of scanning every candidate - #95
sampling: index penalties by token id instead of scanning every candidate#95danielhanchen wants to merge 2 commits into
Conversation
…idate llama_sampler_penalties_apply probed token_count once per candidate, so a model with a large vocab paid n_vocab hash lookups per token while at most penalty_last_n of them can ever match. Walk token_count instead and index cur_p directly. The direct index needs a token id to be its own index, so it is guarded and falls back to the old scan when an earlier sampler has dropped or reordered candidates. Qwen3.6-27B (248320 vocab, RTX/B200, presence_penalty 1.5): 130.3 -> 146.7 t/s. Output is unchanged: same arithmetic, applied to the same tokens.
c26185b to
e2e842a
Compare
The first pin for #95 was cut from fork master, whose merge base with the current base tag is from June 10, so merging it dragged in the whole fork/upstream divergence and conflicted across ~1600 files. Rebased the change onto b10359 and repinned to e2e842a. #95 now sits on penalties-upstream-base, matching how #70 and #91 are carried. Verified: all five pins merge onto b10359 in pr-set.json order, #70 via the additive resolver as before, everything else clean.
The first version checked only that each penalized token sat at its own index. A cur_p holding the same id twice then took the by-index path and penalized one of the two entries, where the scan penalized both. llama.cpp's own samplers never build such an array, but the old behaviour should not change for anyone who does. Check the whole array is the identity layout instead, which also rules out duplicates. It is a compare per candidate against a map probe per candidate, and it stops at the first mismatch. Found by a differential fuzz over 33600 cases (vocab sizes, penalty settings, reversed / shuffled / truncated / duplicated / sorted layouts, inf and nan logits): 209 mismatches before, 0 after.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3db8cb5b2e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // Apply frequency and presence penalties to the cur_p | ||
| if (by_index) { | ||
| for (const auto & it : ctx->token_count) { | ||
| penalize(cur_p->data[it.first], it.second); |
There was a problem hiding this comment.
When a caller has accepted LLAMA_TOKEN_NULL or another token outside [0, n_vocab), a full identity-layout candidate array still selects this path, and it.first indexes before or beyond cur_p->data, causing memory corruption. This can occur when integrations feed multimodal prompt histories containing the documented LLAMA_TOKEN_NULL placeholder into the public sampler API. The previous lookup path simply found no matching candidate and ignored the sentinel, so validate each history ID or fall back to lookup before indexing.
Useful? React with 👍 / 👎.
1 participant
The problem
llama_sampler_penalties_applywalks the whole candidate array and probes the frequency map once per candidate:token_countnever holds more thanpenalty_last_nlive entries (64 by default), so at most 64 of those probes can match. Everything else is a wasted hash lookup. The cost scales with the vocab, not with the penalty window, so it lands hardest on the recent large-vocab models: Qwen3.6-27B has 248320 tokens, which is 248320 lookups per token to penalize at most 64.It only shows up when a penalty is actually enabled, so most default presets never see it. Ours do: the
qwen3.5andqwen3.6family profiles in Unsloth Studio setpresence_penalty: 1.5, and that alone was costing about 11% of generation throughput.The fix
Walk
token_countand indexcur_pdirectly. That turns the loop fromn_vocabintopenalty_last_n.Indexing by id is only valid while
cur_pis still the untouched candidate array, so the fast path is guarded on the array being the identity layout,data[i].id == ifor everyi. Anything else, a truncated, reordered or hand-built array, falls back to the original scan. The check is one compare per candidate against a map probe per candidate, and it stops at the first mismatch.The identity layout also guarantees each id appears exactly once, which matters: the first version of this patch checked only that each penalized token sat at its own index, and an array holding the same id twice then got one of the two entries penalized instead of both. The fuzz below caught it.
The arithmetic itself is untouched and each token is still penalized exactly once, so the result does not depend on which path runs or on map iteration order.
Measured
Both builds from
b10359, single GPU,--parallel 4 --kv-unified --flash-attn on, 400 generated tokens,temperature 0, median of 5. Throughput is llama.cpp's own/metricscounters, not client wall clock.Qwen3-0.6B-Q4_K_M, 151936 vocab:
presence_penalty 1.5frequency_penalty 0.8repeat_penalty 1.15Qwen3.6-27B-UD-Q4_K_XL, 248320 vocab,
--spec-type draft-mtp,presence_penalty 1.5: 130.29 to 144.59 t/s, up 11.0%, against 146.63 with the penalty switched off. So the penalty goes from costing 11% of the decode rate to costing about 1%.The "none" row is noise.
is_disabled()returns before any of this, so the patched path cannot execute there; over 9 runs each the medians were 600.60 and 598.80 with fully overlapping spreads.The win grows with vocab size and shrinks with model size, which is why it has gone unnoticed: it is invisible on small-vocab models and largest exactly where it hurts, small fast models with modern tokenizers.
Output is unchanged
60 comparisons, stock server against patched server, both built from
b10359, 4 prompts x 5 penalty combinations x greedy and two seeded sampled settings:presence 1.5frequency 0.8repeat 1.1560/60 byte identical. Repeating a generation five times on each build gave one distinct output per build, so both remain deterministic.
Tests
tests/test-samplingpasses. The existingtest_penaltiescases already pin the expected probabilities for repeat, frequency and presence penalties and are unchanged. Two cases are added:test_penalties_reorderedapplies the sampler to a normal candidate array and to a reversed one, which forces the fallback, and asserts the resulting logits match exactly.test_penalties_duplicate_idsbuilds an array holding one id twice and asserts both entries end up with the same logit, which is what the scan does and what the first version of this patch got wrong.Beyond the suite, a differential fuzz compared the shipped sampler against a reference implementation of the old scan over 33600 cases: 10 vocab sizes from 1 to 151936, 6 values of
penalty_last_nincluding 0, 5 repeat x 4 frequency x 4 presence settings including negative ones, and 7 candidate layouts (canonical, reversed, shuffled, truncated, duplicated id, sorted by logit, single entry), with logits including 0, -0, inf, -inf, nan, denormals andFLT_MAX. Comparison is on the float bit patterns, so a nan or a signed zero landing in the wrong place fails.g++ -std=c++17and-std=c++20, both with-Wall -Wextra -Werror, are clean, and the file produces no warnings the stock file does not.