Repository navigation
sampling: backend samplers match the CPU samplers, also under draft verification; ggml-alloc re-plans when a tensor becomes an output - #316
Open
danielhanchen wants to merge 37 commits into
Conversation
ggml_gallocr_alloc_graph() reused the previous plan whenever the graph had the same number of nodes and leafs and every tensor still fit its slot. The plan also depends on tensor flags: ggml_gallocr_alloc_graph_impl() allocates inputs before any node runs, never frees an output and never lets a later tensor reuse an output (or a view of one). A graph of the same shape in which a tensor became an output therefore kept a plan where that tensor shares memory with a later tensor, and its data is overwritten before it is read back. The backend sampler hits this: it builds the same graph for every step and marks as outputs only the chains of the rows the step samples, so the outputs change with the number of draft rows. The plan now records, for every tensor slot it was computed from (each node, its sources, the leafs), the flags that decide a lifetime: input, output, view of an output. ggml_gallocr_needs_realloc() requests a new plan when a tensor has a flag the plan was not made with. A tensor that lost a flag keeps the plan: it is only allocated earlier or kept longer than needed, so a graph whose outputs shrink is not re-planned, and once the plan covers the largest set of outputs it holds in steady state. test-alloc: test_plan_follows_output_flags plans a chain while its middle tensor is not an output, then the same shape with it marked as an output, and checks that it no longer overlaps the tensors after it; and that the plan made with the output is kept when the flag is dropped.
…rification Backend (GPU) sampler fixes, each to the CPU sampler's semantics: - top-k: clamp k to the candidates (k > n_vocab, k > 131072 on the 128k probe graph, or a temperature-0 argmax before top-k aborted in ggml_top_k), and sort the k results by descending logit (ggml_top_k order is unspecified, on CUDA not even deterministic, so a fixed seed neither reproduced nor matched the CPU). - top-p: honour min_keep; keep the first max(n, min_keep) candidates where n is the first index whose cumulative probability reaches p. - min-p: keep logit >= max + log(p) (ggml_step is > 0, so p = 1 dropped every token); honour min_keep (the n-th best logit, selected exactly), and sort when fewer than min_keep reach the threshold, as the CPU does. - temp-ext: a positive delta is dynamic temperature also at temp <= 0; count only the kept (finite) candidates for the maximum entropy; no 1e-10 floor on masked probabilities; exponent 0; sort like the CPU. - logit-bias: a token listed several times gets the sum of its biases; candidates already reordered by an earlier sampler are biased by id. - dist: scale the uniform by the cumulative total and clamp the index to the candidates with a non-zero probability (never pick a masked candidate); pick the first candidate whose running sum reaches u, as the CPU. - greedy / temperature 0: the first maximum on ties, as the CPU. - hand-off to CPU samplers after a backend prefix (common_sampler and llama_sampler_sample): drop the candidates the backend masked to -inf, the CPU samplers would have removed them (XTC at threshold 0 picked a token top-p had excluded). - CPU softmax and top-p accumulate in double: a sequential f32 sum over a 150k+ vocabulary is off by ~1e-3, which moved the CPU top-p cut by dozens of candidates away from the backend and from an exact softmax. Draft verification (several outputs per sequence, any speculative decoding): - penalties run on the backend with n_outputs_max_per_seq > 1: chain g of a sequence uses the accepted history plus the input tokens of rows 1..g of this decode (the tokens the server accepts before it samples row g in common_sampler_sample_and_accept_n), with the CPU accept rule (ring buffer, counts). The rows' tokens reach the samplers through a new internal llama_sampler_backend_set_rows() called from the sampling graph input. - dist and penalties track the rows of a sequence across ubatches, so a draft split over two ubatches uses the right draw and history. - one input tensor per sampler for all chains (uniforms, penalty ids and counts) instead of one per chain: a single host-to-device copy each. Tests (test-backend-sampler): test_cpu_parity runs 29 whole chains (Studio's defaults and every edge value above) through common_sampler on the backend and on the CPU with identical logits and seeds, and requires the same token for every draw and the same post-sampling candidates and probabilities; test_cpu_parity_multi_output does the same for draft verification rows, with n_ubatch 512 and 8 (rows split across ubatches). test_temp_ext expects dynamic temperature at temp 0.
…of a mismatch full_nofilter (no sampler before dist) and full_nofilter_k (top-k past the vocabulary, so a full sort before dist) draw over the whole vocabulary, where a CPU and a backend cumulative sum differ most.
test_multi_output_ubatch_split decodes two sequences whose draft verification rows straddle ubatches (n_ubatch 8) with a dist-only chain and checks that every row's sampled probabilities are the softmax of its own logits (CPU context, identical batch). It needs no other sampler fix, so it runs against older trees as well.
…rity tests The CPU dist sampler accumulates in double, the backend one on an fp32 cumulative sum. On an unfiltered 248k-token vocabulary the two can pick different tokens when u * total falls on the boundary between two candidates (Qwen3.5-0.8B: 1 of 96 draws in two cases at n_ubatch 8). A mismatch is now tolerated only when it is a provable near-tie: with u replayed from the seed, u * total lies within eps = n * FLT_EPSILON * total (the fp32 summation error bound over the n candidates) of the CPU cumulative interval of both tokens, and both tokens' probabilities agree within eps. Every other mismatch still fails. Each mismatch prints u, the two intervals, the distances and eps, and every case prints its count of tolerated near-ties.
The worst-case bound n * FLT_EPSILON * total is 3e-2 at a 248k vocabulary: it tolerated draws of a deliberately broken backend dist (index + 1) whose token had less mass than the bound. Use the expected error of a sum of n floats, sqrt(n) * FLT_EPSILON * total (6e-5 at 248k), which still covers the observed fp32 ties (1e-8 to 9e-8).
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
|
The account paying for this security review has reached its Codex usage limits. The payer can check the Codex usage dashboard. For personal accounts, using credits requires enabling “Use credits for security reviews” in Code review settings. If you do not manage the paying account, contact this repository's admins. |
![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: 640a4ff6f7
ℹ️ 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".
Comment on lines
+4101
to
+4105
| if (it == pos.end()) { | ||
| pos[lb.token] = res.size(); | ||
| res.push_back(lb); | ||
| } else { | ||
| res[it->second].bias += lb.bias; |
There was a problem hiding this comment.
When a custom sampler chain places logit bias after a reordering sampler such as top-k or top-p, the CPU path at llama_sampler_logit_bias_apply stops after the first matching bias, while this backend merge sums every duplicate entry. A repeated token therefore receives different logits and can produce different sampled tokens depending on whether backend sampling is enabled; either make the reordered CPU path accumulate all matching entries or preserve its behavior here.
Useful? React with 👍 / 👎.
Comment on lines
+1175
to
+1176
| // draws of the current transaction: the output row g of the sequence in this decode uses draws[g], whatever | ||
| // ubatch the row lands in, so the draws match the CPU sampler that samples the rows one after the other |
There was a problem hiding this comment.
This explanatory sentence is split mid-sentence across fixed-width lines, a pattern repeated throughout the added sampler and allocator code. Consolidate or shorten these comments so the patch follows the repository's explicit requirement that comments usually remain one or two lines and are not hard-wrapped.
AGENTS.md reference: AGENTS.md:L77-L81
Useful? React with 👍 / 👎.
Comment on lines
+1937
to
+1939
| struct ggml_tensor * pair = ggml_concat(ctx, ggml_reshape_1d(ctx, threshold, 1), ggml_reshape_1d(ctx, kth_logit, 1), 0); | ||
| struct ggml_tensor * pick = ggml_cast(ctx, ggml_step(ctx, ggml_sub(ctx, threshold, kth_logit)), GGML_TYPE_I32); | ||
| threshold_keep = ggml_get_rows(ctx, ggml_reshape_2d(ctx, pair, 1, 2), ggml_reshape_1d(ctx, pick, 1)); |
There was a problem hiding this comment.
When fewer than min_keep candidates satisfy the min-p threshold and multiple logits tie at the min_keep boundary, the CPU fallback sorts and stops after exactly the first min_keep entries, but lowering the backend threshold to the kth logit keeps every candidate tied at that value. This changes the candidate set and can change the sampled token; forced retention needs to select the first positions rather than widening the scalar threshold across all ties.
Useful? React with 👍 / 👎.
Comment on lines
+2409
to
+2411
| // a zero temperature is greedy on the CPU: it only happens when the entropy is 0, i.e. one candidate holds | ||
| // all the probability in f32, and dividing by a tiny temperature keeps exactly that; keep -INF / inf finite | ||
| dyn_temp = ggml_clamp(ctx, dyn_temp, 1e-20f, 1e30f); |
There was a problem hiding this comment.
For the public llama_sampler_init_temp_ext API, a caller can pass temp < 0 with a positive delta; when the resulting dynamic temperature is nonpositive, the CPU path invokes llama_sampler_temp_impl and uses first-maximum greedy semantics. Clamping that value to a tiny positive number instead leaves tied maxima equiprobable, so the backend distribution can randomly choose a later tied token. The backend needs to mask to the first maximum whenever the computed dynamic temperature is nonpositive rather than merely clamping it.
Useful? React with 👍 / 👎.
…ous plan covered A plan made again because a tensor became an output covered only the current outputs, so outputs that alternate between two sets (the rows of the sequences of one decode split across ubatches) made a new plan on every alternation. When the graph has the same tensor slots as the previous plan, the tensors that plan had as outputs are marked as outputs while planning (the graph gets its flags back afterwards), so the plan grows to the union of the outputs seen and then holds. test-alloc: test_plan_keeps_union_of_outputs alternates the output between two tensors of one chain and checks that neither shares memory with a later tensor once both were outputs.
A running f32 sum over a large row (the probabilities of a 150k+ vocabulary) drifts by up to ~n * FLT_EPSILON. The backend dist sampler compares it with u * total, and on the CPU backend picked a neighbour of the CPU sampler's token (which sums in double): test_cpu_parity with --device cpu failed full_nofilter_k (46 / 48) on Qwen3-0.6B.
A decode outputs the results of the chains of the rows it samples, which can be any chain of a sequence, but the reserve ubatch has fewer rows than n_seqs * n_outputs_max_per_seq when n_ubatch is small, so it planned only some chains as outputs and the allocator planned again on the first decode that sampled another chain (an unexpected reallocation with GGML_SCHED_DEBUG_REALLOC=1 in test_multi_output_ubatch_split). The reserve graph now marks the results of every chain as outputs; the decode graphs keep marking only the sampled rows, which the reserve plan covers.
- candidates handed to the CPU samplers: a row the backend masked entirely keeps this row's candidates as they are (it kept the resized, default or previous-row entries); no dereference of a missing logits or ids array (common/sampling.cpp and llama_sampler_sample). - logit bias (CPU): the search path after a reordering sampler sums the biases of a token listed several times, as the direct path and the backend do (it stopped at the first one). - min-p: in the min_keep fallback keep exactly min_keep candidates by rank; the threshold lowered to the min_keep-th logit kept its ties. - dynamic temperature: a range that is <= 0 everywhere is greedy, as on the CPU; one candidate left keeps its logit; the CPU counts only the candidates that are not -INF for the maximum entropy, as the backend. - top-k exports 1-D logits. - penalties: clone copies the row state and the chains per sequence, reset clears the row state; dist reset clears its draws. - default member initializers for the new members (-Werror with -Wmissing-field-initializers, LLAMA_FATAL_WARNINGS). - llama.h: document that the output rows of a sequence are verification rows for the backend samplers (dist draws, penalty history). tests: test_synthetic_parity (exact logits on the GPU and CPU backends: min-p tie, negative dynamic temperature range, one candidate, masked candidate entropy, duplicate biases after top-k, 1-D top-k logits), test_all_masked_row, and the parity comparator checks both candidate lists alike (NaN, missing candidates at equal sizes).
…he lifetime bits The outputs of the previous plan were kept for any graph with the same number of nodes and leafs. A signature of the graph structure (op and type of every node and leaf, op of every source; not the sizes, which change with the number of tokens) now has to match too, so an unrelated graph never inherits outputs. No union for a plan without tensor slots (malloc(0) may return NULL). The lifetime bits get names.
…emperature, 1-D candidates - penalties: several outputs per sequence run on the backend only when the chain declares them draft-verification rows (llama_sampler_chain_params.backend_draft_rows, common_params_sampling .backend_draft_rows); otherwise the penalties stay on the CPU, as before this branch, so prompt outputs of which only some are sampled and rows without token ids keep the history the caller accepted. The server declares it for the slots that speculate. Embedding rows of a mixed batch are not read as token ids. - dynamic temperature: a computed temperature <= 0 is greedy (only the first maximum is kept) and a positive one divides the logits without a clamp, as llama_sampler_temp_impl. - top-k and greedy temperature export 1-D candidate ids and logits (a 2-D layout aborted the gather of a following top-k or logit bias); the logit bias gather checks the candidate count. - top-p (CPU): the double running sum is compared with p at float precision, as the backend's cumulative sum.
- test_synthetic_parity: zero and tiny dynamic temperatures, candidate layouts through min-p / top-k / top-k and dynamic temperature / top-k / min-p / logit bias, greedy after top-k, top-p at a rounding boundary; honours --device cpu and skips a device that lacks an op of the chain (the CPU backend must run every case). - test_multi_output_prompt_penalties: several prompt outputs, only the last sampled, without draft rows: the backend candidates equal the CPU. - the multi-output parity cases declare draft rows, as the server does. - near-ties never accept a token without probability; the split-ubatch oracle rejects NaN probabilities; the parity contexts fit their batches; main fails cleanly when the model does not load.
The previous commit passed common_sampler_init a copy of the task's
sampling params with backend_draft_rows set. common_sampler_init turns
backend_sampling off in the params it gets when the chain has a grammar
or a reasoning budget, and the server reads task.params.sampling
.backend_sampling afterwards to decide whether the slot gets a backend
sampler: with the copy it still read true, attached the backend chain
to a slot with a grammar, and common_sampler_sample aborted ("using
grammar in combination with backend sampling is not supported"). Set
the flag on the task's params as before.
test: test_backend_sampling_then_grammar_same_slot runs a backend
sampling request on a speculating slot, then a grammar request and
another plain one on the same slot.
…sked rows, merged biases on the CPU - dynamic temperature: the result row (scaled, greedy, or unchanged for one candidate or none) is selected exactly; the arithmetic gate multiplied an unused infinite temperature by 0 (NaN, and a masked token could be sampled). A row with fewer than two candidates computes its (unused) entropy on zeros, so a fully masked row stays as it is. - argmax_first: the index of ggml_argmax is clamped into the row before the gather (CUDA returns -1 for a row of -INF). - logit bias: the CPU sampler adds the merged bias of each token once, as the backend (a token listed several times rounded differently); the merged biases are computed once; backend_init checks the device like the other backend samplers. - backend_draft_rows reaches the penalty samplers of nested chains. - the sampling input collects the row tokens of every sequence in one pass over the ubatch. - ggml-alloc: lifetime_bits keeps its capacity.
…atch test test_synthetic_parity: one live candidate with a negative exponent, a fully masked row through dynamic temperature and through greedy, and duplicate biases that cancel before greedy. The split-ubatch oracle also rejects an invalid CPU reference.
… the CPU as a whole, min-p argmax bound - a sampler that sets logits to -INF and keeps the candidates on the CPU (logit bias of -INF, temperature <= 0, a dynamic temperature that can reach 0) cannot be told apart from a backend filter in the candidates handed to the CPU samplers, which drop the -INF ones; a chain with such a sampler in its backend prefix and a CPU suffix now runs on the CPU, so CPU samplers that count candidates (XTC min_keep) see the same array. - a nested chain that is not a backend unit has every sampler (and nested chain) marked CPU: its own backend prefix would otherwise be skipped on the CPU too and run nowhere (penalties lost with draft rows). - min-p clamps the argmax index into the row before the gather, as greedy. - the CPU logit bias search stops at the first match again (the merged biases list a token once).
…cter candidate comparison - test_masking_bias_cpu_suffix: -INF biases keep three tokens, XTC with min_keep 3 follows: the backend chain gives the CPU token and probabilities. - test_nested_chain_cpu_fallback: a nested chain with penalties and XTC under draft rows applies the penalties. - synthetic min-p on a fully masked row. - compare_candidates: each token once and a total of 1 on both sides, and the total probability that differs is bounded, not only each candidate.
…ered candidates at infinite temperature, double entropy), logit bias clears sorted
# Conflicts: # tools/server/tests/unit/test_speculative.py
A source is a node or a leaf of the same graph, so its flags are checked at its own entry. Checking every source slot as well re-planned a graph whose consumer switched between a computed tensor and an input leaf (llama's input embeddings built with ggml_build_forward_select: token rows at reserve, raw embeddings at decode), which aborts with GGML_SCHED_NO_REALLOC (test-llama-archs, models-check).
The zero-probability bounds took seven extra passes over the row (two step + sum pairs and their subtractions), which on a full-vocabulary chain (top_k 0) cost +0.2 ms per step at 1 stream and +0.9 ms at 4 on an L4 with Qwen3.5-0.8B. A target clamped to at least FLT_MIN already never lands on a zero-probability candidate and u < 1 keeps it at or below the total, so one comparison against the float just below the target (sub, step, sum) gives the first candidate whose running sum reaches it, as the CPU sampler.
common_sampler::set_logits copies the backend candidates every step, also when the backend already sampled the token and no CPU sampler reads them. The compaction (drop -INF) stored every candidate twice, which on a full-vocabulary chain (top_k 0, ~250k candidates) cost about 0.17 ms per sequence and step on the CPU (+0.2 ms at 1 stream, +0.8 ms at 4 on an L4). Copy them as base does when the token is sampled; the compaction for a CPU suffix stores each kept candidate once and refills an all-masked row.
danielhanchen
changed the base branch from
base/upstream-2639a3b4e
to
base/upstream-ead9f3205
October 22, 2026 21:44
…masks, unsigned min_keep clamp - a backend sampler implemented outside llama.cpp may set logits to -INF and keep the candidates: treat it as masking, so its chain with a CPU suffix runs on the CPU instead of losing candidates at the hand-off - an infinite temperature (dynamic with entropy 0 and a negative exponent, or plain) keeps -INF logits at -INF on the CPU and on the backend (-INF / INF and -INF * 0 are NaN) - top-p and min-p clamp min_keep as size_t before the signed conversion (SIZE_MAX turned into -1 and aborted in ggml_view_1d)
…ure after a -INF bias, min_keep SIZE_MAX
…efix stay on the CPU A backend sampler masks without removing when it can leave a candidate at -INF that the CPU sampler keeps, also by overflowing a finite logit (any finite input, an earlier sampler can set it): - logit bias <= -2^103 (half an ulp of FLT_MAX), -INF included - temperature < 1 (greedy at <= 0), for temp-ext the lowest it applies - penalties with repeat > 1, or last_n * freq + present >= 2^103 A chain with such a sampler in its backend prefix and a CPU suffix runs on the CPU, as before for explicit -INF biases and greedy. set_sampler does not attach a chain whose first sampler does not run on the backend: its CPU samplers get the model logits exactly as without backend sampling, instead of the exported logits with -INF candidates dropped.
…hed, all-masked row with greedy (Debug)
A -INF logit in the row can be a candidate the CPU keeps (the model itself outputs -INF for a reduced vocabulary, eagle3 / dflash) or one a CPU filter removes, so dropping every -INF at the hand-off changed CPU samplers that count candidates. The hand-off now copies the row unchanged, and a chain with a CPU suffix keeps a backend prefix only when the prefix leaves the CPU's candidates exactly: top-p and min-p (which write -INF where the CPU removes) join the inexact samplers that send such a chain to the CPU; top-k shrinks the row as the CPU does. Fully backend chains are unchanged (the token is sampled on the backend and the row is only copied).
…t draft rows, repeated biases on the CPU, no reserve for inactive chains - backend dynamic temperature sorts the row only when it has two live candidates or more, as the CPU sampler, which returns before sorting otherwise (a later greedy picked another token on a fully masked row) - a cloned penalties sampler does not copy backend_draft_rows: the opt-in belongs to the chain the clone is added to - a logit bias listing a token more than once stays on the CPU, which adds the biases one at a time as base (their float sum can overflow where the sequence does not); lists without repeats run on the backend as before - chains at or past n_ubatch can never be active: their results are not reserved as outputs
…natemp + greedy, repeated logit biases, penalties clone
Leaves are allocated before any node whatever their flags, so the input flag does not change the plan of a leaf; only becoming an output does. A reserve of a graph whose leaf is set as input only later no longer makes a new plan (GGML_SCHED_NO_REALLOC aborted there).
…atch, every repeated bias applies - llama_sampler_chain_params is unchanged from upstream: the draft rows opt-in is llama_sampler_chain_set_backend_draft_rows, kept by clones - the context builds min(n_outputs_max_per_seq, n_ubatch) chains per sampler and the penalties and dist inputs have that many columns; the validation and the per-row draws keep n_outputs_max_per_seq - the CPU logit bias adds every entry of a repeated token also when the candidates are not in vocabulary order - tests: reordered repeated biases, chain input columns at a small n_ubatch, the synthetic parity requires both sides to select, the candidate ids are checked against n_vocab, the all-masked row comes from the model output after an exact backend prefix
… remove drops the draft rows opt-in - one physical chain (n_ubatch 1) no longer means one row: the rows of a decode come one ubatch at a time and each gets the draft inputs of the rows before it, as with several chains - a sampler taken out of a chain with llama_sampler_chain_remove does not keep the chain's draft rows opt-in, as for a clone - tests: chain columns at n_ubatch 1 and 3 against one ubatch, penalties removed from an opted-in chain stay on the CPU, the multi-sequence dist reference sums in double as the samplers do
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
The backend (GPU) samplers behind
--backend-samplingnow pick the same token as the CPU samplers for the same logits and seed, and penalties run on the GPU when a step verifies several draft tokens per sequence (speculative decoding). On base b11642 (2639a3b) several samplers crash the server, drop every token, or diverge from the CPU, which is why Studio has to keep the CPU sampler for many chains. Not added toscripts/unsloth/pr-set.jsonin this PR.Changes
ggml-alloc.c(85d7920):ggml_gallocr_alloc_graph()reused its last plan whenever the graph had the same nodes, without checking which tensors are inputs or outputs. The sampling graph marks only the active chains' results as outputs, so when the rows a step samples change (draft rows split across ubatches, several sequences in a ubatch), a later tensor could reuse an active chain's output memory and overwrite its sampled token or probabilities. The plan now records the input / output / view-of-output flags and is recomputed only when a tensor gains one of them; losing a flag only means a tensor is allocated earlier or freed later, which is safe, so the plan settles once it covers the largest output set.llama-sampler.cpp,llama-graph.cpp,common/sampling.cpp(28dd780), each to the CPU sampler's semantics:ggml_top_k), and sort the k results by logit (theggml_top_korder is unspecified, so a fixed seed neither reproduced nor matched the CPU);logit >= max + log(p)(p = 1dropped every token); min-p and top-p honourmin_keep;test_penaltiesfailed on base);test-backend-samplerruns 31 whole chains (Studio defaults and every edge value above) throughcommon_sampleron the backend and on the CPU with identical logits and seeds, single draws and draft-verification rows at n_ubatch 512 and 8, and a split-ubatch case;test-allocadds an output-flag case. In the unfiltered full-vocabulary cases a mismatch is accepted only as a near-tie of the running sum: u withinsqrt(n) * FLT_EPSILON * totalof the CPU boundary of both tokens and both probabilities equal within that bound, counted and printed.Review fixes (0c71fe1, 30d3ac1, 60c4fe2, c165be6)
Each fix below has a test that fails before it and passes after it, unless noted.
common/sampling.cppset_logits,llama_sampler_sample) kept default-initialised entries; a fully masked row now keeps its own candidates at -INF, as on base. The logits/ids null checks are added too.min_keepor beyond are dropped, so a tie at the cut no longer keeps extra candidates.llama_sampler_distfields and the test structs; a CPULLAMA_FATAL_WARNINGS=ONbuild of the library, tests and server passes.llama_sampler_penalties_clonecopies the per-row state and chain count, and reset clears it; dist reset clears its draws. The row rule itself is unchanged: the output rows of a sequence are treated as draft-verification rows, the same rule the base dist sampler uses for its draws, now documented next ton_outputs_max_per_seqininclude/llama.h.ggml-cpucumsum accumulates in double (ggml_vec_cumsum_f32): the f32 sum moved CPU-backend draws over a full vocabulary (test_cpu_parity --device cpuon Qwen3-0.6B: 46/48 before, 1488/1488 after).After the fixes:
test-backend-sampler26/26 on Qwen3-0.6B and Qwen3.5-0.8B, 25/26 on gemma-4-E2B (test_logit_bias, same on base); parity 1488 single and 5952 multi draws per model, 0 failing (2 tolerated near-ties on Qwen3.5); the same parity tests pass on the CPU backend; a broken dist still fails them on all three models;test-alloc34/34;test-backend-opson CUDA0 all pass (MUL_MAT + MUL_MAT_ID 3380, FLASH_ATTN_EXT 3993, the other 123 ops 8464); server 4800/4800 draws and 112/112 generations identical with and without--spec-default. Timing was not rerun after these fixes.Review round 2 (a606099, 992d747, 4df4462)
llama_sampler_chain_params.backend_draft_rows/common_params_sampling.backend_draft_rows. Without it, penalties withn_outputs_max_per_seq > 1stay on the CPU as on base (ordinary multi-output prompt sampling, embedding rows). The server sets it only for slots that speculate (slot.can_speculate()), so draft verification keeps backend penalties. Embedding rows in a mixed batch are no longer read as token ids.test_multi_output_prompt_penaltiesdiffered before (candidate set diff 1, probability diff 2.9e-2) and matches after.ggml_nelementsand exports 1-D candidates; temperature sampling's greedy path is 1-D; the logit-bias gather checks the candidate count (min_p_top_k_top_kaborted before).--device cpuhonoured, near-ties never accept a zero-probability token, NaN rejected, clean exit on a failed model load.After round 2:
test-backend-sampler27/27 on Qwen3 and Qwen3.5, 26/27 on gemma-4 (test_logit_bias, same on base); parity 1488 + 5952 draws per model, 0 failing; synthetic parity 0 failing on CUDA and CPU;test-alloc34/34;test-backend-opsCUDA all pass; server 4800/4800 draws and 112/112 generations identical with and without--spec-default;LLAMA_FATAL_WARNINGS=ONbuild passes. On the CPU backend, Qwen3.5 parity fails because the plain CPU reference context gives NaN logits above about 4 sequences (same with baseggml-alloc.candvec.h, so not this PR).Review round 3 and a CI fix (fb5fc0d, e3aa756, c649dae)
common/sampling.cpp:639): round 2 setbackend_draft_rowson a copy of the request's sampling params, so the grammar check disabled backend sampling in the copy while the server still attached the backend chain. The flag is now set on the request's own params. Newtest_backend_sampling_then_grammar_same_slotaborted before and passes;test_synth_ignores_target_tokensandtest_completion_with_response_formatpass withLLAMA_ARG_BACKEND_SAMPLING=1.dynatemp_one_cand_neg_exp,dynatemp_all_maskedfailed before).backend_draft_rowsreaches penalty samplers in nested chains; the logit-bias sampler runs the same device-support check as the others; the argmax index is clamped before the gather (CUDA returns -1 for an all -INF row).logit_bias_dup_canceldiffered before); the split-ubatch check also rejects an invalid CPU reference; row-token discovery is one pass over the ubatch;ggml-allockeeps its lifetime-bit buffer capacity.After round 3:
test-backend-sampler27/27 on Qwen3 and Qwen3.5, 26/27 on gemma-4 (test_logit_bias, as on base); parity 0 failing on all three; 16 synthetic cases on CUDA and CPU; broken dist still fails;test-alloc34/34;test-backend-opsCUDA all pass; server 4800/4800 draws and 112/112 generations identical with and without--spec-default;-Werrorbuild passes. A CPU server suite run on the build host (load 580-1550) had only timeouts and one pre-existing failure (test_systemone_shared_prompt, which fails the same way on base there).Review round 4 (167e12c, 5e8b7b4)
test_nested_chain_cpu_fallback: before, the backend picked the penalised top token).test_masking_bias_cpu_suffix(-INF biases keeping three tokens, then XTC withmin_keep3) differed before (probability difference 0.459) and matches now.After round 4:
test-backend-sampler29/29 on Qwen3 and Qwen3.5, 28/29 on gemma-4 (test_logit_bias, as on base); parity 0 failing; broken dist still fails;test-alloc34/34;test-backend-opsCUDA all pass; server 4800/4800 draws and 112/112 generations identical with and without--spec-default;-Werrorbuild passes; server sampling tests (test_completion.py,test_chat_completion.py,test_speculative.py) in both CI modes pass apart from a load timeout that passes alone and the slow Phi-3.5 test whose model did not download here.Review round 5 and audit (afa2856 .. b8fea57)
sorted, so a later top-p does not trust a stale order. Four synthetic cases fail before (8 of 8 on CUDA0 and CPU) and pass after;PARITY_SEEDS=1no longer builds an empty batch.test_speculative.pynext to the new test): upstream's test is kept, the new one follows it. The pin-set sim (master pr-set plus this PR at the new head, aged and newest tags: b11804 and b11815 at the merge, b11819 and b11827 at b8fea57) merges this PR clean on every tag; the refusals left are other pins (Model arch: Flint #314, ggml-org#25731), as on master.GGML_SCHED_NO_REALLOC=ON: the plan check compared the flags of every source slot, so a consumer switching between a computed tensor and an input leaf (the input embeddings built withggml_build_forward_select: token rows at reserve, raw embeddings at decode) re-planned and abortedtest-llama-archson the first architecture (CPU and CUDA). A source is a node or a leaf of the same graph, so the check now uses only node and leaf flags.test_plan_holds_for_selected_sourcefails before and passes after;test-llama-archs -s 1with NO_REALLOC passes 139/139 on the CPU, as base.top_k 0, Studio's gpt-oss / llama-3 / deepseek / mistral / glm / phi defaults) were slower than base on an L4: +0.2 ms per step at 1 stream and +0.9 ms at 4 on Qwen3.5-0.8B. Two causes, both fixed: the dist zero-probability bounds added seven passes over the row (now one comparison against the float just below a target floored at FLT_MIN, the passes of base), andcommon_sampler::set_logitscompacted the ~250k candidates every step even after the backend sampled the token, storing each twice (about 0.17 ms per sequence and step on the CPU). It now copies them as base does when the token is already sampled; the compaction for a CPU suffix stores each kept candidate once.--device cpuparity fails again on Qwen3-0.6B (1 single-output and 3 multi-output cases).Review round 6 (b8fea57..805fa03)
test_custom_masking_cpu_suffix(custom sampler keeping tokens 10-12, then XTC, then dist, 8 seeds): 5 of 8 differ before, 8 of 8 match afterdynatemp_inf_biasandtemp_inf_bias: fail before, pass on CUDA and CPU aftermin_keepwas converted to a signed dimension before clamping, soSIZE_MAXaborted inggml_view_1dsize_tfirst (min-p and top-p)min_p(1.1, SIZE_MAX): aborted before, keeps all candidates after, as the CPU doesOn 805fa03:
test-backend-samplerQwen3 30/30, Qwen3.5 30/30, gemma-4 29/30 (test_logit_bias, which also fails on base); GPU vs CPU parity 0 failing;--device cpuparity passes; a deliberately broken dist still fails every model; server parity 4800/4800 draws and 112/112 generations; -Werror CPU build passes; CPU server sampling tests 104 passed in both modes.Review round 7 (805fa03..bf75187)
test_masking_classes_cpu_suffix(two -FLT_MAX biases, temperature 1e-37, FLT_MAX frequency penalty, each before XTC): 3 of 6 seeds differ per class before, 0 afterset_samplerattaches only a chain whose first sampler runs on the backend; such chains (and those forced to the CPU above) get the model logits exactly as without backend samplingtest_cpu_only_chain_not_exported(XTC then dist): exported before, not afterassert(found)The masking rule is conservative: with a CPU-only sampler (XTC, DRY, typical, ...) after it, a temperature below 1 or a repeat penalty above 1 now sends the whole chain to the CPU. Chains that run fully on the backend are unaffected, and the server's Studio chains (presence 1.5, repeat 1.0) still run penalties on the backend.
On bf75187:
test-backend-samplerQwen3 32/32, Qwen3.5 32/32, gemma-4 31/32 (test_logit_bias, which also fails on base); GPU vs CPU parity 0 failing;--device cpuparity passes; a deliberately broken dist fails every model; server parity 4800/4800 draws and 112/112 generations; -Werror CPU build passes; CPU server sampling tests 104 passed in both modes.Review round 8 (bf75187..77e244e)
test_model_inf_logits_cpu_suffix(model output forced to 0.3 / 0.2 / -1 on three tokens and -INF elsewhere, temperature(1.5) then XTC(1, 0.2, min_keep 3) then dist, 8 seeds): 5 of 8 differ before, 0 afterFully-backend chains are unchanged (they have only copied the row since b8fea57, so the L4 numbers above still apply; not re-timed). Chains with top-p or min-p before a CPU-only sampler (XTC, DRY, typical, ...) now sample on the CPU at base speed.
On 77e244e:
test-backend-samplerQwen3 33/33, Qwen3.5 33/33, gemma-4 32/33 (test_logit_bias, which also fails on base); GPU vs CPU parity 0 failing;--device cpuparity passes; a deliberately broken dist fails every model; server parity 4800/4800 draws and 112/112 generations; -Werror CPU build passes; CPU server sampling tests (GGML_SCHED_NO_REALLOCbuild) 104 passed in both modes.Review round 9 (77e244e..5b14d53)
dynatemp_all_masked_greedy(32 -INF logits,temp_ext(0.5, 0.5, 1), greedy): 0 vs 16 on the CPU backend before, equal after (CUDA0 already matched)draft_rows; the chain it is added to decides (cloning a whole opted-in chain still offloads, the chain reapplies the opt-in)test_penalties_clone_draft_rows: penalties cloned into a default chain with 4 outputs per sequence was offloaded before, on the CPU aftertest_logit_bias_duplicates(logits -3e38 / -2e38 / -1e38, bias 2e38 on token 0 twice): +INF before, about 1e38 after, not offloaded; the two synthetic repeated-bias cases are removed (that sampler no longer runs on the backend)On 5b14d53:
test-backend-samplerQwen3 35/35, Qwen3.5 35/35, gemma-4 34/35 (test_logit_bias, which also fails on base); GPU vs CPU parity 0 failing (one tolerated near-tie each on Qwen3 and Qwen3.5);--device cpuparity passes; a deliberately broken dist fails every model (25 single-output and 50 multi-output cases); server parity 4800/4800 draws and 112/112 generations; -Werror CPU build passes; CPU server sampling tests (GGML_SCHED_NO_REALLOCbuild) 104 passed in both modes (the Phi-3.5 test whose model cannot load here excluded, as in every round).Review round 10 (5b14d53..230c8b5)
Round 10 on 5b14d53: 8 APPROVE, 4 REQUEST_CHANGES.
test_logit_bias_duplicatesgains a reordered case (temp_ext(1, 0.1, 0), biases {1: +5, 1: -6}, greedy on [2, 1, 3]): token 1 before, 2 aftertest_chain_columns_n_ubatch(64 outputs per sequence, n_ubatch 8) reads the column count from the graph: 64 before, 8 after; rows sampled at n_ubatch 8 match those at 512llama_sampler_chain_params, which is passed by value, so a client built against the old header could turn it on with garbage register bitsllama_sampler_chain_set_backend_draft_rows(chain, bool), called before the chain is attached (asserted) and copied by a chain clone;common_sampler_initand the tests use itggml_gallocr_needs_realloctreated a leaf that only became an input as a changed plan, so a scheduler withGGML_SCHED_NO_REALLOCaborted; leaves are allocated before every node whatever their flagstest_plan_holds_when_leaf_becomes_input: abort before, pass after; test-alloc 36/36 in the CUDA build and aGGML_SCHED_NO_REALLOC=ONCPU buildgreedy_finitecase); ids >= n_vocab are rejected; the all-masked row comes from the model output through an eval callback, with a backend temperature before XTC, and both rows are asserted to export backend logitsOn 230c8b5:
test-backend-samplerQwen3 36/36, Qwen3.5 36/36, gemma-4 35/36 (test_logit_bias, which also fails on base); GPU vs CPU parity 0 failing on all 3 models;--device cpuparity 0 failing (2 seeds); a deliberately broken dist fails every model (25 single-output and 50 multi-output cases); server parity 4800/4800 draws and 112/112 generations; -Werror CPU build passes; CPU server sampling tests (GGML_SCHED_NO_REALLOCbuild) 104 passed in both modes (the Phi-3.5 test whose model cannot load here excluded).Review round 11 (230c8b5..3817569)
Round 11 on 230c8b5 (4 reviewers): 4 REQUEST_CHANGES.
n_chains == 1: later draft rows of a sequence lost the draft inputs of the rows before them (presence penalty 10 on the next draft token: base and CPU 43, backend 42)test_chain_columns_n_ubatchalso runs n_ubatch 1 and 3 against one ubatch on fixed model logits: 7 of 14 rows differed at n_ubatch 1 before, 0 rows at 1, 3 and 8 on all 3 models afterllama_sampler_chain_remove()handed a penalties sampler over with the draft-row opt-in of the chain it came from (round 9 made clone drop it; remove was missed)test_penalties_clone_draft_rows(penalties removed from an opted-in [XTC, penalties] chain into a default chain with 4 outputs per sequence): offloaded before, on the CPU aftertest_multi_output_multi_sequence_dist's reference summed in float while the kernel sums in double, so it aborted on the CPU backend with stories15M although the token was rightOn 3817569 (one GPU job per card, each under its own timeout and teardown check):
test-backend-samplerQwen3 36/36, Qwen3.5 36/36, gemma-4 35/36 (test_logit_bias, which also fails on base); GPU vs CPU parity 0 failing on all 3 models;--device cpuparity 0 failing (2 seeds); a deliberately broken dist fails every model (25 single-output and 50 multi-output cases); server parity 4800/4800 draws and 112/112 generations. The -Werror CPU build and the CPUGGML_SCHED_NO_REALLOCserver tests were not rerun this round.Timing and memory at b8fea57 (Colab L4, one session, 3 interleaved rounds, base = b11776, -np 4, ms)
--spec-defaultBackend sampling without speculation (what Studio turns on) is within the head A/A spread of base on every model (largest delta +0.06 ms, Qwen3-0.6B at 4 streams). The reserve that plans every chain's result costs 2-8 MiB of peak VRAM with
--spec-default -np 4on a 248k vocabulary, not hundreds of MB. With speculation on the full-vocabulary chain the 4-stream TTFT is about +12 ms (248 vs 260, the same on the A/A arm).Tests (CUDA, B200 unless noted)
test-alloc: base aborts at!memory_overlap, head 33/33.test-backend-opson CUDA0: MUL_MAT + MUL_MAT_ID 3380/3380, FLASH_ATTN_EXT 3993/3993, the other 125 ops 8464/8464.test-backend-samplerparity on Qwen3-0.6B, Qwen3.5-0.8B and gemma-4-E2B: single-output 1488 draws per model, 0 failing; multi-output 0 failing, with 2 tolerated near-ties on Qwen3.5 (distances 9.8e-9 and 8.6e-8, probability differences below 1e-13) and 0 elsewhere. The upstream suite passes on all three excepttest_logit_biason gemma-4, which fails the same way on base (it assumes a +10 bias wins).--spec-default: 4800/4800 first-token draws and 112/112 generations identical between backend and CPU sampling; the server log shows+penaltieson the backend under speculation (base:-penalties). Every crash and divergence repro passes; on base the three top-k cases crash the server.git merge-treeis clean against ggml-org master (9b56ccf) and against the composed Pin unslothai#269 (server-side preemption and exact concurrency) at 18365c117 #211 pin-set tree.Timing (Colab L4, 3 interleaved rounds, 4 streams, ms per step)
Without speculation backend sampling wins on all three. With speculation it pays for chains with penalties (Qwen3.5, -34%) and costs 14-19% on Qwen3 and gemma-4, most likely from the per-slot sampler chains built under speculation (not profiled).
Not run
llama_pinset.py verifyof a pin set containing this PR (the merge sim is in the round 5 section); CUDAGGML_SCHED_NO_REALLOCtest-llama-archsand server suite (host congestion); multi-GPU legs.