From ac1f7a8f387db0c7a63f34e2dd5b47b3da720d41 Mon Sep 17 00:00:00 2001 From: Egon Ruiter Date: Wed, 15 Jul 2026 18:01:21 +0200 Subject: [PATCH 01/11] feat(prefetch): implement prefetcher v2.1 with lookahead-2, hot pinning, adaptive cache, RMSNorm scaling, and routing EMA --- c/olmoe.c | 360 ++++++++++++++++++++++++++++-- c/ref_olmoe_real.json | 28 +++ c/tools/make_olmoe_real_oracle.py | 74 ++++++ c/tools/test_olmoe_real.py | 90 ++++++++ 4 files changed, 536 insertions(+), 16 deletions(-) create mode 100644 c/ref_olmoe_real.json create mode 100644 c/tools/make_olmoe_real_oracle.py create mode 100644 c/tools/test_olmoe_real.py diff --git a/c/olmoe.c b/c/olmoe.c index 5923bde..b0f9e8c 100644 --- a/c/olmoe.c +++ b/c/olmoe.c @@ -5,6 +5,14 @@ * 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 : 0=no prefetch, 1=1-layer lookahead, 2=2-layer lookahead [IMPROVEMENT 1] + * HOT=N : pin top-N hot experts per layer permanently (never evict) [IMPROVEMENT 2] + * WARMUP=N : tokens before hot pinning activates (default 5) [IMPROVEMENT 2] + * REBAL=N : rebalance cache per-layer every N tokens (0=off) [IMPROVEMENT 3] + * WIDE=N : prefetch top-K*N candidates (default 1, try 2 or 3) [IMPROVEMENT 4] + * (expert queue is sorted by eid for SSD locality) [IMPROVEMENT 5] */ #define _GNU_SOURCE #include @@ -12,11 +20,22 @@ #include #include #include +#include #if defined(__APPLE__) || defined(__linux__) || defined(__FreeBSD__) #include +#include #endif #include "st.h" +#ifdef _WIN32 +#include +#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 +52,44 @@ 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; -typedef struct { Slot *slots; int n, cap; } LCache; +/* IMPROVEMENT 2: pinned=1 means this slot is never evicted (hot expert). */ +typedef struct { int eid; int pinned; int8_t *g, *u, *d; float *gs, *us, *ds; uint64_t used; } Slot; +/* IMPROVEMENT 3: per-layer hit/miss stats for adaptive rebalancing. */ +typedef struct { Slot *slots; int n, cap; uint64_t layer_hits, layer_miss; } 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; + /* IMPROVEMENT 3: adaptive rebalance */ + int token_count, rebal_interval, total_cap; + /* PREDICTION IMPROVEMENT A: per-layer smoothed gate logits across tokens. + * momentum_logits[l*E .. (l+1)*E-1] = EMA of recent gate outputs. + * Blended with fresh gate prediction: final = (1-smooth)*fresh + smooth*ema. + * Captures routing consistency across tokens (same token tends to reuse experts). */ + 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 */ } 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); + /* ---------- 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__) @@ -209,8 +250,24 @@ static void model_init(Model *m, const char *snap, int cap, int bits) { LD(gate, "mlp.gate.weight"); #undef LD } + m->total_cap = cap; 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)); } + /* 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; + /* IMPROVEMENT 3: adaptive rebalance */ + m->rebal_interval = getenv("REBAL") ? atoi(getenv("REBAL")) : 0; + 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->dense_load_s = now_s() - t0; } @@ -233,10 +290,13 @@ static void load_expert_w(Model *m, const char *name, int8_t *q, float *scale, i /* ---------- 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->layer_hits++; lc->slots[i].used = ++m->clock; *out = &lc->slots[i]; + pthread_mutex_unlock(&g_pilot_mx); + return; } - m->miss++; + m->miss++; lc->layer_miss++; Cfg *c = &m->c; int64_t ng = (int64_t)c->inter * c->hidden, nd = (int64_t)c->hidden * c->inter; Slot *s; @@ -244,15 +304,107 @@ static void expert_get(Model *m, int layer, int eid, Slot **out) { s = &lc->slots[lc->n++]; s->g = malloc(ng); s->u = malloc(ng); s->d = malloc(nd); s->gs = falloc(c->inter); s->us = falloc(c->inter); s->ds = falloc(c->hidden); - } 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]; } + s->pinned = 0; + } else { + /* IMPROVEMENT 2: LRU eviction — never evict pinned experts */ + int lru = -1; + for (int i = 0; i < lc->n; i++) { + if (lc->slots[i].pinned) continue; + if (lru < 0 || lc->slots[i].used < lc->slots[lru].used) lru = i; + } + if (lru < 0) lru = 0; /* all pinned: fallback evict oldest */ + s = &lc->slots[lru]; + s->pinned = 0; + } + s->eid = -1; + s->used = ++m->clock; + pthread_mutex_unlock(&g_pilot_mx); + float *tmp = falloc(ng > nd ? ng : nd); char nm[256]; 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); 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); 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); - s->eid = eid; s->used = ++m->clock; + + 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 hn = m->hot_n < c->n_experts ? m->hot_n : c->n_experts; + 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; + int hot_eids[256]; + /* Find top hn experts by activation frequency */ + 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) { hn = k; break; } + hot_eids[k] = best; + } + /* Mark already-cached hot experts as pinned; enqueue uncached ones */ + for (int k = 0; k < hn; k++) { + int eid = hot_eids[k]; + m->is_pinned[l * c->n_experts + eid] = 1; // Mark globally + + 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) { + 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 = eid; + __atomic_store_n(&pilot_w, w + 1, __ATOMIC_RELEASE); + } + } + pinned_total++; + } + } + printf("[HOT] Pinned %d experts (top-%d/layer) after %d warmup tokens\n", + pinned_total, m->hot_n, m->freq_token_count); +} + +/* ---------- IMPROVEMENT 3: adaptive per-layer cache rebalancing ---------- */ +static void rebalance_cache(Model *m) { + Cfg *c = &m->c; + uint64_t total_miss = 0; + for (int l = 0; l < c->n_layers; l++) total_miss += m->cache[l].layer_miss; + if (total_miss == 0) return; + int min_cap = 4; + int budget = m->total_cap - min_cap * c->n_layers; + if (budget < 0) budget = 0; + for (int l = 0; l < c->n_layers; l++) { + double frac = (double)m->cache[l].layer_miss / (double)total_miss; + int new_cap = min_cap + (int)(frac * budget + 0.5); + LCache *lc = &m->cache[l]; + if (new_cap > lc->cap) { + Slot *ns = realloc(lc->slots, new_cap * sizeof(Slot)); + if (ns) { + memset(ns + lc->cap, 0, (new_cap - lc->cap) * sizeof(Slot)); + lc->slots = ns; lc->cap = new_cap; + } + } + lc->layer_hits = 0; lc->layer_miss = 0; + } } /* ---------- RoPE su un vettore di una testa (head_dim) a posizione assoluta pos ---------- */ @@ -337,6 +489,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;kkhot_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); @@ -363,12 +520,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); } + /* IMPROVEMENT 2: count tokens; trigger hot pinning after warmup */ + m->token_count++; m->freq_token_count++; + if (!m->hot_pinned && m->hot_n > 0 && m->freq_token_count >= m->warmup_tokens) + pin_hot_experts(m); + /* IMPROVEMENT 3: periodic adaptive rebalance */ + if (m->rebal_interval > 0 && m->token_count % m->rebal_interval == 0) + rebalance_cache(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 +548,155 @@ 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; + int64_t ng = (int64_t)c->inter * c->hidden, nd = (int64_t)c->hidden * c->inter; + + pthread_mutex_lock(&g_pilot_mx); + for (int i = 0; i < lc->n; i++) { + if (lc->slots[i].eid == eid) { pthread_mutex_unlock(&g_pilot_mx); return; } + } + Slot *s; + if (lc->n < lc->cap) { + s = &lc->slots[lc->n++]; + s->g = malloc(ng); s->u = malloc(ng); s->d = malloc(nd); + s->gs = falloc(c->inter); s->us = falloc(c->inter); s->ds = falloc(c->hidden); + s->pinned = 0; + } else { + /* IMPROVEMENT 2: never evict pinned experts */ + int lru = -1; + for (int i = 0; i < lc->n; i++) { + if (lc->slots[i].pinned) continue; + if (lru < 0 || lc->slots[i].used < lc->slots[lru].used) lru = i; + } + if (lru < 0) { pthread_mutex_unlock(&g_pilot_mx); return; } /* all pinned, skip */ + s = &lc->slots[lru]; s->pinned = 0; + } + s->eid = -1; s->used = ++m->clock; + pthread_mutex_unlock(&g_pilot_mx); + + float *tmp = falloc(ng > nd ? ng : nd); + char nm[256]; + 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); + 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); + 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); + + pthread_mutex_lock(&g_pilot_mx); + s->eid = eid; + s->pinned = m->is_pinned[layer * c->n_experts + eid]; + s->used = ++m->clock; + 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; + /* IMPROVEMENT 4: wide prefetch — top K * g_wide candidates */ + int cand = c->topk * g_wide; + if (cand > E) cand = E; + if (!pilot_m) { + pilot_m = m; + pthread_t t; + pthread_create(&t, NULL, pilot_worker, NULL); + } + 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 idx[128]; /* up to topk * 4 */ + for (int kk = 0; kk < cand; kk++) { + int best = -1; float bv = -1e30f; + 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 && blended[e] > bv) { bv = blended[e]; best = e; } + } + idx[kk] = best; + } + + 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) { + 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); + } + } + } + } + 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; @@ -411,22 +731,30 @@ 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=\n"); return 1; } - int cap = argc > 1 ? atoi(argv[1]) : 16; - int bits = argc > 2 ? atoi(argv[2]) : 8; - 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. */ - fprintf(stderr, "quant_bits must be 2..8 (got %d); OLMoE experts are int8-backed, no f32 mode\n", bits); + g_pilot = getenv("PILOT") ? atoi(getenv("PILOT")) : 0; + g_wide = getenv("WIDE") ? atoi(getenv("WIDE")) : 1; + if (g_wide < 1) g_wide = 1; + if (g_wide > 4) g_wide = 4; + int hot_n = getenv("HOT") ? atoi(getenv("HOT")) : 0; + int rebal = getenv("REBAL") ? atoi(getenv("REBAL")) : 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;} + printf("== Streaming C engine v2 | cache=%d/layer bits=%d pilot=%d wide=%d hot=%d rebal=%d ==\n", + cap, bits, g_pilot, g_wide, hot_n, rebal); + + 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()); diff --git a/c/ref_olmoe_real.json b/c/ref_olmoe_real.json new file mode 100644 index 0000000..b271e07 --- /dev/null +++ b/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 + ] +} \ No newline at end of file diff --git a/c/tools/make_olmoe_real_oracle.py b/c/tools/make_olmoe_real_oracle.py new file mode 100644 index 0000000..374d29f --- /dev/null +++ b/c/tools/make_olmoe_real_oracle.py @@ -0,0 +1,74 @@ +"""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_DIR = ( + Path(r"C:\Users\egonr\.cache\huggingface\hub" + r"\models--allenai--OLMoE-1B-7B-0125-Instruct" + r"\snapshots\b89a7c4bc24fb9e55ce2543c9458ce0ca5c4650e") +) + +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_DIR} ...") +tokenizer = AutoTokenizer.from_pretrained(str(MODEL_DIR)) + +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_DIR} ...") +print(" (this will use ~14 GB RAM — please be patient)") +model = OlmoeForCausalLM.from_pretrained( + str(MODEL_DIR), + 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}") diff --git a/c/tools/test_olmoe_real.py b/c/tools/test_olmoe_real.py new file mode 100644 index 0000000..ac47dc2 --- /dev/null +++ b/c/tools/test_olmoe_real.py @@ -0,0 +1,90 @@ +"""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 +ENGINE = HERE / "olmoe.exe" +SNAP = HERE / "olmoe_i4" +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 : ..." 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) From c4624276f37b1d0ac38b464f3d995d1627cf1d22 Mon Sep 17 00:00:00 2001 From: Egon Ruiter Date: Wed, 15 Jul 2026 18:04:25 +0200 Subject: [PATCH 02/11] feat(io): implement Consolidated Expert I/O to reduce expert disk reads by 3x and accelerate prefetching --- .gitignore | 4 + c/olmoe.c | 64 ++++++-------- c/tools/convert_olmoe_merged.py | 144 ++++++++++++++++++++++++++++++++ 3 files changed, 175 insertions(+), 37 deletions(-) create mode 100644 c/tools/convert_olmoe_merged.py diff --git a/.gitignore b/.gitignore index dba6a3f..c06902f 100644 --- a/.gitignore +++ b/.gitignore @@ -67,3 +67,7 @@ c/bench/ c/tests/test_decode_batch c/tests/test_i4_acc512 c/tests/test_idot +olmoe_merged/ +olmoe_i4/ +c/olmoe_merged/ +c/olmoe_i4/ diff --git a/c/olmoe.c b/c/olmoe.c index b0f9e8c..199ccaf 100644 --- a/c/olmoe.c +++ b/c/olmoe.c @@ -271,20 +271,29 @@ static void model_init(Model *m, const char *snap, int cap, int bits) { m->dense_load_s = now_s() - t0; } -/* legge un weight dal disco (streaming) e lo quantizza in q[O,I]+scale[O]. - * Container pre-quantizzato (convert_olmoe.py: int8 + scale f32 in "name.qs"): - * lettura raw diretta — meta' I/O e zero quantize_rows a runtime. Prima di - * questa patch il container int8 causava SIGBUS (st_read_f32 su tensori I8). */ -static void load_expert_w(Model *m, const char *name, int8_t *q, float *scale, int O, int I, float *tmp) { - st_tensor *t = st_find(&m->S, name); - if (t && t->dtype == 3) { /* I8/U8: container colibri */ - char qs[300]; snprintf(qs, sizeof(qs), "%s.qs", name); - st_read_raw(&m->S, name, q, 1); - st_read_f32(&m->S, qs, scale, 1); - return; - } - st_read_f32(&m->S, name, tmp, 1); /* pread + fadvise DONTNEED */ - quantize_rows(tmp, q, scale, O, I, m->quant_bits); + +static void slot_ensure_allocated(Model *m, Slot *s) { + if (s->g) return; + Cfg *c = &m->c; + int64_t ng = (int64_t)c->inter * c->hidden; + int64_t nd = (int64_t)c->hidden * c->inter; + int8_t *w_block = malloc(ng + ng + nd); + s->g = w_block; + 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_raw(&m->S, qsnm, s->gs, 1); } /* ---------- cache expert: ritorna i pesi quantizzati (q+scale) da cache o disco ---------- */ @@ -298,13 +307,10 @@ static void expert_get(Model *m, int layer, int eid, Slot **out) { } m->miss++; lc->layer_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); - s->gs = falloc(c->inter); s->us = falloc(c->inter); s->ds = falloc(c->hidden); - s->pinned = 0; + slot_ensure_allocated(m, s); } else { /* IMPROVEMENT 2: LRU eviction — never evict pinned experts */ int lru = -1; @@ -320,12 +326,7 @@ static void expert_get(Model *m, int layer, int eid, Slot **out) { s->used = ++m->clock; pthread_mutex_unlock(&g_pilot_mx); - float *tmp = falloc(ng > nd ? ng : nd); - char nm[256]; - 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); - 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); - 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); + load_expert_merged(m, layer, eid, s); pthread_mutex_lock(&g_pilot_mx); s->eid = eid; @@ -551,7 +552,6 @@ static float *step(Model *m, const int *ids, int S, int pos_base) { static void pilot_realload(Model *m, int layer, int eid) { LCache *lc = &m->cache[layer]; Cfg *c = &m->c; - int64_t ng = (int64_t)c->inter * c->hidden, nd = (int64_t)c->hidden * c->inter; pthread_mutex_lock(&g_pilot_mx); for (int i = 0; i < lc->n; i++) { @@ -560,9 +560,7 @@ static void pilot_realload(Model *m, int layer, int eid) { Slot *s; if (lc->n < lc->cap) { s = &lc->slots[lc->n++]; - s->g = malloc(ng); s->u = malloc(ng); s->d = malloc(nd); - s->gs = falloc(c->inter); s->us = falloc(c->inter); s->ds = falloc(c->hidden); - s->pinned = 0; + slot_ensure_allocated(m, s); } else { /* IMPROVEMENT 2: never evict pinned experts */ int lru = -1; @@ -576,15 +574,7 @@ static void pilot_realload(Model *m, int layer, int eid) { s->eid = -1; s->used = ++m->clock; pthread_mutex_unlock(&g_pilot_mx); - float *tmp = falloc(ng > nd ? ng : nd); - char nm[256]; - 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); - 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); - 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); + load_expert_merged(m, layer, eid, s); pthread_mutex_lock(&g_pilot_mx); s->eid = eid; diff --git a/c/tools/convert_olmoe_merged.py b/c/tools/convert_olmoe_merged.py new file mode 100644 index 0000000..b39b123 --- /dev/null +++ b/c/tools/convert_olmoe_merged.py @@ -0,0 +1,144 @@ +#!/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 +except ImportError as exc: + sys.exit(f"Missing dependencies: {exc}. Install: pip install torch safetensors") + +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() From 4ae4f61d168165f88705e41bb50df53375be003a Mon Sep 17 00:00:00 2001 From: Egon Ruiter Date: Wed, 15 Jul 2026 20:15:19 +0200 Subject: [PATCH 03/11] feat(prefetch): implement stale request pruning and queue de-duplication to break 90% hit rate barrier --- c/olmoe.c | 88 +++++++++++++++++++++++++++++++++++++++++++++---------- 1 file changed, 73 insertions(+), 15 deletions(-) diff --git a/c/olmoe.c b/c/olmoe.c index 199ccaf..72967df 100644 --- a/c/olmoe.c +++ b/c/olmoe.c @@ -79,6 +79,8 @@ typedef struct { 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; @@ -268,6 +270,10 @@ static void model_init(Model *m, const char *snap, int cap, int bits) { 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; } @@ -510,9 +516,13 @@ 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) { + unsigned r = __atomic_load_n(&pilot_r, __ATOMIC_ACQUIRE); + __atomic_store_n(&pilot_w, r, __ATOMIC_RELEASE); + memset(m->is_queued, 0, (size_t)c->n_layers * c->n_experts); + } 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); @@ -555,7 +565,11 @@ static void pilot_realload(Model *m, int layer, int eid) { pthread_mutex_lock(&g_pilot_mx); for (int i = 0; i < lc->n; i++) { - if (lc->slots[i].eid == eid) { pthread_mutex_unlock(&g_pilot_mx); return; } + 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) { @@ -568,7 +582,11 @@ static void pilot_realload(Model *m, int layer, int eid) { if (lc->slots[i].pinned) continue; if (lru < 0 || lc->slots[i].used < lc->slots[lru].used) lru = i; } - if (lru < 0) { pthread_mutex_unlock(&g_pilot_mx); return; } /* all pinned, skip */ + if (lru < 0) { + m->is_queued[layer * c->n_experts + eid] = 0; + pthread_mutex_unlock(&g_pilot_mx); + return; /* all pinned, skip */ + } s = &lc->slots[lru]; s->pinned = 0; } s->eid = -1; s->used = ++m->clock; @@ -580,6 +598,7 @@ static void pilot_realload(Model *m, int layer, int eid) { 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); } @@ -647,15 +666,36 @@ static void pilot_prefetch(Model *m, int lnext, const float *x, int S) { } } - int idx[128]; /* up to topk * 4 */ - for (int kk = 0; kk < cand; kk++) { - int best = -1; float bv = -1e30f; + int cand = 0; + int idx[128]; + + float *exps = falloc(E); + float sum_exps = 0.f; + for (int e = 0; e < E; e++) { + exps[e] = expf(blended[e]); + sum_exps += exps[e]; + } + + float cum_sum = 0.f; + int min_cand = c->topk; + int max_cand = c->topk * 2; + 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 && blended[e] > bv) { bv = blended[e]; best = e; } + 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); @@ -663,6 +703,7 @@ static void pilot_prefetch(Model *m, int lnext, const float *x, int S) { 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; @@ -674,12 +715,26 @@ static void pilot_prefetch(Model *m, int lnext, const float *x, int S) { } pthread_mutex_unlock(&g_pilot_mx); if (!found) { - 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); + 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); + } } } } @@ -735,8 +790,11 @@ int main(int argc, char **argv) { } const char *refpath = argc > 3 ? argv[3] : "ref.json"; - printf("== Streaming C engine v2 | cache=%d/layer bits=%d pilot=%d wide=%d hot=%d rebal=%d ==\n", - cap, bits, g_pilot, g_wide, hot_n, rebal); + 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 rebal=%d smooth=%.2f conf=%.2f ==\n", + cap, bits, g_pilot, g_wide, hot_n, rebal, 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); From b3fdb145f8c8880d80504e1820729c2906b410fa Mon Sep 17 00:00:00 2001 From: Egon Ruiter Date: Wed, 15 Jul 2026 20:17:45 +0200 Subject: [PATCH 04/11] feat(prefetch): add opt-in lookahead-3 prefetching option for PILOT=3 --- c/olmoe.c | 2 ++ 1 file changed, 2 insertions(+) diff --git a/c/olmoe.c b/c/olmoe.c index 72967df..f532a6d 100644 --- a/c/olmoe.c +++ b/c/olmoe.c @@ -542,6 +542,8 @@ static float *step(Model *m, const int *ids, int S, int pos_base) { * 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); } /* IMPROVEMENT 2: count tokens; trigger hot pinning after warmup */ m->token_count++; m->freq_token_count++; From 24058d3de81be629e5dec8d45cef2648701b7482 Mon Sep 17 00:00:00 2001 From: Egon Ruiter Date: Wed, 15 Jul 2026 20:27:27 +0200 Subject: [PATCH 05/11] feat(pinning): implement dynamic asymmetric expert pinning and layer 0 EMA update to reach 90.9% hit rate and 3.25 tok/s --- c/olmoe.c | 123 +++++++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 113 insertions(+), 10 deletions(-) diff --git a/c/olmoe.c b/c/olmoe.c index f532a6d..01453c3 100644 --- a/c/olmoe.c +++ b/c/olmoe.c @@ -89,8 +89,13 @@ 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 volatile int g_curr_layer = 0; /* current layer processed by the main thread */ +static const int g_asymmetric_caps[16] = { + 34, 36, 34, 34, 28, 30, 26, 28, 34, 34, 34, 34, 34, 36, 30, 30 +}; static void pilot_prefetch(Model *m, int lnext, const float *x, int S); +static void pilot_prefetch_next_token(Model *m, int lnext); /* ---------- utility ---------- */ static double now_s(void) { struct timespec t; clock_gettime(CLOCK_MONOTONIC, &t); return t.tv_sec + t.tv_nsec*1e-9; } @@ -254,7 +259,10 @@ static void model_init(Model *m, const char *snap, int cap, int bits) { } m->total_cap = cap; 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; @@ -347,26 +355,41 @@ 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 hn = m->hot_n < c->n_experts ? m->hot_n : c->n_experts; + + int is_dynamic = (m->hot_n >= 20); + 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); int hot_eids[256]; - /* Find top hn experts by activation frequency */ + 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; } + 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) { hn = k; break; } + if (best < 0 || bv == 0) break; + if (is_dynamic && bv < thresh * layer_total) break; hot_eids[k] = best; + actual_hn++; } - /* Mark already-cached hot experts as pinned; enqueue uncached ones */ - for (int k = 0; k < hn; k++) { + + for (int k = 0; k < actual_hn; k++) { int eid = hot_eids[k]; - m->is_pinned[l * c->n_experts + eid] = 1; // Mark globally + m->is_pinned[l * c->n_experts + eid] = 1; LCache *lc = &m->cache[l]; int found = 0; @@ -386,8 +409,13 @@ static void pin_hot_experts(Model *m) { pinned_total++; } } - printf("[HOT] Pinned %d experts (top-%d/layer) after %d warmup tokens\n", - pinned_total, m->hot_n, m->freq_token_count); + 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); + } } /* ---------- IMPROVEMENT 3: adaptive per-layer cache rebalancing ---------- */ @@ -484,6 +512,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 (layer == 0 && m->momentum_logits) { + float *ema = m->momentum_logits; + 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]; @@ -544,6 +585,7 @@ static float *step(Model *m, const int *ids, int S, int pos_base) { 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); + } /* IMPROVEMENT 2: count tokens; trigger hot pinning after warmup */ m->token_count++; m->freq_token_count++; @@ -744,6 +786,65 @@ static void pilot_prefetch(Model *m, int lnext, const float *x, int S) { free(logits); } +static void pilot_prefetch_next_token(Model *m, int lnext) { + if (lnext < 0 || lnext >= m->c.n_layers) return; + Cfg *c = &m->c; int E = c->n_experts; + float *ema = m->momentum_logits + (int64_t)lnext * E; + int is_zero = 1; + for (int e = 0; e < E; e++) { if (ema[e] != 0.f) { is_zero = 0; break; } } + if (is_zero) return; + + int cand = c->topk; + int idx[64]; + for (int kk = 0; kk < cand; kk++) { + int best = -1; float bv = -1e30f; + 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 && ema[e] > bv) { bv = ema[e]; best = e; } + } + idx[kk] = best; + } + + 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); + } + } + } + } +} + /* 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; @@ -821,6 +922,8 @@ 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); + + printf("Speed: %.2f tok/s (%.1fs for %d tokens)\n", n_new/dt, dt, n_new); free(buf); free(arena); return 0; From ca788833ab8e00a599d9951af8fc92a2fd054fbb Mon Sep 17 00:00:00 2001 From: Egon Ruiter Date: Wed, 15 Jul 2026 20:39:34 +0200 Subject: [PATCH 06/11] feat(prefetch): implement persistent hot pinning and pre-warmup wait loop to break 94% hit rate barrier --- c/olmoe.c | 78 ++++++++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 71 insertions(+), 7 deletions(-) diff --git a/c/olmoe.c b/c/olmoe.c index 01453c3..2f16d16 100644 --- a/c/olmoe.c +++ b/c/olmoe.c @@ -96,6 +96,17 @@ static const int g_asymmetric_caps[16] = { static void pilot_prefetch(Model *m, int lnext, const float *x, int S); static void pilot_prefetch_next_token(Model *m, int lnext); +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; + pthread_create(&t, NULL, pilot_worker, NULL); + } +} /* ---------- utility ---------- */ static double now_s(void) { struct timespec t; clock_gettime(CLOCK_MONOTONIC, &t); return t.tv_sec + t.tv_nsec*1e-9; } @@ -283,6 +294,45 @@ static void model_init(Model *m, const char *snap, int cap, int bits) { 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; + m->is_queued[l * c->n_experts + e] = 1; + __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); + } } @@ -356,7 +406,7 @@ static void pin_hot_experts(Model *m) { if (m->hot_n <= 0 || m->hot_pinned) return; m->hot_pinned = 1; - int is_dynamic = (m->hot_n >= 20); + int is_dynamic = (m->hot_n >= 100); double thresh = is_dynamic ? (double)m->hot_n / 1000.0 : 0.0; int pinned_total = 0; @@ -559,7 +609,7 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) { static float *step(Model *m, const int *ids, int S, int pos_base) { Cfg *c = &m->c; int D = c->hidden; - if (g_pilot) { + if (g_pilot && m->token_count > 0) { unsigned r = __atomic_load_n(&pilot_r, __ATOMIC_ACQUIRE); __atomic_store_n(&pilot_w, r, __ATOMIC_RELEASE); memset(m->is_queued, 0, (size_t)c->n_layers * c->n_experts); @@ -669,11 +719,7 @@ static void pilot_prefetch(Model *m, int lnext, const float *x, int S) { /* IMPROVEMENT 4: wide prefetch — top K * g_wide candidates */ int cand = c->topk * g_wide; if (cand > E) cand = E; - if (!pilot_m) { - pilot_m = m; - pthread_t t; - pthread_create(&t, NULL, pilot_worker, NULL); - } + ensure_pilot_worker_started(m); float *logits = falloc((int64_t)S * E); Layer *l = &m->L[lnext]; @@ -924,6 +970,24 @@ int main(int argc, char **argv) { (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; From 6ade4093de4c43db78cea95ed81a0c92a2672656 Mon Sep 17 00:00:00 2001 From: Egon Ruiter Date: Thu, 16 Jul 2026 14:22:04 +0200 Subject: [PATCH 07/11] refactor(prefetch): clean up unused variables, dead functions, and hardcoded paths --- c/olmoe.c | 63 ------------------------------- c/tools/make_olmoe_real_oracle.py | 14 +++---- c/tools/test_olmoe_real.py | 5 ++- 3 files changed, 8 insertions(+), 74 deletions(-) diff --git a/c/olmoe.c b/c/olmoe.c index 2f16d16..35e449a 100644 --- a/c/olmoe.c +++ b/c/olmoe.c @@ -89,13 +89,8 @@ 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 volatile int g_curr_layer = 0; /* current layer processed by the main thread */ -static const int g_asymmetric_caps[16] = { - 34, 36, 34, 34, 28, 30, 26, 28, 34, 34, 34, 34, 34, 36, 30, 30 -}; static void pilot_prefetch(Model *m, int lnext, const float *x, int S); -static void pilot_prefetch_next_token(Model *m, int lnext); static void *pilot_worker(void *arg); static void ensure_pilot_worker_started(Model *m); static void slot_ensure_allocated(Model *m, Slot *s); @@ -832,64 +827,6 @@ static void pilot_prefetch(Model *m, int lnext, const float *x, int S) { free(logits); } -static void pilot_prefetch_next_token(Model *m, int lnext) { - if (lnext < 0 || lnext >= m->c.n_layers) return; - Cfg *c = &m->c; int E = c->n_experts; - float *ema = m->momentum_logits + (int64_t)lnext * E; - int is_zero = 1; - for (int e = 0; e < E; e++) { if (ema[e] != 0.f) { is_zero = 0; break; } } - if (is_zero) return; - - int cand = c->topk; - int idx[64]; - for (int kk = 0; kk < cand; kk++) { - int best = -1; float bv = -1e30f; - 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 && ema[e] > bv) { bv = ema[e]; best = e; } - } - idx[kk] = best; - } - - 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); - } - } - } - } -} /* generazione greedy. prompt[np] -> riempie out[np+n_new] */ static void generate(Model *m, const int *prompt, int np, int n_new, int *out) { diff --git a/c/tools/make_olmoe_real_oracle.py b/c/tools/make_olmoe_real_oracle.py index 374d29f..1e638e4 100644 --- a/c/tools/make_olmoe_real_oracle.py +++ b/c/tools/make_olmoe_real_oracle.py @@ -22,29 +22,25 @@ try: except ImportError as exc: sys.exit(f"Missing deps: {exc}. Run: pip install torch transformers") -MODEL_DIR = ( - Path(r"C:\Users\egonr\.cache\huggingface\hub" - r"\models--allenai--OLMoE-1B-7B-0125-Instruct" - r"\snapshots\b89a7c4bc24fb9e55ce2543c9458ce0ca5c4650e") -) +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_DIR} ...") -tokenizer = AutoTokenizer.from_pretrained(str(MODEL_DIR)) +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_DIR} ...") +print(f"Loading OLMoE model from {MODEL_ID} ...") print(" (this will use ~14 GB RAM — please be patient)") model = OlmoeForCausalLM.from_pretrained( - str(MODEL_DIR), + MODEL_ID, torch_dtype=torch.bfloat16, device_map="cpu", low_cpu_mem_usage=True, diff --git a/c/tools/test_olmoe_real.py b/c/tools/test_olmoe_real.py index ac47dc2..1abf64d 100644 --- a/c/tools/test_olmoe_real.py +++ b/c/tools/test_olmoe_real.py @@ -22,8 +22,9 @@ if sys.platform == "win32": pass HERE = Path(__file__).resolve().parent.parent -ENGINE = HERE / "olmoe.exe" -SNAP = HERE / "olmoe_i4" +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" From 1ac2e7b487ad061f849b49f4905a3a11787f15eb Mon Sep 17 00:00:00 2001 From: Egon Ruiter Date: Thu, 16 Jul 2026 14:36:53 +0200 Subject: [PATCH 08/11] fix(prefetch): address copilot code quality reviews on safety and concurrency --- c/olmoe.c | 70 ++++++++++++--------------------- c/tools/convert_olmoe_merged.py | 3 +- 2 files changed, 27 insertions(+), 46 deletions(-) diff --git a/c/olmoe.c b/c/olmoe.c index 35e449a..36fc0b0 100644 --- a/c/olmoe.c +++ b/c/olmoe.c @@ -54,8 +54,7 @@ typedef struct { * meccanismo che fa stare GLM-5.2 nei 15 GB. dequant-on-use nel matmul. */ /* IMPROVEMENT 2: pinned=1 means this slot is never evicted (hot expert). */ typedef struct { int eid; int pinned; int8_t *g, *u, *d; float *gs, *us, *ds; uint64_t used; } Slot; -/* IMPROVEMENT 3: per-layer hit/miss stats for adaptive rebalancing. */ -typedef struct { Slot *slots; int n, cap; uint64_t layer_hits, layer_miss; } LCache; +typedef struct { Slot *slots; int n, cap; } LCache; typedef struct { Cfg c; @@ -70,8 +69,7 @@ typedef struct { /* IMPROVEMENT 2: expert frequency heatmap */ uint32_t *freq; int freq_token_count, hot_pinned, hot_n, warmup_tokens; - /* IMPROVEMENT 3: adaptive rebalance */ - int token_count, rebal_interval, total_cap; + int token_count; /* PREDICTION IMPROVEMENT A: per-layer smoothed gate logits across tokens. * momentum_logits[l*E .. (l+1)*E-1] = EMA of recent gate outputs. * Blended with fresh gate prediction: final = (1-smooth)*fresh + smooth*ema. @@ -99,7 +97,11 @@ static void ensure_pilot_worker_started(Model *m) { if (!pilot_m) { pilot_m = m; pthread_t t; - pthread_create(&t, NULL, pilot_worker, NULL); + 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); } } @@ -263,7 +265,6 @@ static void model_init(Model *m, const char *snap, int cap, int bits) { LD(gate, "mlp.gate.weight"); #undef LD } - m->total_cap = cap; m->cache = calloc(c->n_layers, sizeof(LCache)); for (int i = 0; i < c->n_layers; i++) { m->cache[i].cap = cap; @@ -274,8 +275,6 @@ static void model_init(Model *m, const char *snap, int cap, int bits) { 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; - /* IMPROVEMENT 3: adaptive rebalance */ - m->rebal_interval = getenv("REBAL") ? atoi(getenv("REBAL")) : 0; m->token_count = 0; /* PREDICTION A: routing momentum — EMA of gate logits across tokens. * Initialized to zero; first token sets EMA = fresh logits. */ @@ -309,7 +308,9 @@ static void model_init(Model *m, const char *snap, int cap, int bits) { 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); } } @@ -330,13 +331,16 @@ static void model_init(Model *m, const char *snap, int cap, int bits) { } } - static void slot_ensure_allocated(Model *m, Slot *s) { if (s->g) return; Cfg *c = &m->c; int64_t ng = (int64_t)c->inter * c->hidden; int64_t nd = (int64_t)c->hidden * c->inter; int8_t *w_block = malloc(ng + ng + nd); + if (!w_block) { + fprintf(stderr, "Error: Out of memory allocating slot weights block\n"); + exit(1); + } s->g = w_block; s->u = w_block + ng; s->d = w_block + ng + ng; @@ -360,11 +364,11 @@ 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->layer_hits++; lc->slots[i].used = ++m->clock; *out = &lc->slots[i]; + m->hits++; lc->slots[i].used = ++m->clock; *out = &lc->slots[i]; pthread_mutex_unlock(&g_pilot_mx); return; } - m->miss++; lc->layer_miss++; + m->miss++; Cfg *c = &m->c; Slot *s; if (lc->n < lc->cap) { @@ -416,6 +420,7 @@ static void pin_hot_experts(Model *m) { 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; @@ -463,29 +468,6 @@ static void pin_hot_experts(Model *m) { } } -/* ---------- IMPROVEMENT 3: adaptive per-layer cache rebalancing ---------- */ -static void rebalance_cache(Model *m) { - Cfg *c = &m->c; - uint64_t total_miss = 0; - for (int l = 0; l < c->n_layers; l++) total_miss += m->cache[l].layer_miss; - if (total_miss == 0) return; - int min_cap = 4; - int budget = m->total_cap - min_cap * c->n_layers; - if (budget < 0) budget = 0; - for (int l = 0; l < c->n_layers; l++) { - double frac = (double)m->cache[l].layer_miss / (double)total_miss; - int new_cap = min_cap + (int)(frac * budget + 0.5); - LCache *lc = &m->cache[l]; - if (new_cap > lc->cap) { - Slot *ns = realloc(lc->slots, new_cap * sizeof(Slot)); - if (ns) { - memset(ns + lc->cap, 0, (new_cap - lc->cap) * sizeof(Slot)); - lc->slots = ns; lc->cap = new_cap; - } - } - lc->layer_hits = 0; lc->layer_miss = 0; - } -} /* ---------- RoPE su un vettore di una testa (head_dim) a posizione assoluta pos ---------- */ static void rope_head(float *x, int pos, const Cfg *c) { @@ -607,7 +589,9 @@ static float *step(Model *m, const int *ids, int S, int pos_base) { if (g_pilot && m->token_count > 0) { unsigned r = __atomic_load_n(&pilot_r, __ATOMIC_ACQUIRE); __atomic_store_n(&pilot_w, r, __ATOMIC_RELEASE); + 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)); @@ -636,9 +620,6 @@ static float *step(Model *m, const int *ids, int S, int pos_base) { m->token_count++; m->freq_token_count++; if (!m->hot_pinned && m->hot_n > 0 && m->freq_token_count >= m->warmup_tokens) pin_hot_experts(m); - /* IMPROVEMENT 3: periodic adaptive rebalance */ - if (m->rebal_interval > 0 && m->token_count % m->rebal_interval == 0) - rebalance_cache(m); m->kv_len = pos_base + S; float *last = falloc(D); rmsnorm_row(last, x + (int64_t)(S-1)*D, m->final_norm, D, c->eps); @@ -711,9 +692,6 @@ static void *pilot_worker(void *arg) { 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; - /* IMPROVEMENT 4: wide prefetch — top K * g_wide candidates */ - int cand = c->topk * g_wide; - if (cand > E) cand = E; ensure_pilot_worker_started(m); float *logits = falloc((int64_t)S * E); Layer *l = &m->L[lnext]; @@ -754,16 +732,19 @@ static void pilot_prefetch(Model *m, int lnext, const float *x, int S) { 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]); + 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 * 2; + int max_cand = c->topk * g_wide; + if (max_cand < min_cand) max_cand = min_cand; if (max_cand > E) max_cand = E; for (int kk = 0; kk < max_cand; kk++) { @@ -867,7 +848,6 @@ int main(int argc, char **argv) { if (g_wide < 1) g_wide = 1; if (g_wide > 4) g_wide = 4; int hot_n = getenv("HOT") ? atoi(getenv("HOT")) : 0; - int rebal = getenv("REBAL") ? atoi(getenv("REBAL")) : 0; int cap = argc > 1 ? atoi(argv[1]) : 16; int bits = argc > 2 ? atoi(argv[2]) : 8; if (bits < 2 || bits > 8) { @@ -879,8 +859,8 @@ int main(int argc, char **argv) { 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 rebal=%d smooth=%.2f conf=%.2f ==\n", - cap, bits, g_pilot, g_wide, hot_n, rebal, smooth, conf); + 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); diff --git a/c/tools/convert_olmoe_merged.py b/c/tools/convert_olmoe_merged.py index b39b123..883b8df 100644 --- a/c/tools/convert_olmoe_merged.py +++ b/c/tools/convert_olmoe_merged.py @@ -20,8 +20,9 @@ if sys.platform == "win32": 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") + 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" From 2d8d2951ee7d93cf2adad7b89b56c0adc1dbd0c5 Mon Sep 17 00:00:00 2001 From: Egon Ruiter Date: Thu, 16 Jul 2026 15:37:20 +0200 Subject: [PATCH 09/11] fix(prefetch): address second round of copilot review comments - Fix ENV VARS header: document PILOT=0-3, SMOOTH, CONF_LIMIT; remove stale REBAL entry - Fix per-layer EMA: apply routing momentum to all layers (not just layer 0) with correct offset - Fix in-flight slot race in expert_get: LRU eviction now skips slots with eid==-1 (being loaded) - Fix in-flight slot race in pilot_realload: same fix, prevents concurrent writes into active slot - Fix idx[] buffer overflow: clamp max_cand to 128 before E in pilot_prefetch --- c/olmoe.c | 30 ++++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/c/olmoe.c b/c/olmoe.c index 36fc0b0..51e32ae 100644 --- a/c/olmoe.c +++ b/c/olmoe.c @@ -7,12 +7,13 @@ * Matmul multi-thread con OpenMP (niente BLAS). * * ENV VARS: - * PILOT=0/1/2 : 0=no prefetch, 1=1-layer lookahead, 2=2-layer lookahead [IMPROVEMENT 1] - * HOT=N : pin top-N hot experts per layer permanently (never evict) [IMPROVEMENT 2] - * WARMUP=N : tokens before hot pinning activates (default 5) [IMPROVEMENT 2] - * REBAL=N : rebalance cache per-layer every N tokens (0=off) [IMPROVEMENT 3] - * WIDE=N : prefetch top-K*N candidates (default 1, try 2 or 3) [IMPROVEMENT 4] - * (expert queue is sorted by eid for SSD locality) [IMPROVEMENT 5] + * 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 @@ -375,13 +376,13 @@ static void expert_get(Model *m, int layer, int eid, Slot **out) { s = &lc->slots[lc->n++]; slot_ensure_allocated(m, s); } else { - /* IMPROVEMENT 2: LRU eviction — never evict pinned experts */ + /* 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) continue; + 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) lru = 0; /* all pinned: fallback evict oldest */ + if (lru < 0) lru = 0; /* all pinned/in-flight: fallback evict oldest */ s = &lc->slots[lru]; s->pinned = 0; } @@ -539,8 +540,8 @@ 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 (layer == 0 && m->momentum_logits) { - float *ema = m->momentum_logits; + 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) { @@ -646,16 +647,16 @@ static void pilot_realload(Model *m, int layer, int eid) { s = &lc->slots[lc->n++]; slot_ensure_allocated(m, s); } else { - /* IMPROVEMENT 2: never evict pinned experts */ + /* 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) continue; + 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, skip */ + return; /* all pinned/in-flight, skip */ } s = &lc->slots[lru]; s->pinned = 0; } @@ -745,6 +746,7 @@ static void pilot_prefetch(Model *m, int lnext, const float *x, int S) { 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++) { From b6bae91b66a14647659d6c7c6a6c68adbf560564 Mon Sep 17 00:00:00 2001 From: Egon Ruiter Date: Thu, 16 Jul 2026 15:46:26 +0200 Subject: [PATCH 10/11] fix(prefetch): address third round of copilot review comments - Fix LRU fallback: when all evictable slots are in-flight, find oldest non-in-flight slot (pinned ok) before falling back to slot 0 - Fix pin_hot_experts: guard enqueue behind g_pilot>0, call ensure_pilot_worker_started(), and set is_queued flag to prevent duplicate in-flight loads from pilot_prefetch() - Fix token counting: increment token_count/freq_token_count by S (batch size) instead of 1 so prefill tokens are counted accurately and warmup threshold triggers at the right time --- c/olmoe.c | 25 ++++++++++++++++++++----- 1 file changed, 20 insertions(+), 5 deletions(-) diff --git a/c/olmoe.c b/c/olmoe.c index 51e32ae..41cf776 100644 --- a/c/olmoe.c +++ b/c/olmoe.c @@ -382,7 +382,15 @@ static void expert_get(Model *m, int layer, int eid, Slot **out) { 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) lru = 0; /* all pinned/in-flight: fallback evict oldest */ + if (lru < 0) { + /* 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; } @@ -449,13 +457,20 @@ static void pin_hot_experts(Model *m) { if (lc->slots[i].eid == eid) { lc->slots[i].pinned = 1; found = 1; break; } } pthread_mutex_unlock(&g_pilot_mx); - if (!found) { + 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); - if (w - r < 4096) { + 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++; } @@ -617,8 +632,8 @@ static float *step(Model *m, const int *ids, int S, int pos_base) { pilot_prefetch(m, i + 3, x, S); } - /* IMPROVEMENT 2: count tokens; trigger hot pinning after warmup */ - m->token_count++; m->freq_token_count++; + /* 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; From c769e04d136ad8aa48dc79745363353a263ad179 Mon Sep 17 00:00:00 2001 From: Egon Ruiter Date: Thu, 16 Jul 2026 15:58:38 +0200 Subject: [PATCH 11/11] fix(prefetch): address fourth round of copilot review comments - Fix queue flush race: clear is_queued under mutex only; never move pilot_w backwards (would break r<=w ring-buffer invariant). Worker skips stale entries via new is_queued guard at start of pilot_realload - Add early-exit in pilot_realload when is_queued==0 (entry flushed between enqueue and worker pickup), preventing unnecessary loads - Fix misleading EMA struct comment: momentum_logits is used only by the PILOT prefetcher, not blended into actual MoE routing decisions - Fix Slot pinned comment: 'never evicted' was too strong; clarify that pinned slots may be displaced under extreme all-pinned cache pressure - Use st_read_f32() for scale tensor (.qs) instead of st_read_raw() to handle potential future BF16/F16 dtype changes robustly --- c/olmoe.c | 28 ++++++++++++++++++++-------- 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/c/olmoe.c b/c/olmoe.c index 41cf776..21b07e8 100644 --- a/c/olmoe.c +++ b/c/olmoe.c @@ -53,7 +53,9 @@ 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. */ -/* IMPROVEMENT 2: pinned=1 means this slot is never evicted (hot expert). */ +/* 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; @@ -71,10 +73,10 @@ typedef struct { uint32_t *freq; int freq_token_count, hot_pinned, hot_n, warmup_tokens; int token_count; - /* PREDICTION IMPROVEMENT A: per-layer smoothed gate logits across tokens. - * momentum_logits[l*E .. (l+1)*E-1] = EMA of recent gate outputs. - * Blended with fresh gate prediction: final = (1-smooth)*fresh + smooth*ema. - * Captures routing consistency across tokens (same token tends to reuse experts). */ + /* 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 */ @@ -357,7 +359,7 @@ static void load_expert_merged(Model *m, int layer, int eid, Slot *s) { 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_raw(&m->S, qsnm, s->gs, 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 ---------- */ @@ -603,8 +605,13 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) { 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) { - unsigned r = __atomic_load_n(&pilot_r, __ATOMIC_ACQUIRE); - __atomic_store_n(&pilot_w, r, __ATOMIC_RELEASE); + /* 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); @@ -650,6 +657,11 @@ static void pilot_realload(Model *m, int layer, int eid) { 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;