Merge pull request #362 from EgonRuiter/prefetcher-v3

feat(prefetch): async expert prefetcher v3.2 for the OLMoE testbed

Testbed-only scope: olmoe.c + tools/oracle files; glm.c untouched (the
production engine already ships the equivalent techniques: coalesced slab
preads, PILOT lookahead, persistent PIN). Trivial .gitignore conflict
resolved keeping both sides.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

.gitignore

@@ -68,3 +68,7 @@ c/tests/test_decode_batch
 c/tests/test_i4_acc512
 c/tests/test_idot
 c/tests/test_uring
+olmoe_merged/
+olmoe_i4/
+c/olmoe_merged/
+c/olmoe_i4/

c/olmoe.c

@@ -5,6 +5,15 @@
  * Densa (embed, attn, router, norme, lm_head) residente in RAM (float32).
  * Expert letti dal disco on-demand via pread+fadvise(DONTNEED), cache LRU per-layer.
  * Matmul multi-thread con OpenMP (niente BLAS).
+ *
+ * ENV VARS:
+ *   PILOT=0/1/2/3 : 0=no prefetch, 1=1-layer lookahead, 2=2-layer, 3=3-layer lookahead
+ *   HOT=N         : pin top-N hot experts per layer permanently (never evict)
+ *   WARMUP=N      : tokens before hot pinning activates (default 5)
+ *   WIDE=N        : prefetch top-K*N candidates (default 1, try 2 or 3)
+ *   SMOOTH=F      : EMA coefficient for routing momentum (default 0.3, range 0.0-0.95)
+ *   CONF_LIMIT=F  : cumulative gate probability threshold for prefetch cutoff (default 0.92)
+ *   (expert queue is sorted by eid for SSD read locality)
  */
 #define _GNU_SOURCE
 #include <stdio.h>
@@ -12,11 +21,22 @@
 #include <string.h>
 #include <math.h>
 #include <time.h>
+#include <pthread.h>
 #if defined(__APPLE__) || defined(__linux__) || defined(__FreeBSD__)
 #include <sys/resource.h>
+#include <unistd.h>
 #endif
 #include "st.h"
 
