David
/
Pantheon
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Pantheon/tests/test_prefix_cache_qwen35.c

330 lines
14 KiB
C
Raw Blame History

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/* Prefix-cache gate + latency bench for qwen35moe on a K80 host.
*
* The agentic-turn win: across turns the context is mostly unchanged (a small
* delta appended to a fixed prefix). With caching ON we keep the session's
* model state and prefill ONLY the delta; with caching OFF we re-prefill the
* whole context every turn. The two MUST generate identical replies.
*
* Phase 1 (single session, THE win): a ~4k-token base context, then several
* turns each appending a ~200-token delta and greedily generating a short
* reply. Run twice — cached (pn_prefix_cache: delta-prefill only) vs uncached
* (fresh pn_state, full re-prefill each turn) — over identical per-turn
* contexts, comparing the generated tokens argmax-for-argmax and reporting the
* per-turn prefill latency of each.
*
* Phase 2 (multi-session save/restore): two sessions interleaved through ONE
* cache (they share the GPU's SSM buffers, so each acquire snapshots the other
* off-GPU and restores this one). Each session's interleaved generation must
* match that session decoded alone (uncached).
*
* Build (CUDA host): make USE_CUDA=1 tests/test_prefix_cache_qwen35
* Run: QWEN35_GPU_SPLIT_LAYER=20 QWEN35_GPU_DIE_A=0 QWEN35_GPU_DIE_B=1 \
* QWEN35_FUSE=0x37 ./tests/test_prefix_cache_qwen35 <model.gguf> \
* [base] [delta] [turns] [gen]
* base=4000 delta=200 turns=4 gen=8 by default.
* QWEN35_FUSE=0x37 makes the single-token decode path float-identical to the
* batched-prefill kernels (~1-ulp fused qknorm+rope contraction), the same
* convention as flash_spec_gate.sh / test_spec_qwen35 — the strict argmax gate
* runs under it.
* NOT part of the CPU test suite (needs the model + a CUDA GPU). */
#include "../include/pantheon.h"
#include "../src/core/engine.h"
#include "../src/core/prefix_cache.h"
#include "../src/models/qwen35moe.h"
#include "tensor/tensor.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#ifndef USE_CUDA
int main(void)
{
fprintf(stderr, "built without USE_CUDA\n");
return 1;
}
#else
static double now_s(void)
{
struct timespec ts;
clock_gettime(CLOCK_MONOTONIC, &ts);
return (double)ts.tv_sec + (double)ts.tv_nsec * 1e-9;
}
static int argmax(const float *x, int n)
{
int best = 0;
for (int i = 1; i < n; i++)
if (x[i] > x[best]) best = i;
return best;
}
/* Deterministic in-vocab pseudo-random token stream (reproducible contexts). */
static void fill_tokens(uint32_t *dst, int n, uint32_t seed, int n_vocab)
{
uint32_t s = seed * 2654435761u + 12345u;
for (int i = 0; i < n; i++) {
s = s * 1664525u + 1013904223u;
dst[i] = (s >> 8) % (uint32_t)n_vocab;
}
}
/* Batched-prefill tokens[start..n) from the state's current pos, in chunks of
* PREFILL_CHUNK (keeps the batch scratch bounded for multi-thousand-token
* contexts; chunking is argmax-identical — the KV/SSM state is built
* incrementally exactly as one big batch would). Returns the last-position
* logits tensor (caller frees), or NULL. */
static int prefill_chunk(void)
{
const char *v = getenv("QWEN35_PC_CHUNK");
int c = v ? atoi(v) : 256;
return c > 0 ? c : 256;
}
static tm_tensor *prefill_delta(pn_model *model, pn_state *st,
const uint32_t *tokens, int start, int n)
{
if (n <= start)
return NULL;
int chunk = prefill_chunk();
tm_tensor *o = NULL;
for (int p = start; p < n; p += chunk) {
int m = n - p < chunk ? n - p : chunk;
if (o) tm_free(o);
o = pn_qwen35moe_forward_batch_cuda(model, tokens + p, m, st, NULL);
if (!o) return NULL;
}
return o;
}
/* Greedy-generate `gen` tokens from `first` logits (the prefill's last-position
* logits), writing them to out and advancing the state. The first token comes
* straight from `first`; the next gen-1 from gen-1 forwards; then ONE more
* forward feeds the final token so the state ends at pos += gen with every
* generated token committed to the KV/SSM state (state->pos == committed
* context length) — exactly a real serving turn, and what lets the next turn
* delta-prefill only the truly new tokens. Returns 0 / -1. */
static int generate(pn_model *model, pn_state *st, tm_tensor *first,
uint32_t *out, int gen, int n_vocab)
{
uint32_t tok = (uint32_t)argmax(tm_data_f32(first), n_vocab);
out[0] = tok;
for (int i = 1; i < gen; i++) {
tm_tensor *o = model->forward(model, tok, st);
if (!o) return -1;
tok = (uint32_t)argmax(tm_data_f32(o), n_vocab);
out[i] = tok;
tm_free(o);
}
tm_tensor *last = model->forward(model, out[gen - 1], st); /* commit last */
if (!last) return -1;
tm_free(last);
return 0;
}
int main(int argc, char **argv)
{
if (argc < 2) {
fprintf(stderr, "usage: %s <model.gguf> [base] [delta] [turns] [gen]\n",
argv[0]);
return 1;
}
pn_model *model = pn_load(argv[1]);
if (!model) { fprintf(stderr, "load failed\n"); return 1; }
int n_vocab = model->config.n_vocab;
int base = argc > 2 ? atoi(argv[2]) : 4000;
int delta = argc > 3 ? atoi(argv[3]) : 200;
int turns = argc > 4 ? atoi(argv[4]) : 4;
int gen = argc > 5 ? atoi(argv[5]) : 8;
if (base < 1 || delta < 1 || turns < 1 || gen < 1) {
fprintf(stderr, "bad params\n"); return 1;
}
int max_ctx = base + turns * (delta + gen) + 8;
uint32_t *ctx = malloc((size_t)max_ctx * sizeof(uint32_t));
/* per-turn snapshot of the context fed before generation, and each run's
* generated reply. */
uint32_t **snap = malloc((size_t)turns * sizeof(uint32_t *));
int *snap_n = malloc((size_t)turns * sizeof(int));
uint32_t *gc = malloc((size_t)turns * gen * sizeof(uint32_t)); /* cached */
uint32_t *gu = malloc((size_t)turns * gen * sizeof(uint32_t)); /* uncached */
double *t_cache = calloc((size_t)turns, sizeof(double));
double *t_uncab = calloc((size_t)turns, sizeof(double));
int *d_delta = calloc((size_t)turns, sizeof(int));
if (!ctx || !snap || !snap_n || !gc || !gu || !t_cache || !t_uncab || !d_delta) {
fprintf(stderr, "oom\n"); return 1;
}
fprintf(stderr, "prefix-cache gate: base=%d delta=%d turns=%d gen=%d\n",
base, delta, turns, gen);
/* ================= Phase 1a: CACHED run ================= */
pn_prefix_cache *cache = pn_prefix_cache_new(model, 4);
if (!cache) { fprintf(stderr, "cache new failed\n"); return 1; }
int n_ctx = 0;
fill_tokens(ctx, base, 1u, n_vocab); /* the fixed base prefix */
n_ctx = base;
for (int t = 0; t < turns; t++) {
if (t > 0) { /* append this turn's delta */
fill_tokens(ctx + n_ctx, delta, 100u + (uint32_t)t, n_vocab);
n_ctx += delta;
}
/* snapshot the context fed before generation (for the reference run) */
snap_n[t] = n_ctx;
snap[t] = malloc((size_t)n_ctx * sizeof(uint32_t));
memcpy(snap[t], ctx, (size_t)n_ctx * sizeof(uint32_t));
int ds = 0;
pn_state *st = pn_prefix_cache_acquire(cache, "agent", ctx, n_ctx, &ds);
if (!st) { fprintf(stderr, "acquire failed t=%d\n", t); return 1; }
d_delta[t] = n_ctx - ds;
double t0 = now_s();
tm_tensor *o = prefill_delta(model, st, ctx, ds, n_ctx);
t_cache[t] = now_s() - t0;
if (!o) { fprintf(stderr, "cached prefill failed t=%d\n", t); return 1; }
if (generate(model, st, o, gc + (size_t)t * gen, gen, n_vocab) != 0) {
fprintf(stderr, "cached gen failed t=%d\n", t); return 1;
}
tm_free(o);
/* commit the reply into the running context + the session history */
memcpy(ctx + n_ctx, gc + (size_t)t * gen, (size_t)gen * sizeof(uint32_t));
n_ctx += gen;
if (pn_prefix_cache_commit(cache, "agent", ctx, n_ctx) != 0) {
fprintf(stderr, "commit failed t=%d\n", t); return 1;
}
}
pn_prefix_cache_free(cache);
/* ================= Phase 1b: UNCACHED reference ================= */
for (int t = 0; t < turns; t++) {
pn_state *st = pn_state_new(model);
if (!st) { fprintf(stderr, "state failed t=%d\n", t); return 1; }
pn_state_reset(st);
double t0 = now_s();
tm_tensor *o = prefill_delta(model, st, snap[t], 0, snap_n[t]);
t_uncab[t] = now_s() - t0;
if (!o) { fprintf(stderr, "uncached prefill failed t=%d\n", t); return 1; }
if (generate(model, st, o, gu + (size_t)t * gen, gen, n_vocab) != 0) {
fprintf(stderr, "uncached gen failed t=%d\n", t); return 1;
}
tm_free(o);
pn_state_free(st);
}
/* ================= Gate + latency report ================= */
int fail_turn = -1, fail_pos = -1;
for (int t = 0; t < turns && fail_turn < 0; t++)
for (int i = 0; i < gen; i++)
if (gc[(size_t)t * gen + i] != gu[(size_t)t * gen + i]) {
fail_turn = t; fail_pos = i; break;
}
fprintf(stderr, "\nper-turn prefill latency (cached delta vs uncached full):\n");
for (int t = 0; t < turns; t++)
fprintf(stderr,
" turn %d: ctx=%d delta=%d | cached %.3fs | uncached %.3fs | "
"speedup %.1fx\n",
t, snap_n[t], d_delta[t], t_cache[t], t_uncab[t],
t_cache[t] > 0 ? t_uncab[t] / t_cache[t] : 0.0);
if (fail_turn < 0)
fprintf(stderr, "GATE 1 (single session): ARGMAX-EXACT PASS\n");
else
fprintf(stderr,
"GATE 1: FAIL turn=%d pos=%d (cached=%u uncached=%u)\n",
fail_turn, fail_pos,
gc[(size_t)fail_turn * gen + fail_pos],
gu[(size_t)fail_turn * gen + fail_pos]);
/* ================= Phase 2: multi-session interleave ================= */
/* Two sessions, short contexts, advanced turn-by-turn through ONE cache in
* interleaved order — each acquire snapshots the other session off the
* shared GPU SSM buffers and restores this one. Each session's interleaved
* reply must match that session decoded alone. */
int gate2_fail = 0;
{
const int b2 = 96, d2 = 24, T2 = 3;
int cap2 = b2 + T2 * (d2 + gen) + 8;
uint32_t *cx[2] = { malloc((size_t)cap2 * sizeof(uint32_t)),
malloc((size_t)cap2 * sizeof(uint32_t)) };
int nx[2] = { 0, 0 };
uint32_t *g2 = malloc((size_t)2 * T2 * gen * sizeof(uint32_t));
uint32_t **sn = malloc((size_t)2 * T2 * sizeof(uint32_t *));
int *snn = malloc((size_t)2 * T2 * sizeof(int));
if (!cx[0] || !cx[1] || !g2 || !sn || !snn) { fprintf(stderr, "oom2\n"); return 1; }
fill_tokens(cx[0], b2, 7u, n_vocab); nx[0] = b2;
fill_tokens(cx[1], b2, 9u, n_vocab); nx[1] = b2;
pn_prefix_cache *c2 = pn_prefix_cache_new(model, 2);
if (!c2) { fprintf(stderr, "cache2 failed\n"); return 1; }
const char *sid[2] = { "A", "B" };
for (int t = 0; t < T2; t++) {
for (int u = 0; u < 2; u++) {
if (t > 0) {
fill_tokens(cx[u] + nx[u], d2, 200u + (uint32_t)(t * 2 + u),
n_vocab);
nx[u] += d2;
}
int si = (t * 2 + u);
snn[si] = nx[u];
sn[si] = malloc((size_t)nx[u] * sizeof(uint32_t));
memcpy(sn[si], cx[u], (size_t)nx[u] * sizeof(uint32_t));
int ds = 0;
pn_state *st = pn_prefix_cache_acquire(c2, sid[u], cx[u], nx[u],
&ds);
if (!st) { fprintf(stderr, "acq2 failed\n"); return 1; }
tm_tensor *o = prefill_delta(model, st, cx[u], ds, nx[u]);
if (!o) { fprintf(stderr, "pf2 failed\n"); return 1; }
if (generate(model, st, o, g2 + (size_t)si * gen, gen,
n_vocab) != 0) { fprintf(stderr, "gen2\n"); return 1; }
tm_free(o);
memcpy(cx[u] + nx[u], g2 + (size_t)si * gen,
(size_t)gen * sizeof(uint32_t));
nx[u] += gen;
pn_prefix_cache_commit(c2, sid[u], cx[u], nx[u]);
}
}
pn_prefix_cache_free(c2);
/* Reference: each snapshot decoded alone on a fresh state. */
for (int si = 0; si < 2 * T2; si++) {
pn_state *st = pn_state_new(model);
pn_state_reset(st);
tm_tensor *o = prefill_delta(model, st, sn[si], 0, snn[si]);
uint32_t ref[512];
if (gen > 512) { fprintf(stderr, "gen too large for phase 2\n"); return 1; }
if (!o || generate(model, st, o, ref, gen, n_vocab) != 0) {
fprintf(stderr, "ref2 failed si=%d\n", si); return 1;
}
tm_free(o);
pn_state_free(st);
for (int i = 0; i < gen; i++)
if (ref[i] != g2[(size_t)si * gen + i]) {
fprintf(stderr,
"GATE 2: FAIL si=%d pos=%d (interleaved=%u alone=%u)\n",
si, i, g2[(size_t)si * gen + i], ref[i]);
gate2_fail = 1; break;
}
if (gate2_fail) break;
}
if (!gate2_fail)
fprintf(stderr, "GATE 2 (multi-session save/restore): ARGMAX-EXACT PASS\n");
}
fflush(stderr);
int ok = (fail_turn < 0) && !gate2_fail;
if (!getenv("QWEN35_EXIT_CLEAN"))
_Exit(ok ? 0 : 1);
return ok ? 0 : 1;
}
#endif
Reference in New Issue View Git Blame Copy Permalink