David
/
Pantheon
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Pantheon/tests/ngram_accept_stream.c

410 lines
14 KiB
C
Raw Permalink Blame History

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/**
* @file ngram_accept_stream.c
* @brief Implementation of the recorded-stream dump and the speculative-round simulator
* declared in ngram_accept_stream.h.
*
* The drafter itself is NOT reimplemented here: `pn_ngram_draft` from src/core/ngram_draft.c
* is the one that ships, and a second copy written for a measurement is a measurement of the
* copy.
*
* @status Working
* @issues None
* @todo None
*/
/* 1. This module's own header */
#include "ngram_accept_stream.h"
/* 2. Standard library */
#include <stdlib.h>
#include <string.h>
/* 4. Project headers */
#include "../src/core/ngram_draft.h"
#define NGRAM_MAX_HISTORY 65536
#define NGRAM_MAX_DRAFT 32
const char *const ngram_length_bucket_name[NGRAM_LEN_BUCKETS] = {
"<512", "512-1k", "1k-2k", "2k-4k", "4k-8k", ">=8k"
};
/** Which context-length bucket a history of `n` tokens falls in. */
static int ngram_length_bucket(int n)
{
if (n < 512) return 0;
if (n < 1024) return 1;
if (n < 2048) return 2;
if (n < 4096) return 3;
if (n < 8192) return 4;
return 5;
}
/** splitmix64 — a seeded stream so a control arm repeats exactly. */
static uint64_t ngram_rng_next(uint64_t *state)
{
uint64_t z = (*state += 0x9E3779B97F4A7C15ULL);
z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL;
z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL;
return z ^ (z >> 31);
}
/* ================================================================
* The dump file
* ================================================================ */
/** Write one little-endian uint32. */
static int ngram_put_u32(FILE *f, uint32_t v)
{
unsigned char b[4] = { (unsigned char)(v & 0xFF), (unsigned char)((v >> 8) & 0xFF),
(unsigned char)((v >> 16) & 0xFF), (unsigned char)((v >> 24) & 0xFF) };
return fwrite(b, 1, sizeof(b), f) == sizeof(b) ? 0 : -1;
}
/** Read one little-endian uint32. */
static int ngram_get_u32(FILE *f, uint32_t *out)
{
unsigned char b[4];
if (fread(b, 1, sizeof(b), f) != sizeof(b))
return -1;
*out = (uint32_t)b[0] | ((uint32_t)b[1] << 8) | ((uint32_t)b[2] << 16) |
((uint32_t)b[3] << 24);
return 0;
}
/** Write a fixed-width name field, zero-padded, never truncated mid-run without notice. */
static int ngram_put_name(FILE *f, const char *s)
{
char buf[NGRAM_STREAM_NAME_LEN] = {0};
snprintf(buf, sizeof(buf), "%s", s ? s : "");
return fwrite(buf, 1, sizeof(buf), f) == sizeof(buf) ? 0 : -1;
}
int ngram_stream_open_write(const char *path, uint32_t n_vocab, uint32_t arch, uint32_t eos_id,
uint32_t think_end_id, const char *arm, const char *tag,
FILE **out)
{
FILE *f = fopen(path, "wb");
if (!f) {
fprintf(stderr, "ngram_stream: cannot create %s\n", path);
return -1;
}
if (fwrite(NGRAM_STREAM_MAGIC, 1, 8, f) != 8 ||
ngram_put_u32(f, NGRAM_STREAM_ENDIAN) != 0 || ngram_put_u32(f, n_vocab) != 0 ||
ngram_put_u32(f, arch) != 0 || ngram_put_u32(f, eos_id) != 0 ||
ngram_put_u32(f, think_end_id) != 0 || ngram_put_name(f, arm) != 0 ||
ngram_put_name(f, tag) != 0) {
fprintf(stderr, "ngram_stream: header write failed for %s\n", path);
fclose(f);
return -1;
}
*out = f;
return 0;
}
int ngram_stream_write_document(FILE *f, const ngram_document *doc)
{
uint64_t bits;
if (ngram_put_u32(f, (uint32_t)doc->n_prompt) != 0 ||
ngram_put_u32(f, (uint32_t)doc->n_gen) != 0 ||
ngram_put_u32(f, doc->think_end_pos < 0 ? NGRAM_STREAM_NONE
: (uint32_t)doc->think_end_pos) != 0)
return -1;
memcpy(&bits, &doc->decode_s, sizeof(bits));
if (ngram_put_u32(f, (uint32_t)(bits & 0xFFFFFFFFu)) != 0 ||
ngram_put_u32(f, (uint32_t)(bits >> 32)) != 0)
return -1;
for (int i = 0; i < doc->n_prompt; i++)
if (ngram_put_u32(f, doc->prompt[i]) != 0)
return -1;
for (int i = 0; i < doc->n_gen; i++)
if (ngram_put_u32(f, doc->gen[i]) != 0)
return -1;
return fflush(f);
}
int ngram_stream_close_write(FILE *f)
{
int rc = ngram_put_u32(f, NGRAM_STREAM_END);
if (fclose(f) != 0)
rc = -1;
return rc;
}
/** Read one document body; returns 0 on success, 1 at the end marker, -1 on damage. */
static int ngram_read_document(FILE *f, ngram_document *doc)
{
uint32_t first, n_gen, think, lo, hi;
uint64_t bits;
int rc = -1;
if (ngram_get_u32(f, &first) != 0)
return -1;
if (first == NGRAM_STREAM_END)
return 1;
if (ngram_get_u32(f, &n_gen) != 0 || ngram_get_u32(f, &think) != 0 ||
ngram_get_u32(f, &lo) != 0 || ngram_get_u32(f, &hi) != 0)
return -1;
if (first > NGRAM_MAX_HISTORY || n_gen > NGRAM_MAX_HISTORY)
return -1;
bits = (uint64_t)lo | ((uint64_t)hi << 32);
memcpy(&doc->decode_s, &bits, sizeof(doc->decode_s));
doc->n_prompt = (int)first;
doc->n_gen = (int)n_gen;
doc->think_end_pos = think == NGRAM_STREAM_NONE ? -1 : (int)think;
doc->prompt = calloc((size_t)first + 1, sizeof(*doc->prompt));
doc->gen = calloc((size_t)n_gen + 1, sizeof(*doc->gen));
if (!doc->prompt || !doc->gen)
goto cleanup;
for (uint32_t i = 0; i < first; i++)
if (ngram_get_u32(f, &doc->prompt[i]) != 0)
goto cleanup;
for (uint32_t i = 0; i < n_gen; i++)
if (ngram_get_u32(f, &doc->gen[i]) != 0)
goto cleanup;
rc = 0;
cleanup:
if (rc != 0) {
free(doc->prompt);
free(doc->gen);
doc->prompt = NULL;
doc->gen = NULL;
}
return rc;
}
int ngram_stream_read(const char *path, ngram_stream *out)
{
char magic[8];
uint32_t endian;
FILE *f = fopen(path, "rb");
int rc = 0;
memset(out, 0, sizeof(*out));
if (!f) {
fprintf(stderr, "ngram_stream: cannot open %s\n", path);
return -1;
}
if (fread(magic, 1, sizeof(magic), f) != sizeof(magic) ||
memcmp(magic, NGRAM_STREAM_MAGIC, sizeof(magic)) != 0) {
fprintf(stderr, "ngram_stream: %s is not a stream dump\n", path);
fclose(f);
return -1;
}
if (ngram_get_u32(f, &endian) != 0 || endian != NGRAM_STREAM_ENDIAN ||
ngram_get_u32(f, &out->n_vocab) != 0 || ngram_get_u32(f, &out->arch) != 0 ||
ngram_get_u32(f, &out->eos_id) != 0 || ngram_get_u32(f, &out->think_end_id) != 0 ||
fread(out->arm, 1, sizeof(out->arm), f) != sizeof(out->arm) ||
fread(out->tag, 1, sizeof(out->tag), f) != sizeof(out->tag)) {
fprintf(stderr, "ngram_stream: %s has a damaged header\n", path);
fclose(f);
return -1;
}
out->arm[NGRAM_STREAM_NAME_LEN - 1] = '\0';
out->tag[NGRAM_STREAM_NAME_LEN - 1] = '\0';
while (out->n_docs < NGRAM_STREAM_MAX_DOCS) {
rc = ngram_read_document(f, &out->docs[out->n_docs]);
if (rc != 0)
break;
out->n_docs++;
}
fclose(f);
/* A dump whose writer died mid-run has no end marker. Refuse it rather than reporting an
* acceptance rate over however many documents happened to land — the eighth vacuous-gate
* class in this project was a harness scoring bytes it should have refused. */
if (rc != 1) {
fprintf(stderr, "ngram_stream: %s has no end marker — the writer did not finish. "
"Refusing rather than scoring %d partial document(s).\n",
path, out->n_docs);
ngram_stream_free(out);
return -1;
}
return 0;
}
void ngram_stream_free(ngram_stream *s)
{
for (int i = 0; i < s->n_docs; i++) {
free(s->docs[i].prompt);
free(s->docs[i].gen);
}
memset(s, 0, sizeof(*s));
}
/* ================================================================
* The simulator
* ================================================================ */
/**
* Present the drafter's search space for one round, applying the control.
*
* NGRAM_CONTROL_SHUFFLE permutes everything EXCEPT the trailing `ngram_max` tokens, so the
* KEY is still the model's own real suffix and only the structure it is searching for is
* destroyed. Shuffling the key too would confound "the history has no structure" with "the
* query is nonsense".
*
* @return the number of tokens the drafter may see; *view points at them
*/
static int ngram_present_history(const uint32_t *hist, int nh, int window, int ngram_max,
ngram_control control, uint64_t *rng, uint32_t **scratch,
const uint32_t **view)
{
int n = nh;
const uint32_t *base = hist;
if (window > 0 && n > window) {
base = hist + (n - window);
n = window;
}
if (control != NGRAM_CONTROL_SHUFFLE) {
*view = base;
return n;
}
memcpy(*scratch, base, (size_t)n * sizeof(**scratch));
for (int i = n - ngram_max - 1; i > 0; i--) {
int j = (int)(ngram_rng_next(rng) % (uint64_t)(i + 1));
uint32_t t = (*scratch)[i];
(*scratch)[i] = (*scratch)[j];
(*scratch)[j] = t;
}
*view = *scratch;
return n;
}
int ngram_accept_replay(const ngram_document *doc, int max_k, int ngram_max, int ngram_min,
int window, ngram_control control, uint32_t n_vocab, uint64_t seed,
ngram_accept_stats *out)
{
static uint32_t hist[NGRAM_MAX_HISTORY];
static uint32_t scratch[NGRAM_MAX_HISTORY];
uint32_t *scratch_p = scratch;
uint32_t drafts[NGRAM_MAX_DRAFT];
uint64_t rng = seed;
int nh, p;
if (max_k < 1 || max_k > NGRAM_MAX_DRAFT || doc->n_gen < 2)
return -1;
if (doc->n_prompt + doc->n_gen > NGRAM_MAX_HISTORY)
return -1;
memcpy(hist, doc->prompt, (size_t)doc->n_prompt * sizeof(*hist));
nh = doc->n_prompt;
hist[nh++] = doc->gen[0];
for (p = 0; p < doc->n_gen - 1; ) {
int phase = (doc->think_end_pos >= 0 && p < doc->think_end_pos) ? NGRAM_PHASE_REASONING
: NGRAM_PHASE_ANSWER;
ngram_accept_stats *s = &out[phase];
int bucket = ngram_length_bucket(nh);
const uint32_t *view = NULL;
int nv = ngram_present_history(hist, nh, window, ngram_max, control, &rng, &scratch_p,
&view);
int k = pn_ngram_draft(view, nv, ngram_max, ngram_min, max_k, drafts);
int delivered, m = 0;
if (k > 0 && control == NGRAM_CONTROL_RANDOM)
for (int i = 0; i < k; i++)
drafts[i] = (uint32_t)(ngram_rng_next(&rng) % (uint64_t)n_vocab);
if (k > 0) {
/* The round's free token is the argmax after the last ACCEPTED draft, so
* accepting m drafts delivers m+1 tokens and needs gen[p+m+1] to exist. */
int lim = doc->n_gen - p - 2;
if (lim > k)
lim = k;
while (m < lim && drafts[m] == doc->gen[p + 1 + m])
m++;
delivered = m + 1;
s->rounds++;
s->drafted += k;
s->accepted += m;
s->hist[delivered]++;
} else {
delivered = 1;
s->plain_rounds++;
}
s->forwards++;
s->emitted += delivered;
s->bucket_forwards[bucket]++;
s->bucket_emitted[bucket] += delivered;
for (int i = 0; i < delivered; i++)
hist[nh++] = doc->gen[p + 1 + i];
p += delivered;
}
return 0;
}
void ngram_accept_stats_add(ngram_accept_stats *a, const ngram_accept_stats *b)
{
a->rounds += b->rounds;
a->plain_rounds += b->plain_rounds;
a->drafted += b->drafted;
a->accepted += b->accepted;
a->forwards += b->forwards;
a->emitted += b->emitted;
for (int i = 0; i < (int)(sizeof(a->hist) / sizeof(a->hist[0])); i++)
a->hist[i] += b->hist[i];
for (int i = 0; i < NGRAM_LEN_BUCKETS; i++) {
a->bucket_emitted[i] += b->bucket_emitted[i];
a->bucket_forwards[i] += b->bucket_forwards[i];
}
}
double ngram_tokens_per_forward(const ngram_accept_stats *s)
{
return s->forwards ? (double)s->emitted / (double)s->forwards : 0.0;
}
/* ================================================================
* Known-answer self-tests
* ================================================================ */
void ngram_selftest_periodic(ngram_document *doc, int n_prompt, int n_gen, int period,
uint32_t n_vocab)
{
memset(doc, 0, sizeof(*doc));
doc->prompt = calloc((size_t)n_prompt, sizeof(*doc->prompt));
doc->gen = calloc((size_t)n_gen, sizeof(*doc->gen));
doc->n_prompt = n_prompt;
doc->n_gen = n_gen;
doc->think_end_pos = -1;
for (int i = 0; i < n_prompt; i++)
doc->prompt[i] = (uint32_t)((i % period) + 1) % n_vocab;
for (int i = 0; i < n_gen; i++)
doc->gen[i] = (uint32_t)(((n_prompt + i) % period) + 1) % n_vocab;
}
void ngram_selftest_random(ngram_document *doc, int n_prompt, int n_gen, uint32_t n_vocab,
uint64_t seed)
{
uint64_t rng = seed;
memset(doc, 0, sizeof(*doc));
doc->prompt = calloc((size_t)n_prompt, sizeof(*doc->prompt));
doc->gen = calloc((size_t)n_gen, sizeof(*doc->gen));
doc->n_prompt = n_prompt;
doc->n_gen = n_gen;
doc->think_end_pos = -1;
for (int i = 0; i < n_prompt; i++)
doc->prompt[i] = (uint32_t)(ngram_rng_next(&rng) % (uint64_t)n_vocab);
for (int i = 0; i < n_gen; i++)
doc->gen[i] = (uint32_t)(ngram_rng_next(&rng) % (uint64_t)n_vocab);
}
void ngram_selftest_free(ngram_document *doc)
{
free(doc->prompt);
free(doc->gen);
memset(doc, 0, sizeof(*doc));
}
Reference in New Issue View Git Blame Copy Permalink