Commit 3642238cc for llama.cpp
commit 3642238cc12050c8ab70afdf4b3121ea365c7579
Author: Daniel Bevenius <daniel.bevenius@gmail.com>
Date: Sat Oct 24 11:47:49 2026 +0200
tests : add speculative rejection draft sampling test (#30213)
* tests : add speculative rejection draft sampling test
This commit adds a test for speculative rejection sampling test that
excercises the speculative rejection sampling mechanism in
common/speculative.cpp.
The motivation for this is have a test that verifies that the
speculative draft models emitts tokens that match the target models
distribution. So when rejection sampling occurs it should still
produce the same probability distribution as if we directly sampled
from the target model.
Running the test:
```console
$ ctest --test-dir build -R "^test-speculative-rejection-sampling$" \
--n_trials 5000 -V
```
This will generate two models, one target and one draft model.
For subsequent runs use the "fixture exclude any" option to avoid
recreating the models:
```console
$ ctest --test-dir build-cpu-debug -R \
"^test-speculative-rejection-sampling$" \
-FA "^speculative-models$" --n_trails 5000 -V
```
* remove draft model
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index 763154972..36703127a 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -240,6 +240,26 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
FIXTURES_SETUP generate-models
)
+ # Target model for test-speculative-rejection-sampling
+ set(SPEC_TARGET_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-speculative-models/target")
+ file(MAKE_DIRECTORY "${SPEC_TARGET_DIR}")
+ llama_test(
+ test-llama-archs
+ NAME test-generate-speculative-target
+ LABEL main
+ ARGS -a llama -s 999 -o "${SPEC_TARGET_DIR}"
+ )
+ set_tests_properties(
+ test-generate-speculative-target
+ PROPERTIES FIXTURES_SETUP speculative-model
+ )
+
+ # To avoid recreating the model use the fixture exclude any option:
+ # -FA "^speculative-model$"
+ llama_build_and_test(test-speculative-rejection-sampling.cpp ARGS
+ -m "${SPEC_TARGET_DIR}/llama-dense.gguf")
+ set_tests_properties(test-speculative-rejection-sampling PROPERTIES FIXTURES_REQUIRED speculative-model)
+
# Test context state handling across all architectures, using the generated dummy models
llama_test(
test-llama-context
diff --git a/tests/test-speculative-rejection-sampling.cpp b/tests/test-speculative-rejection-sampling.cpp
new file mode 100644
index 000000000..523436028
--- /dev/null
+++ b/tests/test-speculative-rejection-sampling.cpp
@@ -0,0 +1,229 @@
+#include "sampling.h"
+#include "llama.h"
+#include "common.h"
+#include "arg.h"
+#include "log.h"
+
+#include <cmath>
+#include <cstdlib>
+#include <cstring>
+#include <optional>
+#include <random>
+#include <unordered_map>
+
+// Tests that common_sampler_sample_and_accept_n_rejection preserves the target
+// distribution regardless of what the draft distribution Q looks like.
+
+static int test_speculative_rejection_sampling(common_params & params, int n_trials) {
+ params.sampling.backend_sampling = false;
+ auto llama_init = common_init_from_params(params);
+
+ auto * model_tgt = llama_init->model();
+ auto * ctx_tgt = llama_init->context();
+ if (!model_tgt || !ctx_tgt) {
+ LOG_ERR("failed to initialize target model and context\n");
+ return 1;
+ }
+
+ const llama_tokens prompt = { 1, 2, 3, 1, 2 };
+
+ // Process the prompt with the target model to get its true distribution for
+ // the position right after the prompt.
+ common_batch batch = common_batch_get_one(ctx_tgt, prompt);
+ if (llama_process(ctx_tgt, LLAMA_PROCESS_TYPE_DECODE, batch.get()) != 0) {
+ LOG_ERR("failed to decode the test prompt\n");
+ return 1;
+ }
+
+ common_params_sampling sampling;
+ sampling.seed = 18;
+ sampling.samplers = { COMMON_SAMPLER_TYPE_TEMPERATURE };
+ sampling.temp = 1.0f;
+
+ // Initialize the sampler for the target model.
+ common_sampler_ptr smpl(common_sampler_init(model_tgt, sampling));
+ if (!smpl) {
+ LOG_ERR("failed to initialize target sampler\n");
+ return 1;
+ }
+
+ // Sample once so the sampler chain runs and populates cur_p. We don't need
+ // the returned token id itself so the return value is ignored.
+ common_sampler_sample(smpl.get(), ctx_tgt, -1);
+
+ // Get the target candidates distribution in sorted order.
+ const llama_token_data_array * dist_tgt = common_sampler_get_candidates(smpl.get(), true);
+
+ // Clone smpl before the trial loop starts. We need the dist_tgt pointer to stay valid.
+ common_sampler_ptr smpl_verify(common_sampler_clone(smpl.get()));
+ if (!smpl_verify) {
+ LOG_ERR("failed to clone target sampler\n");
+ return 1;
+ }
+
+ const llama_vocab * vocab = llama_model_get_vocab(model_tgt);
+ const int n_vocab = llama_vocab_n_tokens(vocab);
+
+ std::mt19937 rng(18);
+ std::uniform_real_distribution<float> uni(0.0f, 1.0f);
+
+ std::unordered_map<llama_token, int> token_counts;
+
+ std::vector<llama_token_data> q(n_vocab);
+ for (int i = 0; i < n_vocab; ++i) {
+ q[i].id = i;
+ q[i].logit = 0.0f;
+ }
+
+ for (int trial = 0; trial < n_trials; ++trial) {
+ float q_sum = 0.0f;
+
+ for (auto & td : q) {
+ td.p = uni(rng);
+ q_sum += td.p;
+ }
+
+ const float inv_sum = 1.0f / q_sum;
+ for (auto & td : q) {
+ td.p *= inv_sum;
+ }
+
+ float u = uni(rng);
+
+ // Fallback to the last bucket to absorb residual probability from FP precision loss
+ llama_token draft_tok = q.back().id;
+ for (const auto & e : q) {
+ u -= e.p;
+ if (u <= 0.0f) {
+ draft_tok = e.id;
+ break;
+ }
+ }
+
+ // Both indices reuse the prompt's last output logits. The second entry
+ // is the bonus draw slot; it is always stale and its result is ignored.
+ const llama_tokens draft = { draft_tok };
+ const std::vector<std::vector<llama_token_data>> draft_q = { q };
+ const std::vector<int> idxs = { -1, -1 };
+
+ const std::vector<llama_token> result =
+ common_sampler_sample_and_accept_n_rejection(smpl_verify.get(), ctx_tgt, idxs, draft, draft_q, false);
+
+ token_counts[result[0]]++;
+ }
+
+ LOG("\n");
+ LOG("Empirical vs target distribution after %d trials:\n", n_trials);
+
+ bool ok = true;
+ size_t covered = 0;
+
+ auto get_count_for = [&](llama_token id) -> int {
+ const auto it = token_counts.find(id);
+ return it != token_counts.end() ? it->second : 0;
+ };
+
+ // Iterate over the target distribution
+ for (size_t i = 0; i < dist_tgt->size; ++i) {
+ const llama_token id = dist_tgt->data[i].id;
+ const float p_t = dist_tgt->data[i].p;
+
+ const int count = get_count_for(id);
+ const float p_e = (float) count / n_trials;
+ const float diff = std::fabs(p_t - p_e);
+
+ // We only deal with ids in dist_tgt, and if each trial result was one
+ // of dist_tgt's candidates then covered ends up equal to the number of
+ // trials. If not this inconsistency tells us how many trials produced
+ // a token outside of the target's distribution.
+ covered += count;
+
+ const float standard_error = std::sqrt(p_t * (1.0f - p_t) / n_trials);
+ const float tolerance = 6.0f * standard_error;
+
+ LOG("token %3d : target = %f, empirical = %f, diff = %f, tol = %f\n", id, p_t, p_e, diff, tolerance);
+
+ if (diff > tolerance) {
+ LOG_ERR("token %d: empirical probability %f too far from target %f (diff %f > tol %f)\n",
+ id, p_e, p_t, diff, tolerance);
+ ok = false;
+ }
+ }
+
+ if (covered != (size_t) n_trials) {
+ LOG_ERR("%zu of %d trials produced a token outside dist_tgt's candidates\n",
+ (size_t) n_trials - covered, n_trials);
+ ok = false;
+ }
+
+ if (!ok) {
+ return 1;
+ }
+
+ return 0;
+}
+
+struct parsed_args {
+ int n_trials;
+ std::vector<char *> argv;
+};
+
+static std::optional<parsed_args> parse_n_trials_arg(int argc, char ** argv, int default_n_trials) {
+ parsed_args result;
+ result.n_trials = default_n_trials;
+ result.argv.push_back(argv[0]);
+
+ for (int i = 1; i < argc; ++i) {
+ if (strcmp(argv[i], "--n-trials") == 0) {
+ if (i + 1 >= argc) {
+ LOG_ERR("%s: --n-trials requires a number argument\n", __func__);
+ return std::nullopt;
+ }
+ int value = std::atoi(argv[i + 1]);
+ if (value <= 0) {
+ LOG_ERR("%s: --n-trials requires a positive number argument\n", __func__);
+ return std::nullopt;
+ }
+ result.n_trials = value;
+ i++;
+ } else {
+ result.argv.push_back(argv[i]);
+ }
+ }
+ result.argv.push_back(nullptr);
+
+ return result;
+}
+
+int main(int argc, char ** argv) {
+ std::setlocale(LC_NUMERIC, "C");
+
+ common_params params;
+ params.n_ctx = 256;
+
+ common_init();
+
+ auto parsed = parse_n_trials_arg(argc, argv, 5000);
+ if (!parsed) {
+ return 1;
+ }
+ const int fargc = (int) parsed->argv.size() - 1;
+
+ if (!common_params_parse(fargc, parsed->argv.data(), params, LLAMA_EXAMPLE_COMMON)) {
+ return 1;
+ }
+
+ LOG("Using target model '%s'\n", params.model.path.c_str());
+
+ llama_backend_init();
+
+ const int result = test_speculative_rejection_sampling(params, parsed->n_trials);
+ if (result == 0) {
+ LOG("\nAll tests passed.\n");
+ } else {
+ LOG("\nSome tests failed.\n");
+ }
+
+ llama_backend_free();
+ return result;
+}