330 lines
14 KiB
C
330 lines
14 KiB
C
/* 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
|