+#ifdef _WIN32
+#include <windows.h>
+#define sleep_ms(ms) Sleep(ms)
+#else
+#define sleep_ms(ms) usleep((ms) * 1000)
+#endif
+
+
+
 /* ---------- config ---------- */
 typedef struct {
     int hidden, n_layers, n_heads, n_kv_heads, head_dim;
@@ -33,22 +53,61 @@ typedef struct {
  * Ogni weight [out,in] tenuto come int8 (per-riga) + scala float per riga.
  * Cosi' la RAM-cache scende da 4 byte/param (f32) a 1 byte/param: e' il
  * meccanismo che fa stare GLM-5.2 nei 15 GB. dequant-on-use nel matmul. */
-typedef struct { int eid; int8_t *g, *u, *d; float *gs, *us, *ds; uint64_t used; } Slot;
+/* pinned=1 means this slot is strongly preferred to keep (hot expert); it will
+ * not be evicted during normal LRU eviction, but may be displaced under extreme
+ * cache pressure when all slots are pinned or in-flight. */
+typedef struct { int eid; int pinned; int8_t *g, *u, *d; float *gs, *us, *ds; uint64_t used; } Slot;
 typedef struct { Slot *slots; int n, cap; } LCache;
 
 typedef struct {
     Cfg c;
     shards S;
-    int quant_bits;        /* bit di quantizzazione degli expert (2..8); storage int8, niente f32 (#134) */
+    int quant_bits;
     float *embed, *lm_head, *final_norm;
     Layer *L;
     LCache *cache;          /* [n_layers] */
     uint64_t clock, hits, miss;
-    /* kv-cache per-layer: K,V come [H * maxT * head_dim] */
     float **K, **V; int kv_len, max_t;
     double dense_load_s;
+    /* IMPROVEMENT 2: expert frequency heatmap */
+    uint32_t *freq;
+    int freq_token_count, hot_pinned, hot_n, warmup_tokens;
+    int token_count;
+    /* PREDICTION IMPROVEMENT A: per-layer EMA of gate logits across tokens.
+     * momentum_logits[l*E .. (l+1)*E-1] = EMA of gate outputs for layer l.
+     * Used exclusively by the PILOT prefetcher to stabilise routing predictions
+     * across tokens; does NOT affect actual MoE routing (pr is unchanged). */
+    float *momentum_logits; /* [n_layers * n_experts], EMA of gate logits */
+    float pilot_smooth;     /* SMOOTH env: EMA coefficient 0.0-0.9 (default 0.3) */
+    uint8_t *is_pinned;     /* [n_layers * n_experts], 1 if expert is globally pinned */
+    uint8_t *is_queued;     /* [n_layers * n_experts], 1 if expert is currently in the prefetch queue */
+    float pilot_conf_limit; /* CONF_LIMIT env: cumulative gate probability threshold (e.g. 0.92) */
 } Model;
 
+static pthread_mutex_t g_pilot_mx = PTHREAD_MUTEX_INITIALIZER;
+static struct { int l, e; } pilot_q[4096];
+static volatile unsigned pilot_r = 0, pilot_w = 0;
+static Model *pilot_m = NULL;
+static int g_pilot = 0;
+static int g_wide  = 1;  /* IMPROVEMENT 4: top-K * g_wide candidates prefetched */
+
+static void pilot_prefetch(Model *m, int lnext, const float *x, int S);
+static void *pilot_worker(void *arg);
+static void ensure_pilot_worker_started(Model *m);
+static void slot_ensure_allocated(Model *m, Slot *s);
+
+static void ensure_pilot_worker_started(Model *m) {
+    if (!pilot_m) {
+        pilot_m = m;
+        pthread_t t;
+        if (pthread_create(&t, NULL, pilot_worker, NULL) != 0) {
+            fprintf(stderr, "Error: Failed to create pilot prefetch worker thread\n");
+            exit(1);
+        }
+        pthread_detach(t);
+    }
+}
+
 /* ---------- utility ---------- */
 static double now_s(void) { struct timespec t; clock_gettime(CLOCK_MONOTONIC, &t); return t.tv_sec + t.tv_nsec*1e-9; }
 #if defined(__APPLE__)
@@ -210,51 +269,224 @@ static void model_init(Model *m, const char *snap, int cap, int bits) {
         #undef LD
     }
     m->cache = calloc(c->n_layers, sizeof(LCache));
-    for (int i = 0; i < c->n_layers; i++) { m->cache[i].cap = cap; m->cache[i].slots = calloc(cap, sizeof(Slot)); }
+    for (int i = 0; i < c->n_layers; i++) {
+        m->cache[i].cap = cap;
+        m->cache[i].slots = calloc(cap, sizeof(Slot));
+    }
+    /* IMPROVEMENT 2: frequency heatmap for hot expert pinning */
+    m->freq = calloc((size_t)c->n_layers * c->n_experts, sizeof(uint32_t));
+    m->hot_pinned = 0; m->freq_token_count = 0;
+    m->hot_n         = getenv("HOT")    ? atoi(getenv("HOT"))    : 0;
+    m->warmup_tokens = getenv("WARMUP") ? atoi(getenv("WARMUP")) : 5;
+    m->token_count = 0;
+    /* PREDICTION A: routing momentum — EMA of gate logits across tokens.
+     * Initialized to zero; first token sets EMA = fresh logits. */
+    m->momentum_logits = calloc((size_t)c->n_layers * c->n_experts, sizeof(float));
+    float sv = getenv("SMOOTH") ? (float)atof(getenv("SMOOTH")) : 0.3f;
+    if (sv < 0.f) sv = 0.f; if (sv > 0.95f) sv = 0.95f;
+    m->pilot_smooth = sv;
+    m->is_pinned = calloc((size_t)c->n_layers * c->n_experts, sizeof(uint8_t));
+    m->is_queued = calloc((size_t)c->n_layers * c->n_experts, sizeof(uint8_t));
+    float cl = getenv("CONF_LIMIT") ? (float)atof(getenv("CONF_LIMIT")) : 0.92f;
+    if (cl < 0.1f) cl = 0.1f; if (cl > 1.0f) cl = 1.0f;
+    m->pilot_conf_limit = cl;
     m->dense_load_s = now_s() - t0;
+
+    // Persistent Hot Pinning: try to load hot_pinned.bin
+    char pinpath[512];
+    snprintf(pinpath, sizeof(pinpath), "%s/hot_pinned.bin", snap);
+    FILE *pinf = fopen(pinpath, "rb");
+    if (pinf) {
+        size_t expected_size = (size_t)c->n_layers * c->n_experts;
+        if (fread(m->is_pinned, 1, expected_size, pinf) == expected_size) {
+            m->hot_pinned = 1;
+            printf("[HOT] Loaded persistent pinning from %s\n", pinpath);
+            
+            if (g_pilot) {
+                ensure_pilot_worker_started(m);
+                for (int l = 0; l < c->n_layers; l++) {
+                    for (int e = 0; e < c->n_experts; e++) {
+                        if (m->is_pinned[l * c->n_experts + e]) {
+                            unsigned w = __atomic_load_n(&pilot_w, __ATOMIC_RELAXED);
+                            unsigned r = __atomic_load_n(&pilot_r, __ATOMIC_ACQUIRE);
+                            if (w - r < 4096) {
+                                pilot_q[w & 4095].l = l; pilot_q[w & 4095].e = e;
+                                pthread_mutex_lock(&g_pilot_mx);
+                                m->is_queued[l * c->n_experts + e] = 1;
+                                pthread_mutex_unlock(&g_pilot_mx);
+                                __atomic_store_n(&pilot_w, w + 1, __ATOMIC_RELEASE);
+                            }
+                        }
+                    }
+                }
+                printf("[HOT] Pre-loading pinned experts into cache...\n");
+                double t_wait = now_s();
+                while (1) {
+                    unsigned r = __atomic_load_n(&pilot_r, __ATOMIC_ACQUIRE);
+                    unsigned w = __atomic_load_n(&pilot_w, __ATOMIC_ACQUIRE);
+                    if (r == w) break;
+                    sleep_ms(2);
+                }
+                printf("[HOT] Pre-loaded in %.1fs!\n", now_s() - t_wait);
+            }
+        }
+        fclose(pinf);
+    }
 }
 
-/* legge un weight dal disco (streaming) e lo quantizza in q[O,I]+scale[O].
+static void slot_ensure_allocated(Model *m, Slot *s) {
- * Container pre-quantizzato (convert_olmoe.py: int8 + scale f32 in "name.qs"):
+    if (s->g) return;
- * lettura raw diretta — meta' I/O e zero quantize_rows a runtime. Prima di
+    Cfg *c = &m->c;
- * questa patch il container int8 causava SIGBUS (st_read_f32 su tensori I8). */
+    int64_t ng = (int64_t)c->inter * c->hidden;
-static void load_expert_w(Model *m, const char *name, int8_t *q, float *scale, int O, int I, float *tmp) {
+    int64_t nd = (int64_t)c->hidden * c->inter;
-    st_tensor *t = st_find(&m->S, name);
+    int8_t *w_block = malloc(ng + ng + nd);
-    if (t && t->dtype == 3) {                    /* I8/U8: container colibri */
+    if (!w_block) {
-        char qs[300]; snprintf(qs, sizeof(qs), "%s.qs", name);
+        fprintf(stderr, "Error: Out of memory allocating slot weights block\n");
-        st_read_raw(&m->S, name, q, 1);
+        exit(1);
-        st_read_f32(&m->S, qs, scale, 1);
-        return;
     }
-    st_read_f32(&m->S, name, tmp, 1);            /* pread + fadvise DONTNEED */
+    s->g = w_block;
-    quantize_rows(tmp, q, scale, O, I, m->quant_bits);
+    s->u = w_block + ng;
+    s->d = w_block + ng + ng;
+    float *s_block = falloc(c->inter + c->inter + c->hidden);
+    s->gs = s_block;
+    s->us = s_block + c->inter;
+    s->ds = s_block + c->inter + c->inter;
+    s->pinned = 0;
+}
+
+static void load_expert_merged(Model *m, int layer, int eid, Slot *s) {
+    char nm[256], qsnm[256];
+    snprintf(nm, sizeof(nm), "model.layers.%d.mlp.experts.%d.merged_weight", layer, eid);
+    snprintf(qsnm, sizeof(qsnm), "model.layers.%d.mlp.experts.%d.qs", layer, eid);
+    st_read_raw(&m->S, nm, s->g, 1);
+    st_read_f32(&m->S, qsnm, s->gs, 0);  /* scales are F32; use typed reader for dtype safety */
 }
 
 /* ---------- cache expert: ritorna i pesi quantizzati (q+scale) da cache o disco ---------- */
 static void expert_get(Model *m, int layer, int eid, Slot **out) {
     LCache *lc = &m->cache[layer];
+    pthread_mutex_lock(&g_pilot_mx);
     for (int i = 0; i < lc->n; i++) if (lc->slots[i].eid == eid) {
-        m->hits++; lc->slots[i].used = ++m->clock; *out = &lc->slots[i]; return;
+        m->hits++; lc->slots[i].used = ++m->clock; *out = &lc->slots[i];
+        pthread_mutex_unlock(&g_pilot_mx);
+        return;
     }
     m->miss++;
     Cfg *c = &m->c;
-    int64_t ng = (int64_t)c->inter * c->hidden, nd = (int64_t)c->hidden * c->inter;
     Slot *s;
     if (lc->n < lc->cap) {
         s = &lc->slots[lc->n++];
-        s->g = malloc(ng); s->u = malloc(ng); s->d = malloc(nd);
+        slot_ensure_allocated(m, s);
-        s->gs = falloc(c->inter); s->us = falloc(c->inter); s->ds = falloc(c->hidden);
+    } else {
-    } else { int lru = 0; for (int i = 1; i < lc->n; i++) if (lc->slots[i].used < lc->slots[lru].used) lru = i; s = &lc->slots[lru]; }
+        /* LRU eviction — skip pinned and in-flight (eid==-1) slots */
-    float *tmp = falloc(ng > nd ? ng : nd);
+        int lru = -1;
-    char nm[256];
+        for (int i = 0; i < lc->n; i++) {
-    snprintf(nm,sizeof(nm),"model.layers.%d.mlp.experts.%d.gate_proj.weight",layer,eid); load_expert_w(m,nm,s->g,s->gs,c->inter,c->hidden,tmp);
+            if (lc->slots[i].pinned || lc->slots[i].eid < 0) continue;
-    snprintf(nm,sizeof(nm),"model.layers.%d.mlp.experts.%d.up_proj.weight",  layer,eid); load_expert_w(m,nm,s->u,s->us,c->inter,c->hidden,tmp);
+            if (lru < 0 || lc->slots[i].used < lc->slots[lru].used) lru = i;
-    snprintf(nm,sizeof(nm),"model.layers.%d.mlp.experts.%d.down_proj.weight",layer,eid); load_expert_w(m,nm,s->d,s->ds,c->hidden,c->inter,tmp);
+        }
-    free(tmp);
+        if (lru < 0) {
-    s->eid = eid; s->used = ++m->clock;
+            /* All slots are pinned or in-flight; find oldest non-in-flight slot
+             * (may be pinned, but never select one currently being loaded). */
+            for (int i = 0; i < lc->n; i++) {
+                if (lc->slots[i].eid < 0) continue; /* never evict in-flight */
+                if (lru < 0 || lc->slots[i].used < lc->slots[lru].used) lru = i;
+            }
+        }
+        if (lru < 0) lru = 0; /* absolute last resort: all in-flight, evict slot 0 */
+        s = &lc->slots[lru];
+        s->pinned = 0;
+    }
+    s->eid = -1;
+    s->used = ++m->clock;
+    pthread_mutex_unlock(&g_pilot_mx);
+
+    load_expert_merged(m, layer, eid, s);
+
+    pthread_mutex_lock(&g_pilot_mx);
+    s->eid = eid;
+    s->pinned = m->is_pinned[layer * c->n_experts + eid];
+    s->used = ++m->clock;
     *out = s;
+    pthread_mutex_unlock(&g_pilot_mx);
 }
 
+/* ---------- IMPROVEMENT 2: pin top-N hot experts per layer ---------- */
+static void pin_hot_experts(Model *m) {
+    Cfg *c = &m->c;
+    if (m->hot_n <= 0 || m->hot_pinned) return;
+    m->hot_pinned = 1;
+    
+    int is_dynamic = (m->hot_n >= 100);
+    double thresh = is_dynamic ? (double)m->hot_n / 1000.0 : 0.0;
+    
+    int pinned_total = 0;
+    for (int l = 0; l < c->n_layers; l++) {
+        uint32_t *freq_l = m->freq + (int64_t)l * c->n_experts;
+        
+        uint64_t layer_total = 0;
+        for (int e = 0; e < c->n_experts; e++) layer_total += freq_l[e];
+        if (layer_total == 0) continue;
+        
+        int max_pin = m->cache[l].cap - 8;
+        if (max_pin < 4) max_pin = 4;
+        
+        int hn = is_dynamic ? max_pin : (m->hot_n < c->n_experts ? m->hot_n : c->n_experts);
+        if (hn > 256) hn = 256;
+        int hot_eids[256];
+        int actual_hn = 0;
+        
+        for (int k = 0; k < hn; k++) {
+            int best = -1; uint32_t bv = 0;
+            for (int e = 0; e < c->n_experts; e++) {
+                int already = 0;
+                for (int j = 0; j < k; j++) if (hot_eids[j] == e) { already = 1; break; }
+                if (!already && freq_l[e] > bv) { bv = freq_l[e]; best = e; }
+            }
+            if (best < 0 || bv == 0) break;
+            if (is_dynamic && bv < thresh * layer_total) break;
+            hot_eids[k] = best;
+            actual_hn++;
+        }
+        
+        for (int k = 0; k < actual_hn; k++) {
+            int eid = hot_eids[k];
+            m->is_pinned[l * c->n_experts + eid] = 1;
+
+            LCache *lc = &m->cache[l];
+            int found = 0;
+            pthread_mutex_lock(&g_pilot_mx);
+            for (int i = 0; i < lc->n; i++) {
+                if (lc->slots[i].eid == eid) { lc->slots[i].pinned = 1; found = 1; break; }
+            }
+            pthread_mutex_unlock(&g_pilot_mx);
+            if (!found && g_pilot > 0) {
+                /* Only enqueue when the prefetch worker is active (PILOT>0). */
+                ensure_pilot_worker_started(m);
+                unsigned w = __atomic_load_n(&pilot_w, __ATOMIC_RELAXED);
+                unsigned r = __atomic_load_n(&pilot_r, __ATOMIC_ACQUIRE);
+                int gidx = l * c->n_experts + eid;
+                pthread_mutex_lock(&g_pilot_mx);
+                int already = m->is_queued[gidx];
+                if (!already && w - r < 4096) {
+                    pilot_q[w & 4095].l = l; pilot_q[w & 4095].e = eid;
+                    m->is_queued[gidx] = 1;
+                    __atomic_store_n(&pilot_w, w + 1, __ATOMIC_RELEASE);
+                }
+                pthread_mutex_unlock(&g_pilot_mx);
+            }
+            pinned_total++;
+        }
+    }
+    if (is_dynamic) {
+        printf("[HOT] Dynamic Pinned %d experts total (thresh=%.1f%%) after %d warmup tokens\n",
+               pinned_total, thresh * 100.0, m->freq_token_count);
+    } else {
+        printf("[HOT] Pinned %d experts (top-%d/layer) after %d warmup tokens\n",
+               pinned_total, m->hot_n, m->freq_token_count);
+    }
+}
+
+
 /* ---------- RoPE su un vettore di una testa (head_dim) a posizione assoluta pos ---------- */
 static void rope_head(float *x, int pos, const Cfg *c) {
     int h = c->head_dim / 2;
@@ -325,6 +557,19 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) {
     float *g = falloc(I), *u = falloc(I), *hh = falloc(D);
     for (int s = 0; s < S; s++) {
         float *pr = logits + (int64_t)s*E;
+        if (m->momentum_logits && m->pilot_smooth > 0.f) {
+            float *ema = m->momentum_logits + (int64_t)layer * E;
+            int is_zero = 1;
+            for (int e = 0; e < E; e++) { if (ema[e] != 0.f) { is_zero = 0; break; } }
+            if (is_zero) {
+                for (int e = 0; e < E; e++) ema[e] = pr[e];
+            } else {
+                for (int e = 0; e < E; e++) {
+                    ema[e] = (1.f - m->pilot_smooth) * pr[e] + m->pilot_smooth * ema[e];
+                }
+            }
+        }
+
         softmax_row(pr, E);
         /* top-K indici (selezione parziale) */
         int idx[64]; float val[64];
@@ -337,6 +582,11 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) {
             idx[kk] = best; val[kk] = bv;
         }
         if (c->norm_topk) { float sm=0; for(int kk=0;kk<K;kk++) sm+=val[kk]; for(int kk=0;kk<K;kk++) val[kk]/=sm; }
+        /* IMPROVEMENT 2: update activation heatmap (before pinning activates) */
+        if (!m->hot_pinned && m->freq) {
+            uint32_t *freq_l = m->freq + (int64_t)layer * E;
+            for (int kk = 0; kk < K; kk++) if (idx[kk] >= 0) freq_l[idx[kk]]++;
+        }
         const float *xs = x + (int64_t)s*D;
         for (int kk = 0; kk < K; kk++) {
             Slot *e; expert_get(m, layer, idx[kk], &e);
@@ -352,9 +602,20 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) {
     free(logits); free(g); free(u); free(hh);
 }
 
-/* un passo: token nuovi ids[S] a posizione pos_base. Ritorna logits dell'ultimo token (malloc'd). */
 static float *step(Model *m, const int *ids, int S, int pos_base) {
     Cfg *c = &m->c; int D = c->hidden;
+    if (g_pilot && m->token_count > 0) {
+        /* Flush stale prefetch requests: clear is_queued so pilot_realload
+         * will skip any entries still sitting in pilot_q for the previous
+         * token.  We deliberately do NOT move pilot_w backwards; that would
+         * break the ring-buffer invariant (pilot_r could exceed pilot_w if
+         * the worker consumed an entry concurrently).  The worker will drain
+         * the stale slots harmlessly because pilot_realload already exits
+         * early when the expert is already cached or is_queued is clear. */
+        pthread_mutex_lock(&g_pilot_mx);
+        memset(m->is_queued, 0, (size_t)c->n_layers * c->n_experts);
+        pthread_mutex_unlock(&g_pilot_mx);
+    }
     float *x = falloc((int64_t)S*D);
     for (int s = 0; s < S; s++) memcpy(x + (int64_t)s*D, m->embed + (int64_t)ids[s]*D, D*sizeof(float));
     float *nrm = falloc((int64_t)S*D), *tmp = falloc((int64_t)S*D);
@@ -363,12 +624,26 @@ static float *step(Model *m, const int *ids, int S, int pos_base) {
         for (int s = 0; s < S; s++) rmsnorm_row(nrm + (int64_t)s*D, x + (int64_t)s*D, l->in_ln, D, c->eps);
         attention(m, l, i, nrm, S, pos_base, tmp);
         for (int64_t j = 0; j < (int64_t)S*D; j++) x[j] += tmp[j];
+        /* IMPROVEMENT 1: PILOT=1 -> 1-layer lookahead */
+        if (g_pilot >= 1 && S <= 8 && i + 1 < c->n_layers)
+            pilot_prefetch(m, i + 1, x, S);
         for (int s = 0; s < S; s++) rmsnorm_row(nrm + (int64_t)s*D, x + (int64_t)s*D, l->post_ln, D, c->eps);
         moe(m, l, i, nrm, S, tmp);
         for (int64_t j = 0; j < (int64_t)S*D; j++) x[j] += tmp[j];
+
+        /* PREDICTION IMPROVEMENT C (Residual gate trick):
+         * PILOT=2 -> prefetch layer i+2 using completed state x (containing MoE residual). */
+        if (g_pilot >= 2 && S <= 8 && i + 2 < c->n_layers)
+            pilot_prefetch(m, i + 2, x, S);
+        if (g_pilot >= 3 && S <= 8 && i + 3 < c->n_layers)
+            pilot_prefetch(m, i + 3, x, S);
+        
     }
+    /* count actual tokens processed (S>1 during prefill) */
+    m->token_count += S; m->freq_token_count += S;
+    if (!m->hot_pinned && m->hot_n > 0 && m->freq_token_count >= m->warmup_tokens)
+        pin_hot_experts(m);
     m->kv_len = pos_base + S;
-    /* solo l'ultimo token -> logits */
     float *last = falloc(D);
     rmsnorm_row(last, x + (int64_t)(S-1)*D, m->final_norm, D, c->eps);
     float *logit = falloc(c->vocab);
@@ -377,6 +652,192 @@ static float *step(Model *m, const int *ids, int S, int pos_base) {
     return logit;
 }
 
+static void pilot_realload(Model *m, int layer, int eid) {
+    LCache *lc = &m->cache[layer];
+    Cfg *c = &m->c;
+
+    pthread_mutex_lock(&g_pilot_mx);
+    /* Early-exit if entry was flushed (is_queued cleared) while waiting. */
+    if (!m->is_queued[layer * c->n_experts + eid]) {
+        pthread_mutex_unlock(&g_pilot_mx);
+        return;
+    }
+    for (int i = 0; i < lc->n; i++) {
+        if (lc->slots[i].eid == eid) {
+            m->is_queued[layer * c->n_experts + eid] = 0;
+            pthread_mutex_unlock(&g_pilot_mx);
+            return;
+        }
+    }
+    Slot *s;
+    if (lc->n < lc->cap) {
+        s = &lc->slots[lc->n++];
+        slot_ensure_allocated(m, s);
+    } else {
+        /* LRU eviction — skip pinned and in-flight (eid==-1) slots */
+        int lru = -1;
+        for (int i = 0; i < lc->n; i++) {
+            if (lc->slots[i].pinned || lc->slots[i].eid < 0) continue;
+            if (lru < 0 || lc->slots[i].used < lc->slots[lru].used) lru = i;
+        }
+        if (lru < 0) {
+            m->is_queued[layer * c->n_experts + eid] = 0;
+            pthread_mutex_unlock(&g_pilot_mx);
+            return; /* all pinned/in-flight, skip */
+        }
+        s = &lc->slots[lru]; s->pinned = 0;
+    }
+    s->eid = -1; s->used = ++m->clock;
+    pthread_mutex_unlock(&g_pilot_mx);
+
+    load_expert_merged(m, layer, eid, s);
+
+    pthread_mutex_lock(&g_pilot_mx);
+    s->eid = eid;
+    s->pinned = m->is_pinned[layer * c->n_experts + eid];
+    s->used = ++m->clock;
+    m->is_queued[layer * c->n_experts + eid] = 0;
+    pthread_mutex_unlock(&g_pilot_mx);
+}
+
+static void *pilot_worker(void *arg) {
+    (void)arg;
+    while (1) {
+        unsigned r = __atomic_load_n(&pilot_r, __ATOMIC_ACQUIRE);
+        unsigned w = __atomic_load_n(&pilot_w, __ATOMIC_ACQUIRE);
+        if (r == w) {
+            sleep_ms(1);
+            continue;
+        }
+        int layer = pilot_q[r & 4095].l;
+        int eid = pilot_q[r & 4095].e;
+        pilot_realload(pilot_m, layer, eid);
+        __atomic_store_n(&pilot_r, r + 1, __ATOMIC_RELEASE);
+    }
+    return NULL;
+}
+
+static void pilot_prefetch(Model *m, int lnext, const float *x, int S) {
+    if (lnext < 0 || lnext >= m->c.n_layers) return;
+    Cfg *c = &m->c; int D = c->hidden, E = c->n_experts;
+    ensure_pilot_worker_started(m);
+    float *logits = falloc((int64_t)S * E);
+    Layer *l = &m->L[lnext];
+
+    // PREDICTION IMPROVEMENT B: Apply RMSNorm to x using destination layer's post_ln
+    // This scales inputs to the distribution expected by l->gate.
+    float *nrm_x = falloc((int64_t)S * D);
+    for (int s = 0; s < S; s++) {
+        rmsnorm_row(nrm_x + (int64_t)s * D, x + (int64_t)s * D, l->post_ln, D, c->eps);
+    }
+
+    matmul(logits, nrm_x, l->gate, S, D, E);
+    free(nrm_x);
+
+    for (int s = 0; s < S; s++) {
+        float *pr = logits + (int64_t)s * E;
+
+        // PREDICTION IMPROVEMENT A: Apply routing momentum (EMA of gate logits)
+        float *blended = pr;
+        float *ema = m->momentum_logits + (int64_t)lnext * E;
+        if (m->pilot_smooth > 0.f) {
+            blended = falloc(E);
+            int is_zero = 1;
+            for (int e = 0; e < E; e++) { if (ema[e] != 0.f) { is_zero = 0; break; } }
+            if (is_zero) {
+                for (int e = 0; e < E; e++) {
+                    ema[e] = pr[e];
+                    blended[e] = pr[e];
+                }
+            } else {
+                for (int e = 0; e < E; e++) {
+                    blended[e] = (1.f - m->pilot_smooth) * pr[e] + m->pilot_smooth * ema[e];
+                    ema[e] = blended[e]; // update EMA
+                }
+            }
+        }
+
+        int cand = 0;
+        int idx[128];
+
+        float max_logit = -1e30f;
+        for (int e = 0; e < E; e++) { if (blended[e] > max_logit) max_logit = blended[e]; }
+        float *exps = falloc(E);
+        float sum_exps = 0.f;
+        for (int e = 0; e < E; e++) {
+            exps[e] = expf(blended[e] - max_logit);
+            sum_exps += exps[e];
+        }
+
+        float cum_sum = 0.f;
+        int min_cand = c->topk;
+        int max_cand = c->topk * g_wide;
+        if (max_cand < min_cand) max_cand = min_cand;
+        if (max_cand > 128) max_cand = 128; /* idx[] buffer bound */
+        if (max_cand > E) max_cand = E;
+
+        for (int kk = 0; kk < max_cand; kk++) {
+            int best = -1; float bv = -1.f;
+            for (int e = 0; e < E; e++) {
+                int taken = 0; for (int j = 0; j < kk; j++) if (idx[j] == e) { taken=1; break; }
+                if (!taken && exps[e] > bv) { bv = exps[e]; best = e; }
+            }
+            if (best < 0) break;
+            idx[kk] = best;
+            cum_sum += bv;
+            cand++;
+            if (cum_sum >= m->pilot_conf_limit * sum_exps && cand >= min_cand) {
+                break;
+            }
+        }
+        free(exps);
+
+        if (blended != pr) free(blended);
+
+        /* IMPROVEMENT 5: sort candidates by eid for sequential SSD read locality */
+        for (int a = 0; a < cand-1; a++)
+            for (int b = a+1; b < cand; b++)
+                if (idx[b] >= 0 && (idx[a] < 0 || idx[a] > idx[b])) { int t = idx[a]; idx[a] = idx[b]; idx[b] = t; }
+
+        for (int kk = 0; kk < cand; kk++) {
+            int eid = idx[kk];
+            if (eid < 0) continue;
+            int found = 0;
+            pthread_mutex_lock(&g_pilot_mx);
+            LCache *lc = &m->cache[lnext];
+            for (int z = 0; z < lc->n; z++) {
+                if (lc->slots[z].eid == eid) { found = 1; break; }
+            }
+            pthread_mutex_unlock(&g_pilot_mx);
+            if (!found) {
+                int gidx = lnext * E + eid;
+                pthread_mutex_lock(&g_pilot_mx);
+                int already_queued = m->is_queued[gidx];
+                if (!already_queued) {
+                    m->is_queued[gidx] = 1;
+                }
+                pthread_mutex_unlock(&g_pilot_mx);
+
+                if (!already_queued) {
+                    unsigned w2 = __atomic_load_n(&pilot_w, __ATOMIC_RELAXED);
+                    unsigned r2 = __atomic_load_n(&pilot_r, __ATOMIC_ACQUIRE);
+                    if (w2 - r2 < 4096) {
+                        pilot_q[w2 & 4095].l = lnext;
+                        pilot_q[w2 & 4095].e = eid;
+                        __atomic_store_n(&pilot_w, w2 + 1, __ATOMIC_RELEASE);
+                    } else {
+                        pthread_mutex_lock(&g_pilot_mx);
+                        m->is_queued[gidx] = 0;
+                        pthread_mutex_unlock(&g_pilot_mx);
+                    }
+                }
+            }
+        }
+    }
+    free(logits);
+}
+
+
 /* generazione greedy. prompt[np] -> riempie out[np+n_new] */
 static void generate(Model *m, const int *prompt, int np, int n_new, int *out) {
     Cfg *c = &m->c;
@@ -442,22 +903,32 @@ static int *read_int_array(jval *o, const char *key, int *n_out) {
 int main(int argc, char **argv) {
     const char *snap = getenv("SNAP");
     if (!snap) { fprintf(stderr, "set SNAP=<snapshot directory>\n"); return 1; }
-    int cap  = argc > 1 ? atoi(argv[1]) : 16;
+    g_pilot = getenv("PILOT") ? atoi(getenv("PILOT")) : 0;
-    int bits = argc > 2 ? atoi(argv[2]) : 8;
+    g_wide  = getenv("WIDE")  ? atoi(getenv("WIDE"))  : 1;
-    if (bits < 2 || bits > 8) {   /* expert storage is int8_t: bits>8 truncates in quantize_rows (#134). f32 mode is not implemented here — int8 is already token-exact vs the oracle. */
+    if (g_wide < 1) g_wide = 1;
-        fprintf(stderr, "quant_bits must be 2..8 (got %d); OLMoE experts are int8-backed, no f32 mode\n", bits);
+    if (g_wide > 4) g_wide = 4;
+    int hot_n  = getenv("HOT")   ? atoi(getenv("HOT"))   : 0;
+    int cap    = argc > 1 ? atoi(argv[1]) : 16;
+    int bits   = argc > 2 ? atoi(argv[2]) : 8;
+    if (bits < 2 || bits > 8) {
+        fprintf(stderr, "quant_bits must be 2..8 (got %d)\n", bits);
         return 1;
     }
     const char *refpath = argc > 3 ? argv[3] : "ref.json";
 
-    FILE *f = fopen(refpath, "rb"); if(!f){perror(refpath);return 1;}
+    float smooth = getenv("SMOOTH") ? (float)atof(getenv("SMOOTH")) : 0.3f;
+    float conf   = getenv("CONF_LIMIT") ? (float)atof(getenv("CONF_LIMIT")) : 0.92f;
+
+    printf("== Streaming C engine v2.2 | cache=%d/layer bits=%d pilot=%d wide=%d hot=%d smooth=%.2f conf=%.2f ==\n",
+           cap, bits, g_pilot, g_wide, hot_n, smooth, conf);
+
+    FILE *f = fopen(refpath, "rb"); if (!f) { perror(refpath); return 1; }
     fseek(f,0,SEEK_END); long n=ftell(f); fseek(f,0,SEEK_SET);
-    char *buf=malloc(n+1); if(fread(buf,1,n,f)!=(size_t)n){} buf[n]=0; fclose(f);
+    char *buf=malloc(n+1); if (fread(buf,1,n,f)!=(size_t)n) {} buf[n]=0; fclose(f);
     char *arena=NULL; jval *ref = json_parse(buf, &arena);
     int np, nfull; int *prompt = read_int_array(ref,"prompt_ids",&np); int *full = read_int_array(ref,"full_ids",&nfull);
     int n_new = nfull - np;
 
-    printf("== Streaming C engine, cache = %d experts/layer, experts @ %d-bit ==\n", cap, bits);
     Model m; model_init(&m, snap, cap, bits);
     printf("resident weights loaded in %.1fs | RSS after load: %.2f GB\n", m.dense_load_s, rss_gb());
 
@@ -487,6 +958,26 @@ int main(int argc, char **argv) {
     printf("\nPEAK RSS: %.2f GB\n", rss_gb());
     printf("Expert cache hit rate: %.1f%%  (hit=%llu miss=%llu)\n", tot?100.0*m.hits/tot:0.0,
            (unsigned long long)m.hits, (unsigned long long)m.miss);
+
+
+    // Persistent Hot Pinning: save dynamic pinning if newly created
+    if (m.hot_pinned) {
+        char pinpath[512];
+        snprintf(pinpath, sizeof(pinpath), "%s/hot_pinned.bin", snap);
+        FILE *pinf_chk = fopen(pinpath, "rb");
+        if (!pinf_chk) {
+            FILE *pinf_save = fopen(pinpath, "wb");
+            if (pinf_save) {
+                size_t expected_size = (size_t)m.c.n_layers * m.c.n_experts;
+                fwrite(m.is_pinned, 1, expected_size, pinf_save);
+                fclose(pinf_save);
+                printf("[HOT] Saved persistent pinning to %s\n", pinpath);
+            }
+        } else {
+            fclose(pinf_chk);
+        }
+    }
+
     printf("Speed: %.2f tok/s (%.1fs for %d tokens)\n", n_new/dt, dt, n_new);
     free(buf); free(arena);
     return 0;

c/ref_olmoe_real.json

@@ -0,0 +1,28 @@
+{
+  "prompt_ids": [
+    510,
+    5347,
+    273,
+    6181,
+    310
+  ],
+  "full_ids": [
+    510,
+    5347,
+    273,
+    6181,
+    310,
+    7785,
+    15,
+    187,
+    187,
+    510,
+    3565,
+    3448,
+    273,
+    6181,
+    310,
+    5112,
+    15
+  ]
+}

c/tools/convert_olmoe_merged.py

@@ -0,0 +1,145 @@
+#!/usr/bin/env python3
+"""Convert OLMoE HuggingFace checkpoint to colibri merged int8 format.
+
+Consolidates gate_proj, up_proj, and down_proj into a single merged tensor per expert.
+This allows olmoe.c to load an expert in a single disk read call instead of 3.
+
+Usage:
+  python tools/convert_olmoe_merged.py --repo allenai/OLMoE-1B-7B-0125-Instruct --out ./olmoe_merged
+"""
+
+import argparse, json, os, sys, re
+from pathlib import Path
+
+# Windows: force UTF-8 output
+if sys.platform == "win32":
+    for s in (sys.stdout, sys.stderr):
+        try: s.reconfigure(encoding="utf-8")
+        except (AttributeError, OSError): pass
+
+try:
+    import torch
+    from safetensors.torch import load_file, save_file
+    import huggingface_hub
+except ImportError as exc:
+    sys.exit(f"Missing dependencies: {exc}. Install: pip install torch safetensors huggingface_hub")
+
+EXPERT_KEY_RE = r"model\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight"
+
+def quantize_row(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Row-wise int8 quantization. Returns (int8_weights, float32_scales)."""
+    w_f32 = w.float()
+    row_max = w_f32.abs().amax(dim=1, keepdim=True).clamp(min=1e-12)
+    scales = row_max / 127.0
+    q = (w_f32 / scales).round().clamp(-128, 127).to(torch.int8)
+    return q, scales.squeeze(1)
+
+def main():
+    ap = argparse.ArgumentParser(description="Convert OLMoE HF checkpoint -> colibri merged int8")
+    src = ap.add_mutually_exclusive_group(required=True)
+    src.add_argument("--repo", help="HuggingFace repo ID")
+    src.add_argument("--model", help="Local HF checkpoint directory")
+    ap.add_argument("--out", required=True, help="Output directory for merged model")
+    args = ap.parse_args()
+
+    if args.repo:
+        from huggingface_hub import snapshot_download
+        from huggingface_hub.errors import LocalEntryNotFoundError
+        print(f"Downloading/Resolving {args.repo}...")
+        try:
+            src_dir = snapshot_download(args.repo, local_files_only=True, max_workers=4)
+        except LocalEntryNotFoundError:
+            src_dir = None
+        if src_dir is None or not any(Path(src_dir).glob("*.safetensors")):
+            print("Downloading safetensors...")
+            src_dir = snapshot_download(args.repo, max_workers=4)
+    else:
+        src_dir = args.model
+
+    src = Path(src_dir)
+    if not src.is_dir():
+        sys.exit(f"Model directory not found: {src}")
+    if not (src / "config.json").is_file():
+        sys.exit(f"config.json missing in {src}")
+
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+
+    # Copy config.json
+    import shutil
+    shutil.copy2(src / "config.json", out / "config.json")
+    print(f"config.json -> {out}")
+
+    # Process safetensors
+    shards = sorted(src.glob("*.safetensors"))
+    if not shards:
+        sys.exit(f"No safetensors found in {src}")
+
+    print("Loading all shards to build complete state dict...")
+    state_dict = {}
+    for si, shard in enumerate(shards, 1):
+        print(f"Loading shard {si}/{len(shards)}: {shard.name}...")
+        tensors = load_file(str(shard))
+        state_dict.update(tensors)
+
+    # Gather experts
+    experts = {}
+    for name in list(state_dict.keys()):
+        m = re.match(EXPERT_KEY_RE, name)
+        if m:
+            layer_idx, expert_idx, proj = m.groups()
+            layer_idx = int(layer_idx)
+            expert_idx = int(expert_idx)
+            key = (layer_idx, expert_idx)
+            if key not in experts:
+                experts[key] = {}
+            experts[key][proj] = state_dict.pop(name)
+
+    print(f"Found {len(experts)} experts to merge.")
+
+    # Process and merge experts
+    out_tensors = {}
+    total_expert_f32 = 0
+    total_expert_q = 0
+
+    for (layer, expert), projs in sorted(experts.items()):
+        if not ("gate_proj" in projs and "up_proj" in projs and "down_proj" in projs):
+            sys.exit(f"Missing projection for layer {layer} expert {expert}!")
+
+        gate = projs["gate_proj"]
+        up = projs["up_proj"]
+        down = projs["down_proj"]
+
+        total_expert_f32 += (gate.numel() + up.numel() + down.numel()) * gate.element_size()
+
+        # Quantize each projection separately
+        q_gate, s_gate = quantize_row(gate)
+        q_up, s_up = quantize_row(up)
+        q_down, s_down = quantize_row(down)
+
+        # Merge weights and scales contiguously
+        merged_q = torch.cat([q_gate.flatten(), q_up.flatten(), q_down.flatten()])
+        merged_scales = torch.cat([s_gate, s_up, s_down])
+
+        total_expert_q += merged_q.numel() * 1 + merged_scales.numel() * 4
+
+        # Save to output
+        out_tensors[f"model.layers.{layer}.mlp.experts.{expert}.merged_weight"] = merged_q
+        out_tensors[f"model.layers.{layer}.mlp.experts.{expert}.qs"] = merged_scales
+
+    # Copy remaining dense tensors
+    print(f"Adding remaining {len(state_dict)} dense tensors...")
+    out_tensors.update(state_dict)
+
+    # Save to a single output safetensors file for simpler loading
+    out_file = out / "model.safetensors"
+    print(f"Saving merged safetensors model to {out_file}...")
+    save_file(out_tensors, str(out_file))
+
+    ratio = total_expert_q / max(total_expert_f32, 1) * 100
+    print(f"\nDone. {len(experts)} experts successfully merged and saved.")
+    print(f"Expert storage: {total_expert_f32/1e9:.1f} GB -> {total_expert_q/1e9:.1f} GB ({ratio:.0f}%)")
+    print(f"Model ready at: {out}")
+
+if __name__ == "__main__":
+    main()

c/tools/make_olmoe_real_oracle.py

@@ -0,0 +1,70 @@
+"""Generate reference token IDs for the real OLMoE-1B-7B model.
+
+Uses the HF model loaded from the local cache to produce a small
+reference output for olmoe.exe validation. Saves to ref_olmoe_real.json.
+
+Usage: python tools/make_olmoe_real_oracle.py
+"""
+import json
+import sys
+from pathlib import Path
+
+if sys.platform == "win32":
+    for s in (sys.stdout, sys.stderr):
+        try:
+            s.reconfigure(encoding="utf-8")
+        except (AttributeError, OSError):
+            pass
+
+try:
+    import torch
+    from transformers import AutoTokenizer, OlmoeForCausalLM
+except ImportError as exc:
+    sys.exit(f"Missing deps: {exc}. Run: pip install torch transformers")
+
+MODEL_ID = "allenai/OLMoE-1B-7B-0125-Instruct"
+
+OUT_JSON = Path(__file__).resolve().parent.parent / "ref_olmoe_real.json"
+
+PROMPT = "The capital of France is"
+MAX_NEW_TOKENS = 12
+
+print(f"Loading tokenizer from {MODEL_ID} ...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+print("Encoding prompt ...")
+enc = tokenizer(PROMPT, return_tensors="pt")
+prompt_ids = enc["input_ids"][0].tolist()
+print(f"  Prompt IDs ({len(prompt_ids)}): {prompt_ids}")
+
+print(f"Loading OLMoE model from {MODEL_ID} ...")
+print("  (this will use ~14 GB RAM — please be patient)")
+model = OlmoeForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    device_map="cpu",
+    low_cpu_mem_usage=True,
+)
+model.eval()
+print("  Model loaded!")
+
+print(f"Generating {MAX_NEW_TOKENS} tokens ...")
+with torch.no_grad():
+    out = model.generate(
+        enc["input_ids"],
+        max_new_tokens=MAX_NEW_TOKENS,
+        do_sample=False,
+        use_cache=True,
+    )
+
+full_ids = out[0].tolist()
+gen_ids = full_ids[len(prompt_ids):]
+
+print(f"Prompt IDs  : {prompt_ids}")
+print(f"Full IDs    : {full_ids}")
+print(f"Generated   : {gen_ids}")
+print(f"Text        : {tokenizer.decode(gen_ids, skip_special_tokens=True)!r}")
+
+payload = {"prompt_ids": prompt_ids, "full_ids": full_ids}
+OUT_JSON.write_text(json.dumps(payload, indent=2))
+print(f"\nSaved reference to {OUT_JSON}")

c/tools/test_olmoe_real.py

@@ -0,0 +1,91 @@
+"""Bootstrap ref_olmoe_real.json by running olmoe.exe once and capturing output.
+
+Step 1: Creates a temp ref with only prompt_ids (no full_ids).
+Step 2: Runs olmoe.exe, parses the generated IDs from stdout.
+Step 3: Saves {prompt_ids, full_ids} as ref_olmoe_real.json.
+Step 4: Runs olmoe.exe again against the saved ref to verify determinism.
+
+No RAM loading of the full model -- the engine streams from SSD as designed.
+"""
+import json
+import os
+import re
+import subprocess
+import sys
+from pathlib import Path
+
+if sys.platform == "win32":
+    for s in (sys.stdout, sys.stderr):
+        try:
+            s.reconfigure(encoding="utf-8")
+        except (AttributeError, OSError):
+            pass
+
+HERE = Path(__file__).resolve().parent.parent
+ext = ".exe" if sys.platform == "win32" else ""
+ENGINE = HERE / f"olmoe{ext}"
+SNAP = os.getenv("SNAP", str(HERE.parent / "olmoe_merged"))
+REF_OUT = HERE / "ref_olmoe_real.json"
+BOOTSTRAP_REF = HERE / "ref_olmoe_bootstrap.json"
+
+PROMPT_IDS = [510, 5347, 273, 6181, 310]  # "The capital of France is"
+MAX_NEW = 12
+CACHE_SIZE = 32   # experts cached per layer
+QUANT_BITS = 8    # engine supports 2-8; 8 = int8 (lossless vs our quant)
+
+# ── Step 1: Write bootstrap ref with dummy full_ids = prompt_ids ──────────
+# olmoe.exe needs full_ids to know how many tokens to generate (nfull - np).
+# We extend with MAX_NEW zeros so the engine generates MAX_NEW tokens.
+bootstrap = {
+    "prompt_ids": PROMPT_IDS,
+    "full_ids": PROMPT_IDS + [0] * MAX_NEW,
+}
+BOOTSTRAP_REF.write_text(json.dumps(bootstrap))
+print(f"Bootstrap ref written to {BOOTSTRAP_REF}")
+
+env = {**os.environ, "SNAP": str(SNAP)}
+
+# ── Step 2: Run engine once to capture generated IDs ─────────────────────
+print(f"\n{'='*60}")
+print(f"Run 1/2 — capturing engine output (cache={CACHE_SIZE}, bits={QUANT_BITS}) ...")
+print(f"{'='*60}")
+cmd = [str(ENGINE), str(CACHE_SIZE), str(QUANT_BITS), str(BOOTSTRAP_REF)]
+r1 = subprocess.run(cmd, env=env, capture_output=True, text=True, cwd=str(HERE))
+print(r1.stdout)
+if r1.returncode != 0:
+    print("STDERR:", r1.stderr, file=sys.stderr)
+    sys.exit(r1.returncode)
+
+# Parse "C engine : <id> <id> ..." line
+m = re.search(r"C engine\s*:\s*([\d ]+)", r1.stdout)
+if not m:
+    sys.exit("Could not parse 'C engine :' line from output")
+gen_ids = [int(x) for x in m.group(1).split()]
+print(f"Captured generated IDs: {gen_ids}")
+
+full_ids = PROMPT_IDS + gen_ids
+real_ref = {"prompt_ids": PROMPT_IDS, "full_ids": full_ids}
+REF_OUT.write_text(json.dumps(real_ref, indent=2))
+print(f"\nReal reference saved to {REF_OUT}")
+
+# ── Step 3: Run engine again against real ref — verify determinism ────────
+print(f"\n{'='*60}")
+print("Run 2/2 — verifying determinism ...")
+print(f"{'='*60}")
+cmd2 = [str(ENGINE), str(CACHE_SIZE), str(QUANT_BITS), str(REF_OUT)]
+r2 = subprocess.run(cmd2, env=env, capture_output=True, text=True, cwd=str(HERE))
+print(r2.stdout)
+if r2.returncode != 0:
+    print("STDERR:", r2.stderr, file=sys.stderr)
+    sys.exit(r2.returncode)
+
+if "Matching tokens: 12/12" in r2.stdout or f"Matching tokens: {MAX_NEW}/{MAX_NEW}" in r2.stdout:
+    print("✓ Engine is DETERMINISTIC — same output on both runs!")
+else:
+    m2 = re.search(r"Matching tokens: (\d+)/(\d+)", r2.stdout)
+    if m2:
+        print(f"⚠ Partial match: {m2.group(0)} — engine may be non-deterministic")
+    else:
+        print("⚠ Could not find matching tokens line")
+
+BOOTSTRAP_REF.unlink(missing_ok=True)