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

448 lines
18 KiB
C
Raw Blame History

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/*
* Unit tests for speculative decoding under the tool-call grammar
* (src/serve/grammar_spec_chooser.c) with the real ornith JSON grammar in its
* tool_choice "required" think-gated mode.
*
* The central test is an EQUIVALENCE SIMULATION: a deterministic toy model
* (logits = a hash of the token prefix, over a small vocab of real ornith
* fragments) is decoded twice —
* plain : one token at a time, grammar mask on the dense row, argmax;
* spec : seed + drafts verified in one cycle exactly as
* pn_qwen35moe_tt_spec_step_ex accepts them (a draft is accepted iff
* the chooser picks it; DEFER ends the cycle at the previous row),
* rows handed over as top-N pairs, drafts mixing the true
* continuation with grammar-breaking junk —
* and the two token streams must be identical for every trial. Every
* completed call must parse with the serve's own parser (tc_parse_tool_calls).
*
* The simulation is SELF-FALSIFYING: the same spec loop with the grammar
* ignored (mask never active) must DIVERGE from plain decoding in some trial,
* or the test exits 2 (VACUOUS) — a comparison that cannot fail proves nothing.
*
* Pure C, libc-only — no engine/model/device.
*/
#include "lib/test.h"
#include "../src/serve/grammar_spec_chooser.h"
#include "../src/core/grammar_mask_shortlist.h"
#include "../src/serve/tool_grammar_json.h"
#include "../src/serve/tool_calling.h"
#include "../src/serve/json_reader.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
static const char *TOOLS_JSON =
"[{\"type\":\"function\",\"function\":{\"name\":\"read\","
"\"parameters\":{\"type\":\"object\",\"properties\":{"
"\"path\":{\"type\":\"string\",\"minLength\":1}},\"required\":[\"path\"]}}},"
"{\"type\":\"function\",\"function\":{\"name\":\"add\","
"\"parameters\":{\"type\":\"object\",\"properties\":{"
"\"a\":{\"type\":\"integer\"},\"b\":{\"type\":\"integer\"}},"
"\"required\":[\"a\",\"b\"]}}}]";
/* Toy vocab: real ornith fragments, reasoning text, and junk the grammar must
* refuse once the call is compelled. Id 0 is the end token (no text). */
static const char *VOCAB[] = {
"", "Let", " me", " think", "{", "</think>", "\n\n", "```json\n",
"{\"function\": \"read\"", ", \"arguments\": {", "\"path\": \"", "src",
"/main.c", "\"}}", "\n```", "Sure, here it is", " ", "\"", "```", "}",
"{\"function\": \"add\"", "\"a\": 1, \"b\": 2}}", "ok", "\n",
};
#define VOCAB_N ((int)(sizeof(VOCAB) / sizeof(VOCAB[0])))
#define END_ID 0
#define THINK_CLOSE_ID 5
#define TOP_N 4
#define MAX_SEQ 128
/* ---- grammar ops over tgj_state ---- */
static void ops_copy(void *d, const void *s) { tgj_state_copy(d, s); }
static bool ops_allowed(const void *st, const char *t) { return tgj_token_allowed(st, t); }
static void ops_advance(void *st, const char *t) { tgj_advance(st, t); }
static bool ops_mask(const void *st) { return tgj_mask_active(st); }
static bool ops_mask_never(const void *st) { (void)st; return false; }
static bool ops_complete(const void *st) { return tgj_complete(st); }
static const gsc_grammar_ops OPS = { ops_copy, ops_allowed, ops_advance, ops_mask, ops_complete };
static const gsc_grammar_ops OPS_IGNORING_GRAMMAR = {
ops_copy, ops_allowed, ops_advance, ops_mask_never, ops_complete };
static const char *vocab_text(void *ctx, uint32_t id) {
(void)ctx;
return id < (uint32_t)VOCAB_N ? VOCAB[id] : NULL;
}
/* ---- the toy model: logits are a hash of (salt, prefix) ---- */
static uint64_t mix64(uint64_t x) {
x += 0x9E3779B97F4A7C15ull;
x = (x ^ (x >> 30)) * 0xBF58476D1CE4E5B9ull;
x = (x ^ (x >> 27)) * 0x94D049BB133111EBull;
return x ^ (x >> 31);
}
static void toy_logits(uint64_t salt, const uint32_t *seq, int n, float *out) {
uint64_t h = mix64(salt);
for (int i = 0; i < n; i++) h = mix64(h ^ seq[i]);
for (int v = 0; v < VOCAB_N; v++)
out[v] = (float)(mix64(h + (uint64_t)v) >> 40) / (float)(1u << 24);
/* Close the think block after a while, and favour the fragments that finish
* a call, so most trials reach a complete call inside MAX_SEQ. */
if (n >= 6 + (int)(h % 8)) out[THINK_CLOSE_ID] += 0.6f;
static const int CALL_FRAGMENTS[] = { 7, 8, 9, 10, 13, 14 };
for (int i = 0; i < (int)(sizeof(CALL_FRAGMENTS) / sizeof(CALL_FRAGMENTS[0])); i++)
out[CALL_FRAGMENTS[i]] += 0.3f;
out[END_ID] -= 0.5f; /* an early end token would make trials trivial */
}
/* Argmax, lowest id on a tie; under `walk`'s active mask only allowed ids. */
static uint32_t masked_argmax(const float *logits, const tgj_state *walk) {
bool masking = walk && tgj_mask_active(walk);
int best = -1;
for (int v = 0; v < VOCAB_N; v++) {
if (masking && VOCAB[v][0] && !tgj_token_allowed(walk, VOCAB[v])) continue;
if (best < 0 || logits[v] > logits[best]) best = v;
}
return (uint32_t)best;
}
/* The card's top-N (index, value) pairs of a dense row, in id order. */
static uint32_t top_n_pairs(const float *logits, uint32_t *index, float *value) {
bool taken[VOCAB_N] = { false };
for (int k = 0; k < TOP_N; k++) {
int best = -1;
for (int v = 0; v < VOCAB_N; v++)
if (!taken[v] && (best < 0 || logits[v] > logits[best])) best = v;
taken[best] = true;
}
uint32_t n = 0;
for (int v = 0; v < VOCAB_N; v++)
if (taken[v]) { index[n] = (uint32_t)v; value[n] = logits[v]; n++; }
return n;
}
/* Plain constrained decode. Returns the token count. */
static int decode_plain(const tgj_grammar *g, uint64_t salt, uint32_t *seq) {
tgj_state *walk = tgj_state_new_required(g, true);
float logits[VOCAB_N];
int n = 0;
while (n < MAX_SEQ) {
toy_logits(salt, seq, n, logits);
uint32_t tok = masked_argmax(logits, walk);
if (tok == END_ID) break;
seq[n++] = tok;
tgj_advance(walk, VOCAB[tok]);
if (tgj_complete(walk)) break;
}
tgj_state_free(walk);
return n;
}
/* A row under an active mask, decided as the serve does: from the row's top-N
* by the shortlist rule (greedy), else — row 0 only — the dense row at the
* cycle's start, masked; a later undecidable row defers. */
typedef struct {
uint64_t salt;
const uint32_t *seq;
int n;
unsigned long dense_rows;
} masked_ctx;
static bool walk_allows(void *ctx, uint32_t id) {
return !VOCAB[id][0] || tgj_token_allowed((const tgj_state *)ctx, VOCAB[id]);
}
static uint32_t masked_pick(void *ctx, void *walk, int row, const uint32_t *index,
const float *value, uint32_t count) {
masked_ctx *m = (masked_ctx *)ctx;
float sparse[VOCAB_N], logits[VOCAB_N];
uint32_t ids[GMS_MAX];
float bound;
gms_sampling greedy = { .temperature = 0.0f, .top_k = 1, .repeat_penalty = 1.0f };
uint32_t neg = 0xff800000u;
for (int v = 0; v < VOCAB_N; v++) memcpy(&sparse[v], &neg, sizeof(float));
for (uint32_t i = 0; i < count; i++) sparse[index[i]] = value[i];
int n = gms_select_shortlist(sparse, VOCAB_N, ids, &bound);
if (gms_mask_from_shortlist(sparse, VOCAB_N, ids, n, bound, &greedy, walk_allows, walk)) {
int best = -1;
for (int v = 0; v < VOCAB_N; v++)
if (sparse[v] > -1e30f && (best < 0 || sparse[v] > sparse[best])) best = v;
return (uint32_t)best;
}
if (row > 0) return GSC_DEFER;
m->dense_rows++;
toy_logits(m->salt, m->seq, m->n, logits);
return masked_argmax(logits, (const tgj_state *)walk);
}
/* Drafts: the true plain continuation, with junk mixed in by the hash. */
static int make_drafts(const tgj_grammar *g, uint64_t salt, const uint32_t *seq, int n,
uint32_t *drafts) {
uint32_t truth[MAX_SEQ + 4];
int nt = decode_plain(g, salt, truth);
int nd = 0;
uint64_t h = mix64(salt ^ (uint64_t)n * 7919u);
for (int j = 0; j < 3; j++) {
h = mix64(h);
uint32_t d = (h % 3 == 0 || n + j >= nt) ? (uint32_t)(h >> 33) % VOCAB_N
: truth[n + j];
if (d == END_ID) break; /* the serve cuts drafts at an end token */
drafts[nd++] = d;
}
(void)seq;
return nd;
}
/* One speculative verify cycle, accepting exactly as spec_step_ex does. */
static int verify_cycle(gsc_chooser *c, uint64_t salt, uint32_t *seq, int n,
const uint32_t *drafts, int nd, uint32_t *emit) {
uint32_t index[TOP_N];
float value[TOP_N], logits[VOCAB_N];
uint32_t row_seq[MAX_SEQ + 8];
int n_accept = 0;
uint32_t next;
memcpy(row_seq, seq, (size_t)n * sizeof(*seq));
gsc_begin_cycle(c);
for (;;) {
for (int i = 0; i < n_accept; i++) row_seq[n + i] = drafts[i];
toy_logits(salt, row_seq, n + n_accept, logits);
uint32_t cnt = top_n_pairs(logits, index, value);
uint32_t tok = gsc_choose(c, index, value, cnt);
if (tok == GSC_DEFER) {
if (n_accept == 0) return -1; /* spec_step_ex fails the cycle */
n_accept--;
next = drafts[n_accept];
break;
}
if (n_accept < nd && tok == drafts[n_accept]) { n_accept++; continue; }
next = tok;
break;
}
for (int i = 0; i < n_accept; i++) emit[i] = drafts[i];
emit[n_accept] = next;
return n_accept + 1;
}
/* Speculative constrained decode through the chooser. Returns the count. */
static int decode_spec(const tgj_grammar *g, uint64_t salt, const gsc_grammar_ops *ops,
uint32_t *seq, unsigned long *cycles) {
tgj_state *live = tgj_state_new_required(g, true);
tgj_state *walk = tgj_state_clone(live);
masked_ctx dctx = { salt, seq, 0, 0 };
gsc_chooser c = { .ops = ops, .live = live, .walk = walk, .text = vocab_text,
.masked = masked_pick, .masked_ctx = &dctx };
float logits[VOCAB_N];
int n = 0;
bool stop = false;
while (n < MAX_SEQ && !stop) {
uint32_t drafts[4], emit[8];
int nd = make_drafts(g, salt, seq, n, drafts);
nd = gsc_truncate_drafts(&c, drafts, nd);
int ne;
if (n == 0 || nd < 1) { /* no cycle: one plain step */
toy_logits(salt, seq, n, logits);
emit[0] = masked_argmax(logits, ops == &OPS ? live : NULL);
ne = 1;
} else {
dctx.n = n;
ne = verify_cycle(&c, salt, seq, n, drafts, nd, emit);
(*cycles)++;
if (ne < 0) { /* a refused cycle: report a divergence */
seq[n++] = END_ID;
break;
}
}
for (int i = 0; i < ne && !stop; i++) {
if (emit[i] == END_ID || n >= MAX_SEQ) { stop = true; break; }
seq[n++] = emit[i];
tgj_advance(live, VOCAB[emit[i]]);
if (tgj_complete(live)) stop = true;
}
}
tgj_state_free(walk);
tgj_state_free(live);
return n;
}
/* The text the model emitted after its think block. */
static void content_after_think(const uint32_t *seq, int n, char *out, size_t cap) {
size_t len = 0;
int start = 0;
for (int i = 0; i < n; i++)
if (seq[i] == THINK_CLOSE_ID) { start = i + 1; break; }
out[0] = '\0';
for (int i = start; i < n; i++) {
size_t t = strlen(VOCAB[seq[i]]);
if (len + t + 1 >= cap) break;
memcpy(out + len, VOCAB[seq[i]], t);
len += t;
out[len] = '\0';
}
}
/* ============================================================
* Spec == plain over many trials; completed calls parse.
* ============================================================ */
static void test_spec_equals_plain_constrained(void) {
jr_value *tools = jr_parse(TOOLS_JSON);
tgj_grammar *g = tgj_build(tools);
int mismatches = 0, completed = 0, parsed = 0, ignoring_mismatches = 0;
unsigned long cycles = 0;
for (uint64_t salt = 1; salt <= 400; salt++) {
uint32_t plain[MAX_SEQ + 8], spec[MAX_SEQ + 8], loose[MAX_SEQ + 8];
int np = decode_plain(g, salt, plain);
int ns = decode_spec(g, salt, &OPS, spec, &cycles);
int nl = decode_spec(g, salt, &OPS_IGNORING_GRAMMAR, loose, &cycles);
if (np != ns || memcmp(plain, spec, (size_t)np * sizeof(*plain)) != 0) mismatches++;
if (np != nl || memcmp(plain, loose, (size_t)np * sizeof(*plain)) != 0)
ignoring_mismatches++;
tgj_state *w = tgj_state_new_required(g, true);
for (int i = 0; i < np; i++) tgj_advance(w, VOCAB[plain[i]]);
if (tgj_complete(w)) {
char text[4096];
tc_tool_call *calls = NULL;
int n_calls = 0;
completed++;
content_after_think(plain, np, text, sizeof(text));
if (tc_parse_tool_calls(text, tools, TC_DIALECT_ORNITH_JSON, &calls, &n_calls) == 0 &&
n_calls == 1 && tc_tool_name_known(tools, calls[0].name))
parsed++;
tc_tool_calls_free(calls, n_calls);
}
tgj_state_free(w);
}
printf(" trials=400 cycles=%lu completed=%d parsed=%d mismatches=%d "
"grammar-ignoring mismatches=%d\n", cycles, completed, parsed, mismatches,
ignoring_mismatches);
if (ignoring_mismatches == 0) {
printf("VACUOUS: ignoring the grammar never changed the output\n");
exit(2);
}
ASSERT_INT_EQ(mismatches, 0);
ASSERT_TRUE(completed > 100);
ASSERT_INT_EQ(parsed, completed);
tgj_free(g);
jr_free(tools);
}
/* ============================================================
* Truncation: a grammar-breaking draft is cut, and nothing follows completion.
* ============================================================ */
static void test_truncate_drafts(void) {
jr_value *tools = jr_parse(TOOLS_JSON);
tgj_grammar *g = tgj_build(tools);
tgj_state *live = tgj_state_new_required(g, true);
tgj_state *walk = tgj_state_clone(live);
gsc_chooser c = { .ops = &OPS, .live = live, .walk = walk, .text = vocab_text };
/* Reasoning: anything goes, '{' included. */
uint32_t think[3] = { 4, 15, 18 };
ASSERT_INT_EQ(gsc_truncate_drafts(&c, think, 3), 3);
tgj_advance(live, "</think>");
tgj_advance(live, "\n\n");
uint32_t prose[3] = { 7, 15, 8 }; /* fence, then junk */
ASSERT_INT_EQ(gsc_truncate_drafts(&c, prose, 3), 1);
uint32_t good[3] = { 7, 8, 9 };
ASSERT_INT_EQ(gsc_truncate_drafts(&c, good, 3), 3);
tgj_advance(live, "```json\n");
tgj_advance(live, "{\"function\": \"read\"");
tgj_advance(live, ", \"arguments\": {");
uint32_t finish[4] = { 10, 11, 13, 14 }; /* path, value, close, fence close */
ASSERT_INT_EQ(gsc_truncate_drafts(&c, finish, 4), 4);
uint32_t past[3] = { 10, 13, 14 }; /* "": an empty path is refused */
ASSERT_INT_EQ(gsc_truncate_drafts(&c, past, 3), 1);
tgj_advance(live, "\"path\": \"");
tgj_advance(live, "src");
uint32_t over[3] = { 13, 14, 22 }; /* complete after 14 → 22 cut */
ASSERT_INT_EQ(gsc_truncate_drafts(&c, over, 3), 2);
ASSERT_TRUE(c.drafts_cut >= 4);
tgj_state_free(walk);
tgj_state_free(live);
tgj_free(g);
jr_free(tools);
}
/* Routing: an unmasked row goes to the base chooser, a masked row to the masked
* pick (with the private walk and the row index), and the masked pick's DEFER
* ends the cycle. */
static int g_base_calls, g_masked_calls, g_masked_last_row;
static uint32_t base_first(void *ctx, int row, const uint32_t *index, const float *value,
uint32_t count) {
(void)ctx; (void)row; (void)value; (void)count;
g_base_calls++;
return index[0];
}
static uint32_t masked_fence_then_defer(void *ctx, void *walk, int row, const uint32_t *index,
const float *value, uint32_t count) {
(void)ctx; (void)index; (void)value; (void)count;
g_masked_calls++;
g_masked_last_row = row;
return tgj_in_call((const tgj_state *)walk) ? GSC_DEFER : 7u; /* fence, then defer */
}
static void test_rows_route_by_mask(void) {
jr_value *tools = jr_parse(TOOLS_JSON);
tgj_grammar *g = tgj_build(tools);
tgj_state *live = tgj_state_new_required(g, true);
tgj_state *walk = tgj_state_clone(live);
gsc_chooser c = { .ops = &OPS, .live = live, .walk = walk, .text = vocab_text,
.base = base_first, .masked = masked_fence_then_defer };
uint32_t idx_close[2] = { 5, 22 }; /* "</think>" first */
uint32_t idx_any[2] = { 15, 22 };
float val[2] = { 1.0f, 0.5f };
g_base_calls = g_masked_calls = 0;
gsc_begin_cycle(&c);
ASSERT_TRUE(gsc_choose(&c, idx_close, val, 2) == 5u); /* reasoning: base */
ASSERT_INT_EQ(g_base_calls, 1);
ASSERT_TRUE(gsc_choose(&c, idx_any, val, 2) == 7u); /* compelled: masked pick */
ASSERT_INT_EQ(g_masked_calls, 1);
ASSERT_INT_EQ(g_masked_last_row, 1);
ASSERT_TRUE(gsc_choose(&c, idx_any, val, 2) == GSC_DEFER);
ASSERT_INT_EQ(g_masked_last_row, 2);
ASSERT_INT_EQ(g_base_calls, 1);
tgj_state_free(walk);
tgj_state_free(live);
tgj_free(g);
jr_free(tools);
}
/* Nothing past the call is committed: the row after the one that completed
* the call defers, so the completing token ends the cycle uncommitted. */
static void test_defer_after_completion(void) {
jr_value *tools = jr_parse(TOOLS_JSON);
tgj_grammar *g = tgj_build(tools);
tgj_state *live = tgj_state_new_required(g, true);
const char *frags[] = { "</think>", "```json\n", "{\"function\": \"read\"",
", \"arguments\": {", "\"path\": \"", "src", "\"}}" };
for (int i = 0; i < 7; i++) tgj_advance(live, frags[i]);
tgj_state *walk = tgj_state_clone(live);
masked_ctx m = { 0, NULL, 0, 0 };
gsc_chooser c = { .ops = &OPS, .live = live, .walk = walk, .text = vocab_text,
.masked = masked_pick, .masked_ctx = &m };
uint32_t idx[2] = { 14, 22 };
float val[2] = { 2.0f, 1.0f };
gsc_begin_cycle(&c);
ASSERT_TRUE(gsc_choose(&c, idx, val, 2) == 14u); /* the fence close completes it */
ASSERT_TRUE(c.completed);
ASSERT_TRUE(gsc_choose(&c, idx, val, 2) == GSC_DEFER);
tgj_state_free(walk);
tgj_state_free(live);
tgj_free(g);
jr_free(tools);
}
int main(void) {
printf("\n=== grammar_spec_chooser (speculative decoding under the grammar) ===\n\n");
RUN_TEST(truncate_drafts);
RUN_TEST(rows_route_by_mask);
RUN_TEST(defer_after_completion);
RUN_TEST(spec_equals_plain_constrained);
TEST_REPORT();
return __ft_failed == 0 ? 0 : 1;
}
Reference in New Issue View Git Blame Copy Permalink