459 lines
15 KiB
C
459 lines
15 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/* Differential parity test for the optimized token sampler.
|
|
*
|
|
* (2026-10-07: both copies take the llama.cpp/HF order — top-k renormalised before top-p.)
|
|
* This test keeps a VERBATIM copy of the ORIGINAL sample_token implementation
|
|
* (sample_token_ref, the pre-optimization slow path) and a copy of the new
|
|
* OPTIMIZED implementation (sample_token_opt, kept textually in sync with
|
|
* src/core/engine.c). For thousands of random logit vectors and random
|
|
* parameter combinations (temperature including 0 for greedy, top_k, top_p,
|
|
* repeat_penalty, and random recent-token windows), it feeds BOTH samplers the
|
|
* exact same seeded RNG state and asserts they return the identical token id.
|
|
*
|
|
* The two samplers are self-contained (no engine/model dependencies), so this
|
|
* builds and runs anywhere with just a C compiler and libm.
|
|
*
|
|
* Determinism: each sampler calls the global rand() at most once (the final
|
|
* draw). Before each call we srand() to the same per-trial seed, so both draws
|
|
* see the identical random value r. Test data is produced by a private xorshift
|
|
* PRNG so that reseeding rand() never perturbs the generated logits/params.
|
|
*/
|
|
|
|
#include <stdint.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <math.h>
|
|
|
|
/* Mirror of the real sampling params struct (see include/pantheon.h). */
|
|
typedef struct {
|
|
float temperature;
|
|
float top_p;
|
|
int top_k;
|
|
float repeat_penalty;
|
|
int repeat_last_n;
|
|
} pn_sampling_params;
|
|
|
|
static float rand_float(void) {
|
|
return (float)rand() / ((float)RAND_MAX + 1.0f);
|
|
}
|
|
|
|
/* ==================================================================
|
|
* REFERENCE: verbatim copy of the ORIGINAL sample_token (slow path).
|
|
* ================================================================== */
|
|
static uint32_t sample_token_ref(float *logits, int n_vocab, pn_sampling_params *params,
|
|
uint32_t *recent_tokens, int n_recent) {
|
|
if (params->repeat_penalty != 1.0f && recent_tokens && n_recent > 0) {
|
|
for (int i = 0; i < n_recent; i++) {
|
|
uint32_t tok = recent_tokens[i];
|
|
if (tok < (uint32_t)n_vocab) {
|
|
if (logits[tok] > 0.0f)
|
|
logits[tok] /= params->repeat_penalty;
|
|
else
|
|
logits[tok] *= params->repeat_penalty;
|
|
}
|
|
}
|
|
}
|
|
|
|
if (params->temperature <= 0.0f) {
|
|
int best = 0;
|
|
for (int i = 1; i < n_vocab; i++) {
|
|
if (logits[i] > logits[best]) best = i;
|
|
}
|
|
return (uint32_t)best;
|
|
}
|
|
|
|
float inv_temp = 1.0f / params->temperature;
|
|
for (int i = 0; i < n_vocab; i++) {
|
|
logits[i] *= inv_temp;
|
|
}
|
|
|
|
int top_k = params->top_k;
|
|
if (top_k <= 0 || top_k > n_vocab) top_k = n_vocab;
|
|
|
|
float max_logit = logits[0];
|
|
for (int i = 1; i < n_vocab; i++) {
|
|
if (logits[i] > max_logit) max_logit = logits[i];
|
|
}
|
|
|
|
float *probs = malloc(n_vocab * sizeof(float));
|
|
if (!probs) return 0;
|
|
|
|
float sum = 0.0f;
|
|
for (int i = 0; i < n_vocab; i++) {
|
|
probs[i] = expf(logits[i] - max_logit);
|
|
sum += probs[i];
|
|
}
|
|
float inv_sum = 1.0f / sum;
|
|
for (int i = 0; i < n_vocab; i++) {
|
|
probs[i] *= inv_sum;
|
|
}
|
|
|
|
if (top_k < n_vocab) {
|
|
float *sorted = malloc(n_vocab * sizeof(float));
|
|
memcpy(sorted, probs, n_vocab * sizeof(float));
|
|
|
|
for (int k = 0; k < top_k; k++) {
|
|
float best_val = sorted[k];
|
|
int best_idx = k;
|
|
for (int j = k + 1; j < n_vocab; j++) {
|
|
if (sorted[j] > best_val) {
|
|
best_val = sorted[j];
|
|
best_idx = j;
|
|
}
|
|
}
|
|
float tmp = sorted[k];
|
|
sorted[k] = sorted[best_idx];
|
|
sorted[best_idx] = tmp;
|
|
}
|
|
float threshold = sorted[top_k - 1];
|
|
free(sorted);
|
|
|
|
for (int i = 0; i < n_vocab; i++) {
|
|
if (probs[i] < threshold) probs[i] = 0.0f;
|
|
}
|
|
/* llama.cpp / HF order (2026-10-07): renormalise within the top-k set. */
|
|
sum = 0.0f;
|
|
for (int i = 0; i < n_vocab; i++) sum += probs[i];
|
|
if (sum > 0.0f) {
|
|
inv_sum = 1.0f / sum;
|
|
for (int i = 0; i < n_vocab; i++) probs[i] *= inv_sum;
|
|
}
|
|
}
|
|
|
|
if (params->top_p < 1.0f) {
|
|
float cum_prob = 0.0f;
|
|
float top_p_threshold = 0.0f;
|
|
|
|
float *probs_copy = malloc(n_vocab * sizeof(float));
|
|
memcpy(probs_copy, probs, n_vocab * sizeof(float));
|
|
|
|
while (cum_prob < params->top_p) {
|
|
int best = 0;
|
|
for (int i = 1; i < n_vocab; i++) {
|
|
if (probs_copy[i] > probs_copy[best]) best = i;
|
|
}
|
|
if (probs_copy[best] <= 0.0f) break;
|
|
cum_prob += probs_copy[best];
|
|
top_p_threshold = probs_copy[best];
|
|
probs_copy[best] = 0.0f;
|
|
}
|
|
free(probs_copy);
|
|
|
|
for (int i = 0; i < n_vocab; i++) {
|
|
if (probs[i] < top_p_threshold) probs[i] = 0.0f;
|
|
}
|
|
}
|
|
|
|
sum = 0.0f;
|
|
for (int i = 0; i < n_vocab; i++) sum += probs[i];
|
|
if (sum > 0.0f) {
|
|
inv_sum = 1.0f / sum;
|
|
for (int i = 0; i < n_vocab; i++) probs[i] *= inv_sum;
|
|
}
|
|
|
|
float r = rand_float();
|
|
float cum = 0.0f;
|
|
for (int i = 0; i < n_vocab; i++) {
|
|
cum += probs[i];
|
|
if (r < cum) {
|
|
free(probs);
|
|
return (uint32_t)i;
|
|
}
|
|
}
|
|
|
|
for (int i = n_vocab - 1; i >= 0; i--) {
|
|
if (probs[i] > 0.0f) {
|
|
free(probs);
|
|
return (uint32_t)i;
|
|
}
|
|
}
|
|
|
|
free(probs);
|
|
return 0;
|
|
}
|
|
|
|
/* ==================================================================
|
|
* OPTIMIZED: copy kept textually in sync with src/core/engine.c.
|
|
* ================================================================== */
|
|
typedef struct {
|
|
float *probs;
|
|
float *work;
|
|
int cap;
|
|
} pn_sampler_scratch;
|
|
|
|
static _Thread_local pn_sampler_scratch g_sampler_scratch = {0};
|
|
|
|
static pn_sampler_scratch *sampler_scratch_get(int n) {
|
|
pn_sampler_scratch *s = &g_sampler_scratch;
|
|
if (s->cap < n) {
|
|
free(s->probs); free(s->work);
|
|
s->probs = malloc((size_t)n * sizeof(float));
|
|
s->work = malloc((size_t)n * sizeof(float));
|
|
if (!s->probs || !s->work) {
|
|
free(s->probs); free(s->work);
|
|
s->probs = NULL; s->work = NULL; s->cap = 0;
|
|
return NULL;
|
|
}
|
|
s->cap = n;
|
|
}
|
|
return s;
|
|
}
|
|
|
|
static float sampler_select_kth(float *a, int n, int k) {
|
|
int lo = 0, hi = n - 1;
|
|
while (lo < hi) {
|
|
int i = lo, j = hi;
|
|
float pivot = a[lo + ((hi - lo) >> 1)];
|
|
while (i <= j) {
|
|
while (a[i] < pivot) i++;
|
|
while (a[j] > pivot) j--;
|
|
if (i <= j) {
|
|
float t = a[i]; a[i] = a[j]; a[j] = t;
|
|
i++; j--;
|
|
}
|
|
}
|
|
if (k <= j) hi = j;
|
|
else if (k >= i) lo = i;
|
|
else break;
|
|
}
|
|
return a[k];
|
|
}
|
|
|
|
static int sampler_cmp_desc(const void *pa, const void *pb) {
|
|
float x = *(const float *)pa, y = *(const float *)pb;
|
|
return (x < y) - (x > y);
|
|
}
|
|
|
|
static uint32_t sample_token_opt(float *logits, int n_vocab, pn_sampling_params *params,
|
|
uint32_t *recent_tokens, int n_recent) {
|
|
if (params->repeat_penalty != 1.0f && recent_tokens && n_recent > 0) {
|
|
for (int i = 0; i < n_recent; i++) {
|
|
uint32_t tok = recent_tokens[i];
|
|
if (tok < (uint32_t)n_vocab) {
|
|
if (logits[tok] > 0.0f)
|
|
logits[tok] /= params->repeat_penalty;
|
|
else
|
|
logits[tok] *= params->repeat_penalty;
|
|
}
|
|
}
|
|
}
|
|
|
|
if (params->temperature <= 0.0f) {
|
|
int best = 0;
|
|
for (int i = 1; i < n_vocab; i++) {
|
|
if (logits[i] > logits[best]) best = i;
|
|
}
|
|
return (uint32_t)best;
|
|
}
|
|
|
|
float inv_temp = 1.0f / params->temperature;
|
|
for (int i = 0; i < n_vocab; i++) {
|
|
logits[i] *= inv_temp;
|
|
}
|
|
|
|
int top_k = params->top_k;
|
|
if (top_k <= 0 || top_k > n_vocab) top_k = n_vocab;
|
|
|
|
float max_logit = logits[0];
|
|
for (int i = 1; i < n_vocab; i++) {
|
|
if (logits[i] > max_logit) max_logit = logits[i];
|
|
}
|
|
|
|
pn_sampler_scratch *sc = sampler_scratch_get(n_vocab);
|
|
if (!sc) return 0;
|
|
float *probs = sc->probs;
|
|
|
|
float sum = 0.0f;
|
|
for (int i = 0; i < n_vocab; i++) {
|
|
probs[i] = expf(logits[i] - max_logit);
|
|
sum += probs[i];
|
|
}
|
|
float inv_sum = 1.0f / sum;
|
|
for (int i = 0; i < n_vocab; i++) {
|
|
probs[i] *= inv_sum;
|
|
}
|
|
|
|
if (top_k < n_vocab) {
|
|
memcpy(sc->work, probs, (size_t)n_vocab * sizeof(float));
|
|
float threshold = sampler_select_kth(sc->work, n_vocab, n_vocab - top_k);
|
|
for (int i = 0; i < n_vocab; i++) {
|
|
if (probs[i] < threshold) probs[i] = 0.0f;
|
|
}
|
|
/* llama.cpp / HF order (2026-10-07): renormalise within the top-k set. */
|
|
sum = 0.0f;
|
|
for (int i = 0; i < n_vocab; i++) sum += probs[i];
|
|
if (sum > 0.0f) {
|
|
inv_sum = 1.0f / sum;
|
|
for (int i = 0; i < n_vocab; i++) probs[i] *= inv_sum;
|
|
}
|
|
}
|
|
|
|
if (params->top_p < 1.0f) {
|
|
float *cand = sc->work;
|
|
int nc = 0;
|
|
for (int i = 0; i < n_vocab; i++) {
|
|
if (probs[i] > 0.0f) cand[nc++] = probs[i];
|
|
}
|
|
qsort(cand, (size_t)nc, sizeof(float), sampler_cmp_desc);
|
|
|
|
float cum_prob = 0.0f;
|
|
float top_p_threshold = 0.0f;
|
|
for (int t = 0; t < nc; t++) {
|
|
if (!(cum_prob < params->top_p)) break;
|
|
if (cand[t] <= 0.0f) break;
|
|
cum_prob += cand[t];
|
|
top_p_threshold = cand[t];
|
|
}
|
|
for (int i = 0; i < n_vocab; i++) {
|
|
if (probs[i] < top_p_threshold) probs[i] = 0.0f;
|
|
}
|
|
}
|
|
|
|
sum = 0.0f;
|
|
for (int i = 0; i < n_vocab; i++) sum += probs[i];
|
|
if (sum > 0.0f) {
|
|
inv_sum = 1.0f / sum;
|
|
for (int i = 0; i < n_vocab; i++) probs[i] *= inv_sum;
|
|
}
|
|
|
|
float r = rand_float();
|
|
float cum = 0.0f;
|
|
for (int i = 0; i < n_vocab; i++) {
|
|
cum += probs[i];
|
|
if (r < cum) return (uint32_t)i;
|
|
}
|
|
for (int i = n_vocab - 1; i >= 0; i--) {
|
|
if (probs[i] > 0.0f) return (uint32_t)i;
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
/* ==================================================================
|
|
* Test driver.
|
|
* ================================================================== */
|
|
|
|
/* Private xorshift PRNG for generating test data (kept independent of the
|
|
* global rand() used inside the samplers). */
|
|
static uint64_t g_rng = 0x123456789abcdef0ull;
|
|
static uint32_t xr(void) {
|
|
uint64_t x = g_rng;
|
|
x ^= x << 13; x ^= x >> 7; x ^= x << 17;
|
|
g_rng = x;
|
|
return (uint32_t)(x >> 32);
|
|
}
|
|
static float xr_uniform(float lo, float hi) {
|
|
return lo + (hi - lo) * ((float)xr() / 4294967296.0f);
|
|
}
|
|
static int xr_range(int lo, int hi) { /* inclusive */
|
|
return lo + (int)(xr() % (uint32_t)(hi - lo + 1));
|
|
}
|
|
|
|
int main(void) {
|
|
const int TRIALS = 20000;
|
|
int mismatches = 0;
|
|
int greedy_cases = 0, topk_cases = 0, topp_cases = 0, penalty_cases = 0, tie_heavy = 0;
|
|
|
|
int max_vocab = 4096;
|
|
float *base = malloc(max_vocab * sizeof(float));
|
|
float *a = malloc(max_vocab * sizeof(float));
|
|
float *b = malloc(max_vocab * sizeof(float));
|
|
uint32_t *recent = malloc(max_vocab * sizeof(uint32_t));
|
|
|
|
for (int trial = 0; trial < TRIALS; trial++) {
|
|
/* Vocab size: mix of tiny, small and large. */
|
|
int nv;
|
|
int roll = xr_range(0, 9);
|
|
if (roll < 2) nv = xr_range(1, 8); /* tiny edge cases */
|
|
else if (roll < 7) nv = xr_range(16, 512); /* small */
|
|
else nv = xr_range(1000, max_vocab);
|
|
|
|
/* Decide whether this trial is "tie heavy": draw logits from a small
|
|
* discrete set so lots of tokens share the exact same value. This
|
|
* stresses the top-k threshold and top-p tie handling. */
|
|
int ties = (xr_range(0, 2) == 0);
|
|
if (ties) tie_heavy++;
|
|
for (int i = 0; i < nv; i++) {
|
|
if (ties) {
|
|
int levels = xr_range(2, 6);
|
|
base[i] = (float)xr_range(0, levels) * xr_uniform(0.5f, 3.0f) - 4.0f;
|
|
} else {
|
|
base[i] = xr_uniform(-25.0f, 25.0f);
|
|
}
|
|
}
|
|
|
|
/* Random params. Temperature 0 ~1/6 of the time (greedy). */
|
|
pn_sampling_params p;
|
|
int troll = xr_range(0, 5);
|
|
if (troll == 0) { p.temperature = 0.0f; greedy_cases++; }
|
|
else if (troll == 1) p.temperature = 0.2f; /* production OCR */
|
|
else p.temperature = xr_uniform(0.05f, 2.0f);
|
|
|
|
int kroll = xr_range(0, 4);
|
|
if (kroll == 0) p.top_k = 0; /* -> all */
|
|
else if (kroll == 1) p.top_k = xr_range(nv + 1, nv + 100); /* > vocab */
|
|
else if (kroll == 2) p.top_k = 40; /* production OCR */
|
|
else p.top_k = xr_range(1, nv);
|
|
if (p.top_k > 0 && p.top_k <= nv) topk_cases++;
|
|
|
|
int proll = xr_range(0, 4);
|
|
if (proll == 0) p.top_p = 1.0f; /* off */
|
|
else if (proll == 1) p.top_p = 0.9f; /* production OCR */
|
|
else p.top_p = xr_uniform(0.05f, 1.0f);
|
|
if (p.top_p < 1.0f) topp_cases++;
|
|
|
|
int rpr = xr_range(0, 3);
|
|
if (rpr == 0) p.repeat_penalty = 1.0f; /* off */
|
|
else if (rpr == 1) p.repeat_penalty = 1.15f; /* production OCR */
|
|
else p.repeat_penalty = xr_uniform(1.0f, 2.0f);
|
|
|
|
p.repeat_last_n = 64;
|
|
|
|
/* Random recent-token window. */
|
|
int n_recent = 0;
|
|
if (p.repeat_penalty != 1.0f && xr_range(0, 1)) {
|
|
n_recent = xr_range(0, nv < 64 ? nv : 64);
|
|
for (int i = 0; i < n_recent; i++) {
|
|
/* Sometimes emit out-of-range ids to exercise the guard. */
|
|
recent[i] = (xr_range(0, 9) == 0) ? (uint32_t)(nv + xr_range(0, 5))
|
|
: (uint32_t)xr_range(0, nv - 1);
|
|
}
|
|
if (n_recent > 0) penalty_cases++;
|
|
}
|
|
|
|
memcpy(a, base, nv * sizeof(float));
|
|
memcpy(b, base, nv * sizeof(float));
|
|
|
|
unsigned seed = 0x9E3779B9u ^ (unsigned)(trial * 2654435761u);
|
|
|
|
srand(seed);
|
|
uint32_t tok_ref = sample_token_ref(a, nv, &p, n_recent ? recent : NULL, n_recent);
|
|
srand(seed);
|
|
uint32_t tok_opt = sample_token_opt(b, nv, &p, n_recent ? recent : NULL, n_recent);
|
|
|
|
if (tok_ref != tok_opt) {
|
|
if (mismatches < 20) {
|
|
fprintf(stderr,
|
|
"MISMATCH trial=%d nv=%d temp=%.4f top_k=%d top_p=%.4f rp=%.4f "
|
|
"n_recent=%d ties=%d : ref=%u opt=%u\n",
|
|
trial, nv, p.temperature, p.top_k, p.top_p, p.repeat_penalty,
|
|
n_recent, ties, tok_ref, tok_opt);
|
|
}
|
|
mismatches++;
|
|
}
|
|
}
|
|
|
|
free(base); free(a); free(b); free(recent);
|
|
|
|
printf("sampler parity: %d trials\n", TRIALS);
|
|
printf(" greedy=%d top_k=%d top_p=%d penalty=%d tie_heavy=%d\n",
|
|
greedy_cases, topk_cases, topp_cases, penalty_cases, tie_heavy);
|
|
if (mismatches == 0) {
|
|
printf("PASS: 100%% token-id match (0 mismatches)\n");
|
|
return 0;
|
|
}
|
|
printf("FAIL: %d mismatches out of %d\n", mismatches, TRIALS);
|
|
return 1;
|
|
}
|