93 lines
3.5 KiB
C
93 lines
3.5 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/* bench_verify_rows.c - M1 milestone bench for batched speculative decoding.
|
|
*
|
|
* Measures the cost of one pn_llama_verify pass as a function of the number of
|
|
* rows M (candidate tokens verified in one shot). The verify pass is the kernel
|
|
* the batched speculative decoder runs once per step over the concatenation of
|
|
* every slot's [pending, drafts...] block, so this table confirms the payoff
|
|
* (verify of M rows costs far less than M separate single-token forwards) and
|
|
* picks a sane default draft width K.
|
|
*
|
|
* The pass runs on a state prefilled to a realistic context length so the
|
|
* attention cost (which grows with committed length) is representative. pos is
|
|
* never advanced by verify, so the same committed cache is reused for every M
|
|
* and every repetition.
|
|
*
|
|
* Usage: bench_verify_rows [model.gguf] [ctx_len] [reps]
|
|
*/
|
|
|
|
#include "../include/pantheon.h"
|
|
#include "../src/core/engine.h"
|
|
#include "../src/models/llama.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 * 1e3 + ts.tv_nsec / 1e6;
|
|
}
|
|
|
|
int main(int argc, char **argv) {
|
|
const char *path = (argc > 1) ? argv[1] : "models/lightonocr-lm-Q4_K_M.gguf";
|
|
int ctx_len = (argc > 2) ? atoi(argv[2]) : 256;
|
|
int reps = (argc > 3) ? atoi(argv[3]) : 20;
|
|
|
|
pn_model *model = pn_load(path);
|
|
if (!model) { fprintf(stderr, "load failed: %s\n", path); return 1; }
|
|
pn_state *state = pn_state_new(model);
|
|
if (!state) { fprintf(stderr, "state failed\n"); return 1; }
|
|
int nv = model->config.n_vocab;
|
|
|
|
/* Prefill a deterministic context so verify runs against a warm cache. */
|
|
uint32_t *prompt = malloc((size_t)ctx_len * sizeof(uint32_t));
|
|
uint32_t x = 0x9e3779b9u;
|
|
for (int i = 0; i < ctx_len; i++) {
|
|
x ^= x << 13; x ^= x >> 17; x ^= x << 5;
|
|
prompt[i] = (x % (uint32_t)(nv - 4)) + 2;
|
|
}
|
|
pn_state_reset(state);
|
|
tm_tensor *pl = pn_llama_prefill(model, prompt, NULL, ctx_len, state);
|
|
if (!pl) {
|
|
printf("[SKIP] model not on the quantized non-MoE verify path\n");
|
|
free(prompt); pn_state_free(state); pn_model_free(model);
|
|
return 0;
|
|
}
|
|
tm_free(pl);
|
|
printf("model: %s n_vocab=%d n_layer=%d n_embd=%d ctx_len=%d reps=%d\n",
|
|
path, nv, model->config.n_layer, model->config.n_embd, ctx_len, reps);
|
|
|
|
const int Ms[] = { 6, 8, 12, 16, 24, 32, 48, 64 };
|
|
const int nM = (int)(sizeof(Ms) / sizeof(Ms[0]));
|
|
|
|
uint32_t cand[64];
|
|
for (int i = 0; i < 64; i++) {
|
|
x ^= x << 13; x ^= x >> 17; x ^= x << 5;
|
|
cand[i] = (x % (uint32_t)(nv - 4)) + 2;
|
|
}
|
|
float *logits = malloc((size_t)64 * (size_t)nv * sizeof(float));
|
|
if (!logits) { fprintf(stderr, "oom\n"); return 1; }
|
|
|
|
printf("\n M ms/call ms/row rows/s\n");
|
|
printf(" -- -------- ------- --------\n");
|
|
for (int mi = 0; mi < nM; mi++) {
|
|
int M = Ms[mi];
|
|
/* one warm-up */
|
|
if (pn_llama_verify(model, cand, NULL, M, state, logits) != 0) {
|
|
fprintf(stderr, "verify failed at M=%d\n", M); return 1;
|
|
}
|
|
double t0 = now_ms();
|
|
for (int r = 0; r < reps; r++)
|
|
pn_llama_verify(model, cand, NULL, M, state, logits);
|
|
double dt = (now_ms() - t0) / reps;
|
|
printf(" %2d %8.3f %7.4f %8.1f\n", M, dt, dt / M, 1000.0 * M / dt);
|
|
}
|
|
|
|
free(logits); free(prompt);
|
|
pn_state_free(state); pn_model_free(model);
|
|
return 0;
|
|
}
|