From ba5c44a35f68704f0b12ab12a169b31be5b1a57f Mon Sep 17 00:00:00 2001 From: Mario Gutierrez Date: Sat, 1 Aug 2026 08:24:14 -0500 Subject: [PATCH 1/2] =?UTF-8?q?olmoe:=20adopt=20route=5Ftrace.h=20?= =?UTF-8?q?=E2=80=94=20the=20last=20engine=20on=20the=20shared=20history?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fourth and final engine for #700, after colibri.c (#716), kimi_k3.c (#719) and inkling.c (#742). Every engine that routes experts now writes the same format. m->freq moves from one flat [n_layers * n_experts] block to route_trace.h's row per layer, so the three read sites -- the hot-pin ranking, the bump in moe() and the LFRU eviction guard -- change shape but not meaning. last_access stays flat; it is a separate array route_trace.h does not own. Every olmoe layer routes and there is no MTP layer, so the only row dropped is the spare one rt_init leaves. The history is opt-in behind COLI_USAGE, exactly as kimi_k3 does it: with the variable unset nothing is loaded and nothing is written, so the default path is unchanged. PPL=1 deliberately does not save, so a loss sweep cannot fold its own tokens into the persisted ranking. Verified against the real OLMoE-1B-7B-0125-Instruct converted with tools/convert_olmoe_merged.py. Output is byte-identical to dev on the same run, down to hit=1216 miss=832 -- the eviction path decides exactly as before, which is the property the layout change could have broken. The COLI_USAGE round trip writes the -1/-2 header with the olmoe identity, leaves no row for the dropped layer, and accumulates 2048 -> 4096 across two runs. make check green. Co-Authored-By: Claude Opus 5 --- c/olmoe.c | 32 +++++++++++++++++++++++--------- 1 file changed, 23 insertions(+), 9 deletions(-) diff --git a/c/olmoe.c b/c/olmoe.c index bb7ac1b81..325a8b5d5 100644 --- a/c/olmoe.c +++ b/c/olmoe.c @@ -36,6 +36,7 @@ #include /* omp_set_num_threads/omp_get_max_threads per omp_tune.h */ #endif #include "omp_tune.h" +#include "route_trace.h" /* shared routing telemetry (#700) */ #ifdef _WIN32 #include @@ -82,7 +83,7 @@ typedef struct { float **K, **V; int kv_len, max_t; double dense_load_s; /* IMPROVEMENT 2: expert frequency heatmap */ - uint32_t *freq; + uint32_t **freq; /* per-layer expert counts, owned by route_trace.h */ int freq_token_count, hot_pinned, hot_n, warmup_tokens; int token_count; /* PREDICTION IMPROVEMENT A: per-layer EMA of gate logits across tokens. @@ -322,7 +323,13 @@ static void model_init(Model *m, const char *snap, int cap, int bits) { 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)); + rt_init("olmoe", c->n_layers, c->n_experts); + rt_drop_row(c->n_layers); /* every olmoe layer routes; no MTP row */ + m->freq = rt_counts_all(); /* alias: the read sites keep their shape */ + { const char *up = getenv("COLI_USAGE"); /* optional history to seed from */ + if (up && *up) { int64_t h = rt_load(up); + if (h > 0) fprintf(stderr, "[USAGE] expert history: %lld selections (%s)\n", + (long long)h, up); } } 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; @@ -491,7 +498,8 @@ static void pin_hot_experts(Model *m) { 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; + uint32_t *freq_l = m->freq[l]; + if (!freq_l) continue; /* a layer with no row cannot be ranked */ uint64_t layer_total = 0; for (int e = 0; e < c->n_experts; e++) layer_total += freq_l[e]; @@ -654,8 +662,8 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out) { 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]]++; + uint32_t *freq_l = m->freq[layer]; + if (freq_l) 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++) { @@ -757,10 +765,11 @@ static void pilot_realload(Model *m, int layer, int eid) { } /* LFRU eviction guard: don't displace a warm resident expert with a speculation */ - if (g_pilot_evict_guard && m->freq && m->last_access && lc->slots[lru].eid >= 0) { + if (g_pilot_evict_guard && m->freq && m->freq[layer] && m->last_access && + lc->slots[lru].eid >= 0) { int vid = lc->slots[lru].eid; - uint64_t vs = lfru_score(m->freq[layer * c->n_experts + vid], m->last_access[layer * c->n_experts + vid], m->clock); - uint64_t cs = lfru_score(m->freq[layer * c->n_experts + eid], m->last_access[layer * c->n_experts + eid], m->clock); + uint64_t vs = lfru_score(m->freq[layer][vid], m->last_access[layer * c->n_experts + vid], m->clock); + uint64_t cs = lfru_score(m->freq[layer][eid], m->last_access[layer * c->n_experts + eid], m->clock); if (cs <= vs + (vs >> 2) + (4u << 8)) { m->is_queued[layer * c->n_experts + eid] = 0; pthread_mutex_unlock(&g_pilot_mx); @@ -1120,6 +1129,8 @@ int main(int argc, char **argv) { char tokpath[2048]; snprintf(tokpath, sizeof(tokpath), "%s/tokenizer.json", snap); tok_load(&T, tokpath); run_chat(&m, &T, ctx_cap); + { const char *up = getenv("COLI_USAGE"); + if (up && *up) rt_save(up, 0); } return 0; } @@ -1151,7 +1162,8 @@ int main(int argc, char **argv) { (unsigned long long)m.hits, (unsigned long long)m.miss); printf("Speed: %.2f tok/s (%.1fs for %d tokens) | PEAK RSS: %.2f GB\n", scored/dt, dt, scored, rss_gb()); free(buf); free(arena); - return 0; + return 0; /* PPL is a measurement run: no rt_save on purpose, so a loss + * sweep cannot fold its own tokens into the persisted ranking */ } int *out = malloc((np + n_new) * sizeof(int)); @@ -1187,6 +1199,8 @@ int main(int argc, char **argv) { } } + { const char *up = getenv("COLI_USAGE"); + if (up && *up) rt_save(up, 0); } /* same bytes as every other engine */ printf("Speed: %.2f tok/s (%.1fs for %d tokens)\n", n_new/dt, dt, n_new); free(buf); free(arena); return 0; From 8e13a0cfe706834925426058bff4f12b1900483d Mon Sep 17 00:00:00 2001 From: Mario Gutierrez Date: Sat, 1 Aug 2026 08:24:16 -0500 Subject: [PATCH 2/2] =?UTF-8?q?olmoe:=20adopt=20route=5Ftrace.h=20?= =?UTF-8?q?=E2=80=94=20the=20last=20engine=20on=20the=20shared=20history?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fourth and final engine for #700, after colibri.c (#716), kimi_k3.c (#719) and inkling.c (#742). Every engine that routes experts now writes the same format. m->freq moves from one flat [n_layers * n_experts] block to route_trace.h's row per layer, so the three read sites -- the hot-pin ranking, the bump in moe() and the LFRU eviction guard -- change shape but not meaning. last_access stays flat; it is a separate array route_trace.h does not own. Every olmoe layer routes and there is no MTP layer, so the only row dropped is the spare one rt_init leaves. The history is opt-in behind COLI_USAGE, exactly as kimi_k3 does it: with the variable unset nothing is loaded and nothing is written, so the default path is unchanged. PPL=1 deliberately does not save, so a loss sweep cannot fold its own tokens into the persisted ranking. Verified against the real OLMoE-1B-7B-0125-Instruct converted with tools/convert_olmoe_merged.py. Output is byte-identical to dev on the same run, down to hit=1216 miss=832 -- the eviction path decides exactly as before, which is the property the layout change could have broken. The COLI_USAGE round trip writes the -1/-2 header with the olmoe identity, leaves no row for the dropped layer, and accumulates 2048 -> 4096 across two runs. make check green. Co-Authored-By: Claude Opus 5 --- c/Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/c/Makefile b/c/Makefile index 1f51f4387..bf11ba243 100644 --- a/c/Makefile +++ b/c/Makefile @@ -458,7 +458,7 @@ cuda-bench: backend_cuda.cu backend_cuda.h backend_gpu_compat.h tests/bench_tens "$(GPUCC)" $(GPUFLAGS) backend_cuda.cu tests/bench_tensor_core.cu -o backend_cuda_bench$(EXE) $(ANS_NVCC_LIBS) ./backend_cuda_bench$(EXE) -olmoe$(EXE): olmoe.c st.h json.h compat.h sample.h tok.h tok_unicode.h tok_unicode_o200k.h omp_tune.h +olmoe$(EXE): olmoe.c st.h json.h compat.h sample.h tok.h tok_unicode.h tok_unicode_o200k.h omp_tune.h route_trace.h $(CC) $(CFLAGS) olmoe.c -o olmoe$(EXE) $(LDFLAGS) inkling$(EXE): inkling.c st.h json.h compat.h $(INK_CUDA_OBJ) $(METAL_OBJ)