Pantheon/tests/test_tt_prefill_parity.c

413 lines
19 KiB
C

/* SPDX-License-Identifier: AGPL-3.0-or-later */
/* Copyright (C) 2026 David (Corey) Herrera */
/**
* @file test_tt_prefill_parity.c
* @brief BYTE-EXACTNESS of the batched TT prefill against the one-at-a-time path.
*
* ---------------------------------------------------------------------------
* WHY BYTES AND NOT PCC
* ---------------------------------------------------------------------------
* A batched prefill changes the SHAPE of the pass, not the arithmetic: 32 prompt positions
* ride rows 0..31 of tiles whose rows 1..31 were previously zero padding the FPU multiplied
* anyway. Every claim that makes this safe is a claim about INDEPENDENCE — that row m of a
* matmul reads only row m of A, that a reduce block reduces each of its 32 rows over itself,
* that an elementwise op is per element. Independence is exactly what a correlation cannot
* see: PCC saturated at 1.000000000 on a tensor with relRMS 0.235 in this project, and it
* caught none of four kernel mutations on norm-shaped ops (foundry/tt/DEV_NOTES.md §3). If
* the two paths are the same arithmetic they are the same BYTES, and anything less than that
* is a defect however small it scores.
*
* ---------------------------------------------------------------------------
* WHAT IS COMPARED, AND WHY IT IS SUFFICIENT
* ---------------------------------------------------------------------------
* Two arms over the same prompt on the same card in the same process:
*
* A. SEQUENTIAL — `pn_llama_forward_tt` once per prompt token, the shipping path.
* B. BATCHED — `pn_llama_prefill_tt`, which runs tokens 0..P-2 in passes of up to 32
* rows and hands token P-1 to `pn_llama_forward_tt`.
*
* and it compares:
*
* 1. **the whole KV cache**, every layer's K and V, byte for byte over the P positions
* actually written. This is the primary instrument: the KV cache is the ENTIRE state a
* prefill produces, so two prefills that agree on it are the same prefill.
* 2. **the last position's logits**, byte for byte. This is a consequence rather than an
* independent check — arm B's final token goes through the same function arm A's does —
* and it is compared anyway because a consequence that fails is evidence the premise is
* wrong.
*
* The KV comparison is the strict one and it is what a divergence will show up in first: the
* logits of the last position depend on the whole cache, but they are one row of one softmax
* and could in principle absorb a small error in an early position. The cache cannot absorb
* anything.
*
* ---------------------------------------------------------------------------
* PROVING IT CAN FAIL
* ---------------------------------------------------------------------------
* `PN_TT_RED=<name>` breaks the batched arm in one named way and the harness must then
* report DIFFER. Four vacuous gates have been caught in this project, so a green run of this
* file means nothing on its own; `tests/run_tt_prefill_red.sh` is the arm that makes it mean
* something. The mutations are documented at `llama_tt_red` in llama_forward_tt.c:
*
* prefill_pos every row rotates and appends at the pass's FIRST position
* prefill_causal the per-row loop runs in reverse, so a row reads later rows' K/V
* prefill_rowfill q/k/v rows 1..M-1 are overwritten with row 0 — the bug of a transfer
* that still only moves logical row 0
*
* A fourth control needs no mutation at all: `PN_TT_PREFILL_ROWS=1` makes arm B one row per
* pass, which is the sequential computation reached through the batched code. It must be
* byte-identical, and if it is not, the divergence is in the batched PLUMBING rather than in
* the batching.
*
* ---------------------------------------------------------------------------
* USAGE
* ---------------------------------------------------------------------------
* PN_TT_ARTIFACT=<...fttw> PN_TT_ONCARD=norm,swiglu,residual,topk \
* ./tests/test_tt_prefill_parity <model.gguf> [--tokens N] [--seed S]
*
* ON-CARD ATTENTION IS DELIBERATELY NOT SUPPORTED ON THE LLAMA ARM and the batched path
* refuses it by name: `pn_tt_rope`, `pn_tt_kv_append` and the flash decode each take ONE
* position, so a 32-row pass has no way to express 32 of them. The harness says so and exits
* non-zero rather than quietly scoring a refusal as a pass — a refusal that reads as green is
* how a gate becomes vacuous.
*
* ---------------------------------------------------------------------------
* THE ORNITH ARM, AND WHY IT COMPARES A DIFFERENT THING
* ---------------------------------------------------------------------------
* Ornith (`qwen35`) reaches this file through `pn_qwen35moe_prefill_tt`, and the state a
* prefill produces there is NOT a host KV cache. It is 24 layers of Gated-Delta-Net
* recurrence carried in card DRAM plus 8 layers of K/V in card DRAM — **120 MiB of GDN state
* alone**, and the card->host direction measures **5.3 MB/s flat from 16 KB to 786 KB**
* (`bench_tt_prefill_quantum`, the TT host, 2026-08-21), so capturing it twice would cost ~46 s of
* link per comparison and would itself perturb the residencies.
*
* So the Ornith arm compares **the last prompt position's logit row, byte for byte**. That is
* a consequence rather than an independent check, and it is stated as one — but it is a
* consequence of the WHOLE card state: the row is produced by the ordinary single-token
* forward reading every one of the 24 recurrent states and all 8 K/V regions the prefill
* left. A batched pass that wrote one position's state to the wrong slot, ran the recurrence
* out of order, or served row 0's projection to all 32 rows cannot leave that row intact.
* The three red arms in `qwen35moe_prefill_tt.c` are what makes that a measurement.
*
* **It REFUSES `PN_TT_SPARSE_LOGITS`.** That knob fills 248,312 of 248,320 entries with
* `-INFINITY`, so two rows would agree everywhere the card did not name — the eighth
* vacuous-gate class recorded in this project, and the reason it is a refusal and not a
* warning.
*
* @status Working. The llama arm compares the KV cache; the Ornith arm compares the last
* logit row and says why.
* @issues None
* @todo None
*/
#include "../include/pantheon.h"
#include "../src/core/engine.h"
#include "../src/models/llama.h"
#include "../src/models/qwen35moe.h"
#include "tensor/tensor.h"
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#ifndef USE_TT
int main(void)
{
fprintf(stderr, "test_tt_prefill_parity: built without USE_TT — nothing to compare.\n");
return 77;
}
#else
/** Default prompt length. 100 crosses three 32-row passes AND leaves a ragged last one
* (32 + 32 + 32 + 3 + the unbatched final token), so the remainder path is exercised by
* the default rather than only by an option nobody passes. */
#define PARITY_DEFAULT_TOKENS 100
/** A deterministic pseudo-prompt. Real ids matter less here than COVERAGE — every row of
* every pass must carry a different vector, or a bug that copies row 0 everywhere would be
* invisible. A fixed LCG gives that and makes the run reproducible. */
static double parity_now(void)
{
struct timespec ts;
clock_gettime(CLOCK_MONOTONIC, &ts);
return (double)ts.tv_sec + (double)ts.tv_nsec * 1e-9;
}
static uint32_t parity_token(uint32_t *rng, int n_vocab)
{
*rng = *rng * 1664525u + 1013904223u;
return (uint32_t)((*rng >> 8) % (uint32_t)n_vocab);
}
/** Copy every layer's K and V for positions [0, positions) into one packed fp32 buffer.
* Packed rather than compared in place because the two arms need two snapshots alive at
* once, and because a single buffer makes the comparison one memcmp per layer-half. */
static float *parity_capture_kv(pn_model *model, pn_state *state, int positions,
size_t *out_floats)
{
pn_llama_kv_cache *kv = (pn_llama_kv_cache *)state->kv_cache;
pn_config *cfg = &model->config;
int head_dim = pn_head_dim(cfg);
int n_kv_head = cfg->n_kv_head;
size_t per_half = (size_t)n_kv_head * (size_t)positions * (size_t)head_dim;
size_t total = per_half * 2u * (size_t)cfg->n_layer;
float *out = malloc(total * sizeof(float));
size_t at = 0;
if (!out)
return NULL;
for (int layer = 0; layer < cfg->n_layer; layer++) {
const float *k = (const float *)kv->k_cache[layer]->data;
const float *v = (const float *)kv->v_cache[layer]->data;
for (int h = 0; h < n_kv_head; h++) {
memcpy(out + at, k + (size_t)h * kv->max_seq * head_dim,
(size_t)positions * head_dim * sizeof(float));
at += (size_t)positions * head_dim;
}
for (int h = 0; h < n_kv_head; h++) {
memcpy(out + at, v + (size_t)h * kv->max_seq * head_dim,
(size_t)positions * head_dim * sizeof(float));
at += (size_t)positions * head_dim;
}
}
*out_floats = total;
return out;
}
/** First differing float index, or -1. Reported as an INDEX and a pair of values, because
* "they differ" without a where is not a diagnosis. */
static long parity_first_diff(const float *a, const float *b, size_t n)
{
const unsigned char *pa = (const unsigned char *)a;
const unsigned char *pb = (const unsigned char *)b;
size_t bytes = n * sizeof(float);
for (size_t i = 0; i < bytes; i++)
if (pa[i] != pb[i])
return (long)(i / sizeof(float));
return -1;
}
/**
* WHEN THE TWO ROWS DIFFER, SAY BY HOW MUCH — because one arm this harness now scores is
* KNOWN not to be byte-identical, and printing only "DIFFER" makes that arm unreadable.
*
* The batched ATTENTION chunk is BIT-EXACT against the decode descriptors it replaces
* (max-abs 0.000e+00, row for row), which is why the batched arm alone is byte-identical and
* must stay so. The two CHUNKWISE GDN mixers are not: they are the same recurrence with FEWER
* ROUNDINGS — PHASE CONV crosses one more unpack stage than the decode form, and the
* chunkwise delta rule measures y 8.6e-03 against 32 sequential decode steps' 2.1e-02 through
* the same Tf32 unpack. Byte-identity against a per-position path is therefore unreachable BY
* CONSTRUCTION for that arm, and the honest question is a magnitude.
*
* MAX-NORM RELATIVE, `max|a-b| / max|a|`, which is the metric every recorded figure for these
* two ops is in. The ARGMAX is reported beside it because it is what a greedy completion
* actually consumes.
*/
static void parity_magnitude(const float *a, const float *b, size_t n, const char *label)
{
double max_abs = 0.0, max_ref = 0.0, sum_sq = 0.0, ref_sq = 0.0;
size_t at = 0, over = 0, finite = 0, argmax_a = 0, argmax_b = 0;
for (size_t i = 0; i < n; i++) {
const double x = (double)a[i], y = (double)b[i];
double gap;
/* A SPARSE ROW IS -INFINITY ALMOST EVERYWHERE and this harness refuses one, but a
* -INFINITY can still reach a dense row — and `inf - inf` is NaN, which would make
* every figure below NaN with nothing saying so. */
if (!(x > -1e30 && x < 1e30) || !(y > -1e30 && y < 1e30))
continue;
finite++;
gap = x > y ? x - y : y - x;
if (gap > max_abs) { max_abs = gap; at = i; }
if ((x < 0 ? -x : x) > max_ref) max_ref = (x < 0 ? -x : x);
sum_sq += gap * gap;
ref_sq += x * x;
if (gap > 1e-2) over++;
if (x > (double)a[argmax_a]) argmax_a = i;
if (y > (double)b[argmax_b]) argmax_b = i;
}
printf("PARITY %s max-norm relative %.3e (max|delta| %.4g at float %zu of %zu finite), "
"relRMS %.3e, %zu past 1e-2\n",
label, max_ref > 0.0 ? max_abs / max_ref : 0.0, max_abs, at, finite,
ref_sq > 0.0 ? sqrt(sum_sq / ref_sq) : 0.0, over);
printf("PARITY %s ARGMAX sequential %zu (%.6g) vs batched %zu (%.6g) — %s\n", label,
argmax_a, (double)a[argmax_a], argmax_b, (double)b[argmax_b],
argmax_a == argmax_b ? "THE SAME TOKEN" : "A DIFFERENT TOKEN");
}
int main(int argc, char **argv)
{
const char *model_path = NULL;
int want_tokens = PARITY_DEFAULT_TOKENS;
uint32_t rng = 0x5eed1234u;
pn_model *model;
pn_state *sa, *sb;
uint32_t *prompt;
float *kv_a = NULL, *kv_b = NULL;
float *log_a = NULL, *log_b = NULL;
size_t kv_floats_a = 0, kv_floats_b = 0;
tm_tensor *la = NULL, *lb = NULL;
long kv_diff = -1, logit_diff;
double t_seq = 0.0, t_batch = 0.0, mark;
int ornith;
int rc = 1;
for (int i = 1; i < argc; i++) {
if (strcmp(argv[i], "--tokens") == 0 && i + 1 < argc)
want_tokens = atoi(argv[++i]);
else if (strcmp(argv[i], "--seed") == 0 && i + 1 < argc)
rng = (uint32_t)strtoul(argv[++i], NULL, 0);
else if (!model_path)
model_path = argv[i];
}
if (!model_path || want_tokens < 2) {
fprintf(stderr, "usage: %s <model.gguf> [--tokens N>=2] [--seed S]\n", argv[0]);
return 2;
}
model = pn_load(model_path);
if (!model) {
fprintf(stderr, "PARITY REFUSED: cannot load %s\n", model_path);
return 2;
}
/* THE VACUOUS-RUN GUARD. Without the card there is no batched path to test and the two
* arms would be the same CPU code compared against itself — a green run that proves
* nothing, which is the exact failure this project has paid for four times. */
ornith = (model->forward == pn_qwen35moe_forward_tt);
if (model->forward != pn_llama_forward_tt && !ornith) {
fprintf(stderr, "PARITY REFUSED: the engine did not select a Tenstorrent forward. "
"Set PN_TT_ARTIFACT and check the card. Comparing the CPU path "
"against itself would be vacuous.\n");
pn_model_free(model);
return 2;
}
/* A SPARSE ROW IS -INFINITY IN 248,312 OF 248,320 ENTRIES, so two of them agree
* everywhere the card did not name. Scoring that is the eighth vacuous-gate class in
* this project, and the Ornith arm has nothing to compare BUT that row. */
if (ornith && getenv("PN_TT_SPARSE_LOGITS")) {
fprintf(stderr, "PARITY REFUSED: PN_TT_SPARSE_LOGITS is set. The Ornith arm compares "
"the last position's LOGIT ROW, and a sparse row is -INFINITY almost "
"everywhere — two of them would agree for a reason that is not "
"parity. Unset it.\n");
pn_model_free(model);
return 2;
}
prompt = malloc((size_t)want_tokens * sizeof(*prompt));
sa = pn_state_new(model);
sb = pn_state_new(model);
if (!prompt || !sa || !sb) {
fprintf(stderr, "PARITY REFUSED: allocation failed\n");
goto done;
}
for (int i = 0; i < want_tokens; i++)
prompt[i] = parity_token(&rng, model->config.n_vocab);
/* ── ARM A: the shipping path, one forward per prompt token. */
mark = parity_now();
for (int i = 0; i < want_tokens; i++) {
if (la) tm_free(la);
la = model->forward(model, prompt[i], sa);
if (!la) {
fprintf(stderr, "PARITY REFUSED: sequential forward failed at token %d\n", i);
goto done;
}
}
t_seq = parity_now() - mark;
/* ORNITH'S PREFILL STATE IS IN CARD DRAM, not in a host KV cache: 120 MiB of carried
* Gated-Delta-Net state plus 8 layers of K/V, over a link that reads at 5.3 MB/s. The
* logit row below is the instrument there, and the header says why. */
if (!ornith)
kv_a = parity_capture_kv(model, sa, want_tokens, &kv_floats_a);
/* ── ARM B: the batched path over the same ids from a fresh state. */
mark = parity_now();
lb = ornith ? pn_qwen35moe_prefill_tt(model, prompt, want_tokens, sb)
: pn_llama_prefill_tt(model, prompt, NULL, want_tokens, sb);
t_batch = parity_now() - mark;
if (!lb) {
fprintf(stderr, "PARITY REFUSED: the batched prefill REFUSED this configuration. "
"The llama arm refuses on-card attention (its rotary, KV append and "
"flash decode each take ONE position) and the diagnostic arms; the "
"Ornith arm defaults OFF and needs PN_TT_PREFILL=1, both on-card "
"mixers, and a sequence that starts at position 0. Scoring a refusal "
"as a pass is how a gate goes vacuous, so this is an error and not a "
"skip.\n");
goto done;
}
if (!ornith) {
kv_b = parity_capture_kv(model, sb, want_tokens, &kv_floats_b);
if (!kv_a || !kv_b || kv_floats_a != kv_floats_b) {
fprintf(stderr, "PARITY REFUSED: KV capture failed or shapes disagree\n");
goto done;
}
}
if (sa->pos != sb->pos) {
fprintf(stderr, "PARITY FAIL position: sequential pos=%d batched pos=%d\n",
sa->pos, sb->pos);
goto done;
}
if (la->numel != lb->numel) {
fprintf(stderr, "PARITY REFUSED: logit rows differ in length\n");
goto done;
}
log_a = (float *)la->data;
log_b = (float *)lb->data;
if (!ornith)
kv_diff = parity_first_diff(kv_a, kv_b, kv_floats_a);
logit_diff = parity_first_diff(log_a, log_b, (size_t)la->numel);
printf("PARITY family=%s tokens=%d rows=%s kv_floats=%zu vocab=%d pos=%d\n",
ornith ? "ornith/qwen35" : "llama", want_tokens,
getenv("PN_TT_PREFILL_ROWS") ? getenv("PN_TT_PREFILL_ROWS") : "32(default)",
kv_floats_a, model->config.n_vocab, sa->pos);
if (ornith)
printf("PARITY kv_cache NOT COMPARED — Ornith's prefill state is 120 MiB of "
"card-resident GDN recurrence plus card K/V, and the link reads at 5.3 MB/s. "
"The logit row below is the instrument.\n");
else if (kv_diff < 0)
printf("PARITY kv_cache IDENTICAL over %zu floats\n", kv_floats_a);
else
printf("PARITY kv_cache DIFFER first float %ld sequential=%.9g batched=%.9g\n",
kv_diff, (double)kv_a[kv_diff], (double)kv_b[kv_diff]);
if (logit_diff < 0)
printf("PARITY last_logits IDENTICAL over %lld floats\n", (long long)la->numel);
else {
printf("PARITY last_logits DIFFER first float %ld sequential=%.9g batched=%.9g\n",
logit_diff, (double)log_a[logit_diff], (double)log_b[logit_diff]);
parity_magnitude(log_a, log_b, (size_t)la->numel, "last_logits");
}
/* PREFILL TOK/S, and it is COMPUTED tokens over the wall clock of the prefill alone —
* both arms start from a fresh state and neither has a prefix cache in front of it, so
* there is nothing here that a warm cache could inflate. Both arms include the same
* unbatched final token, which is why the batched rate is a lower bound rather than a
* best case. */
printf("PARITY rate sequential %.3f s = %.2f tok/s | batched %.3f s = %.2f tok/s | "
"speedup %.2fx\n", t_seq, (double)want_tokens / t_seq, t_batch,
(double)want_tokens / t_batch, t_seq / t_batch);
rc = (kv_diff < 0 && logit_diff < 0) ? 0 : 1;
printf("PARITY VERDICT %s\n", rc == 0 ? "BYTE-IDENTICAL" : "DIFFER");
done:
if (la) tm_free(la);
if (lb) tm_free(lb);
free(kv_a);
free(kv_b);
free(prompt);
if (sa) pn_state_free(sa);
if (sb) pn_state_free(sb);
pn_model_free(model);
return rc;
}
#endif /* USE_TT */