83 lines
3.0 KiB
C
83 lines
3.0 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/* Correctness check for pn_llama_forward_batch.
|
|
*
|
|
* (1) A single sequence decoded via forward_batch(B=1) must produce the SAME
|
|
* greedy token stream as the sequential pn_llama_forward (pn_infer).
|
|
* (2) Two sequences with different seeds, decoded TOGETHER in one batch, must
|
|
* each match their own independent single-stream decode — proving the B
|
|
* caches don't cross-contaminate.
|
|
*/
|
|
|
|
#include "pantheon.h"
|
|
#include "src/core/engine.h"
|
|
#include "src/models/llama.h"
|
|
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
|
|
static uint32_t argmax(const float *v, int n) {
|
|
int best = 0; float bv = v[0];
|
|
for (int i = 1; i < n; i++) if (v[i] > bv) { bv = v[i]; best = i; }
|
|
return (uint32_t)best;
|
|
}
|
|
|
|
/* Reference: sequential decode of `n` greedy tokens from a seed. */
|
|
static void decode_seq(pn_model *m, uint32_t seed, int n, uint32_t *out) {
|
|
pn_state *st = pn_state_new(m);
|
|
uint32_t tok = seed;
|
|
for (int i = 0; i < n; i++) {
|
|
pn_output *o = pn_infer(m, st, tok);
|
|
tok = argmax(o->data, o->n);
|
|
out[i] = tok;
|
|
pn_output_free(o);
|
|
}
|
|
pn_state_free(st);
|
|
}
|
|
|
|
int main(int argc, char **argv) {
|
|
const char *path = argc > 1 ? argv[1] : "models/lightonocr-lm-Q4_K_M.gguf";
|
|
int n = argc > 2 ? atoi(argv[2]) : 24;
|
|
pn_model *m = pn_load(path);
|
|
if (!m) { fprintf(stderr, "load failed\n"); return 1; }
|
|
int nv = m->config.n_vocab;
|
|
|
|
uint32_t seedA = 5, seedB = 99;
|
|
uint32_t *refA = malloc(n * sizeof(uint32_t));
|
|
uint32_t *refB = malloc(n * sizeof(uint32_t));
|
|
decode_seq(m, seedA, n, refA);
|
|
decode_seq(m, seedB, n, refB);
|
|
|
|
/* Batched decode of both sequences together. */
|
|
pn_state *st[2] = { pn_state_new(m), pn_state_new(m) };
|
|
uint32_t tok[2] = { seedA, seedB };
|
|
float *logits = malloc((size_t)2 * nv * sizeof(float));
|
|
uint32_t *batA = malloc(n * sizeof(uint32_t));
|
|
uint32_t *batB = malloc(n * sizeof(uint32_t));
|
|
for (int i = 0; i < n; i++) {
|
|
if (pn_llama_forward_batch(m, tok, NULL, st, 2, logits) != 0) {
|
|
fprintf(stderr, "forward_batch unsupported\n"); return 1;
|
|
}
|
|
tok[0] = argmax(logits + 0 * nv, nv);
|
|
tok[1] = argmax(logits + 1 * nv, nv);
|
|
batA[i] = tok[0];
|
|
batB[i] = tok[1];
|
|
}
|
|
|
|
int mismatchA = 0, mismatchB = 0;
|
|
for (int i = 0; i < n; i++) {
|
|
if (batA[i] != refA[i]) mismatchA++;
|
|
if (batB[i] != refB[i]) mismatchB++;
|
|
}
|
|
printf("seq A: %d/%d tokens match sequential decode\n", n - mismatchA, n);
|
|
printf("seq B: %d/%d tokens match sequential decode\n", n - mismatchB, n);
|
|
printf("%s\n", (mismatchA == 0 && mismatchB == 0) ? "PASS: batched == sequential, caches independent"
|
|
: "FAIL: batched decode diverges");
|
|
|
|
pn_state_free(st[0]); pn_state_free(st[1]);
|
|
free(refA); free(refB); free(batA); free(batB); free(logits);
|
|
pn_model_free(m);
|
|
return (mismatchA == 0 && mismatchB == 0) ? 0 : 1;
|
|
}
|