86 lines
3.1 KiB
C
86 lines
3.1 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/* End-to-end per-token generation throughput: single-thread vs pool-parallel
|
|
* pn_qmv on the LightOnOCR LM. Drives the decoder token-by-token (greedy) and
|
|
* verifies the two paths produce the identical token sequence. Bounded run. */
|
|
|
|
#include "pantheon.h"
|
|
#include "src/core/engine.h"
|
|
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <time.h>
|
|
|
|
static double now_ms(void) {
|
|
struct timespec ts;
|
|
clock_gettime(CLOCK_MONOTONIC, &ts);
|
|
return ts.tv_sec * 1000.0 + ts.tv_nsec / 1e6;
|
|
}
|
|
|
|
/* Greedy-generate n_gen tokens starting from a fixed seed token. Records the
|
|
* produced token ids and the mean ms/token. Returns 0 on success. */
|
|
static int run_gen(pn_model *model, int n_gen, uint32_t seed_tok,
|
|
uint32_t *toks_out, double *ms_per_tok) {
|
|
pn_state *st = pn_state_new(model);
|
|
if (!st) return -1;
|
|
uint32_t tok = seed_tok;
|
|
double t0 = now_ms();
|
|
for (int i = 0; i < n_gen; i++) {
|
|
pn_output *o = pn_infer(model, st, tok);
|
|
if (!o || o->type != PN_OUT_LOGITS) { pn_state_free(st); return -1; }
|
|
int best = 0;
|
|
float bv = o->data[0];
|
|
for (int v = 1; v < o->n; v++) {
|
|
if (o->data[v] > bv) { bv = o->data[v]; best = v; }
|
|
}
|
|
tok = (uint32_t)best;
|
|
toks_out[i] = tok;
|
|
pn_output_free(o);
|
|
}
|
|
*ms_per_tok = (now_ms() - t0) / n_gen;
|
|
pn_state_free(st);
|
|
return 0;
|
|
}
|
|
|
|
int main(int argc, char **argv) {
|
|
const char *path = argc > 1 ? argv[1]
|
|
: "models/lightonocr-lm-Q4_K_M.gguf";
|
|
int n_gen = argc > 2 ? atoi(argv[2]) : 32;
|
|
|
|
pn_model *model = pn_load(path);
|
|
if (!model) { fprintf(stderr, "load failed: %s\n", path); return 1; }
|
|
|
|
uint32_t *a = malloc(n_gen * sizeof(uint32_t));
|
|
uint32_t *b = malloc(n_gen * sizeof(uint32_t));
|
|
double ms_mt = 0, ms_1t = 0;
|
|
uint32_t seed = 1; /* arbitrary fixed start token */
|
|
|
|
/* Warm up (page-in weights) */
|
|
double junk; run_gen(model, 2, seed, a, &junk);
|
|
|
|
/* AFTER: multithreaded pn_qmv (model->pool registered each forward) */
|
|
struct fw_pool *saved_pool = model->pool;
|
|
if (run_gen(model, n_gen, seed, a, &ms_mt) != 0) { fprintf(stderr, "gen MT failed\n"); return 1; }
|
|
|
|
/* BEFORE: force single-thread by detaching the pool. pn_llama_forward will
|
|
* register NULL, so pn_qmv runs serially. */
|
|
model->pool = NULL;
|
|
if (run_gen(model, n_gen, seed, b, &ms_1t) != 0) { fprintf(stderr, "gen 1T failed\n"); return 1; }
|
|
model->pool = saved_pool;
|
|
|
|
int mismatch = 0;
|
|
for (int i = 0; i < n_gen; i++) if (a[i] != b[i]) mismatch++;
|
|
|
|
printf("\ngenerate %d tokens on %s\n", n_gen, path);
|
|
printf(" single-thread pn_qmv : %.2f ms/token (%.1f tok/s)\n", ms_1t, 1000.0 / ms_1t);
|
|
printf(" pool-parallel pn_qmv : %.2f ms/token (%.1f tok/s)\n", ms_mt, 1000.0 / ms_mt);
|
|
printf(" speedup : %.2fx\n", ms_1t / ms_mt);
|
|
printf(" token-seq mismatch : %d -> %s\n", mismatch,
|
|
mismatch == 0 ? "IDENTICAL OUTPUT" : "DIVERGED");
|
|
|
|
free(a); free(b);
|
|
pn_model_free(model);
|
|
return mismatch == 0 ? 0 : 1;
|
|
}
|