Something went wrong. Try again.
strata but optimized for m5 (and maybe more later)
Something went wrong. Try again.
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948// metal/src/strata_window_cli.cpp - Direct zero-bloat Metal runner for Strata's Window + ExpertStream + PleTable.// Supports both:// 1. Prefill (PromptGemm) + decode (Window::run + ExpertStream F_NOCACHE pread + speculative router lookahead)// 2. Teacher-forced full-vocab (248,320) FP64 KLD, PPL, and top-1 agreement evaluation (--logits-out / --logits-ref)
#include "strata/core/conversation_cache.hpp"#include "strata/core/layer.hpp"#include "strata/core/session.hpp"#include "strata/kernels/ngram.hpp"#include "strata/spec/suffix_drafter.hpp"#include "strata/spec/spec_gate.hpp"#include "engine/scr_state.hpp"#include "engine/suffix_window.hpp"#include "strata/core/scr.hpp"
#include <algorithm>#include <chrono>#include <cmath>#include <cstdio>#include <cstdlib>#include <cstring>#include <fcntl.h>#include <filesystem>#include <fstream>#include <sstream>#include <string>#include <sys/stat.h>#include <unistd.h>#include <vector>
namespace strata::metal_engine {void set_model_shard(const char* shard1);void* window_create();void window_destroy(void* w);bool window_init(void* w, const strata::core::ModelGeometry& g, strata::core::SessionState& ss, int max_t, int64_t max_cells, int64_t* n_vocab, std::string& err);void window_mrope_host_set(const int32_t* host_table);bool prompt_run(const int64_t* tokens, int64_t n, int64_t pos0, const float* const* embd_rows, const std::function<bool(const float*, int64_t, int64_t, std::string&)>& on_chunk, std::string& err);bool prompt_run_rows(const int64_t* tokens, int64_t n, int64_t pos0, const std::function<bool(const float*, int64_t, int64_t, std::string&)>& on_chunk, std::string& err);bool window_run(void* w, int T, const int32_t* tokens, int64_t pos0, bool head, int32_t* out, std::string& err);bool window_commit(void* w, int n_keep, std::string& err);bool window_copy_logits(void* w, int t, float* host);void window_reset_stats(void* w);void window_reset();bool window_checkpoint_save(strata::core::ConversationCheckpoint& c, std::string& err);bool window_checkpoint_restore(const strata::core::ConversationCheckpoint& c, std::string& err);bool window_session_save(const char* path, const int32_t* ids, int64_t n_cells, std::string& err, int64_t base_cells = 0);bool window_session_load(const char* path, std::vector<int32_t>& ids, int64_t* n_cells, std::string& err, const char* parent_path = nullptr);bool window_scr_relocate(int64_t c_src, int64_t c_dst, int64_t n, std::string& err);bool mtp_load(const std::string& rt_dir, int max_t, std::string& err);bool mtp_prefill(const int32_t* next_tokens, int64_t n, int64_t cell0, std::string& err);bool mtp_draft(int T, const int32_t* tokens, int64_t p, int a, int max_steps, int32_t* drafts, float* probs, float min_p, int* n_drafts, std::string& err);} // namespace strata::metal_engine
static double now_ms() { return std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now().time_since_epoch()).count();}
// One policy instance per decode: a new prompt does not inherit another workload's timings.struct SpecControl { strata::spec::SpecGate gate; bool enabled = true, full = false; int mtp_used = 0;
SpecControl() { const char* mode = std::getenv("STRATA_SPEC_GATE"); enabled = mode && std::strcmp(mode, "0") != 0; full = mode && std::strcmp(mode, "1") == 0; }
int propose(strata::spec::SuffixDrafter& suffix, int32_t* input, const int32_t* mtp, const float* probs, int n_mtp, int suffix_max, int rem, int64_t pos, int64_t cells) { int nd = strata::spec::suffix_limit(n_mtp, rem, pos, cells); if (full) nd = gate.spec_gate_pick(0, nd, probs); mtp_used = nd; for (int j = 0; j < nd; ++j) input[1 + j] = mtp[j]; int limit = strata::spec::suffix_limit(std::max(suffix_max, nd), rem, pos, cells) - nd; if (suffix_max <= 0 || limit <= 0) return nd; if (enabled) limit = gate.spec_gate_pick(1, limit, nullptr, nd, probs); if (limit <= 0) return nd; const double begin = now_ms(); const int sk = nd > 0 ? suffix.propose_after(input + 1, nd, limit, input + 1 + nd) : suffix.propose(limit, input + 1); if (sk > 0) gate.spec_gate_draft_cost(1, (now_ms() - begin) / (1000 * sk)); return nd + sk; }
void result(int nd, int keep) { const int accepted = keep - 1; gate.spec_gate_result(0, std::min(mtp_used, keep), std::min(mtp_used, accepted)); if (accepted >= mtp_used) gate.spec_gate_result(1, std::min(nd - mtp_used, keep - mtp_used), accepted - mtp_used); }
void describe() const { if (enabled) std::fprintf(stderr, "%s\n", gate.spec_gate_describe().c_str()); }};
static double logsumexp_f64(const float* lg, int64_t n, int32_t* argmax_out) { float mx = lg[0]; int32_t am = 0; for (int64_t i = 1; i < n; ++i) { if (lg[i] > mx) { mx = lg[i]; am = (int32_t) i; } } if (argmax_out) *argmax_out = am; double sum = 0.0; for (int64_t i = 0; i < n; ++i) { sum += std::exp((double) lg[i] - (double) mx); } return (double) mx + std::log(sum);}
static std::vector<int32_t> read_ids(const std::string& path) { std::vector<int32_t> v; std::ifstream f(path); int64_t x; while (f >> x) v.push_back((int32_t) x); return v;}
// teacher-forced next-token distributions, as log-probabilities: inputs[i] runs at pos + i and predicts targets[i]struct TfRun { std::vector<float> lp; // [n][n_vocab] double nll = 0; double ms = 0;};static bool tf_run(void* w, const std::vector<int32_t>& inputs, const std::vector<int32_t>& targets, int64_t pos, int64_t n_vocab, TfRun& out, std::string& err) { namespace me = strata::metal_engine; const size_t n = inputs.size(); out.lp.assign(n * (size_t) n_vocab, 0.0f); out.nll = 0; const double t0 = now_ms(); for (size_t i = 0; i < n; ++i) { int32_t tok = inputs[i], next = 0; if (!me::window_run(w, 1, &tok, pos + (int64_t) i, true, &next, err) || !me::window_commit(w, 1, err)) return false; float* lg = out.lp.data() + i * (size_t) n_vocab; if (!me::window_copy_logits(w, 0, lg)) { err = "copy_logits failed"; return false; } const double lse = logsumexp_f64(lg, n_vocab, nullptr); for (int64_t v = 0; v < n_vocab; ++v) lg[v] = (float) ((double) lg[v] - lse); out.nll -= lg[targets[i]]; } out.ms = now_ms() - t0; return true;}// KL(ref || cand) per position, natsstatic void tf_compare(const TfRun& ref, const TfRun& c, int64_t n_vocab, double& mean, double& p90, double& mx, double& top1) { const size_t n = ref.lp.size() / (size_t) n_vocab; std::vector<double> k(n); int agree = 0; for (size_t i = 0; i < n; ++i) { const float* a = ref.lp.data() + i * (size_t) n_vocab; const float* b = c.lp.data() + i * (size_t) n_vocab; double s = 0; int64_t ia = 0, ib = 0; for (int64_t v = 0; v < n_vocab; ++v) { s += std::exp((double) a[v]) * ((double) a[v] - (double) b[v]); if (a[v] > a[ia]) ia = v; if (b[v] > b[ib]) ib = v; } k[i] = std::max(0.0, s); agree += ia == ib; } mean = 0; for (double x : k) mean += x; mean /= (double) n; std::vector<double> s = k; std::sort(s.begin(), s.end()); p90 = s[(size_t) (0.9 * (double) (n - 1))]; mx = s.back(); top1 = 100.0 * agree / (double) n;}
// --scr-test: the SCR splice [A B C] -> [A B' C] against a fresh prefill of [A B' C], scored teacher-forced on D.// C's last token is D's first input (so every variant's state is "A, the middle span, C minus its last token").// Variants, each from the saved [A B C] session:// orig no edit: the old context (how much B -> B' matters at all)// keep J C's cells relocated, B' + C[:J] prefilled from the state at |A|, then the running state of [A B C]// restored (the CLM rule: the post-edit GDN keeps having read B)// gdnex J as keep, but the running state of the fresh [A B' C] restored: only the KV staleness is left// lossy J as keep, without the restore: the GDN state has read A B' C[:J] only// J is raised to the junction block's minimum; lossy with J = |C| is the fresh prefill again (a sanity row: KL 0).static int scr_test(void* w, strata::core::SessionState& ss, int64_t n_vocab, const std::string& pa, const std::string& pb, const std::string& pbp, const std::string& pc, const std::string& pd, const std::string& js, int32_t pad, int n_eval, const std::string& tmp) { namespace me = strata::metal_engine; std::string err; const std::vector<int32_t> A = read_ids(pa), B = read_ids(pb), D = read_ids(pd); std::vector<int32_t> Bp = read_ids(pbp), C = read_ids(pc); if (A.size() < 4 || B.empty() || Bp.empty() || C.size() < 2 || D.size() < 2) { std::fprintf(stderr, "scr-test: empty span\n"); return 2; } while ((Bp.size() - B.size()) % 4 != 0) Bp.push_back(pad); // the shift must be a multiple of 4 cells const int32_t c_last = C.back(); C.pop_back(); const int64_t nA = (int64_t) A.size(), nB = (int64_t) B.size(), nBp = (int64_t) Bp.size(), nC = (int64_t) C.size(); std::vector<int32_t> inputs{c_last}, targets; for (size_t i = 0; i + 1 < D.size() && (int) inputs.size() < n_eval; ++i) inputs.push_back(D[i]); targets.assign(D.begin(), D.begin() + (ptrdiff_t) inputs.size()); auto prefill = [&](const std::vector<int32_t>& ids, int64_t pos) { std::vector<int64_t> v(ids.begin(), ids.end()); return v.empty() || me::prompt_run_rows(v.data(), (int64_t) v.size(), pos, nullptr, err); }; auto cat = [](std::vector<int32_t> a, const std::vector<int32_t>& b, size_t nb) { a.insert(a.end(), b.begin(), b.begin() + (ptrdiff_t) nb); return a; }; auto fresh = [&] { me::window_reset(); ss.ple_prev[0] = -1; ss.ple_prev[1] = -1; }; const int32_t prev_c[2] = {C.size() >= 2 ? C[C.size() - 2] : Bp.back(), C.back()}; auto fail = [&](const char* what) { std::fprintf(stderr, "scr-test: %s: %s\n", what, err.c_str()); return 1; };
// the reference: [A B' C] prefilled fresh strata::core::ConversationCheckpoint ckA, ckABC, ckRef; int32_t prevA[2]; fresh(); if (!prefill(A, 0) || !me::window_checkpoint_save(ckA, err)) return fail("prefill A"); prevA[0] = ss.ple_prev[0]; prevA[1] = ss.ple_prev[1]; double t0 = now_ms(); if (!prefill(cat(Bp, C, C.size()), nA)) return fail("prefill B' C"); const double ref_ms = now_ms() - t0; if (!me::window_checkpoint_save(ckRef, err)) return fail("checkpoint"); TfRun ref; if (!tf_run(w, inputs, targets, nA + nBp + nC, n_vocab, ref, err)) return fail("ref eval"); // the old context [A B C], saved as a session fresh(); if (!prefill(cat(A, B, B.size()), 0) || !prefill(C, nA + nB) || !me::window_checkpoint_save(ckABC, err)) return fail("prefill A B C"); const std::vector<int32_t> abc = cat(cat(A, B, B.size()), C, C.size()); t0 = now_ms(); if (!me::window_session_save(tmp.c_str(), abc.data(), (int64_t) abc.size(), err)) return fail("session save"); const double save_ms = now_ms() - t0; TfRun orig; if (!tf_run(w, inputs, targets, nA + nB + nC, n_vocab, orig, err)) return fail("orig eval"); const double n = (double) inputs.size(); std::printf("SCR_SPANS A=%lld B=%lld B'=%lld (padded) C=%lld eval=%zu delta=%lld session_save_ms=%.1f ref_prefill_ms(B'+C)=%.1f ref_ppl=%.4f\n", (long long) nA, (long long) nB, (long long) nBp, (long long) nC, inputs.size(), (long long) (nBp - nB), save_ms, ref_ms, std::exp(ref.nll / n)); double km, kp, kx, t1; tf_compare(ref, orig, n_vocab, km, kp, kx, t1); std::printf("SCR_ROW variant=orig J=- mean_kld=%.6f p90_kld=%.6f max_kld=%.6f top1=%.2f ppl=%.4f\n", km, kp, kx, t1, std::exp(orig.nll / n));
std::vector<int64_t> Js; { std::stringstream s(js); std::string x; while (std::getline(s, x, ',')) Js.push_back(std::atoll(x.c_str())); Js.push_back(nC); } for (const char* mode : {"keep", "gdnex", "lossy"}) for (int64_t J : Js) { if (J == nC && std::string(mode) != "lossy") continue; const int64_t Je = std::min(nC, std::max(J, scr::scr_min_junction(nA + nBp))); std::vector<int32_t> ids; int64_t cells = 0; t0 = now_ms(); if (!me::window_session_load(tmp.c_str(), ids, &cells, err)) return fail("session load"); const double load_ms = now_ms() - t0; t0 = now_ms(); if (!me::window_scr_relocate(nA + nB, nA + nBp, nC, err)) return fail("relocate"); const double rel_ms = now_ms() - t0; t0 = now_ms(); if (!me::window_checkpoint_restore(ckA, err)) return fail("restore A"); ss.ple_prev[0] = prevA[0]; ss.ple_prev[1] = prevA[1]; if (!prefill(cat(Bp, C, (size_t) Je), nA)) return fail("prefill B' C[:J]"); if (std::string(mode) == "keep" && !me::window_checkpoint_restore(ckABC, err)) return fail("restore ABC"); if (std::string(mode) == "gdnex" && !me::window_checkpoint_restore(ckRef, err)) return fail("restore ref"); ss.ple_prev[0] = prev_c[0]; ss.ple_prev[1] = prev_c[1]; // token ids are never stale const double splice_ms = now_ms() - t0; TfRun r; if (!tf_run(w, inputs, targets, nA + nBp + nC, n_vocab, r, err)) return fail("eval"); tf_compare(ref, r, n_vocab, km, kp, kx, t1); std::printf("SCR_ROW variant=%s J=%lld mean_kld=%.6f p90_kld=%.6f max_kld=%.6f top1=%.2f ppl=%.4f load_ms=%.1f relocate_ms=%.1f " "splice_prefill_ms=%.1f speedup_vs_ref=%.2f\n", mode, (long long) Je, km, kp, kx, t1, std::exp(r.nll / n), load_ms, rel_ms, splice_ms, ref_ms / std::max(1e-3, rel_ms + splice_ms)); std::fflush(stdout); if (std::string(mode) == "keep" && (J == 0 || J == 8)) { // Connectome loop check (PLAN.md §2.9): restore spliced state and free-generate 128 tokens greedy if (!me::window_session_load(tmp.c_str(), ids, &cells, err)) return fail("loop session load"); if (!me::window_scr_relocate(nA + nB, nA + nBp, nC, err)) return fail("loop relocate"); if (!me::window_checkpoint_restore(ckA, err)) return fail("loop restore A"); ss.ple_prev[0] = prevA[0]; ss.ple_prev[1] = prevA[1]; if (!prefill(cat(Bp, C, (size_t) Je), nA)) return fail("loop prefill"); if (!me::window_checkpoint_restore(ckABC, err)) return fail("loop restore ABC"); ss.ple_prev[0] = prev_c[0]; ss.ple_prev[1] = prev_c[1]; std::vector<int32_t> gen_toks; int32_t cur = inputs[0], nxt = 0; int64_t gpos = nA + nBp + nC; for (int step = 0; step < 128; ++step) { if (!me::window_run(w, 1, &cur, gpos, true, &nxt, err) || !me::window_commit(w, 1, err)) break; gen_toks.push_back(nxt); cur = nxt; ++gpos; } int rep4 = 0, tot4 = std::max(0, (int) gen_toks.size() - 3); for (int i = 0; i < tot4; ++i) { for (int k = 0; k < i; ++k) { if (gen_toks[i] == gen_toks[k] && gen_toks[i+1] == gen_toks[k+1] && gen_toks[i+2] == gen_toks[k+2] && gen_toks[i+3] == gen_toks[k+3]) { ++rep4; break; } } } std::printf("SCR_LOOP_CHECK variant=keep J=%lld gen_tokens=%zu repeated_4grams=%d/%d (%.2f%%)\n", (long long) Je, gen_toks.size(), rep4, tot4, tot4 > 0 ? 100.0 * rep4 / tot4 : 0.0); std::fflush(stdout); } } std::remove(tmp.c_str()); return 0;}
#include <iostream>#include <random>
// sample_logits' distribution: the top-k logits, the repetition / presence penalties over the last 128 tokens of// `recent`, the temperature, min-p and top-p. tok / p: its tokens and unnormalized probabilities (sum: theirs); empty// for greedy decoding.struct Dist { std::vector<int32_t> tok; std::vector<double> p; double sum = 0.0;};static void sample_dist(const float* lg, int64_t n_vocab, float temp, int top_k, float top_p, float min_p, float pen_rep, float pen_pres, const std::vector<int64_t>& recent, Dist& d) { d.tok.clear(); d.p.clear(); d.sum = 0.0; if (temp <= 0.0f || top_k == 1) return; const int k = top_k > 0 ? std::min<int>((int) n_vocab, top_k) : 64; std::vector<std::pair<float, int32_t>> cand; cand.reserve(256); float cutoff = -INFINITY; for (int64_t i = 0; i < n_vocab; ++i) { float v = lg[i]; if (std::isnan(v) || v <= cutoff) continue; cand.push_back({v, (int32_t) i}); if ((int) cand.size() >= k * 4) { std::nth_element(cand.begin(), cand.begin() + k - 1, cand.end(), [](const auto& a, const auto& b) { return a.first > b.first; }); cand.resize((size_t) k); cutoff = cand.back().first; } } if (cand.empty()) return; if (pen_rep != 1.0f || pen_pres != 0.0f) { const size_t r0 = recent.size() > 128 ? recent.size() - 128 : 0; for (auto& kv : cand) { bool seen = false; for (size_t ri = r0; ri < recent.size(); ++ri) if (recent[ri] == kv.second) { seen = true; break; } if (seen) { if (pen_rep > 1.0f) kv.first = kv.first > 0.0f ? kv.first / pen_rep : kv.first * pen_rep; kv.first -= pen_pres; } } } std::partial_sort(cand.begin(), cand.begin() + std::min<size_t>(cand.size(), (size_t) k), cand.end(), [](const auto& a, const auto& b) { return a.first > b.first; }); if ((int) cand.size() > k) cand.resize((size_t) k); const float mx = cand[0].first; double sum = 0.0; std::vector<double> probs(cand.size()); for (size_t i = 0; i < cand.size(); ++i) { probs[i] = std::exp((double) (cand[i].first - mx) / (double) temp); sum += probs[i]; } if (sum <= 0.0) { d.tok = {cand[0].second}; d.p = {1.0}; d.sum = 1.0; return; } if (min_p > 0.0f && cand.size() > 1) { const double p_best = probs[0] / sum; size_t keep = 1; while (keep < cand.size() && (probs[keep] / sum) >= p_best * (double) min_p) ++keep; cand.resize(keep); probs.resize(keep); sum = 0.0; for (double p : probs) sum += p; } if (top_p > 0.0f && top_p < 1.0f) { double cum = 0.0; size_t keep = cand.size(); for (size_t i = 0; i < cand.size(); ++i) { cum += probs[i] / sum; if (cum >= (double) top_p) { keep = i + 1; break; } } cand.resize(keep); probs.resize(keep); sum = 0.0; for (double p : probs) sum += p; } d.tok.resize(cand.size()); for (size_t i = 0; i < cand.size(); ++i) d.tok[i] = cand[i].second; d.p = std::move(probs); d.sum = sum;}
// a draw from d, token `skip` left out (-1: none)static int32_t dist_draw(const Dist& d, std::mt19937& rng, int32_t skip = -1) { double tot = d.sum; for (size_t i = 0; i < d.tok.size(); ++i) if (d.tok[i] == skip) tot -= d.p[i]; if (tot <= 0.0) return skip; std::uniform_real_distribution<double> dist(0.0, tot); double r = dist(rng); int32_t last = d.tok[0]; for (size_t i = 0; i < d.tok.size(); ++i) { if (d.tok[i] == skip) continue; last = d.tok[i]; r -= d.p[i]; if (r <= 0.0) return d.tok[i]; } return last;}
static int32_t sample_logits(const float* lg, int64_t n_vocab, int32_t greedy_tok, float temp, int top_k, float top_p, float min_p, float pen_rep, float pen_pres, const std::vector<int64_t>& recent, std::mt19937& rng) { Dist d; sample_dist(lg, n_vocab, temp, top_k, top_p, min_p, pen_rep, pen_pres, recent, d); return d.tok.empty() ? greedy_tok : dist_draw(d, rng);}
// Speculative sampling (Leviathan et al. 2023, sec. 2.3) of a verified window: input[0] the last token, input[1..nd] the// drafts, output[i] the window's greedy pick after input[0..i]. The drafts are deterministic (q = one-hot), so draft// x after row i is accepted with probability p_i(x) (p_i: sample_dist of row i's logits, history `hist` + input[0..i]);// a rejected row samples p_i without x (the residual max(p - q, 0) normalized), and a fully accepted window samples its// last row. The tokens emitted are distributed as sample_logits one at a time. Returns keep (output[0..keep-1] the// tokens, input[0..keep-1] appended to hist).static int spec_sample(void* w, const int32_t* input, int32_t* output, int nd, int64_t n_vocab, float temp, int top_k, float top_p, float min_p, float pen_rep, float pen_pres, std::vector<int64_t>& hist, std::vector<float>& lg, std::mt19937& rng) { namespace me = strata::metal_engine; Dist d; std::uniform_real_distribution<double> uni(0.0, 1.0); for (int i = 0;; ++i) { hist.push_back(input[i]); me::window_copy_logits(w, i, lg.data()); sample_dist(lg.data(), n_vocab, temp, top_k, top_p, min_p, pen_rep, pen_pres, hist, d); if (d.tok.empty()) { // greedy (top_k 1): the drafts must be the picks if (i == nd || output[i] != input[i + 1] || strata::spec::suffix_stop(output[i])) return i + 1; continue; } if (i == nd) { output[i] = dist_draw(d, rng); return i + 1; } const int32_t x = input[i + 1]; static const bool typical = std::getenv("STRATA_SPEC_SAMPLE") && std::string(std::getenv("STRATA_SPEC_SAMPLE")) == "typical"; double px = 0.0; for (size_t k = 0; k < d.tok.size(); ++k) if (d.tok[k] == x) px = d.p[k]; const bool accept = typical ? (!d.tok.empty() && (d.tok[0] == x || px >= std::max(0.05 * d.sum, 0.25 * d.p[0]))) : (uni(rng) * d.sum < px); if (accept) { output[i] = x; if (strata::spec::suffix_stop(x)) return i + 1; continue; } output[i] = typical ? dist_draw(d, rng) : dist_draw(d, rng, x); return i + 1; }}
struct DiskCache { struct Entry { std::string path; std::string parent_path; std::vector<int64_t> ids; std::vector<int64_t> pads; int64_t base_cells = 0; size_t bytes = 0; bool pinned = false; };
static constexpr size_t kLearnMinTokens = 32; static constexpr size_t kLearnMaxTokens = 65536; static constexpr size_t kMaxPins = 16; static constexpr size_t kSeenMax = 32;
std::string dir; size_t limit_bytes = 0; std::vector<Entry> entries; std::deque<std::vector<int64_t>> seen_prompts;
bool enabled() const { return !dir.empty(); }
static uint64_t fnv1a(const int64_t* data, size_t n) { uint64_t h = 1469598103934665603ull; const auto* p = reinterpret_cast<const uint8_t*>(data); const size_t nb = n * sizeof(int64_t); for (size_t i = 0; i < nb; ++i) h = (h ^ p[i]) * 1099511628211ull; return h; }
static bool is_prefix(const std::vector<int64_t>& a, const std::vector<int64_t>& b) { return a.size() <= b.size() && std::equal(a.begin(), a.end(), b.begin()); }
static size_t common_prefix(const std::vector<int64_t>& a, const std::vector<int64_t>& b, size_t cap) { const size_t lim = std::min({a.size(), b.size(), cap}); size_t i = 0; while (i < lim && a[i] == b[i]) ++i; return i; }
static void read_pads(const std::string& ppath, size_t n_ids, int32_t pad_tok, const std::vector<int64_t>& ids, std::vector<int64_t>& out_pads) { out_pads.clear(); FILE* pf = std::fopen(ppath.c_str(), "rb"); if (!pf) return; uint64_t np = 0; if (std::fread(&np, sizeof(np), 1, pf) == 1 && np <= n_ids) { std::vector<int64_t> p((size_t) np); if (std::fread(p.data(), sizeof(int64_t), (size_t) np, pf) == (size_t) np) { bool ok = true; for (int64_t at : p) { if (at < 0 || (size_t) at >= n_ids || ids[(size_t) at] != pad_tok) ok = false; } if (ok) out_pads = std::move(p); } } std::fclose(pf); }
static void write_pads(const std::string& ppath, const std::vector<int64_t>& pads) { if (pads.empty()) { std::remove(ppath.c_str()); return; } if (FILE* pf = std::fopen(ppath.c_str(), "wb")) { const uint64_t np = (uint64_t) pads.size(); std::fwrite(&np, sizeof(np), 1, pf); std::fwrite(pads.data(), sizeof(int64_t), pads.size(), pf); std::fclose(pf); } }
static std::string read_parent(const std::string& ppath) { std::ifstream in(ppath, std::ios::binary); if (!in) return ""; std::string s; std::getline(in, s); while (!s.empty() && (s.back() == '\r' || s.back() == '\n')) s.pop_back(); return s; }
static void write_parent(const std::string& ppath, const std::string& parent_path) { if (parent_path.empty()) { std::remove(ppath.c_str()); return; } std::ofstream out(ppath, std::ios::binary | std::ios::trunc); if (out) out << parent_path << "\n"; }
static void remove_files(const std::string& p) { namespace fs = std::filesystem; std::error_code ec; fs::remove(p, ec); fs::remove(p + ".pads", ec); fs::remove(p + ".pin", ec); fs::remove(p + ".parent", ec); }
void load_seen() { seen_prompts.clear(); if (dir.empty()) return; const std::string spath = dir + "/prompts-seen.bin"; FILE* f = std::fopen(spath.c_str(), "rb"); if (!f) return; uint32_t magic = 0, count = 0; if (std::fread(&magic, 4, 1, f) == 1 && magic == 0x5345454Eu && std::fread(&count, 4, 1, f) == 1 && count <= kSeenMax) { for (uint32_t i = 0; i < count; ++i) { uint32_t len = 0; if (std::fread(&len, 4, 1, f) != 1 || len > kLearnMaxTokens) break; std::vector<int32_t> i32(len); if (len > 0 && std::fread(i32.data(), 4, len, f) != len) break; seen_prompts.emplace_back(i32.begin(), i32.end()); } } std::fclose(f); }
void save_seen() const { if (dir.empty()) return; const std::string spath = dir + "/prompts-seen.bin"; const std::string tmp = spath + ".tmp"; FILE* f = std::fopen(tmp.c_str(), "wb"); if (!f) return; const uint32_t magic = 0x5345454Eu; const uint32_t count = (uint32_t) seen_prompts.size(); bool ok = std::fwrite(&magic, 4, 1, f) == 1 && std::fwrite(&count, 4, 1, f) == 1; for (const auto& p : seen_prompts) { const uint32_t len = (uint32_t) std::min(p.size(), kLearnMaxTokens); std::vector<int32_t> i32(p.begin(), p.begin() + len); ok = ok && std::fwrite(&len, 4, 1, f) == 1 && (len == 0 || std::fwrite(i32.data(), 4, len, f) == len); } std::fclose(f); if (ok) std::rename(tmp.c_str(), spath.c_str()); else std::remove(tmp.c_str()); }
void init(const std::string& cache_dir, int64_t limit_mib, int64_t max_cells, int32_t pad_tok) { namespace fs = std::filesystem; if (cache_dir.empty()) return; std::error_code ec; fs::create_directories(cache_dir, ec); if (ec) { std::fprintf(stderr, "strata metal cache: cannot create %s: %s\n", cache_dir.c_str(), ec.message().c_str()); return; } dir = cache_dir; limit_bytes = limit_mib > 0 ? (size_t) limit_mib * 1024ull * 1024ull : 0; std::vector<std::pair<fs::file_time_type, Entry>> found; for (const auto& de : fs::directory_iterator(dir, ec)) { const auto p = de.path(); const auto ext = p.extension().string(); if (ext == ".partial" || ext == ".tmp") { fs::remove(p, ec); continue; } if (ext != ".kvc") continue; FILE* f = std::fopen(p.string().c_str(), "rb"); if (!f) continue; scr::SessionHeader h, want; const bool hdr_ok = std::fread(&h, sizeof(h), 1, f) == 1; const int64_t n_ids = hdr_ok ? (int64_t) (uint32_t) h.n_ids : 0; const int64_t base_cells = (hdr_ok && (h.flags & (1u << 23))) ? (int64_t) ((uint64_t) h.n_ids >> 32) : 0; bool ok = hdr_ok && std::memcmp(h.magic, want.magic, 8) == 0 && h.version == want.version && h.n_cells > 0 && h.n_cells < max_cells && n_ids == h.n_cells && base_cells >= 0 && base_cells < h.n_cells && (base_cells & 3) == 0; std::vector<int32_t> i32; if (ok) { i32.resize((size_t) n_ids); ok = std::fread(i32.data(), 4, (size_t) n_ids, f) == (size_t) n_ids; } std::fclose(f); std::string parent_path = base_cells > 0 ? read_parent(p.string() + ".parent") : ""; if (base_cells > 0 && parent_path.empty()) ok = false; if (!ok) { remove_files(p.string()); continue; } Entry e; e.path = p.string(); e.parent_path = std::move(parent_path); e.ids.assign(i32.begin(), i32.end()); read_pads(e.path + ".pads", e.ids.size(), pad_tok, e.ids, e.pads); e.base_cells = base_cells; e.bytes = (size_t) de.file_size(ec); e.pinned = base_cells == 0 && fs::exists(e.path + ".pin", ec); found.push_back({de.last_write_time(ec), std::move(e)}); } std::sort(found.begin(), found.end(), [](const auto& a, const auto& b) { return a.first < b.first; }); for (auto& kv : found) entries.push_back(std::move(kv.second)); prune_orphans(); enforce_limits(); load_seen(); size_t total = 0, n_pin = 0, n_shared = 0; for (const auto& e : entries) { total += e.bytes; if (e.pinned) ++n_pin; if (e.base_cells > 0) ++n_shared; } std::fprintf(stderr, "strata metal cache: disk tier %s, %zu sessions (%zu pinned, %zu shared), %.2f GiB (limit %.1f GiB)\n", dir.c_str(), entries.size(), n_pin, n_shared, total / 1073741824.0, limit_bytes / 1073741824.0); }
void prune_orphans() { for (auto it = entries.begin(); it != entries.end();) { if (it->base_cells > 0) { bool parent_ok = false; for (const auto& p : entries) { if (p.path == it->parent_path && p.base_cells == 0 && (int64_t) common_prefix(p.ids, it->ids, p.ids.size()) >= it->base_cells) { parent_ok = true; break; } } if (!parent_ok) { remove_files(it->path); it = entries.erase(it); continue; } } ++it; } }
void enforce_limits() { namespace fs = std::filesystem; std::error_code ec; size_t n_pin = 0; for (const auto& e : entries) if (e.pinned) ++n_pin; for (size_t i = 0; n_pin > kMaxPins && i < entries.size(); ++i) { if (entries[i].pinned) { entries[i].pinned = false; fs::remove(entries[i].path + ".pin", ec); --n_pin; } } if (limit_bytes == 0) return; size_t total = 0; for (const auto& e : entries) total += e.bytes; for (auto it = entries.begin(); total > limit_bytes && entries.size() > 1 && it != entries.end();) { if (it->pinned) { ++it; continue; } const std::string gone = it->path; total -= it->bytes; remove_files(gone); it = entries.erase(it); for (auto cit = entries.begin(); cit != entries.end();) { if (cit->parent_path == gone) { if (cit < it) --it; total -= cit->bytes; remove_files(cit->path); cit = entries.erase(cit); } else { ++cit; } } } }
int64_t observe_prefix(const std::vector<int64_t>& ids, int64_t n_pre) { if (!enabled() || n_pre < (int64_t) kLearnMinTokens) return -1; const size_t cap = std::min<size_t>((size_t) n_pre, kLearnMaxTokens); std::vector<int64_t> cur(ids.begin(), ids.begin() + cap); size_t best_lcp = 0; bool added = false; for (auto& s : seen_prompts) { if (is_prefix(s, cur)) { s = cur; added = true; break; } if (is_prefix(cur, s)) { added = true; break; } const size_t lcp = common_prefix(cur, s, kLearnMaxTokens); const auto spans = strata::core::scr_detail::find_spans( s, cur, (int64_t) lcp, strata::core::ScrParams{32, 64, 16, 4}); if (!spans.empty()) { // Same CLM conversation after a mid-prompt edit/compaction: update in place, do not pin the fold point s = cur; added = true; break; } if (lcp >= kLearnMinTokens && lcp < cur.size() && lcp < s.size()) { best_lcp = std::max(best_lcp, lcp); } } if (!added) { if (seen_prompts.size() >= kSeenMax) seen_prompts.pop_front(); seen_prompts.push_back(std::move(cur)); } save_seen(); if (best_lcp >= kLearnMinTokens) { for (const auto& e : entries) { if (e.ids.size() == best_lcp && is_prefix(e.ids, ids)) return -1; } return (int64_t) best_lcp; } return -1; }
void touch(size_t idx) { if (idx >= entries.size()) return; namespace fs = std::filesystem; std::error_code ec; fs::last_write_time(entries[idx].path, fs::file_time_type::clock::now(), ec); if (idx + 1 < entries.size()) { Entry e = std::move(entries[idx]); entries.erase(entries.begin() + (long) idx); entries.push_back(std::move(e)); } }
void save_current(const std::vector<int64_t>& ids, const std::vector<int64_t>& pads, bool pinned, std::string& err, std::string* active_path = nullptr) { if (!enabled() || ids.size() < kLearnMinTokens) return; namespace fs = std::filesystem; std::error_code ec; for (size_t i = 0; i < entries.size(); ++i) { if (is_prefix(ids, entries[i].ids) && entries[i].pads == pads) { if (!pinned || (entries[i].ids.size() == ids.size() && entries[i].base_cells == 0)) { if (pinned && !entries[i].pinned) { entries[i].pinned = true; if (FILE* pf = std::fopen((entries[i].path + ".pin").c_str(), "wb")) std::fclose(pf); } if (!pinned && active_path != nullptr) *active_path = entries[i].path; touch(i); return; } } } std::string parent_path; int64_t base_cells = 0; if (!pinned) { size_t best_len = 0; for (const auto& e : entries) { if (e.pinned && e.base_cells == 0 && e.pads.empty() && e.ids.size() >= kLearnMinTokens && ids.size() >= kLearnMinTokens + 4) { const size_t lcp = common_prefix(e.ids, ids, std::min(e.ids.size(), ids.size() - 4)); const int64_t c0 = ((int64_t) lcp / 4) * 4; if (c0 >= (int64_t) kLearnMinTokens && (size_t) c0 > best_len && (pads.empty() || pads.front() >= c0)) { best_len = (size_t) c0; base_cells = c0; parent_path = e.path; } } } } const uint64_t h = fnv1a(ids.data(), ids.size()); char fname[96]; std::snprintf(fname, sizeof(fname), "/%016llx-%zu.kvc", (unsigned long long) h, ids.size()); const std::string path = dir + fname; std::vector<int32_t> i32(ids.begin(), ids.end()); if (!strata::metal_engine::window_session_save(path.c_str(), i32.data(), (int64_t) i32.size(), err, base_cells)) return; write_pads(path + ".pads", pads); write_parent(path + ".parent", parent_path); if (pinned) { if (FILE* pf = std::fopen((path + ".pin").c_str(), "wb")) std::fclose(pf); } else { fs::remove(path + ".pin", ec); } const size_t sz = (size_t) fs::file_size(path, ec); const std::string prev_active = (!pinned && active_path != nullptr) ? *active_path : ""; for (auto it = entries.begin(); it != entries.end();) { if (it->path == path || (!it->pinned && (is_prefix(it->ids, ids) || (!prev_active.empty() && it->path == prev_active)))) { if (it->path != path) { remove_files(it->path); } it = entries.erase(it); } else { ++it; } } if (!pinned && active_path != nullptr) *active_path = path; entries.push_back({path, parent_path, ids, pads, base_cells, sz, pinned}); prune_orphans(); enforce_limits(); }
int find_best(const std::vector<int64_t>& raw_ids, int64_t max_new, int64_t max_cells, int32_t pad_tok, int64_t& out_reused) const { out_reused = 0; if (!enabled()) return -1; int best_idx = -1; int64_t best_score = 0; for (size_t i = 0; i < entries.size(); ++i) { const auto& e = entries[i]; if (e.ids.empty()) continue; std::vector<int64_t> q_ids = raw_ids; if (!e.pads.empty()) { std::vector<int64_t> carried; strata::core::scr_carry(raw_ids, e.ids, e.pads, pad_tok, carried); if ((int64_t) carried.size() + max_new <= max_cells) q_ids = std::move(carried); } const int64_t q_n_pre = std::max<int64_t>(0, (int64_t) q_ids.size() - 1); int64_t cand = 0; if ((int64_t) e.ids.size() <= q_n_pre && is_prefix(e.ids, q_ids)) { cand = (int64_t) e.ids.size(); } else if (e.ids.size() >= 32 && q_n_pre >= 32) { const strata::core::ScrPlan plan = strata::core::scr_plan( e.ids, q_ids, {(int64_t) e.ids.size()}, strata::core::ScrParams{32, 64, 16, 4}); if (!plan.empty() && (int64_t) q_ids.size() + plan.pads() + max_new <= max_cells) { cand = plan.n_prefix + plan.reused(); } } if (cand >= 32 && cand >= best_score) { best_score = cand; best_idx = (int) i; } } out_reused = best_score; return best_idx; }};
static int serve_loop(void* w, strata::core::SessionState& ss, int64_t n_vocab, int64_t max_cells, int suffix_draft, bool use_mtp, int spec, float spec_min_p, const std::string& cache_disk_dir, int64_t cache_disk_mib) { namespace me = strata::metal_engine; constexpr int64_t kScrPadToken = 198; // '\n' pad for 4-cell QSA block alignment (ported from 01a11408 / scr.hpp) DiskCache disk_cache; disk_cache.init(cache_disk_dir, cache_disk_mib, max_cells, kScrPadToken);
std::printf("INFO kv=%lld\n", (long long) max_cells); std::printf("READY %lld\n", (long long) max_cells); std::fflush(stdout);
std::string line, err; struct CkSlot { std::vector<int64_t> ids; strata::core::ConversationCheckpoint ck; int32_t prev0 = -1, prev1 = -1; }; std::vector<CkSlot> ck_ring; auto save_ck_at = [&](const std::vector<int64_t>& pfx) { if (pfx.empty()) return; for (size_t ci = 0; ci < ck_ring.size(); ++ci) { if (ck_ring[ci].ids == pfx) { CkSlot slot = std::move(ck_ring[ci]); ck_ring.erase(ck_ring.begin() + (long) ci); if (me::window_checkpoint_save(slot.ck, err)) { slot.prev0 = ss.ple_prev[0]; slot.prev1 = ss.ple_prev[1]; ck_ring.push_back(std::move(slot)); } return; } } if (ck_ring.size() >= 2) ck_ring.erase(ck_ring.begin()); CkSlot slot; slot.ids = pfx; if (me::window_checkpoint_save(slot.ck, err)) { slot.prev0 = ss.ple_prev[0]; slot.prev1 = ss.ple_prev[1]; ck_ring.push_back(std::move(slot)); } }; std::vector<int64_t> cached_ids; std::vector<int64_t> cached_pads; std::vector<std::pair<int64_t, uint64_t>> cached_imgs; std::vector<int32_t> mrope_host((size_t) (max_cells + 16) * 3, 0); std::vector<float> img_rows; std::vector<const float*> row_ptr; std::string active_path; bool ram_dirty = false; std::vector<float> lg_buf((size_t) n_vocab); std::mt19937 rng(42); // temp > 0 with drafts: speculative sampling (spec_sample), else STRATA_SPEC_SAMPLE=greedy: the drafts kept while // they are the greedy picks and only the last token sampled (fast, but those tokens are greedy, unpenalized) static const bool spec_exact = std::getenv("STRATA_SPEC_SAMPLE") && (std::string(std::getenv("STRATA_SPEC_SAMPLE")) == "exact" || std::string(std::getenv("STRATA_SPEC_SAMPLE")) == "typical");
while (std::getline(std::cin, line)) { if (line.empty()) continue; if (line == "QUIT" || line.rfind("QUIT ", 0) == 0) { if (ram_dirty && cached_ids.size() >= 32) disk_cache.save_current(cached_ids, cached_pads, false, err, &active_path); break; } if (line == "STOP") continue; if (line == "VRAM" || line.rfind("VRAM ", 0) == 0) { std::printf("VRAM expert_slots=4129 expert_cache_mib=6656 expert_cache_full_mib=6656 vram_free_mib=13312\n"); std::fflush(stdout); continue; } if (line.rfind("SAVE ", 0) == 0) { std::string path = line.substr(5); while (!path.empty() && (path.back() == '\r' || path.back() == ' ')) path.pop_back(); const double t0 = now_ms(); std::vector<int32_t> i32(cached_ids.begin(), cached_ids.end()); if (!me::window_session_save(path.c_str(), i32.data(), (int64_t) i32.size(), err)) { std::printf("SERR io 0 %s\n", err.c_str()); } else { DiskCache::write_pads(path + ".pads", cached_pads); struct stat st{}; const size_t fbytes = ::stat(path.c_str(), &st) == 0 ? (size_t) st.st_size : scr::fixed_bytes() + scr::cell_bytes(i32.size()); std::printf("SAVED %zu %zu %.1f\n", i32.size(), fbytes, now_ms() - t0); } std::fflush(stdout); continue; } if (line.rfind("RESTORE ", 0) == 0) { std::string path = line.substr(8); while (!path.empty() && (path.back() == '\r' || path.back() == ' ')) path.pop_back(); const double t0 = now_ms(); std::vector<int32_t> i32; int64_t cells = 0; if (!me::window_session_load(path.c_str(), i32, &cells, err)) { std::printf("SERR io 0 %s\n", err.c_str()); } else { cached_ids.assign(i32.begin(), i32.end()); DiskCache::read_pads(path + ".pads", cached_ids.size(), kScrPadToken, cached_ids, cached_pads); ram_dirty = false; struct stat st{}; const size_t fbytes = ::stat(path.c_str(), &st) == 0 ? (size_t) st.st_size : scr::fixed_bytes() + scr::cell_bytes((size_t) cells); std::printf("RESTORED %lld %zu %.1f\n", (long long) cells, fbytes, now_ms() - t0); } std::fflush(stdout); continue; } const bool geni = line.rfind("GENI ", 0) == 0; if (!geni && line.rfind("GEN ", 0) != 0) continue;
std::istringstream iss(line.substr(geni ? 5 : 4)); int max_new = 512; if (!(iss >> max_new)) continue; float temp = 0.0f, top_p = 1.0f, min_p = 0.05f, pen_rep = 1.05f, pen_pres = 0.2f; int top_k = 40; std::vector<uint64_t> req_img_keys; std::string tok_str; while (iss >> tok_str) { auto eq = tok_str.find('='); if (eq != std::string::npos) { std::string k = tok_str.substr(0, eq), v = tok_str.substr(eq + 1); if (k == "temperature") temp = std::stof(v); else if (k == "top_p") top_p = std::stof(v); else if (k == "top_k") top_k = std::stoi(v); else if (k == "min_p") min_p = std::stof(v); else if (k == "penalty_repeat") pen_rep = std::stof(v); else if (k == "penalty_present") pen_pres = std::stof(v); else if (k == "seed") rng.seed((uint32_t) std::stoul(v)); else if (k == "img_keys") { req_img_keys.clear(); const char* q = v.c_str(); while (*q != '\0') { char* qe = nullptr; const unsigned long long hk = std::strtoull(q, &qe, 16); if (qe == q) { req_img_keys.clear(); break; } req_img_keys.push_back((uint64_t) hk); q = *qe == ',' ? qe + 1 : qe; if (*qe != ',' && *qe != '\0') { req_img_keys.clear(); break; } } } } else { break; } } std::string emb_path; if (geni) { emb_path = tok_str; tok_str.clear(); iss >> tok_str; } std::vector<int64_t> raw_ids; std::stringstream tss(tok_str); std::string item; while (std::getline(tss, item, ',')) { if (!item.empty()) raw_ids.push_back(std::stoll(item)); } if (raw_ids.empty()) { std::printf("DONE 0 0 0.0 0.0 stop 0 0 0\n"); std::fflush(stdout); continue; } if ((int64_t) raw_ids.size() + max_new > max_cells) { max_new = (int) std::max<int64_t>(1, max_cells - (int64_t) raw_ids.size() - 1); }
const double t_pre0 = now_ms(); std::vector<std::pair<int64_t, uint64_t>> req_imgs; row_ptr.clear(); img_rows.clear(); if (geni) { constexpr int64_t kImagePad = 248056; const int64_t n = (int64_t) raw_ids.size(); row_ptr.assign((size_t) n, nullptr); const int64_t cells = (int64_t) mrope_host.size() / 3; auto put = [&](int64_t c, int64_t t, int64_t h, int64_t w) { mrope_host[(size_t) c * 3] = (int32_t) t; mrope_host[(size_t) c * 3 + 1] = (int32_t) h; mrope_host[(size_t) c * 3 + 2] = (int32_t) w; }; struct Img { int64_t n, nx, ny; size_t off; }; std::vector<Img> imgs; std::string ve; std::FILE* f = std::fopen(emb_path.c_str(), "rb"); if (!f) ve = "cannot open " + emb_path; while (f && ve.empty()) { int32_t hdr[5]; const size_t got = std::fread(hdr, sizeof(int32_t), 5, f); if (got == 0) break; if (got != 5 || hdr[0] != 0x31455653 || hdr[1] < 1 || hdr[2] < 1 || hdr[3] < 1 || (int64_t) hdr[2] * hdr[3] != hdr[1] || hdr[4] != 2560) { ve = "bad embeddings file (expected strata-vision records of width 2560)"; break; } const size_t off = img_rows.size(), cnt = (size_t) hdr[1] * (size_t) hdr[4]; img_rows.resize(off + cnt); if (std::fread(img_rows.data() + off, sizeof(float), cnt, f) != cnt) { ve = "short embeddings file"; break; } imgs.push_back({hdr[1], hdr[2], hdr[3], off}); } if (f) std::fclose(f); int64_t p = 0, i = 0; size_t k = 0; while (ve.empty() && i < n) { if (raw_ids[(size_t) i] != kImagePad) { put(i, p, p, p); ++p; ++i; continue; } if (k >= imgs.size()) { ve = "the prompt has more images than the embeddings file"; break; } const Img& im = imgs[k++]; uint64_t h = DiskCache::fnv1a(&im.n, 3); if (req_img_keys.size() == imgs.size()) { h ^= req_img_keys[k - 1] * 1099511628211ull; } else { h ^= DiskCache::fnv1a(reinterpret_cast<const int64_t*>(img_rows.data() + im.off), (size_t) im.n * (2560 * sizeof(float) / sizeof(int64_t))); } req_imgs.push_back({i, h}); for (int64_t j = 0; j < im.n && ve.empty(); ++j) { if (i + j >= n || raw_ids[(size_t) (i + j)] != kImagePad) { ve = "image has fewer <|image_pad|> tokens than rows"; } } for (int64_t j = 0; j < im.n && ve.empty(); ++j) { const int64_t y = j / im.nx, x = j % im.nx; put(i + j, p, p + y, p + x); row_ptr[(size_t) (i + j)] = img_rows.data() + im.off + (size_t) j * 2560; } i += im.n; p += std::max(im.nx, im.ny); } if (ve.empty() && k != imgs.size()) ve = "the embeddings file has more images than the prompt"; if (ve.empty() && n > 0 && raw_ids[(size_t) (n - 1)] == kImagePad) ve = "the prompt cannot end in an image"; for (int64_t c = n; ve.empty() && c < cells; ++c) put(c, p + (c - n), p + (c - n), p + (c - n)); if (!ve.empty()) { me::window_mrope_host_set(nullptr); std::printf("ERR %s\n", ve.c_str()); std::fflush(stdout); continue; } me::window_mrope_host_set(mrope_host.data()); } else { me::window_mrope_host_set(nullptr); }
std::vector<int64_t> ids = raw_ids; // Carry forward any 0..3 SCR alignment pads from prior turns that still lie in the shared prefix if (!geni && cached_imgs.empty() && !cached_pads.empty() && !cached_ids.empty()) { std::vector<int64_t> carried_ids; cached_pads = strata::core::scr_carry(raw_ids, cached_ids, cached_pads, kScrPadToken, carried_ids); if ((int64_t) carried_ids.size() + max_new <= max_cells) { ids = std::move(carried_ids); } else { cached_pads.clear(); } } int64_t n_pre = std::max<int64_t>(0, (int64_t) ids.size() - 1); size_t lcp = 0; const size_t lim0 = std::min(cached_ids.size(), (size_t) n_pre); while (lcp < lim0 && cached_ids[lcp] == ids[lcp]) ++lcp; if (geni != !cached_imgs.empty()) { // If switching between text and vision (or if image hashes differ inside 0..lcp), clamp lcp before the first differing image size_t img_cap = lim0; const size_t n_im = std::max(cached_imgs.size(), req_imgs.size()); for (size_t ii = 0; ii < n_im; ++ii) { if (ii >= cached_imgs.size() || ii >= req_imgs.size() || cached_imgs[ii] != req_imgs[ii]) { const int64_t p_cut = std::min(ii < cached_imgs.size() ? cached_imgs[ii].first : (int64_t) lim0, ii < req_imgs.size() ? req_imgs[ii].first : (int64_t) lim0); img_cap = (size_t) std::max<int64_t>(0, p_cut); break; } } lcp = std::min(lcp, img_cap); } else if (geni) { for (size_t ii = 0; ii < req_imgs.size(); ++ii) { if (ii >= cached_imgs.size() || cached_imgs[ii] != req_imgs[ii]) { const int64_t p_cut = std::min(ii < cached_imgs.size() ? cached_imgs[ii].first : (int64_t) lcp, req_imgs[ii].first); lcp = std::min(lcp, (size_t) std::max<int64_t>(0, p_cut)); break; } } }
// Estimate how many tokens the live RAM state + ck_ring can reuse before checking the NVMe tier int64_t ram_reused = 0; if (lcp == cached_ids.size() && lcp > 0) { ram_reused = (int64_t) lcp; } else if (!geni && cached_imgs.empty() && cached_ids.size() >= 32 && n_pre >= 32) { std::vector<int64_t> ck_lens = {(int64_t) cached_ids.size()}; for (const auto& s : ck_ring) { if (!s.ids.empty() && s.ids.size() <= cached_ids.size() && std::equal(s.ids.begin(), s.ids.end(), cached_ids.begin())) { ck_lens.push_back((int64_t) s.ids.size()); } } const strata::core::ScrPlan plan = strata::core::scr_plan( cached_ids, ids, ck_lens, strata::core::ScrParams{32, 64, 16, 4}); if (!plan.empty() && (int64_t) ids.size() + plan.pads() + max_new <= max_cells) { ram_reused = plan.n_prefix + plan.reused(); } } if (ram_reused == 0) { for (const auto& s : ck_ring) { if (!s.ids.empty() && (int64_t) s.ids.size() <= n_pre && (int64_t) s.ids.size() > ram_reused && std::equal(s.ids.begin(), s.ids.end(), ids.begin())) { ram_reused = (int64_t) s.ids.size(); } } }
// Check NVMe disk tier if it can beat the live in-RAM reuse by >= 32 tokens (text conversations only) if (!geni && disk_cache.enabled()) { int64_t disk_reused = 0; const int disk_idx = disk_cache.find_best(raw_ids, max_new, max_cells, kScrPadToken, disk_reused); if (disk_idx >= 0 && disk_reused >= ram_reused + 32) { if (ram_dirty && cached_ids.size() >= 32) { disk_cache.save_current(cached_ids, cached_pads, false, err, &active_path); ram_dirty = false; } const int disk_idx2 = disk_cache.find_best(raw_ids, max_new, max_cells, kScrPadToken, disk_reused); if (disk_idx2 >= 0) { const std::string load_path = disk_cache.entries[(size_t) disk_idx2].path; const std::string load_parent = disk_cache.entries[(size_t) disk_idx2].parent_path; const bool load_pinned = disk_cache.entries[(size_t) disk_idx2].pinned; const std::vector<int64_t> load_pads = disk_cache.entries[(size_t) disk_idx2].pads; std::vector<int32_t> i32; int64_t cells = 0; if (me::window_session_load(load_path.c_str(), i32, &cells, err, load_parent.c_str())) { cached_ids.assign(i32.begin(), i32.end()); cached_pads = load_pads; active_path = load_pinned ? "" : load_path; ram_dirty = false; ck_ring.clear(); save_ck_at(cached_ids); disk_cache.touch((size_t) disk_idx2); ids = raw_ids; if (!cached_pads.empty() && !cached_ids.empty()) { std::vector<int64_t> carried_ids; cached_pads = strata::core::scr_carry(raw_ids, cached_ids, cached_pads, kScrPadToken, carried_ids); if ((int64_t) carried_ids.size() + max_new <= max_cells) ids = std::move(carried_ids); else cached_pads.clear(); } n_pre = std::max<int64_t>(0, (int64_t) ids.size() - 1); lcp = 0; const size_t lim1 = std::min(cached_ids.size(), (size_t) n_pre); while (lcp < lim1 && cached_ids[lcp] == ids[lcp]) ++lcp; } } } else if (ram_reused == 0 && ram_dirty && cached_ids.size() >= 32) { disk_cache.save_current(cached_ids, cached_pads, false, err, &active_path); ram_dirty = false; } }
const int64_t learn_at = (!geni && disk_cache.enabled() && cached_pads.empty()) ? disk_cache.observe_prefix(ids, n_pre) : -1; int64_t reused = 0;
auto on_chunk = [&](const float*, int64_t k, int64_t p0, std::string& e) -> bool { if (!use_mtp) return true; std::vector<int32_t> nxt((size_t) k); for (int64_t t = 0; t < k; ++t) nxt[(size_t) t] = (int32_t) ids[(size_t) (p0 + t + 1)]; return me::mtp_prefill(nxt.data(), k, p0, e); }; const float* const* embd_ptr = geni ? row_ptr.data() : nullptr; auto run_prefill_span = [&](int64_t start_p, int64_t end_p) -> bool { if (end_p <= start_p) return true; if (learn_at > start_p && learn_at <= end_p && cached_pads.empty()) { if (!me::prompt_run(ids.data() + start_p, learn_at - start_p, start_p, embd_ptr, on_chunk, err)) return false; std::vector<int64_t> pfx(ids.begin(), ids.begin() + learn_at); save_ck_at(pfx); disk_cache.save_current(pfx, {}, /*pinned=*/true, err); if (end_p > learn_at) { if (!me::prompt_run(ids.data() + learn_at, end_p - learn_at, learn_at, embd_ptr, on_chunk, err)) return false; } return true; } return me::prompt_run(ids.data() + start_p, end_p - start_p, start_p, embd_ptr, on_chunk, err); };
if (lcp == cached_ids.size() && lcp > 0) { reused = (int64_t) lcp; if (reused > 0) { std::printf("RESUME %lld\n", (long long) reused); std::fflush(stdout); } if (n_pre > reused) { if (!run_prefill_span(reused, n_pre)) { std::printf("ERR %s\n", err.c_str()); std::fflush(stdout); cached_ids.clear(); ram_dirty = false; continue; } } cached_ids.assign(ids.begin(), ids.begin() + n_pre); ram_dirty = true; } else { // 2. Check SCR mid-context fold [A B C] -> [A B' C D] (including stripped <think> blocks) via scr_plan if (!geni && cached_imgs.empty() && cached_ids.size() >= 32 && n_pre >= 32) { std::vector<int64_t> ck_lens = {(int64_t) cached_ids.size()}; for (const auto& s : ck_ring) { if (!s.ids.empty() && s.ids.size() <= cached_ids.size() && std::equal(s.ids.begin(), s.ids.end(), cached_ids.begin())) { ck_lens.push_back((int64_t) s.ids.size()); } } const strata::core::ScrPlan plan = strata::core::scr_plan( cached_ids, ids, ck_lens, strata::core::ScrParams{32, 64, 16, 4}); if (!plan.empty() && (int64_t) ids.size() + plan.pads() + max_new <= max_cells) { ids = strata::core::scr_padded(ids, plan, kScrPadToken); n_pre = std::max<int64_t>(0, (int64_t) ids.size() - 1); const auto new_pads = strata::core::scr_pad_positions(plan); cached_pads.insert(cached_pads.end(), new_pads.begin(), new_pads.end()); std::sort(cached_pads.begin(), cached_pads.end()); cached_pads.erase(std::unique(cached_pads.begin(), cached_pads.end()), cached_pads.end()); save_ck_at(cached_ids); bool scr_ok = true; // Pass 1: relocate left-shifting runs (shift < 0) in forward order 0 .. m-1 for (size_t ri = 0; scr_ok && ri < plan.runs.size(); ++ri) { const auto& r = plan.runs[ri]; if (r.shift() < 0 && !me::window_scr_relocate(r.old_start, r.at, r.n, err)) scr_ok = false; } // Pass 2: relocate right-shifting runs (shift > 0) in reverse order m-1 .. 0 for (size_t ri = plan.runs.size(); scr_ok && ri-- > 0;) { const auto& r = plan.runs[ri]; if (r.shift() > 0 && !me::window_scr_relocate(r.old_start, r.at, r.n, err)) scr_ok = false; } int64_t cur_p = plan.restore > 0 ? plan.restore : (int64_t) lcp; if (scr_ok && plan.restore > 0) { for (const auto& s : ck_ring) { if ((int64_t) s.ids.size() == plan.restore) { me::window_checkpoint_restore(s.ck, err); ss.ple_prev[0] = s.prev0; ss.ple_prev[1] = s.prev1; break; } } } // Pass 3: prefill edited gaps [cur_p, r.at) now that all runs sit at [r.at, r.at + r.n) for (const auto& r : plan.runs) { if (!scr_ok) break; if (r.at > cur_p) { if (!me::prompt_run_rows(ids.data() + cur_p, r.at - cur_p, cur_p, on_chunk, err)) { scr_ok = false; break; } } cur_p = r.at + r.n; } if (scr_ok && plan.adopt > 0) { for (const auto& s : ck_ring) { if ((int64_t) s.ids.size() == plan.adopt) { me::window_checkpoint_restore(s.ck, err); ss.ple_prev[0] = s.prev0; ss.ple_prev[1] = s.prev1; break; } } } if (scr_ok && n_pre > cur_p) { if (!me::prompt_run_rows(ids.data() + cur_p, n_pre - cur_p, cur_p, on_chunk, err)) { scr_ok = false; } } if (scr_ok) { reused = plan.n_prefix + plan.reused(); std::printf("RESUME %lld\n", (long long) reused); std::fflush(stdout); cached_ids.assign(ids.begin(), ids.begin() + n_pre); ram_dirty = true; } else { cached_ids.clear(); cached_pads.clear(); ram_dirty = false; } } } if (reused == 0) { int best_idx = -1; size_t best_len = 0; for (size_t ci = 0; ci < ck_ring.size(); ++ci) { const auto& cids = ck_ring[ci].ids; if (!cids.empty() && cids.size() <= (size_t) n_pre && cids.size() > best_len && std::equal(cids.begin(), cids.end(), ids.begin())) { best_len = cids.size(); best_idx = (int) ci; } } if (best_idx >= 0) { me::window_checkpoint_restore(ck_ring[(size_t) best_idx].ck, err); ss.ple_prev[0] = ck_ring[(size_t) best_idx].prev0; ss.ple_prev[1] = ck_ring[(size_t) best_idx].prev1; reused = (int64_t) best_len; std::printf("RESUME %lld\n", (long long) reused); std::fflush(stdout); if (n_pre > reused) { if (!run_prefill_span(reused, n_pre)) { std::printf("ERR %s\n", err.c_str()); std::fflush(stdout); cached_ids.clear(); ram_dirty = false; continue; } } cached_ids.assign(ids.begin(), ids.begin() + n_pre); ram_dirty = true; } else { me::window_reset(); ss.ple_prev[0] = ss.ple_prev[1] = -1; cached_ids.clear(); cached_pads.clear(); active_path.clear(); if (n_pre > 0) { if (!run_prefill_span(0, n_pre)) { std::printf("ERR %s\n", err.c_str()); std::fflush(stdout); ram_dirty = false; continue; } } cached_ids.assign(ids.begin(), ids.begin() + n_pre); ram_dirty = true; } } }
const double prefill_ms = now_ms() - t_pre0; const double t_dec0 = now_ms(); int32_t cur_tok = (int32_t) ids.back(); int64_t pos = n_pre; int gen = 0; const char* finish = "length";
if ((use_mtp || (suffix_draft > 0 && temp == 0.0f))) { strata::spec::SuffixDrafter suffix(4, 32, (size_t) max_cells); for (int64_t id : ids) suffix.append((int32_t) id); int32_t mtp_drafts[8] = {}; float mtp_probs[8] = {}; int n_mtp_drafts = 0; bool first_win = true; SpecControl spec_control; while (gen < max_new) { int32_t input[8] = {cur_tok}, output[8] = {}; int nd = 0; const int rem = max_new - gen; nd = spec_control.propose(suffix, input, mtp_drafts, mtp_probs, !first_win && use_mtp ? n_mtp_drafts : 0, !first_win ? suffix_draft : 0, rem, pos, max_cells); first_win = false; const double verify_begin = now_ms(); if (!me::window_run(w, 1 + nd, input, pos, true, output, err)) { std::printf("ERR %s\n", err.c_str()); std::fflush(stdout); cached_ids.clear(); finish = "error"; break; } spec_control.gate.spec_gate_forward(1 + nd, (now_ms() - verify_begin) / 1000); // the penalties see the window's tokens (cached_ids + input[0..keep-1]), as the one-token path's int keep; if (temp > 0.0f && spec_exact) { keep = spec_sample(w, input, output, nd, n_vocab, temp, top_k, top_p, min_p, pen_rep, pen_pres, cached_ids, lg_buf, rng); } else { keep = strata::spec::suffix_keep(input, output, nd); cached_ids.insert(cached_ids.end(), input, input + keep); if (temp > 0.0f && !strata::spec::suffix_stop(output[keep - 1])) { me::window_copy_logits(w, keep - 1, lg_buf.data()); output[keep - 1] = sample_logits(lg_buf.data(), n_vocab, output[keep - 1], temp, top_k, top_p, min_p, pen_rep, pen_pres, cached_ids, rng); } } spec_control.result(nd, keep); if (!me::window_commit(w, keep, err)) { std::printf("ERR %s\n", err.c_str()); std::fflush(stdout); cached_ids.clear(); finish = "error"; break; } for (int i = 0; i < keep; ++i) { std::printf("T %d\n", (int) output[i]); std::fflush(stdout); suffix.append(output[i]); cur_tok = output[i]; ++gen; if (strata::spec::suffix_stop(cur_tok)) break; } if (!strata::spec::suffix_stop(cur_tok) && gen < max_new && use_mtp) { const double draft_begin = now_ms(); if (!me::mtp_draft(1 + nd, output, pos, keep - 1, spec - 1, mtp_drafts, mtp_probs, spec_min_p, &n_mtp_drafts, err)) { std::printf("ERR %s\n", err.c_str()); std::fflush(stdout); cached_ids.clear(); finish = "error"; break; } spec_control.gate.spec_gate_draft_cost(0, (now_ms() - draft_begin) / (1000 * std::max(1, spec - 1))); } else if (use_mtp) { me::mtp_prefill(output, keep, pos, err); } pos += keep; if (strata::spec::suffix_stop(cur_tok)) { finish = "stop"; break; } } spec_control.describe(); } else { for (int step = 0; step < max_new; ++step) { int32_t greedy_tok = 0; if (!me::window_run(w, 1, &cur_tok, pos, true, &greedy_tok, err) || !me::window_commit(w, 1, err)) { std::printf("ERR %s\n", err.c_str()); std::fflush(stdout); cached_ids.clear(); finish = "error"; break; } cached_ids.push_back(cur_tok); ++pos;
int32_t next_tok = greedy_tok; if (temp > 0.0f) { me::window_copy_logits(w, 0, lg_buf.data()); next_tok = sample_logits(lg_buf.data(), n_vocab, greedy_tok, temp, top_k, top_p, min_p, pen_rep, pen_pres, cached_ids, rng); } std::printf("T %d\n", (int) next_tok); std::fflush(stdout); ++gen;
if (next_tok == 248046 || next_tok == 248044) { finish = "stop"; break; } cur_tok = next_tok; } } cached_imgs = req_imgs; if (!geni) save_ck_at(cached_ids); const double decode_ms = now_ms() - t_dec0; if (std::strcmp(finish, "error") != 0) { std::printf("DONE %d %zu %.1f %.1f %s 0 0 %lld\n", gen, ids.size(), prefill_ms, decode_ms, finish, (long long) reused); std::fflush(stdout); if (!geni && disk_cache.enabled() && cached_ids.size() >= 32) { disk_cache.save_current(cached_ids, cached_pads, false, err, &active_path); ram_dirty = false; } } } return 0;}
int main(int argc, char** argv) { bool serve_mode = false; std::string metal_gguf, ple_gguf, mtp_gguf, tokens_file, logits_out, logits_ref, logits_ref2; int n_expert = 256; int64_t max_cells = 4096; int n_prompt = 22; int n_gen = 64; int suffix_draft = 0; int spec = 1; float spec_min_p = 0.0f; std::string cache_disk_dir = std::getenv("STRATA_METAL_CACHE_DISK") ? std::getenv("STRATA_METAL_CACHE_DISK") : ""; int64_t cache_disk_mib = std::getenv("STRATA_METAL_CACHE_DISK_MIB") ? std::atoll(std::getenv("STRATA_METAL_CACHE_DISK_MIB")) : 16384; std::string save_session, load_session, scr_a, scr_b, scr_bp, scr_c, scr_d, scr_j = "0,8,16", scr_tmp = "/tmp/strata_scr_test.session"; int32_t scr_pad = 198;
for (int i = 1; i < argc; ++i) { std::string a = argv[i]; if (a == "--metal-gguf" && i + 1 < argc) metal_gguf = argv[++i]; else if (a == "--ple-gguf" && i + 1 < argc) ple_gguf = argv[++i]; else if (a == "--tokens-file" && i + 1 < argc) tokens_file = argv[++i]; else if (a == "--n-expert" && i + 1 < argc) n_expert = std::atoi(argv[++i]); else if ((a == "--ctx" || a == "--max-context") && i + 1 < argc) max_cells = std::atoll(argv[++i]); else if (a == "--serve") serve_mode = true; else if (a == "--vision") { /* enabled automatically via GENI */ } else if (a == "--n-prompt" && i + 1 < argc) n_prompt = std::atoi(argv[++i]); else if (a == "--n-gen" && i + 1 < argc) n_gen = std::atoi(argv[++i]); else if (a == "--suffix-draft" && i + 1 < argc) suffix_draft = std::max(0, std::min(7, std::atoi(argv[++i]))); else if (a == "--mtp-gguf" && i + 1 < argc) mtp_gguf = argv[++i]; else if (a == "--spec" && i + 1 < argc) spec = std::max(1, std::min(4, std::atoi(argv[++i]))); else if (a == "--spec-min-p" && i + 1 < argc) spec_min_p = (float) std::atof(argv[++i]); else if (a == "--cache-disk" && i + 1 < argc) cache_disk_dir = argv[++i]; else if (a == "--cache-disk-size" && i + 1 < argc) cache_disk_mib = std::atoll(argv[++i]); else if (a == "--logits-out" && i + 1 < argc) logits_out = argv[++i]; else if (a == "--logits-ref" && i + 1 < argc) logits_ref = argv[++i]; else if (a == "--logits-ref2" && i + 1 < argc) logits_ref2 = argv[++i]; else if (a == "--save-session" && i + 1 < argc) save_session = argv[++i]; else if (a == "--load-session" && i + 1 < argc) load_session = argv[++i]; else if (a == "--scr-a" && i + 1 < argc) scr_a = argv[++i]; else if (a == "--scr-b" && i + 1 < argc) scr_b = argv[++i]; else if (a == "--scr-bp" && i + 1 < argc) scr_bp = argv[++i]; else if (a == "--scr-c" && i + 1 < argc) scr_c = argv[++i]; else if (a == "--scr-d" && i + 1 < argc) scr_d = argv[++i]; else if (a == "--scr-j" && i + 1 < argc) scr_j = argv[++i]; else if (a == "--scr-pad" && i + 1 < argc) scr_pad = std::atoi(argv[++i]); else if (a == "--scr-tmp" && i + 1 < argc) scr_tmp = argv[++i]; } if (metal_gguf.empty()) { std::fprintf(stderr, "usage: strata_window_cli --metal-gguf <K256-metal.gguf> [--ple-gguf <shard1.gguf>] " "[--n-expert 256] [--ctx 4096] [--n-prompt 22] [--n-gen 64] [--tokens-file <ids.txt>] " "[--logits-out <out.bin>] [--logits-ref <ref.bin>] [--save-session <f> | --load-session <f>] " "[--scr-a A --scr-b B --scr-bp B' --scr-c C --scr-d D [--scr-j 0,8,16] [--scr-pad 198]]\n"); return 2; }
const bool tf_mode = !logits_out.empty() || !logits_ref.empty(); std::vector<int64_t> all_ids; if (!tokens_file.empty()) { std::ifstream f(tokens_file); int64_t tid; while (f >> tid) all_ids.push_back(tid); } if (all_ids.empty()) { const int64_t def_ids[] = { 248045, 8678, 198, 2523, 513, 264, 10631, 17313, 13, 248046, 198, 248045, 846, 198, 814, 20139, 1204, 27413, 71284, 3728, 13081, 32907, 12361, 321, 264, 3074, 6700, 11, 321, 948, 264, 28336, 628, 6561, 571, 13, 248046, 198, 248045, 74455, 198, 248068, 198 }; const size_t nd = sizeof(def_ids) / sizeof(def_ids[0]); for (int i = 0; i < n_prompt + (tf_mode ? n_gen : 0); ++i) all_ids.push_back(def_ids[i % nd]); } if (tf_mode && (int) all_ids.size() < n_prompt + n_gen) { std::fprintf(stderr, "strata_window_cli: teacher-forced mode needs at least n_prompt + n_gen (%d) tokens in --tokens-file, got %zu\n", n_prompt + n_gen, all_ids.size()); return 1; } std::vector<int64_t> prompt_ids(all_ids.begin(), all_ids.begin() + std::min<size_t>(all_ids.size(), (size_t) n_prompt));
const double t_init0 = now_ms(); std::string err; strata::kernels::PleTable ple_table; strata::core::SessionState ss{}; ss.ple_prev[0] = -1; ss.ple_prev[1] = -1; ss.ple.consts = strata::kernels::ple_artifact_consts(); if (!ple_gguf.empty()) { strata::kernels::PleIoOptions io{}; io.mode = strata::kernels::PleIo::Direct; io.max_inflight = 256; io.cache_rows = 1u << 18; if (!ple_table.open(ple_gguf, err, io)) { std::fprintf(stderr, "strata_window_cli: PleTable::open failed: %s\n", err.c_str()); return 1; } ss.ple.table = &ple_table; }
strata::core::ModelGeometry g{}; g.n_embd = 2560; g.hc = 4; g.hc_lr = 320; g.n_ff = 640; g.n_expert = n_expert; g.n_layers = 48;
strata::metal_engine::set_model_shard(metal_gguf.c_str()); const bool use_mtp = !mtp_gguf.empty() && spec > 1; const int max_win_t = std::max({1, 1 + suffix_draft, spec}); if (use_mtp) { setenv("STRATA_METAL_MTP_GGUF", mtp_gguf.c_str(), 1); if (!strata::metal_engine::mtp_load(mtp_gguf, max_win_t, err)) { std::fprintf(stderr, "strata_window_cli: mtp_load failed: %s\n", err.c_str()); return 1; } } void* w = strata::metal_engine::window_create(); int64_t n_vocab = 0; if (!strata::metal_engine::window_init(w, g, ss, max_win_t, max_cells, &n_vocab, err)) { std::fprintf(stderr, "strata_window_cli: window_init failed: %s\n", err.c_str()); return 1; } const double init_ms = now_ms() - t_init0; if (serve_mode) return serve_loop(w, ss, n_vocab, max_cells, suffix_draft, use_mtp, spec, spec_min_p, cache_disk_dir, cache_disk_mib); if (!scr_a.empty()) return scr_test(w, ss, n_vocab, scr_a, scr_b, scr_bp, scr_c, scr_d, scr_j, scr_pad, n_gen, scr_tmp);
FILE* fout_lg = nullptr; if (!logits_out.empty()) { fout_lg = std::fopen(logits_out.c_str(), "wb"); if (!fout_lg) { std::perror("fopen logits_out"); return 1; } fcntl(fileno(fout_lg), F_NOCACHE, 1); } FILE* fref_lg = nullptr; if (!logits_ref.empty()) { fref_lg = std::fopen(logits_ref.c_str(), "rb"); if (!fref_lg) { std::perror("fopen logits_ref"); return 1; } fcntl(fileno(fref_lg), F_NOCACHE, 1); } FILE* fref2_lg = nullptr; if (!logits_ref2.empty()) { fref2_lg = std::fopen(logits_ref2.c_str(), "rb"); if (!fref2_lg) { std::perror("fopen logits_ref2"); return 1; } fcntl(fileno(fref2_lg), F_NOCACHE, 1); }
strata::spec::SuffixDrafter suffix(4, 32, (size_t) max_cells); for (int64_t tid : prompt_ids) suffix.append((int32_t) tid);
// Prefill all but the last prompt token via prompt_run_rows const int64_t n_pre = std::max<int64_t>(0, (int64_t) prompt_ids.size() - 1); const double t_pre0 = now_ms(); if (!load_session.empty()) { // the prefix from disk instead of a prefill: it must be this prompt's first n_pre tokens std::vector<int32_t> ids; int64_t cells = 0; if (!strata::metal_engine::window_session_load(load_session.c_str(), ids, &cells, err)) { std::fprintf(stderr, "strata_window_cli: %s\n", err.c_str()); return 1; } bool same = (std::getenv("STRATA_BENCH_EXPAND_CELLS") != nullptr) || (cells == n_pre && (int64_t) ids.size() == n_pre); if (!std::getenv("STRATA_BENCH_EXPAND_CELLS")) { for (int64_t i = 0; same && i < n_pre; ++i) same = ids[(size_t) i] == (int32_t) prompt_ids[(size_t) i]; } if (!same) { std::fprintf(stderr, "strata_window_cli: the session holds %lld cells of another prefix (want this prompt's first %lld)\n", (long long) cells, (long long) n_pre); return 1; } std::printf("STRATA_SESSION load_ms=%.2f cells=%lld bytes=%zu\n", now_ms() - t_pre0, (long long) cells, scr::fixed_bytes() + scr::cell_bytes((size_t) cells)); } else if (n_pre > 0) { auto on_chunk = [&](const float*, int64_t k, int64_t p0, std::string& e) -> bool { if (!use_mtp) return true; std::vector<int32_t> nxt((size_t) k); for (int64_t t = 0; t < k; ++t) nxt[(size_t) t] = (int32_t) prompt_ids[(size_t) (p0 + t + 1)]; return strata::metal_engine::mtp_prefill(nxt.data(), k, p0, e); }; if (!strata::metal_engine::prompt_run_rows(prompt_ids.data(), n_pre, 0, on_chunk, err)) { std::fprintf(stderr, "strata_window_cli: prompt_run_rows failed: %s\n", err.c_str()); return 1; } } if (!save_session.empty()) { std::vector<int32_t> ids(prompt_ids.begin(), prompt_ids.begin() + n_pre); const double t_s = now_ms(); if (!strata::metal_engine::window_session_save(save_session.c_str(), ids.data(), n_pre, err)) { std::fprintf(stderr, "strata_window_cli: %s\n", err.c_str()); return 1; } std::printf("STRATA_SESSION save_ms=%.2f cells=%lld\n", now_ms() - t_s, (long long) n_pre); } // Run the last prompt token with head=true to get the first output distribution int32_t cur_tok = (int32_t) prompt_ids.back(); int32_t next_tok = 0; int64_t pos = std::getenv("STRATA_BENCH_EXPAND_CELLS") ? std::atoll(std::getenv("STRATA_BENCH_EXPAND_CELLS")) : n_pre; if (!strata::metal_engine::window_run(w, 1, &cur_tok, pos, true, &next_tok, err) || !strata::metal_engine::window_commit(w, 1, err)) { std::fprintf(stderr, "strata_window_cli: first window_run failed: %s\n", err.c_str()); return 1; } const double prefill_ms = now_ms() - t_pre0; int32_t mtp_drafts[8] = {}; float mtp_probs[8] = {}; int n_mtp_drafts = 0; double mtp_ms = 0.0; if (use_mtp && !tf_mode && n_gen > 1) { const double tm0 = now_ms(); if (!strata::metal_engine::mtp_draft(1, &next_tok, pos, 0, spec - 1, mtp_drafts, mtp_probs, spec_min_p, &n_mtp_drafts, err)) { std::fprintf(stderr, "strata_window_cli: mtp_draft(0) failed: %s\n", err.c_str()); return 1; } mtp_ms += now_ms() - tm0; } ++pos;
std::vector<float> cur_lg(tf_mode ? (size_t) n_vocab : 0); std::vector<float> ref_lg((fref_lg || fref2_lg) ? (size_t) n_vocab : 0); double sum_nll_cand = 0.0, sum_nll_ref = 0.0, sum_nll_ref2 = 0.0; std::vector<double> klds, klds2; int top1_matches = 0, top1_matches2 = 0;
auto eval_tf_step = [&](int step_idx) -> bool { if (!strata::metal_engine::window_copy_logits(w, 0, cur_lg.data())) return false; if (fout_lg) { if (std::fwrite(cur_lg.data(), sizeof(float), (size_t) n_vocab, fout_lg) != (size_t) n_vocab) return false; } const int32_t tgt = (int32_t) all_ids[(size_t) n_prompt + (size_t) step_idx]; int32_t cand_top1 = 0; const double lse_cand = logsumexp_f64(cur_lg.data(), n_vocab, &cand_top1); sum_nll_cand += -((double) cur_lg[(size_t) tgt] - lse_cand);
if (fref_lg) { if (std::fread(ref_lg.data(), sizeof(float), (size_t) n_vocab, fref_lg) != (size_t) n_vocab) return false; int32_t ref_top1 = 0; const double lse_ref = logsumexp_f64(ref_lg.data(), n_vocab, &ref_top1); sum_nll_ref += -((double) ref_lg[(size_t) tgt] - lse_ref); if (cand_top1 == ref_top1) ++top1_matches; double kld = 0.0; for (int64_t v = 0; v < n_vocab; ++v) { const double log_p = (double) ref_lg[(size_t) v] - lse_ref; const double log_q = (double) cur_lg[(size_t) v] - lse_cand; const double p = std::exp(log_p); if (p > 0.0) kld += p * (log_p - log_q); } klds.push_back(std::max(0.0, kld)); } if (fref2_lg) { if (std::fread(ref_lg.data(), sizeof(float), (size_t) n_vocab, fref2_lg) != (size_t) n_vocab) return false; int32_t ref2_top1 = 0; const double lse_ref2 = logsumexp_f64(ref_lg.data(), n_vocab, &ref2_top1); sum_nll_ref2 += -((double) ref_lg[(size_t) tgt] - lse_ref2); if (cand_top1 == ref2_top1) ++top1_matches2; double kld2 = 0.0; for (int64_t v = 0; v < n_vocab; ++v) { const double log_p = (double) ref_lg[(size_t) v] - lse_ref2; const double log_q = (double) cur_lg[(size_t) v] - lse_cand; const double p = std::exp(log_p); if (p > 0.0) kld2 += p * (log_p - log_q); } klds2.push_back(std::max(0.0, kld2)); } return true; };
if (tf_mode && n_gen > 0) { if (!eval_tf_step(0)) { std::fprintf(stderr, "strata_window_cli: eval_tf_step(0) failed\n"); return 1; } }
std::vector<int32_t> gen_tokens; gen_tokens.push_back(next_tok); suffix.append(next_tok); cur_tok = tf_mode ? (int32_t) all_ids[(size_t) n_prompt] : next_tok;
strata::metal_engine::window_reset_stats(w); const int decode_steps = std::max(0, n_gen - 1); const double t_dec0 = now_ms(); int produced = 0, spec_rounds = 0, spec_drafted = 0, spec_accepted = 0; SpecControl spec_control; while (produced < decode_steps) { int32_t win_in[8] = {cur_tok}; int32_t win_out[8] = {}; int nd = 0; const int rem = decode_steps - produced; nd = spec_control.propose(suffix, win_in, mtp_drafts, mtp_probs, !tf_mode && use_mtp ? n_mtp_drafts : 0, !tf_mode ? suffix_draft : 0, rem, pos, max_cells); const int T = 1 + nd; const double verify_begin = now_ms(); if (!strata::metal_engine::window_run(w, T, win_in, pos, true, win_out, err)) { std::fprintf(stderr, "strata_window_cli: decode step failed: %s\n", err.c_str()); return 1; } spec_control.gate.spec_gate_forward(T, (now_ms() - verify_begin) / 1000); const int n_keep = strata::spec::suffix_keep(win_in, win_out, nd, false); spec_control.result(nd, n_keep); if (!strata::metal_engine::window_commit(w, n_keep, err)) { std::fprintf(stderr, "strata_window_cli: window_commit failed: %s\n", err.c_str()); return 1; } ++spec_rounds; spec_drafted += nd; spec_accepted += (n_keep - 1); if (tf_mode) { if (!eval_tf_step(produced + 1)) { std::fprintf(stderr, "strata_window_cli: eval_tf_step(%d) failed\n", produced + 1); return 1; } gen_tokens.push_back(win_out[0]); ++produced; if (produced < decode_steps) { cur_tok = (int32_t) all_ids[(size_t) n_prompt + (size_t) produced]; } pos += 1; } else { for (int i = 0; i < n_keep && produced < decode_steps; ++i) { gen_tokens.push_back(win_out[i]); suffix.append(win_out[i]); cur_tok = win_out[i]; ++produced; } if (use_mtp && produced < decode_steps) { const double tm0 = now_ms(); if (!strata::metal_engine::mtp_draft(T, win_out, pos, n_keep - 1, spec - 1, mtp_drafts, mtp_probs, spec_min_p, &n_mtp_drafts, err)) { std::fprintf(stderr, "strata_window_cli: mtp_draft failed: %s\n", err.c_str()); return 1; } const double draft_ms = now_ms() - tm0; mtp_ms += draft_ms; spec_control.gate.spec_gate_draft_cost(0, draft_ms / (1000 * std::max(1, spec - 1))); } pos += n_keep; } } if (!tf_mode && (use_mtp || suffix_draft > 0)) spec_control.describe(); const double decode_ms = now_ms() - t_dec0; const double decode_tok_s = decode_steps > 0 ? (decode_steps * 1000.0 / decode_ms) : 0.0; const double decode_ms_tok = decode_steps > 0 ? (decode_ms / decode_steps) : 0.0; const double prefill_tok_s = prompt_ids.size() * 1000.0 / prefill_ms;
if (fout_lg) std::fclose(fout_lg); if (fref_lg) std::fclose(fref_lg); if (fref2_lg) std::fclose(fref2_lg);
std::printf("STRATA_RESULT init_ms=%.2f prompt_tokens=%zu prefill_ms=%.2f prefill_tok_s=%.2f " "decode_tokens=%d decode_ms=%.2f decode_tok_s=%.2f decode_ms_per_tok=%.2f\n", init_ms, prompt_ids.size(), prefill_ms, prefill_tok_s, decode_steps, decode_ms, decode_tok_s, decode_ms_tok); if (spec_rounds > 0 && (use_mtp || suffix_draft > 0)) { std::printf("STRATA_SPEC rounds=%d drafted=%d accepted=%d accept_pct=%.1f tok_per_round=%.2f mtp_ms=%.2f\n", spec_rounds, spec_drafted, spec_accepted, spec_drafted > 0 ? 100.0 * (double) spec_accepted / (double) spec_drafted : 0.0, (double) decode_steps / (double) spec_rounds, mtp_ms); } if (tf_mode && n_gen > 0) { const double cand_ppl = std::exp(sum_nll_cand / n_gen); if (!klds.empty()) { const double ref_ppl = std::exp(sum_nll_ref / n_gen); double mean_kld = 0.0; for (double x : klds) mean_kld += x; mean_kld /= klds.size(); std::vector<double> sk = klds; std::sort(sk.begin(), sk.end()); const double med_kld = sk[sk.size() / 2]; const double p90_kld = sk[(size_t) (0.90 * (sk.size() - 1))]; const double top1_pct = 100.0 * (double) top1_matches / (double) n_gen; std::printf("STRATA_KLD eval_tokens=%d ref_ppl=%.4f cand_ppl=%.4f delta_ppl=%.4f mean_kld=%.6f median_kld=%.6f p90_kld=%.6f top1_agree_pct=%.2f\n", n_gen, ref_ppl, cand_ppl, cand_ppl - ref_ppl, mean_kld, med_kld, p90_kld, top1_pct); if (!klds2.empty()) { const double ref2_ppl = std::exp(sum_nll_ref2 / n_gen); double mean_kld2 = 0.0; for (double x : klds2) mean_kld2 += x; mean_kld2 /= klds2.size(); std::vector<double> sk2 = klds2; std::sort(sk2.begin(), sk2.end()); const double med_kld2 = sk2[sk2.size() / 2]; const double p90_kld2 = sk2[(size_t) (0.90 * (sk2.size() - 1))]; const double top1_pct2 = 100.0 * (double) top1_matches2 / (double) n_gen; std::printf("STRATA_KLD_REF2 eval_tokens=%d ref_ppl=%.4f cand_ppl=%.4f delta_ppl=%.4f mean_kld=%.6f median_kld=%.6f p90_kld=%.6f top1_agree_pct=%.2f\n", n_gen, ref2_ppl, cand_ppl, cand_ppl - ref2_ppl, mean_kld2, med_kld2, p90_kld2, top1_pct2); } } else { std::printf("STRATA_PPL eval_tokens=%d cand_ppl=%.4f\n", n_gen, cand_ppl); } } std::printf("GEN_IDS:"); for (size_t i = 0; i < gen_tokens.size(); ++i) { std::printf(" %d", gen_tokens[i]); } std::printf("\n"); std::fflush(stdout);
strata::metal_engine::window_destroy(w); return 0;}