From fbcef6628e1c6e6e21f253df5cc8e8d1485c7329 Mon Sep 17 00:00:00 2001 From: BE Date: Thu, 23 Jul 2026 11:33:35 +0530 Subject: [PATCH] olmoe: interactive CHAT=1 mode with a growing multi-turn KV cache MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit olmoe.c was benchmark-harness-only (ref.json prompt/full ids in, token ids out) — no tokenizer wiring, no way to actually converse with it. Adds: - CHAT=1 env var on main(), bypassing the ref.json harness entirely - tok.h/sample.h wiring (tokenizer load, temperature/top-p sampling, the tokenizer's own special-token set arms the stop condition automatically) - run_chat(): allocates the KV cache once for CTX tokens (default 4096, capped there since attention()'s per-head score buffer is fixed-size) and never reallocates it -- each turn after the first reuses the previous turns' cached keys/values instead of re-processing them, so this is real multi-turn context, not a fresh generate() call per message - OLMoE-Instruct's actual chat template (<|user|>/<|assistant|>, with bos_token/eos_token both being the tokenizer's "|||IP_ADDRESS|||" token -- a genuine artifact of its tokenizer, not a bug here) - /reset to clear context without restarting the process - chat_olmoe.sh: a memory-safe launcher (systemd-run --scope with MemoryMax + MemorySwapMax=0) so a misbehaving run degrades to a clean self-kill instead of taking the whole machine down Validated with a real multi-turn conversation (two follow-up coding questions) -- second turn correctly used context from the first without re-explaining itself, confirming the KV-cache carry-over works. --- c/Makefile | 2 +- c/chat_olmoe.sh | 33 ++++++++++++ c/olmoe.c | 133 ++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 167 insertions(+), 1 deletion(-) create mode 100755 c/chat_olmoe.sh diff --git a/c/Makefile b/c/Makefile index c7175681..ae4a8e38 100644 --- a/c/Makefile +++ b/c/Makefile @@ -341,7 +341,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) ./backend_cuda_bench$(EXE) -olmoe$(EXE): olmoe.c st.h json.h compat.h +olmoe$(EXE): olmoe.c st.h json.h compat.h sample.h tok.h tok_unicode.h tok_unicode_o200k.h $(CC) $(CFLAGS) olmoe.c -o olmoe$(EXE) $(LDFLAGS) inkling$(EXE): inkling.c st.h json.h compat.h $(INK_CUDA_OBJ) diff --git a/c/chat_olmoe.sh b/c/chat_olmoe.sh new file mode 100755 index 00000000..bcbafdb4 --- /dev/null +++ b/c/chat_olmoe.sh @@ -0,0 +1,33 @@ +#!/bin/bash +# Launch OLMoE-1B-7B chat, memory-safe: hard RAM cap + swap disabled for this +# process (via a systemd-run scope), so a misbehaving run degrades to a clean +# OOM-kill of itself instead of taking the whole machine down. +# +# Usage: +# SNAP_DIR=/path/to/olmoe_merged ./chat_olmoe.sh +# Override any default via env var, e.g.: MAX_NEW=600 ./chat_olmoe.sh +set -e +cd "$(dirname "$0")" + +if [ -z "$SNAP_DIR" ]; then + echo "set SNAP_DIR= (see tools/convert_olmoe_merged.py)" >&2 + exit 1 +fi + +MEMORY_MAX="${MEMORY_MAX:-10G}" +MAX_NEW="${MAX_NEW:-300}" +PILOT="${PILOT:-0}" # measured faster than PILOT=1 once CACHE=64 gives a ~94% hit rate -- + # the prefetch thread has little disk-wait left to hide and just + # competes with compute for CPU on an 8-core CPU-only box +TEMP="${TEMP:-0.7}" +NUCLEUS="${NUCLEUS:-0.95}" +CTX="${CTX:-4096}" +# 64 = OLMoE's total experts/layer, so the cache holds ALL of them -- no eviction, +# no thrashing. This engine doesn't dedupe repeated expert picks within a single +# prefill batch, so a too-small cache causes constant reload storms; caching every +# expert sidesteps that entirely for a model this size (~6GB cache + ~1.8GB dense +# =~ 7.8GB peak, hence MEMORY_MAX=10G above for margin). +CACHE="${CACHE:-64}" + +exec systemd-run --user --scope -p MemoryMax="$MEMORY_MAX" -p MemorySwapMax=0 --collect \ + -- bash -c "SNAP='$SNAP_DIR' CHAT=1 MAX_NEW=$MAX_NEW PILOT=$PILOT TEMP=$TEMP NUCLEUS=$NUCLEUS CTX=$CTX ./olmoe $CACHE 8" diff --git a/c/olmoe.c b/c/olmoe.c index d9460d40..2bf618e9 100644 --- a/c/olmoe.c +++ b/c/olmoe.c @@ -42,6 +42,9 @@ typedef struct { int hidden, n_layers, n_heads, n_kv_heads, head_dim; int n_experts, topk, inter, vocab; float theta, eps; int norm_topk; + int stop_ids[8], n_stop; /* unused (no model-specific hardcoded stops) — chat mode's + * stop set comes entirely from the tokenizer's own special- + * token flags via sample.h's stops_arm_tok */ } Cfg; /* ---------- pesi densi per-layer ---------- */ @@ -117,6 +120,14 @@ static double rss_gb(void) { struct rusage r; getrusage(RUSAGE_SELF, &r); return #endif static float *falloc(int64_t n) { float *p = malloc(n*sizeof(float)); if(!p){fprintf(stderr,"OOM %ld\n",(long)n);exit(1);} return p; } +/* chat mode only (main()'s CHAT=1 path): sampling temperature/top-p and the + * tokenizer + stop-set machinery. g_temp<=0 -> greedy (sample.h's pick_tok). + * Declared before #include "sample.h" — it references these by name, and + * Cfg/falloc above, without its own extern declarations. */ +static float g_temp = 0.7f; /* TEMP env overrides */ +static float g_nuc = 0.95f; /* NUCLEUS env overrides */ +#include "sample.h" + /* y[S,O] = x[S,I] @ W^T, W e' [O,I] row-major */ static void matmul(float *y, const float *x, const float *W, int S, int I, int O) { #pragma omp parallel for schedule(static) @@ -926,6 +937,110 @@ static int tf_nll(Model *m, const int *full, int nfull, int np, double *nll_out) return scored; } +/* ---------- interactive chat mode (CHAT=1) ---------- */ +/* OLMoE-Instruct's real template (tokenizer_config.json's chat_template): + * {{bos_token}}<|user|>\n{msg}\n<|assistant|>\n{reply}{eos_token}\n<|user|>\n... + * bos_token == eos_token == "|||IP_ADDRESS|||" (a genuine OLMoE tokenizer quirk, + * a PII-scrubbing artifact repurposed as the BOS/EOS marker — not a bug here). + * It's an added special token, so tok_encode() tokenizes the literal string as + * one atomic id, same as any other added token. <|user|>/<|assistant|> are NOT + * special tokens in this tokenizer — just plain text the model was trained to + * treat as turn markers. + * Turn 1 gets bos_token with no separator before "<|user|>"; every later turn + * gets eos_token+"\n" first (closing the previous assistant turn we never + * explicitly appended to history, matching the is_stop-skips-append design + * below) before its own "<|user|>...". */ +static int fmt_user_turn(char *out, int cap, const char *msg, int first_turn) { + int n = first_turn + ? snprintf(out, cap, "|||IP_ADDRESS|||<|user|>\n%s\n<|assistant|>\n", msg) + : snprintf(out, cap, "|||IP_ADDRESS|||\n<|user|>\n%s\n<|assistant|>\n", msg); + return (n < 0 || n >= cap) ? -1 : n; +} + +/* KV cache allocated ONCE for ctx_cap and never reallocated — hist_len only + * grows (or resets to 0 on /reset), so every turn after the first reuses the + * previous turns' cached keys/values: real multi-turn context, not a fresh + * generate() call per message. ctx_cap capped at 4096: attention()'s per-head + * score buffer (sc[4096]) is fixed-size, any position >= 4096 would overflow it. */ +static void run_chat(Model *m, Tok *T, int ctx_cap) { + Cfg *c = &m->c; + m->max_t = ctx_cap; + m->K = calloc(c->n_layers, sizeof(float*)); m->V = calloc(c->n_layers, sizeof(float*)); + for (int i = 0; i < c->n_layers; i++) { + m->K[i] = falloc((int64_t)c->n_heads * m->max_t * c->head_dim); + m->V[i] = falloc((int64_t)c->n_heads * m->max_t * c->head_dim); + } + + int tok_eos = tok_id_of(T, "|||IP_ADDRESS|||"); + stops_arm_tok(c, tok_eos, T); + + int max_new = getenv("MAX_NEW") ? atoi(getenv("MAX_NEW")) : 512; + if (max_new < 1) max_new = 1; + + int *hist = malloc((size_t)ctx_cap * sizeof(int)); + int hist_len = 0; + int first_turn = 1; + + char *line = malloc(8192); + char *turn = malloc(8192 + 64); + int *newids = malloc(8192 * sizeof(int)); + int *gen = malloc((size_t)max_new * sizeof(int)); + char *outbuf = malloc(65536); + + fprintf(stderr, "olmoe chat — SNAP=%s, ctx=%d, TEMP=%.2f, NUCLEUS=%.2f\n" + " type a message and press enter; /reset clears context; Ctrl-D exits\n", + getenv("SNAP"), ctx_cap, g_temp, g_nuc); + + for (;;) { + printf("\n> "); fflush(stdout); + if (!fgets(line, 8192, stdin)) break; + size_t L = strlen(line); + while (L > 0 && (line[L-1] == '\n' || line[L-1] == '\r')) line[--L] = 0; + if (L == 0) continue; + if (!strcmp(line, "/reset")) { hist_len = 0; first_turn = 1; fprintf(stderr, "[chat] context reset\n"); continue; } + + int tn = fmt_user_turn(turn, 8192 + 64, line, first_turn); + if (tn < 0) { fprintf(stderr, "[chat] message too long, skipped\n"); continue; } + first_turn = 0; + int nn = tok_encode(T, turn, tn, newids, 8192); + + if (hist_len + nn + max_new > ctx_cap) { + fprintf(stderr, "[chat] context window full (%d/%d tokens) — /reset to start over\n", + hist_len + nn, ctx_cap); + continue; + } + + float *logit = step(m, newids, nn, hist_len); + hist_len += nn; + + /* Every token counted in hist_len must have gone through step() exactly + * once (that's what writes its KV-cache slot) — otherwise a later turn's + * attention would read an uninitialized slot. So even on the turn's LAST + * token we still call step() once to populate its KV before breaking; + * only its returned logit (which nothing will consume) is discarded. */ + int ngen = 0; + for (int s = 0; s < max_new; s++) { + int nt = pick_tok(logit, c->vocab, -1); + free(logit); logit = NULL; + if (is_stop(nt)) break; + hist[hist_len] = nt; gen[ngen++] = nt; hist_len++; + int room_left = (s < max_new - 1) && (hist_len < ctx_cap); + logit = step(m, &nt, 1, hist_len - 1); + if (!room_left) { + if (hist_len >= ctx_cap) fprintf(stderr, "\n[chat] context window full mid-reply — /reset to start over\n"); + free(logit); logit = NULL; + break; + } + } + + int outn = tok_decode(T, gen, ngen, outbuf, 65535); + outbuf[outn] = 0; + printf("%s\n", outbuf); + fflush(stdout); + } + free(line); free(turn); free(newids); free(gen); free(outbuf); free(hist); +} + /* ---------- lettura ref.json ---------- */ static int *read_int_array(jval *o, const char *key, int *n_out) { jval *a = json_get(o, key); @@ -948,6 +1063,24 @@ int main(int argc, char **argv) { fprintf(stderr, "quant_bits must be 2..8 (got %d)\n", bits); return 1; } + + if (getenv("CHAT")) { /* interactive mode: bypasses the ref.json harness entirely */ + g_temp = getenv("TEMP") ? (float)atof(getenv("TEMP")) : g_temp; + g_nuc = getenv("NUCLEUS") ? (float)atof(getenv("NUCLEUS")) : g_nuc; + int ctx_cap = getenv("CTX") ? atoi(getenv("CTX")) : 4096; + if (ctx_cap < 1 || ctx_cap > 4096) { /* attention()'s sc[4096] score buffer hard-caps this */ + fprintf(stderr, "CTX must be 1..4096 (got %d)\n", ctx_cap); + return 1; + } + 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()); + Tok T; + char tokpath[2048]; snprintf(tokpath, sizeof(tokpath), "%s/tokenizer.json", snap); + tok_load(&T, tokpath); + run_chat(&m, &T, ctx_cap); + return 0; + } + const char *refpath = argc > 3 ? argv[3] : "ref.json"; float smooth = getenv("SMOOTH") ? (float)atof(getenv("SMOOTH")) : 0.3f;