401 lines
19 KiB
C
401 lines
19 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/* test_spec_batch.c - the M4/M5 losslessness gate for BATCHED speculative decode.
|
|
*
|
|
* Proves the batched speculative decoder (pn_batch_decoder_step_spec) emits a
|
|
* transcript BIT-IDENTICAL to plain batched greedy decode for ANY draft, across
|
|
* batch sizes, draft widths, draft qualities, mixed spec/plain slots, staggered
|
|
* admits, and stop (EOS / max_tokens) handling. Drafts only change how many
|
|
* tokens each verify pass yields, never which tokens are produced, because
|
|
* step_spec emits only the penalized argmax of the model's own verify logits and
|
|
* accepts a draft solely when it already equals that argmax.
|
|
*
|
|
* References (greedy, temperature 0):
|
|
* A = sequential pn_llama_forward greedy, mirroring the plain per-slot step
|
|
* loop (push pending -> forward -> penalized-argmax next pending), giving
|
|
* both the token stream and its decoded UTF-8 text.
|
|
* B = plain pn_batch_decoder (the shipped batched greedy decoder).
|
|
* C = pn_batch_decoder_step_spec under drafts {none, oracle=replay of A,
|
|
* garbage, truncated}.
|
|
* The shipped decoders (B, C) expose only decoded text, so the end-to-end gate
|
|
* is TEXT byte-for-byte: A-text == B-text == C-text for every page and config.
|
|
* A separate TOKEN-for-token check (sequential A vs a batched forward_batch loop
|
|
* captured in-test) confirms the batched kernel is token-lossless per page.
|
|
*
|
|
* Usage: test_spec_batch [model.gguf]
|
|
*/
|
|
|
|
#include "../include/pantheon.h"
|
|
#include "../src/core/engine.h"
|
|
#include "../src/models/llama.h"
|
|
#include "tokenizer/tokenizer.h"
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
|
|
#define PROMPT_LEN 8
|
|
#define GEN 24
|
|
#define MAXB 6
|
|
|
|
static int g_nv = 0, g_eos = -1, g_imend = -1;
|
|
|
|
/* --- greedy sampler mirror (sample_token's temperature<=0 path) --- */
|
|
static uint32_t pen_argmax(float *logits, int nv, float rp, uint32_t *recent, int nrec) {
|
|
if (rp != 1.0f && recent && nrec > 0) {
|
|
for (int i = 0; i < nrec; i++) {
|
|
uint32_t t = recent[i];
|
|
if (t < (uint32_t)nv) { if (logits[t] > 0.0f) logits[t] /= rp; else logits[t] *= rp; }
|
|
}
|
|
}
|
|
int b = 0;
|
|
for (int i = 1; i < nv; i++) if (logits[i] > logits[b]) b = i;
|
|
return (uint32_t)b;
|
|
}
|
|
static void ring_push(uint32_t *recent, int *nrec, int maxrec, uint32_t tok) {
|
|
if (*nrec < maxrec) recent[(*nrec)++] = tok;
|
|
else { memmove(recent, recent + 1, (maxrec - 1) * sizeof(uint32_t)); recent[maxrec - 1] = tok; }
|
|
}
|
|
|
|
static pn_sampling_params mkparams(float rp) {
|
|
pn_sampling_params p = pn_sampling_defaults();
|
|
p.temperature = 0.0f; p.repeat_penalty = rp; p.repeat_last_n = 64;
|
|
p.top_k = 0; p.top_p = 1.0f;
|
|
return p;
|
|
}
|
|
|
|
static void make_prompt(uint32_t *p, int n, int idx) {
|
|
uint32_t x = 0x9e3779b9u * (uint32_t)(idx + 7);
|
|
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;
|
|
}
|
|
}
|
|
|
|
static int is_stop(uint32_t t) { return ((g_eos >= 0 && (int)t == g_eos) || (int)t == g_imend); }
|
|
|
|
/* Prefill a fresh state, return (state, seed=argmax of prefill logits). */
|
|
static pn_state* prefill_state(pn_model *m, const uint32_t *prompt, int plen, uint32_t *seed) {
|
|
pn_state *s = pn_state_new(m);
|
|
tm_tensor *l = pn_llama_prefill(m, prompt, NULL, plen, s);
|
|
if (!l) { pn_state_free(s); return NULL; }
|
|
*seed = pen_argmax((float *)l->data, g_nv, 1.0f, NULL, 0);
|
|
tm_free(l);
|
|
return s;
|
|
}
|
|
|
|
/* Reference A: sequential greedy via pn_llama_forward, mirroring plain step.
|
|
* Fills tok[0..*ntok-1] (tokens the decoder would output), returns 0/-1. */
|
|
static int gen_seq(pn_model *m, const uint32_t *prompt, int plen, float rp,
|
|
int max_tokens, uint32_t *tok, int *ntok) {
|
|
uint32_t seed; pn_state *s = prefill_state(m, prompt, plen, &seed);
|
|
if (!s) return -1;
|
|
uint32_t recent[64]; int nrec = 0, maxrec = 64;
|
|
uint32_t pending = seed; int no = 0;
|
|
while (!(is_stop(pending) || no >= max_tokens)) {
|
|
tok[no++] = pending;
|
|
ring_push(recent, &nrec, maxrec, pending);
|
|
tm_tensor *l = m->forward(m, pending, s);
|
|
if (!l) { pn_state_free(s); return -1; }
|
|
pending = pen_argmax((float *)l->data, g_nv, rp, recent, nrec);
|
|
tm_free(l);
|
|
}
|
|
*ntok = no;
|
|
pn_state_free(s);
|
|
return 0;
|
|
}
|
|
|
|
/* Batched token reference: decode `np` pages together via pn_llama_forward_batch
|
|
* (argmax + penalty per slot), capturing each page's token stream. This is the
|
|
* plain batched greedy path at the token level (what pn_batch_decoder does
|
|
* internally), used for the token-for-token check vs sequential A. */
|
|
typedef struct { uint32_t tok[GEN]; int no; uint32_t recent[64]; int nrec; uint32_t pending; int done; } bslot;
|
|
static int gen_batch_tokens(pn_model *m, const uint32_t *const *prompts, int np,
|
|
float rp, int max_tokens, bslot *slots) {
|
|
pn_state **st = malloc((size_t)np * sizeof(pn_state *));
|
|
for (int p = 0; p < np; p++) {
|
|
uint32_t seed; st[p] = prefill_state(m, prompts[p], PROMPT_LEN, &seed);
|
|
slots[p].no = 0; slots[p].nrec = 0; slots[p].pending = seed; slots[p].done = 0;
|
|
}
|
|
float *logits = malloc((size_t)np * g_nv * sizeof(float));
|
|
uint32_t *toks = malloc((size_t)np * sizeof(uint32_t));
|
|
pn_state **act = malloc((size_t)np * sizeof(pn_state *));
|
|
int *aslot = malloc((size_t)np * sizeof(int));
|
|
int remaining = np;
|
|
while (remaining > 0) {
|
|
int na = 0;
|
|
for (int p = 0; p < np; p++) {
|
|
bslot *b = &slots[p];
|
|
if (b->done) continue;
|
|
if (is_stop(b->pending) || b->no >= max_tokens) { b->done = 1; remaining--; continue; }
|
|
b->tok[b->no++] = b->pending;
|
|
ring_push(b->recent, &b->nrec, 64, b->pending);
|
|
act[na] = st[p]; toks[na] = b->pending; aslot[na] = p; na++;
|
|
}
|
|
if (na == 0) break;
|
|
if (pn_llama_forward_batch(m, toks, NULL, act, na, logits) != 0) {
|
|
for (int p = 0; p < np; p++) pn_state_free(st[p]);
|
|
free(st); free(logits); free(toks); free(act); free(aslot); return -1;
|
|
}
|
|
for (int j = 0; j < na; j++) {
|
|
bslot *b = &slots[aslot[j]];
|
|
b->pending = pen_argmax(logits + (size_t)j * g_nv, g_nv, rp, b->recent, b->nrec);
|
|
}
|
|
}
|
|
for (int p = 0; p < np; p++) pn_state_free(st[p]);
|
|
free(st); free(logits); free(toks); free(act); free(aslot);
|
|
return 0;
|
|
}
|
|
|
|
static char* decode_text(pn_model *m, const uint32_t *tok, int n) {
|
|
int *ids = malloc((size_t)(n > 0 ? n : 1) * sizeof(int));
|
|
for (int i = 0; i < n; i++) ids[i] = (int)tok[i];
|
|
char *raw = ft_decode(m->tokenizer, ids, n);
|
|
char *text = raw ? gpt2_bytes_decode(raw) : NULL;
|
|
free(raw); free(ids);
|
|
return text ? text : strdup("");
|
|
}
|
|
|
|
/* --- draft providers --- */
|
|
typedef enum { D_NONE, D_ORACLE, D_GARBAGE, D_TRUNC } dkind;
|
|
typedef struct { const uint32_t *ref; int ref_len; dkind kind; unsigned rng; } dctx;
|
|
static int draft_cb(const uint32_t *prefix, int n, uint32_t *out, int max_k, void *ud) {
|
|
dctx *c = (dctx *)ud; (void)prefix;
|
|
if (c->kind == D_NONE || max_k <= 0) return 0;
|
|
int k = max_k;
|
|
if (c->kind == D_TRUNC) k = (max_k > 3) ? 3 : max_k; /* short draft */
|
|
for (int j = 0; j < k; j++) {
|
|
if (c->kind == D_GARBAGE) {
|
|
c->rng ^= c->rng << 13; c->rng ^= c->rng >> 17; c->rng ^= c->rng << 5;
|
|
out[j] = c->rng % (uint32_t)g_nv;
|
|
} else {
|
|
int idx = n + j; /* oracle: the true next tokens */
|
|
out[j] = (idx < c->ref_len) ? c->ref[idx] : 0;
|
|
}
|
|
}
|
|
return k;
|
|
}
|
|
|
|
/* Run the shipped decoder (plain or spec) over `np` pre-supplied (state,seed)
|
|
* pages admitted up front, collect each page's text by page_id. */
|
|
static int run_decoder(pn_model *m, pn_state **states, const uint32_t *seeds,
|
|
int np, float rp, int max_tokens, int max_batch,
|
|
int spec, int max_k, pn_draft_fn *dfns, void **duds,
|
|
char **out_text, long *passes, long *drafted, long *accepted) {
|
|
pn_sampling_params params = mkparams(rp);
|
|
pn_batch_decoder *d = pn_batch_decoder_new(m, max_batch, max_tokens, ¶ms);
|
|
if (!d) return -1;
|
|
if (spec && pn_batch_decoder_spec_config(d, max_k, 0.10, 4) != 0) { pn_batch_decoder_free(d); return -1; }
|
|
for (int p = 0; p < np; p++) {
|
|
int rc = spec ? pn_batch_decoder_admit_spec(d, states[p], p, seeds[p], dfns ? dfns[p] : NULL, duds ? duds[p] : NULL)
|
|
: pn_batch_decoder_admit(d, states[p], p, seeds[p]);
|
|
if (rc != 0) { pn_batch_decoder_free(d); return -1; }
|
|
}
|
|
pn_page_out outbuf[MAXB];
|
|
int rc = 0;
|
|
while (pn_batch_decoder_active(d) > 0) {
|
|
int nf = spec ? pn_batch_decoder_step_spec(d, outbuf, max_batch)
|
|
: pn_batch_decoder_step(d, outbuf, max_batch);
|
|
if (nf < 0) { rc = -1; break; }
|
|
for (int i = 0; i < nf; i++) {
|
|
out_text[outbuf[i].page_id] = outbuf[i].text; /* borrow; caller frees */
|
|
pn_state_free(outbuf[i].state);
|
|
}
|
|
}
|
|
if (spec) pn_batch_decoder_spec_stats(d, passes, drafted, accepted);
|
|
pn_batch_decoder_free(d);
|
|
return rc;
|
|
}
|
|
|
|
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;
|
|
g_eos = ft_eos_id(model->tokenizer);
|
|
g_imend = ft_token_to_id(model->tokenizer, "<|im_end|>");
|
|
printf("model: %s n_vocab=%d n_layer=%d arch=%d eos=%d im_end=%d\n",
|
|
path, g_nv, model->config.n_layer, model->config.arch, g_eos, g_imend);
|
|
|
|
{ /* 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;
|
|
const int Bs[] = { 1, 2, 6 };
|
|
const float RPs[] = { 1.0f, 1.15f };
|
|
const int Ks[] = { 0, 1, 8 };
|
|
const dkind kinds[] = { D_NONE, D_ORACLE, D_GARBAGE, D_TRUNC };
|
|
const char *kname[] = { "none", "oracle", "garbage", "trunc" };
|
|
|
|
uint32_t prompts[MAXB][PROMPT_LEN];
|
|
|
|
for (int bi = 0; bi < 3; bi++) {
|
|
int B = Bs[bi];
|
|
const uint32_t *pp[MAXB];
|
|
for (int p = 0; p < B; p++) { make_prompt(prompts[p], PROMPT_LEN, bi * 10 + p); pp[p] = prompts[p]; }
|
|
|
|
for (int ri = 0; ri < 2; ri++) {
|
|
float rp = RPs[ri];
|
|
printf("\n==== B=%d repeat_penalty=%.2f ====\n", B, rp);
|
|
|
|
/* Reference A: sequential greedy, per page (tokens + text). */
|
|
uint32_t refTok[MAXB][GEN]; int refN[MAXB]; char *refText[MAXB];
|
|
for (int p = 0; p < B; p++) {
|
|
if (gen_seq(model, pp[p], PROMPT_LEN, rp, GEN, refTok[p], &refN[p]) != 0) {
|
|
fprintf(stderr, "gen_seq failed\n"); fail = 1; goto out;
|
|
}
|
|
refText[p] = decode_text(model, refTok[p], refN[p]);
|
|
}
|
|
|
|
/* Token-for-token: batched forward_batch loop must equal sequential A. */
|
|
bslot *bs = malloc((size_t)B * sizeof(bslot));
|
|
if (gen_batch_tokens(model, pp, B, rp, GEN, bs) != 0) { fprintf(stderr, "batch tokens failed\n"); fail = 1; free(bs); goto out; }
|
|
int tokdiff = 0;
|
|
for (int p = 0; p < B; p++) {
|
|
if (bs[p].no != refN[p]) tokdiff++;
|
|
else for (int j = 0; j < refN[p]; j++) if (bs[p].tok[j] != refTok[p][j]) { tokdiff++; break; }
|
|
}
|
|
printf(" [tok] batched forward_batch vs sequential: %s\n", tokdiff ? "DIFFERS" : "token-identical");
|
|
if (tokdiff) fail = 1;
|
|
free(bs);
|
|
|
|
/* B: plain batch decoder. Compare text byte-for-byte to A. */
|
|
{
|
|
pn_state *st[MAXB]; uint32_t seeds[MAXB]; char *txt[MAXB] = {0};
|
|
for (int p = 0; p < B; p++) st[p] = prefill_state(model, pp[p], PROMPT_LEN, &seeds[p]);
|
|
if (run_decoder(model, st, seeds, B, rp, GEN, B, 0, 0, NULL, NULL, txt, NULL, NULL, NULL) != 0) { fail = 1; }
|
|
int bad = 0;
|
|
for (int p = 0; p < B; p++) { if (!txt[p] || strcmp(txt[p], refText[p]) != 0) bad++; free(txt[p]); }
|
|
printf(" [B ] plain batch decoder text: %s\n", bad ? "DIFFERS" : "byte-identical");
|
|
if (bad) fail = 1;
|
|
}
|
|
|
|
/* C: step_spec across K and draft strategies. */
|
|
for (int ki = 0; ki < 3; ki++) {
|
|
int K = Ks[ki];
|
|
int nkind = (K == 0) ? 1 : 4; /* K=0: only "none" */
|
|
for (int kd = 0; kd < nkind; kd++) {
|
|
dkind kind = kinds[kd];
|
|
pn_state *st[MAXB]; uint32_t seeds[MAXB]; char *txt[MAXB] = {0};
|
|
pn_draft_fn dfns[MAXB]; void *duds[MAXB]; dctx *ctxs[MAXB];
|
|
for (int p = 0; p < B; p++) {
|
|
st[p] = prefill_state(model, pp[p], PROMPT_LEN, &seeds[p]);
|
|
ctxs[p] = malloc(sizeof(dctx));
|
|
*ctxs[p] = (dctx){ refTok[p], refN[p], kind, 0xBEEF + p * 131 + kd };
|
|
dfns[p] = (K == 0 || kind == D_NONE) ? NULL : draft_cb;
|
|
duds[p] = ctxs[p];
|
|
}
|
|
long passes = 0, drafted = 0, accepted = 0;
|
|
if (run_decoder(model, st, seeds, B, rp, GEN, B, 1, K, dfns, duds,
|
|
txt, &passes, &drafted, &accepted) != 0) fail = 1;
|
|
int bad = 0;
|
|
for (int p = 0; p < B; p++) { if (!txt[p] || strcmp(txt[p], refText[p]) != 0) bad++; free(txt[p]); free(ctxs[p]); }
|
|
double ar = drafted ? (double)accepted / (double)drafted : 0.0;
|
|
printf(" [C ] K=%d %-8s text %s passes=%ld drafted=%ld accepted=%ld acc=%.2f\n",
|
|
K, kname[kd], bad ? "DIFFERS" : "byte-identical",
|
|
passes, drafted, accepted, ar);
|
|
if (bad) fail = 1;
|
|
}
|
|
}
|
|
|
|
for (int p = 0; p < B; p++) free(refText[p]);
|
|
}
|
|
}
|
|
|
|
/* Mixed spec/plain slots + staggered admit, B=6: pages 0,2,4 speculate
|
|
* (oracle), pages 1,3,5 ride plain; pages 3..5 admitted after two steps. */
|
|
printf("\n==== mixed spec/plain + staggered admit (B=6, rp=1.15) ====\n");
|
|
{
|
|
float rp = 1.15f; int B = 6, max_batch = 6, max_tokens = GEN;
|
|
const uint32_t *pp[6];
|
|
for (int p = 0; p < 6; p++) { make_prompt(prompts[p], PROMPT_LEN, 200 + p); pp[p] = prompts[p]; }
|
|
uint32_t refTok[6][GEN]; int refN[6]; char *refText[6];
|
|
for (int p = 0; p < 6; p++) { gen_seq(model, pp[p], PROMPT_LEN, rp, max_tokens, refTok[p], &refN[p]); refText[p] = decode_text(model, refTok[p], refN[p]); }
|
|
|
|
pn_sampling_params params = mkparams(rp);
|
|
pn_batch_decoder *d = pn_batch_decoder_new(model, max_batch, max_tokens, ¶ms);
|
|
pn_batch_decoder_spec_config(d, 8, 0.10, 4);
|
|
dctx *ctxs[6]; char *txt[6] = {0};
|
|
pn_state *st[6]; uint32_t seeds[6];
|
|
for (int p = 0; p < 6; p++) {
|
|
st[p] = prefill_state(model, pp[p], PROMPT_LEN, &seeds[p]);
|
|
ctxs[p] = malloc(sizeof(dctx));
|
|
*ctxs[p] = (dctx){ refTok[p], refN[p], D_ORACLE, 0u };
|
|
}
|
|
/* admit 0..2 up front; 0 & 2 speculate, 1 plain */
|
|
pn_batch_decoder_admit_spec(d, st[0], 0, seeds[0], draft_cb, ctxs[0]);
|
|
pn_batch_decoder_admit_spec(d, st[1], 1, seeds[1], NULL, NULL); /* plain slot */
|
|
pn_batch_decoder_admit_spec(d, st[2], 2, seeds[2], draft_cb, ctxs[2]);
|
|
pn_page_out ob[6];
|
|
int admitted_more = 0, steps = 0;
|
|
while (pn_batch_decoder_active(d) > 0 || !admitted_more) {
|
|
int nf = pn_batch_decoder_step_spec(d, ob, 6);
|
|
if (nf < 0) { fail = 1; break; }
|
|
for (int i = 0; i < nf; i++) { txt[ob[i].page_id] = ob[i].text; pn_state_free(ob[i].state); }
|
|
steps++;
|
|
if (steps == 2 && !admitted_more) { /* staggered admit of 3..5 */
|
|
pn_batch_decoder_admit_spec(d, st[3], 3, seeds[3], NULL, NULL); /* plain */
|
|
pn_batch_decoder_admit_spec(d, st[4], 4, seeds[4], draft_cb, ctxs[4]);
|
|
pn_batch_decoder_admit_spec(d, st[5], 5, seeds[5], draft_cb, ctxs[5]);
|
|
admitted_more = 1;
|
|
}
|
|
if (admitted_more && pn_batch_decoder_active(d) == 0) break;
|
|
}
|
|
int bad = 0;
|
|
for (int p = 0; p < 6; p++) { if (!txt[p] || strcmp(txt[p], refText[p]) != 0) bad++; free(txt[p]); free(ctxs[p]); free(refText[p]); }
|
|
printf(" mixed/staggered text: %s\n", bad ? "DIFFERS" : "byte-identical (all 6 pages)");
|
|
if (bad) fail = 1;
|
|
pn_batch_decoder_free(d);
|
|
}
|
|
|
|
/* Forced stop cases: (1) seed == EOS -> immediate finalize, empty output;
|
|
* (2) max_tokens truncation -> both plain and spec cut at the same point. */
|
|
printf("\n==== forced-stop handling ====\n");
|
|
if (g_eos >= 0) {
|
|
float rp = 1.0f;
|
|
uint32_t pr[PROMPT_LEN]; make_prompt(pr, PROMPT_LEN, 999);
|
|
/* (1) EOS seed */
|
|
uint32_t junk; pn_state *sa = prefill_state(model, pr, PROMPT_LEN, &junk);
|
|
pn_state *sb = prefill_state(model, pr, PROMPT_LEN, &junk);
|
|
uint32_t seeds[1]; pn_state *st1[1];
|
|
char *tp = NULL, *ts = NULL;
|
|
seeds[0] = (uint32_t)g_eos; st1[0] = sa;
|
|
run_decoder(model, st1, seeds, 1, rp, GEN, 1, 0, 0, NULL, NULL, &tp, NULL, NULL, NULL);
|
|
st1[0] = sb; dctx c = { NULL, 0, D_NONE, 0 }; pn_draft_fn df[1] = { NULL }; void *du[1] = { &c };
|
|
run_decoder(model, st1, seeds, 1, rp, GEN, 1, 1, 8, df, du, &ts, &(long){0}, &(long){0}, &(long){0});
|
|
int ok = tp && ts && tp[0] == '\0' && ts[0] == '\0';
|
|
printf(" EOS seed: plain=%s spec=%s %s\n", tp && !tp[0] ? "empty" : "NONEMPTY",
|
|
ts && !ts[0] ? "empty" : "NONEMPTY", ok ? "OK" : "FAIL");
|
|
if (!ok) fail = 1;
|
|
free(tp); free(ts);
|
|
|
|
/* (2) max_tokens truncation: plain and spec must both equal A[0..N-1]. */
|
|
int N = 5;
|
|
uint32_t rt[GEN]; int rn; gen_seq(model, pr, PROMPT_LEN, rp, N, rt, &rn);
|
|
char *refT = decode_text(model, rt, rn);
|
|
uint32_t s2; pn_state *sc = prefill_state(model, pr, PROMPT_LEN, &s2);
|
|
pn_state *sd = prefill_state(model, pr, PROMPT_LEN, &s2);
|
|
char *tcp = NULL, *tcs = NULL; uint32_t sd2[1] = { s2 }; pn_state *stc[1] = { sc };
|
|
run_decoder(model, stc, sd2, 1, rp, N, 1, 0, 0, NULL, NULL, &tcp, NULL, NULL, NULL);
|
|
stc[0] = sd; dctx c2 = { rt, rn, D_ORACLE, 0 }; pn_draft_fn df2[1] = { draft_cb }; void *du2[1] = { &c2 };
|
|
run_decoder(model, stc, sd2, 1, rp, N, 1, 1, 8, df2, du2, &tcs, &(long){0}, &(long){0}, &(long){0});
|
|
int ok2 = tcp && tcs && strcmp(tcp, refT) == 0 && strcmp(tcs, refT) == 0;
|
|
printf(" max_tokens=%d truncation: plain=%s spec=%s %s\n", N,
|
|
(tcp && !strcmp(tcp, refT)) ? "match" : "DIFF",
|
|
(tcs && !strcmp(tcs, refT)) ? "match" : "DIFF", ok2 ? "OK" : "FAIL");
|
|
if (!ok2) fail = 1;
|
|
free(tcp); free(tcs); free(refT);
|
|
}
|
|
|
|
out:
|
|
printf("\nRESULT: %s\n", fail ? "FAIL"
|
|
: "PASS (batched speculative decode == plain batched greedy, byte for byte)");
|
|
pn_model_free(model);
|
|
return fail ? 1 : 0;
|
|
}
|