278 lines
14 KiB
C
278 lines
14 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/* test_verify_batch.c - the M2/M3 correctness gate for pn_llama_verify_batch.
|
|
*
|
|
* pn_llama_verify_batch verifies several INDEPENDENT sequences in one forward
|
|
* pass (block b = n_tok[b] candidate tokens for sequence b, starting at that
|
|
* sequence's committed position). Each output row must be the logits the model
|
|
* produces for the prefix ending at that candidate, using only that sequence's
|
|
* own cache. Greedy accept/reject over these rows must then reproduce the plain
|
|
* batched-greedy token stream (proven end to end in test_spec_batch.c).
|
|
*
|
|
* A NOTE ON "BIT-IDENTICAL" (read this): with USE_BLAS the projection / FFN /
|
|
* LM-head GEMMs (pn_qmm) lower to Accelerate cblas_sgemm, whose K-dimension
|
|
* summation is reassociated as a function of the batch size M. So the logits of
|
|
* one row computed in an M-row batch are NOT byte-equal to the same row computed
|
|
* in a batch of a different size (e.g. a batch=1 replay) - they differ at the
|
|
* ~5e-5 level, which is far below any argmax gap and never flips a token (this
|
|
* is exactly the reassociation test_spec_decode.c documents). The emitted TOKEN
|
|
* stream is therefore bit-identical (the real losslessness guarantee), while raw
|
|
* logit bytes are equal only when M is held fixed. This gate is built around
|
|
* what is bitwise-achievable AND still pins down every ragged-indexing bug:
|
|
*
|
|
* [A] argmax-exact replay: every verify_batch row's ARGMAX equals an
|
|
* INDEPENDENT per-slot replay (prefill the committed prefix, then run
|
|
* pn_llama_forward_batch batch=1 over the block one token at a time so the
|
|
* block's earlier K/V is built by the SAME pn_qmm path), and the max
|
|
* |logit delta| stays at reassociation noise (<1e-2). A wrong row_pos /
|
|
* cache_pos / seq_len yields a totally different prediction (delta >> 1 and
|
|
* argmax flips), so this catches every within-block position bug.
|
|
* [B] bit-exact cross-cache independence (fixed M): rerun the SAME multi-seq
|
|
* verify_batch with ONE block's prefix+tokens swapped for different ones of
|
|
* the SAME length. Because M is unchanged the reassociation is identical,
|
|
* so every OTHER block's rows must be BYTE-identical and the swapped block
|
|
* must differ. A wrong row_seq index (reading another sequence's cache)
|
|
* would leak the swap into a neighbour's rows and break this exactly.
|
|
* [C] negative control: a reference built against the WRONG sequence's cache
|
|
* must differ grossly, proving the checks can actually see contamination.
|
|
*
|
|
* Usage: test_verify_batch [model.gguf]
|
|
*/
|
|
|
|
#include "../include/pantheon.h"
|
|
#include "../src/core/engine.h"
|
|
#include "../src/models/llama.h"
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <math.h>
|
|
|
|
#define DELTA_TOL 1e-2f /* reassociation noise is ~5e-5; a bug is >> 1 */
|
|
|
|
static int g_nv = 0;
|
|
|
|
static uint32_t argmax(const float *l, int n) {
|
|
int b = 0;
|
|
for (int i = 1; i < n; i++) if (l[i] > l[b]) b = i;
|
|
return (uint32_t)b;
|
|
}
|
|
|
|
static void rng_fill(uint32_t *p, int n, uint32_t seed) {
|
|
uint32_t x = 0x9e3779b9u * (seed + 1);
|
|
for (int i = 0; i < n; i++) {
|
|
x ^= x << 13; x ^= x >> 17; x ^= x << 5;
|
|
p[i] = (x % (uint32_t)(g_nv - 4)) + 2; /* valid, non-special ids */
|
|
}
|
|
}
|
|
|
|
/* Replay reference: fresh state, prefill `prefix`, then run forward_batch
|
|
* (batch=1) over block[0..n_blk-1] one token at a time, writing the logits
|
|
* predicted after block[j] into ref[j*nv]. Builds the block's earlier K/V via
|
|
* the same pn_qmm path verify_batch uses, so the ONLY numeric gap left is the
|
|
* batch-size reassociation. */
|
|
static int build_reference(pn_model *m, const uint32_t *prefix, int prefix_len,
|
|
const uint32_t *block, int n_blk, float *ref) {
|
|
pn_state *st = pn_state_new(m);
|
|
if (!st) return -1;
|
|
tm_tensor *pl = pn_llama_prefill(m, prefix, NULL, prefix_len, st);
|
|
if (!pl) { pn_state_free(st); return -1; }
|
|
tm_free(pl);
|
|
for (int j = 0; j < n_blk; j++) {
|
|
uint32_t tok = block[j];
|
|
pn_state *sts[1] = { st };
|
|
if (pn_llama_forward_batch(m, &tok, NULL, sts, 1, ref + (size_t)j * g_nv) != 0) {
|
|
pn_state_free(st); return -1;
|
|
}
|
|
}
|
|
pn_state_free(st);
|
|
return 0;
|
|
}
|
|
|
|
/* Prefill fresh states for every sequence and run ONE ragged verify_batch,
|
|
* writing [M, nv] logits into out. Returns 0 on success, or asserts pos was not
|
|
* advanced (verify never commits). */
|
|
static int run_verify(pn_model *m, const uint32_t *const *prefixes,
|
|
const int *prefix_len, const uint32_t *const *blocks,
|
|
const int *n_tok, int n_seqs, float *out) {
|
|
int M = 0; for (int b = 0; b < n_seqs; b++) M += n_tok[b];
|
|
pn_state **states = malloc((size_t)n_seqs * sizeof(pn_state *));
|
|
uint32_t *cat = malloc((size_t)M * sizeof(uint32_t));
|
|
int rc = 0, r = 0;
|
|
for (int b = 0; b < n_seqs; b++) {
|
|
states[b] = pn_state_new(m);
|
|
tm_tensor *pl = pn_llama_prefill(m, prefixes[b], NULL, prefix_len[b], states[b]);
|
|
if (!pl) { rc = -1; goto done; }
|
|
tm_free(pl);
|
|
for (int j = 0; j < n_tok[b]; j++) cat[r++] = blocks[b][j];
|
|
}
|
|
if (pn_llama_verify_batch(m, cat, states, n_tok, n_seqs, out) != 0) { rc = -1; goto done; }
|
|
for (int b = 0; b < n_seqs; b++)
|
|
if (states[b]->pos != prefix_len[b]) {
|
|
fprintf(stderr, " seq %d pos advanced (%d != %d) - verify must not commit\n",
|
|
b, states[b]->pos, prefix_len[b]);
|
|
rc = -1;
|
|
}
|
|
done:
|
|
for (int b = 0; b < n_seqs; b++) if (states[b]) pn_state_free(states[b]);
|
|
free(states); free(cat);
|
|
return rc;
|
|
}
|
|
|
|
/* Check A: every row's argmax equals its independent replay, and the logit
|
|
* delta stays at reassociation noise. */
|
|
static int check_argmax_replay(pn_model *m, const char *label,
|
|
const uint32_t *const *prefixes, const int *prefix_len,
|
|
const uint32_t *const *blocks, const int *n_tok, int n_seqs) {
|
|
int M = 0; for (int b = 0; b < n_seqs; b++) M += n_tok[b];
|
|
float *vb = malloc((size_t)M * g_nv * sizeof(float));
|
|
if (run_verify(m, prefixes, prefix_len, blocks, n_tok, n_seqs, vb) != 0) {
|
|
fprintf(stderr, "%s: verify run failed\n", label); free(vb); return 1;
|
|
}
|
|
int argmax_mism = 0; float maxd = 0.0f; int r = 0;
|
|
for (int b = 0; b < n_seqs; b++) {
|
|
float *ref = malloc((size_t)n_tok[b] * g_nv * sizeof(float));
|
|
if (build_reference(m, prefixes[b], prefix_len[b], blocks[b], n_tok[b], ref) != 0) {
|
|
fprintf(stderr, "%s: reference build failed\n", label); free(ref); free(vb); return 1;
|
|
}
|
|
for (int j = 0; j < n_tok[b]; j++) {
|
|
const float *row = vb + (size_t)(r + j) * g_nv;
|
|
const float *rr = ref + (size_t)j * g_nv;
|
|
if (argmax(row, g_nv) != argmax(rr, g_nv)) argmax_mism++;
|
|
for (int i = 0; i < g_nv; i++) {
|
|
float d = fabsf(row[i] - rr[i]); if (d > maxd) maxd = d;
|
|
}
|
|
}
|
|
free(ref); r += n_tok[b];
|
|
}
|
|
int ok = (argmax_mism == 0 && maxd < DELTA_TOL);
|
|
printf(" [A] %-20s seqs=%d M=%2d argmax-mism=%d max|delta|=%.2e %s\n",
|
|
label, n_seqs, M, argmax_mism, maxd, ok ? "OK" : "FAIL");
|
|
free(vb);
|
|
return ok ? 0 : 1;
|
|
}
|
|
|
|
/* Check B: bit-exact cross-cache independence. Swap block `swap`'s prefix and
|
|
* tokens for a different set of the SAME length; M is unchanged so the GEMM
|
|
* reassociation is identical. Every other block's rows must be BYTE-identical,
|
|
* and the swapped block's rows must differ. */
|
|
static int check_independence(pn_model *m, const char *label,
|
|
const uint32_t *const *prefixes, const int *prefix_len,
|
|
const uint32_t *const *blocks, const int *n_tok, int n_seqs,
|
|
int swap, const uint32_t *alt_prefix, const uint32_t *alt_block) {
|
|
int M = 0, off_swap = 0; for (int b = 0; b < n_seqs; b++) { if (b < swap) off_swap += n_tok[b]; M += n_tok[b]; }
|
|
float *base = malloc((size_t)M * g_nv * sizeof(float));
|
|
float *alt = malloc((size_t)M * g_nv * sizeof(float));
|
|
if (run_verify(m, prefixes, prefix_len, blocks, n_tok, n_seqs, base) != 0) {
|
|
free(base); free(alt); return 1;
|
|
}
|
|
const uint32_t *pf2[8]; const uint32_t *bl2[8];
|
|
for (int b = 0; b < n_seqs; b++) { pf2[b] = prefixes[b]; bl2[b] = blocks[b]; }
|
|
pf2[swap] = alt_prefix; bl2[swap] = alt_block; /* same prefix_len[swap], n_tok[swap] */
|
|
if (run_verify(m, pf2, prefix_len, bl2, n_tok, n_seqs, alt) != 0) {
|
|
free(base); free(alt); return 1;
|
|
}
|
|
int other_diff = 0, swapped_same = 0, r = 0;
|
|
for (int b = 0; b < n_seqs; b++) {
|
|
for (int j = 0; j < n_tok[b]; j++) {
|
|
int eq = memcmp(base + (size_t)(r + j) * g_nv, alt + (size_t)(r + j) * g_nv,
|
|
(size_t)g_nv * sizeof(float)) == 0;
|
|
if (b == swap) { if (eq) swapped_same++; }
|
|
else { if (!eq) other_diff++; }
|
|
}
|
|
r += n_tok[b];
|
|
}
|
|
(void)off_swap;
|
|
int ok = (other_diff == 0 && swapped_same == 0);
|
|
printf(" [B] %-20s swap seq %d: other-blocks-changed=%d swapped-unchanged=%d %s\n",
|
|
label, swap, other_diff, swapped_same,
|
|
ok ? "OK (byte-identical, no contamination)" : "FAIL");
|
|
free(base); free(alt);
|
|
return ok ? 0 : 1;
|
|
}
|
|
|
|
int main(int argc, char **argv) {
|
|
const char *path = (argc > 1) ? argv[1] : "models/lightonocr-lm-Q4_K_M.gguf";
|
|
pn_model *model = pn_load(path);
|
|
if (!model) { fprintf(stderr, "load failed: %s\n", path); return 1; }
|
|
g_nv = model->config.n_vocab;
|
|
printf("model: %s n_vocab=%d n_layer=%d n_embd=%d arch=%d\n",
|
|
path, g_nv, model->config.n_layer, model->config.n_embd, model->config.arch);
|
|
|
|
{ /* support gate */
|
|
pn_state *s = pn_state_new(model);
|
|
uint32_t t = 2; int nt = 1;
|
|
float *lg = malloc((size_t)g_nv * sizeof(float));
|
|
int ok = lg && pn_llama_verify_batch(model, &t, &s, &nt, 1, lg) == 0;
|
|
free(lg); pn_state_free(s);
|
|
if (!ok) { printf("[SKIP] model not on the quantized non-MoE verify path\n");
|
|
pn_model_free(model); return 0; }
|
|
}
|
|
|
|
int fail = 0;
|
|
|
|
enum { P0 = 32, P1 = 17, P2 = 48, P3 = 5 };
|
|
uint32_t pre0[P0], pre1[P1], pre2[P2], pre3[P3], preX[P1];
|
|
rng_fill(pre0, P0, 10); rng_fill(pre1, P1, 11);
|
|
rng_fill(pre2, P2, 12); rng_fill(pre3, P3, 13); rng_fill(preX, P1, 77);
|
|
uint32_t blk0[9], blk1[9], blk2[9], blk3[9], blkX[9];
|
|
rng_fill(blk0, 9, 100); rng_fill(blk1, 9, 101);
|
|
rng_fill(blk2, 9, 102); rng_fill(blk3, 9, 103); rng_fill(blkX, 9, 200);
|
|
|
|
printf("\n== M2: single sequence, block [seed, d0..d_{K-1}] ==\n");
|
|
for (int K = 0; K <= 8; K += 2) {
|
|
const uint32_t *pf[] = { pre0 }; int pl[] = { P0 };
|
|
const uint32_t *bl[] = { blk0 }; int nt[] = { K + 1 };
|
|
char lbl[24]; snprintf(lbl, sizeof(lbl), "K=%d", K);
|
|
fail |= check_argmax_replay(model, lbl, pf, pl, bl, nt, 1);
|
|
}
|
|
|
|
printf("\n== M3: ragged multi-sequence (mixed lengths, one plain K=1 block) ==\n");
|
|
{
|
|
const uint32_t *pf[] = { pre0, pre1, pre2 }; int pl[] = { P0, P1, P2 };
|
|
const uint32_t *bl[] = { blk0, blk1, blk2 }; int nt[] = { 5, 1, 8 };
|
|
fail |= check_argmax_replay(model, "3-seq {5,1,8}", pf, pl, bl, nt, 3);
|
|
/* swap the MIDDLE (length-1) block: proves K=1 blocks are isolated too */
|
|
fail |= check_independence(model, "3-seq {5,1,8}", pf, pl, bl, nt, 3, 1, preX, blkX);
|
|
}
|
|
{
|
|
const uint32_t *pf[] = { pre3, pre0, pre1, pre2 }; int pl[] = { P3, P0, P1, P2 };
|
|
const uint32_t *bl[] = { blk3, blk0, blk1, blk2 }; int nt[] = { 4, 4, 2, 6 };
|
|
fail |= check_argmax_replay(model, "4-seq {4,4,2,6}", pf, pl, bl, nt, 4);
|
|
fail |= check_independence(model, "4-seq {4,4,2,6}", pf, pl, bl, nt, 4, 2, preX, blkX);
|
|
}
|
|
{
|
|
const uint32_t *pf[] = { pre0, pre1, pre2, pre3 }; int pl[] = { P0, P1, P2, P3 };
|
|
const uint32_t *bl[] = { blk0, blk1, blk2, blk3 }; int nt[] = { 1, 1, 1, 1 };
|
|
fail |= check_argmax_replay(model, "4-seq all K=1", pf, pl, bl, nt, 4);
|
|
}
|
|
|
|
printf("\n== [C] negative control: reference against the WRONG cache must differ ==\n");
|
|
{
|
|
pn_state *sa = pn_state_new(model), *sb = pn_state_new(model);
|
|
tm_tensor *p;
|
|
p = pn_llama_prefill(model, pre0, NULL, P0, sa); tm_free(p);
|
|
p = pn_llama_prefill(model, pre1, NULL, P1, sb); tm_free(p);
|
|
pn_state *states[2] = { sa, sb }; int nt[2] = { 4, 4 };
|
|
uint32_t cat[8];
|
|
for (int j = 0; j < 4; j++) { cat[j] = blk0[j]; cat[4 + j] = blk1[j]; }
|
|
float *vb = malloc((size_t)8 * g_nv * sizeof(float));
|
|
pn_llama_verify_batch(model, cat, states, nt, 2, vb);
|
|
float *ref = malloc((size_t)4 * g_nv * sizeof(float));
|
|
build_reference(model, pre1, P1, blk0, 4, ref); /* WRONG prefix for seq 0 */
|
|
int gross = 0;
|
|
for (int j = 0; j < 4; j++) {
|
|
float md = 0; for (int i = 0; i < g_nv; i++) { float d = fabsf(vb[(size_t)j*g_nv+i]-ref[(size_t)j*g_nv+i]); if (d>md) md=d; }
|
|
if (md > DELTA_TOL) gross++;
|
|
}
|
|
printf(" wrong-cache reference differs grossly on %d/4 rows %s\n",
|
|
gross, (gross == 4) ? "GOOD" : "SUSPICIOUS");
|
|
if (gross != 4) fail = 1;
|
|
free(vb); free(ref); pn_state_free(sa); pn_state_free(sb);
|
|
}
|
|
|
|
printf("\nRESULT: %s\n", fail ? "FAIL"
|
|
: "PASS (argmax-exact vs replay, bit-exact cross-cache isolation, negative control)");
|
|
pn_model_free(model);
|
|
return fail ? 1 : 0;
|
|
}
|