Something went wrong. Try again.
strata but optimized for m5 (and maybe more later)
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670--- a/metal/src/engine/scr_state.hpp+++ b/metal/src/engine/scr_state.hpp@@ -0,0 +1,143 @@+// /tmp/opus_scr_probes/scr_state.hpp - the Metal window's session state (window.cpp) as bytes, and the SCR+// relocation of one QSA layer's cells (qsa.metal's cache layout). Header-only: shared by the window patch and the+// CPU probe, so the probe tests the code the engine would run.+//+// window.cpp's state, per its init (L554-567) and checkpoint_save (L1688):+// 36 GDN layers (l % 4 != 3): gst_ f32 [128 i][48 h][128 j] = 3 MiB, ghist_[hp_] f32 [10240 ch][3] = 120 KiB+// pleHist_ f32 [10240 ch][9] = 360 KiB, and ss.ple_prev[2] (the PLE n-gram token window)+// 12 QSA layers (l % 4 == 3), per cell c and kv head h (2):+// kq_/vq_ int8 [c][h][256], ksc_/vsc_ half [c][h][4] (one scale per 64 dims)+// pooled_ f32 [b][128] for complete blocks b = c / 4, tail_ f32 [4][128] (raw keys, slot c % 4), dead_ f32 [128]+// qsa_prep (qsa.metal L104-147) applies NEOX RoPE at position = cell (pos_base is 0 on every path) to dims 0..63+// only, pairs (i, i + 32), i < 32, angle cell * base^(-2i/64); k is rotated BEFORE its int8 quantization, so the+// rotated dims are exactly scale group 0 of each row. pooled[b] is rotated at position 4 b, also dims 0..63.+#pragma once+#include <algorithm>+#include <cmath>+#include <cstddef>+#include <cstdint>+#include <cstring>++namespace scr {++constexpr int NGDN = 36, NQSA = 12, QKV = 2, QD = 256, QROT = 64, ID = 128, NG_HIST = 9, CH = 10240;+constexpr size_t GST_BYTES = (size_t) 4 * 128 * 48 * 128; // 3,145,728+constexpr size_t GHIST_BYTES = (size_t) 4 * CH * 3; // 122,880+constexpr size_t PLE_BYTES = (size_t) 4 * NG_HIST * CH; // 368,640+constexpr size_t TAIL_BYTES = (size_t) 4 * 4 * ID, DEAD_BYTES = (size_t) 4 * ID;+constexpr size_t KQ_CELL = (size_t) QKV * QD; // int8 bytes per cell (kq or vq)+constexpr size_t SC_CELL = (size_t) QKV * 4 * 2; // half bytes per cell (ksc or vsc)+constexpr size_t POOL_BLOCK = (size_t) 4 * ID; // f32 bytes per pooled block (4 cells)++// bytes of a saved session: the running state, then the positional cells [0, n) of every QSA layer+constexpr size_t fixed_bytes() { return NGDN * (GST_BYTES + GHIST_BYTES) + PLE_BYTES + NQSA * (TAIL_BYTES + DEAD_BYTES); }+constexpr size_t cell_bytes(size_t n) { return NQSA * (n * (2 * KQ_CELL + 2 * SC_CELL) + (n / 4) * POOL_BLOCK); }++// the file: this header, the token ids, then the sections in a fixed order, each starting on a 16 KiB boundary+// (the page size: F_NOCACHE reads into page-aligned MTLBuffer storage go straight to the device)+constexpr size_t ALIGN = 16384;+inline size_t align_up(size_t x) { return (x + ALIGN - 1) / ALIGN * ALIGN; }+struct SessionHeader {+ char magic[8] = {'S', 'T', 'R', 'S', 'E', 'S', 'S', '1'};+ uint32_t version = 1;+ uint32_t flags = 0; // bit 0: PLE on, bit 1: MTP sections follow+ int64_t n_cells = 0; // positions [0, n_cells) are in the cache; the next token runs at n_cells+ int64_t n_ids = 0; // the token ids of those positions (the harness's prefix check)+ uint64_t model_bytes = 0; // the model shard's size and its tensor count: a session is one model's+ uint64_t model_tensors = 0;+ double rope_base = 0;+ int32_t ple_prev[2] = {-1, -1};+ uint64_t payload_bytes = 0; // everything after the header and ids, for the size check before any copy+};+static_assert(sizeof(SessionHeader) % 8 == 0, "packed as written");++// ---- fp16 and the qsa_prep int8 rule (L136-147): per 64 dims sh = half(amax / 127), q = clamp(rint(x / sf))+inline float h2f(uint16_t h) { _Float16 v; std::memcpy(&v, &h, 2); return (float) v; }+inline uint16_t f2h(float f) { _Float16 v = (_Float16) f; uint16_t h; std::memcpy(&h, &v, 2); return h; }+inline void quant64(const float* x, int8_t* q, uint16_t* sc) {+ float m = 0.0f;+ for (int i = 0; i < 64; ++i) m = std::max(m, std::fabs(x[i]));+ *sc = f2h(m / 127.0f);+ const float sf = h2f(*sc);+ for (int i = 0; i < 64; ++i) q[i] = (int8_t) (sf > 0.0f ? std::clamp((int) std::nearbyint(x[i] / sf), -127, 127) : 0);+}++// the rotation by delta positions: cos/sin of delta * theta_i in double, as window.cpp's table (L545-551)+struct Rot {+ float c[32], s[32];+ Rot(double base, int64_t delta) {+ for (int i = 0; i < 32; ++i) {+ const double th = (double) delta * std::pow(base, -2.0 * i / 64.0);+ c[i] = (float) std::cos(th);+ s[i] = (float) std::sin(th);+ }+ }+};+// the first 64 dims of a row rotated in place (norm_rope's pairs: a = x[i], b = x[i + 32])+inline void rotate64(float* x, const Rot& r) {+ for (int i = 0; i < 32; ++i) {+ const float a = x[i], b = x[i + 32];+ x[i] = a * r.c[i] - b * r.s[i];+ x[i + 32] = a * r.s[i] + b * r.c[i];+ }+}++// One QSA layer's span of n cells moved from c_src to c_dst = c_src + delta, keys re-rotated by delta.+// kq: only group 0 (dims 0..63, scale ksc[.][.][0]) is rotated: dequant, rotate, requant with a fresh scale.+// dims 64..255 and scales 1..3 move as bytes. vq / vsc carry no position: bytes.+// pooled: the complete blocks inside the span, b -> b + delta / 4, rotated by delta (4 b -> 4 b + delta). delta must+// be a multiple of 4, else every block of the span would mix different cells and the raw keys aren't kept.+// The junction block (the one holding c_dst when c_dst % 4 != 0) is not moved: it mixes the new span before it,+// so the caller recomputes the span's first (4 - c_dst % 4) % 4 cells or more (scr_min_junction).+// Cells are walked away from the overlap (ascending when moving down), so source rows are read before they're hit.+inline int64_t scr_min_junction(int64_t c_dst) { return (4 - (c_dst & 3)) & 3; }+inline bool relocate_layer(int8_t* kq, int8_t* vq, uint16_t* ksc, uint16_t* vsc, float* pooled, int64_t c_src,+ int64_t c_dst, int64_t n, const Rot& rot) {+ const int64_t delta = c_dst - c_src;+ if (delta == 0 || n <= 0) return true;+ if ((delta & 3) != 0 || c_dst < 0) return false;+ auto cell = [&](int64_t k) {+ const int64_t s = c_src + k, d = c_dst + k;+ for (int h = 0; h < QKV; ++h) {+ const size_t rs = (size_t) s * QKV + h, rd = (size_t) d * QKV + h;+ float x[64];+ const float sf = h2f(ksc[rs * 4]);+ for (int i = 0; i < 64; ++i) x[i] = (float) kq[rs * QD + i] * sf;+ rotate64(x, rot);+ std::memmove(kq + rd * QD + 64, kq + rs * QD + 64, QD - 64);+ std::memmove(ksc + rd * 4 + 1, ksc + rs * 4 + 1, 3 * 2);+ quant64(x, kq + rd * QD, ksc + rd * 4);+ std::memmove(vq + rd * QD, vq + rs * QD, QD);+ std::memmove(vsc + rd * 4, vsc + rs * 4, 4 * 2);+ }+ };+ const int64_t b0 = (c_src + 3) / 4, b1 = (c_src + n) / 4, db = delta / 4; // complete blocks [b0, b1)+ auto block = [&](int64_t b) {+ float x[ID];+ std::memcpy(x, pooled + (size_t) b * ID, sizeof x);+ rotate64(x, rot);+ std::memcpy(pooled + (size_t) (b + db) * ID, x, sizeof x);+ };+ if (delta < 0) {+ for (int64_t k = 0; k < n; ++k) cell(k);+ for (int64_t b = b0; b < b1; ++b) block(b);+ } else {+ for (int64_t k = n - 1; k >= 0; --k) cell(k);+ for (int64_t b = b1 - 1; b >= b0; --b) block(b);+ }+ return true;+}++// ---- GDN staleness, the part with a closed form. gdn.metal (L89-119): S <- g S + k delta^T, delta = (v - g S^T k) b,+// i.e. S' = g (I - b k k^T) S + b k v^T with |k| <= 1 (k = cv / sqrt(sum cv^2 + eps)) and b = sigmoid in (0, 1). Two+// runs over the same C inputs from states S_1, S_2 differ by D' = g (I - b k k^T) D, and |I - b k k^T|_2 = 1, so+// |D_end|_F <= exp(sum_t log g_t) |D_0|_F, log g_t = softplus(alpha_t + dt_h) * ssm_a_h (<= 0)+// per value head. That bounds the direct term of keeping S_ABC instead of S_AB'C; the indirect term (C's own k, v,+// gate inputs computed under B) has no closed form and is what the splice test measures.+inline double log_gate(float alpha, float dt, float ssm_a) {+ const double v = (double) alpha + dt;+ const double sp = v > 20.0 ? v : std::log1p(std::exp(v));+ return sp * (double) ssm_a;+}++} // namespace scr--- a/metal/src/engine/window.cpp+++ b/metal/src/engine/window.cpp@@ -22,6 +22,7 @@ #include "strata/core/conversation_cache.hpp" #include "strata/core/session.hpp" #include "strata/kernels/ngram.hpp"+#include "scr_state.hpp" #include <algorithm> #include <array>@@ -31,7 +32,10 @@ #include <cstdlib> #include <cstring> #include <cerrno>+#include <atomic>+#include <fcntl.h> #include <filesystem>+#include <sys/stat.h> #include <functional> #include <signal.h> #include <unistd.h>@@ -135,6 +139,12 @@ // PLE history, each QSA layer's indexer tail and tail-block key); the KV and pooled keys are positional bool checkpoint_save(strata::core::ConversationCheckpoint& c, std::string& err); bool checkpoint_restore(const strata::core::ConversationCheckpoint& c, std::string& err);+ // a disk session: the running state of checkpoint_save, ss.ple_prev and the positional cells [0, n_cells) of every+ // QSA layer (and the MTP layer's), read and written straight into the buffers (scr_state.hpp's file layout)+ bool session_save(const std::string& path, const int32_t* ids, int64_t n_cells, std::string& err);+ bool session_load(const std::string& path, std::vector<int32_t>& ids, int64_t& n_cells, std::string& err);+ // SCR: every QSA layer's cells [c_src, c_src + n) moved to c_dst, keys and pooled keys re-rotated by c_dst - c_src+ bool scr_relocate(int64_t c_src, int64_t c_dst, int64_t n, std::string& err); void reset(); // a new sequence (session_zero): the running state zeroed void set_max_t(int t) { max_t_ = std::min(t, MAXT); } // ---- the MTP draft layer (mtp.hpp's math on mac-kern's kernels; the MTP GGUF from tools/mtp_pack.py)@@ -222,6 +232,7 @@ int64_t max_cells_ = 0; int64_t n_vocab_ = 0; float eps_ = 1e-6f;+ double rope_base_ = 1e7; // buffers mtl::Buffer R_, bo_, inj_[2], xn_, lo_, qkv_, z_, al_, be_, yg_, logits_, ids_, wts_, plan_, h_, ye_, gsh_, ush_, ysh_, sg_, qfull_, kcur_, vcur_, iqr_, ikr_, qcur_, qidx_, scores_, qids_, part_, cs_, lm_, apv_, api_, tokout_,@@ -543,6 +554,7 @@ XA_ = act(N); XF_ = act(N); XY_ = act(6144); XQ_ = act(6144); XH_ = act(NFF); XE_ = act(N); // RoPE: cos/sin [pos][32] for the first 64 dims, NEOX pairs (i, i + 32), in double as build_rope_table const double base = f_.meta_num.count("qwen4exp.rope.freq_base") ? f_.meta_num["qwen4exp.rope.freq_base"] : 1e7;+ rope_base_ = base; for (int p = 0; p < MAXC + 16; ++p) for (int i = 0; i < 32; ++i) { const double th = (double) p * std::pow(base, -2.0 * i / 64.0);@@ -1748,8 +1760,174 @@ for (size_t j = 0; j < tail_.size(); ++j) { std::memcpy(tail_[j].ptr, c.tails.data() + j * ts, ts); std::memcpy(dead_[j].ptr, c.dead.data() + j * dsz, dsz);+ }+ return true;+}++namespace {+// one contiguous range of a buffer in the session file+struct SessionSec {+ void* p;+ size_t bytes;+ size_t off = 0;+};+// the sections in file order from `first` on, each on a 16 KiB boundary; returns the end+size_t session_place(std::vector<SessionSec>& secs, size_t first) {+ size_t o = scr::align_up(first);+ for (auto& s : secs) { s.off = o; o = scr::align_up(o + s.bytes); }+ return o;+}+// pread / pwrite of every section, 8 MiB pieces over 4 threads (one stream does not fill the SSD's queue)+bool session_io(int fd, const std::vector<SessionSec>& secs, bool write, std::string& err) {+ struct Piece { uint8_t* p; size_t n; off_t off; };+ std::vector<Piece> ps;+ constexpr size_t PIECE = (size_t) 8 << 20;+ for (const auto& s : secs)+ for (size_t o = 0; o < s.bytes; o += PIECE)+ ps.push_back({(uint8_t*) s.p + o, std::min(PIECE, s.bytes - o), (off_t) (s.off + o)});+ std::atomic<size_t> next{0};+ std::atomic<bool> ok{true};+ auto work = [&] {+ for (size_t i; ok && (i = next++) < ps.size();) {+ const Piece& q = ps[i];+ for (size_t d = 0; d < q.n;) {+ const ssize_t r = write ? ::pwrite(fd, q.p + d, q.n - d, q.off + (off_t) d) : ::pread(fd, q.p + d, q.n - d, q.off + (off_t) d);+ if (r <= 0) { if (r < 0 && errno == EINTR) continue; ok = false; break; }+ d += (size_t) r;+ }+ }+ };+ std::vector<std::thread> th;+ for (int t = 0; t < 4; ++t) th.emplace_back(work);+ for (auto& t : th) t.join();+ if (!ok) err = std::string("metal window: session ") + (write ? "write" : "read") + " failed: " + std::strerror(errno);+ return ok;+}+} // namespace++// the session's buffers in file order; the cells [0, n) of the positional caches+static void session_sections(std::vector<SessionSec>& secs, const std::vector<mtl::Buffer>& gst, const std::vector<mtl::Buffer>& ghist,+ const mtl::Buffer& ple, const std::vector<mtl::Buffer>* const* q, int64_t n) {+ for (size_t i = 0; i < gst.size(); ++i) {+ secs.push_back({gst[i].ptr, scr::GST_BYTES});+ secs.push_back({ghist[i].ptr, scr::GHIST_BYTES});+ }+ secs.push_back({ple.ptr, scr::PLE_BYTES});+ const size_t un = (size_t) n;+ for (int set = 0; set < 2; ++set) { // the 12 QSA layers, then the MTP layer's (when loaded)+ const std::vector<mtl::Buffer>* const* b = q + (set ? 7 : 0);+ if (set && b[0] == nullptr) break;+ for (size_t j = 0; j < b[0]->size(); ++j) {+ secs.push_back({(*b[0])[j].ptr, un * scr::KQ_CELL});+ secs.push_back({(*b[1])[j].ptr, un * scr::KQ_CELL});+ secs.push_back({(*b[2])[j].ptr, un * scr::SC_CELL});+ secs.push_back({(*b[3])[j].ptr, un * scr::SC_CELL});+ secs.push_back({(*b[4])[j].ptr, scr::TAIL_BYTES});+ secs.push_back({(*b[5])[j].ptr, scr::DEAD_BYTES});+ secs.push_back({(*b[6])[j].ptr, un / 4 * scr::POOL_BLOCK});+ }+ }+}++bool Window::session_save(const std::string& path, const int32_t* ids, int64_t n_cells, std::string& err) {+ if (n_cells < 0 || n_cells >= max_cells_) { err = "metal window: session cell count out of range"; return false; }+ if (mrope_on()) { err = "metal window: sessions of M-RoPE (vision) prompts are not supported"; return false; }+ if (pend_fold_ > 0) { // as checkpoint_save: the state after the last commit's fold+ std::vector<mtl::Dispatch> ds;+ nb_ = false;+ fold(ds, pend_fold_, 1 - hp_);+ cx_->run_graph(ds);+ pend_fold_ = 0;+ }+ const std::vector<mtl::Buffer>* q[14] = {&kq_, &vq_, &ksc_, &vsc_, &tail_, &dead_, &pooled_};+ if (mtp_on_) {+ const std::vector<mtl::Buffer>* m[7] = {&mkq_, &mvq_, &mksc_, &mvsc_, &mtail_, &mdead_, &mpooled_};+ std::copy(m, m + 7, q + 7);+ }+ std::vector<SessionSec> secs;+ session_sections(secs, gst_, ghist_[hp_], pleHist_, q, n_cells);+ scr::SessionHeader h;+ h.flags = (ple_on_ ? 1u : 0u) | (mtp_on_ ? 2u : 0u);+ h.n_cells = n_cells;+ h.n_ids = ids ? n_cells : 0;+ h.model_bytes = f_.file_bytes;+ h.model_tensors = f_.tensors.size();+ h.rope_base = rope_base_;+ h.ple_prev[0] = ss_->ple_prev[0];+ h.ple_prev[1] = ss_->ple_prev[1];+ const size_t head = sizeof h + (size_t) h.n_ids * 4, end = session_place(secs, head);+ h.payload_bytes = end - head;+ const std::string tmp = path + ".partial";+ const int fd = ::open(tmp.c_str(), O_WRONLY | O_CREAT | O_TRUNC, 0644);+ if (fd < 0) { err = "metal window: cannot create " + tmp + ": " + std::strerror(errno); return false; }+ bool ok = ::ftruncate(fd, (off_t) end) == 0 && ::pwrite(fd, &h, sizeof h, 0) == (ssize_t) sizeof h &&+ (h.n_ids == 0 || ::pwrite(fd, ids, (size_t) h.n_ids * 4, sizeof h) == (ssize_t) (h.n_ids * 4));+ if (!ok) err = "metal window: session header write failed: " + std::string(std::strerror(errno));+ ok = ok && session_io(fd, secs, true, err);+ ok = (::close(fd) == 0) && ok;+ if (ok && std::rename(tmp.c_str(), path.c_str()) != 0) { err = "metal window: rename to " + path + " failed"; ok = false; }+ if (!ok) ::unlink(tmp.c_str());+ return ok;+}++bool Window::session_load(const std::string& path, std::vector<int32_t>& ids, int64_t& n_cells, std::string& err) {+ const int fd = ::open(path.c_str(), O_RDONLY);+ if (fd < 0) { err = "metal window: cannot open " + path + ": " + std::strerror(errno); return false; }+ ::fcntl(fd, F_NOCACHE, 1);+ scr::SessionHeader h, want;+ struct stat st{};+ bool ok = ::fstat(fd, &st) == 0 && ::pread(fd, &h, sizeof h, 0) == (ssize_t) sizeof h;+ if (!ok || std::memcmp(h.magic, want.magic, 8) != 0 || h.version != want.version) err = "metal window: " + path + " is not a session file";+ else if (h.model_bytes != f_.file_bytes || h.model_tensors != f_.tensors.size() || h.rope_base != rope_base_)+ err = "metal window: the session is another model's";+ else if (h.n_cells < 0 || h.n_cells >= max_cells_) err = "metal window: the session is longer than this context (--ctx)";+ else if ((h.flags & 1u) != (ple_on_ ? 1u : 0u) || ((h.flags & 2u) && !mtp_on_)) err = "metal window: the session's PLE / MTP setup differs";+ if (!err.empty()) { ::close(fd); return false; }+ const std::vector<mtl::Buffer>* q[14] = {&kq_, &vq_, &ksc_, &vsc_, &tail_, &dead_, &pooled_};+ if (h.flags & 2u) {+ const std::vector<mtl::Buffer>* m[7] = {&mkq_, &mvq_, &mksc_, &mvsc_, &mtail_, &mdead_, &mpooled_};+ std::copy(m, m + 7, q + 7); }+ std::vector<SessionSec> secs;+ session_sections(secs, gst_, ghist_[hp_], pleHist_, q, h.n_cells);+ const size_t head = sizeof h + (size_t) h.n_ids * 4, end = session_place(secs, head);+ if (end - head != h.payload_bytes || (size_t) st.st_size != end) { ::close(fd); err = "metal window: the session file is truncated or of another layout"; return false; }+ ids.resize((size_t) h.n_ids);+ ok = h.n_ids == 0 || ::pread(fd, ids.data(), (size_t) h.n_ids * 4, sizeof h) == (ssize_t) (h.n_ids * 4);+ ok = ok && session_io(fd, secs, false, err);+ ::close(fd);+ if (!ok) { if (err.empty()) err = "metal window: session read failed"; reset(); return false; }+ pend_fold_ = 0;+ last_T_ = 0;+ ss_->ple_prev[0] = h.ple_prev[0];+ ss_->ple_prev[1] = h.ple_prev[1];+ n_cells = h.n_cells; return true;+}++bool Window::scr_relocate(int64_t c_src, int64_t c_dst, int64_t n, std::string& err) {+ if (mrope_on()) { err = "metal window: SCR of M-RoPE (vision) cells is not supported"; return false; }+ if (n < 0 || c_src < 0 || c_dst < 0 || c_src + n >= max_cells_ || c_dst + n >= max_cells_) { err = "metal window: SCR span out of range"; return false; }+ if (((c_dst - c_src) & 3) != 0) { err = "metal window: SCR needs a shift by a multiple of 4 cells (pad the new span)"; return false; }+ const scr::Rot rot(rope_base_, c_dst - c_src);+ std::vector<const mtl::Buffer*> L[5];+ auto add = [&](const std::vector<mtl::Buffer>& k, const std::vector<mtl::Buffer>& v, const std::vector<mtl::Buffer>& ks,+ const std::vector<mtl::Buffer>& vs, const std::vector<mtl::Buffer>& p) {+ for (size_t j = 0; j < k.size(); ++j) { L[0].push_back(&k[j]); L[1].push_back(&v[j]); L[2].push_back(&ks[j]); L[3].push_back(&vs[j]); L[4].push_back(&p[j]); }+ };+ add(kq_, vq_, ksc_, vsc_, pooled_);+ if (mtp_on_) add(mkq_, mvq_, mksc_, mvsc_, mpooled_);+ std::atomic<bool> ok{true};+ std::vector<std::thread> th; // the layers are independent: one thread each+ for (size_t j = 0; j < L[0].size(); ++j)+ th.emplace_back([&, j] {+ if (!scr::relocate_layer(L[0][j]->as<int8_t>(), L[1][j]->as<int8_t>(), L[2][j]->as<uint16_t>(), L[3][j]->as<uint16_t>(),+ L[4][j]->as<float>(), c_src, c_dst, n, rot))+ ok = false;+ });+ for (auto& t : th) t.join();+ if (!ok) err = "metal window: SCR relocation refused";+ return ok; } bool Window::commit(int n_keep, std::string& err) {@@ -2977,7 +3155,19 @@ bool window_checkpoint_restore(const strata::core::ConversationCheckpoint& c, std::string& err) { if (g_the_window == nullptr || !g_the_window->ready()) { err = "metal window: not up"; return false; } return g_the_window->checkpoint_restore(c, err);+}+bool window_session_save(const char* path, const int32_t* ids, int64_t n_cells, std::string& err) {+ if (g_the_window == nullptr || !g_the_window->ready()) { err = "metal window: not up"; return false; }+ return g_the_window->session_save(path, ids, n_cells, err); }+bool window_session_load(const char* path, std::vector<int32_t>& ids, int64_t* n_cells, std::string& err) {+ if (g_the_window == nullptr || !g_the_window->ready()) { err = "metal window: not up"; return false; }+ return g_the_window->session_load(path, ids, *n_cells, err);+}+bool window_scr_relocate(int64_t c_src, int64_t c_dst, int64_t n, std::string& err) {+ if (g_the_window == nullptr || !g_the_window->ready()) { err = "metal window: not up"; return false; }+ return g_the_window->scr_relocate(c_src, c_dst, n, err);+} bool window_commit(void* w, int n_keep, std::string& err) { const double t0 = now_ms(); const bool ok = ((Window*) w)->commit(n_keep, err);--- a/metal/src/strata_window_cli.cpp+++ b/metal/src/strata_window_cli.cpp@@ -3,10 +3,12 @@ // 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 "engine/scr_state.hpp" #include <algorithm> #include <chrono>@@ -32,6 +34,12 @@ 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);+bool window_session_load(const char* path, std::vector<int32_t>& ids, int64_t* n_cells, std::string& err);+bool window_scr_relocate(int64_t c_src, int64_t c_dst, int64_t n, std::string& err); } // namespace strata::metal_engine static double now_ms() {@@ -53,6 +61,168 @@ 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, nats+static 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);+ }+ std::remove(tmp.c_str());+ return 0; } int main(int argc, char** argv) {@@ -62,6 +232,8 @@ int n_prompt = 22; int n_gen = 64; int suffix_draft = 0;+ 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];@@ -76,11 +248,22 @@ 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>]\n");+ "[--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; } @@ -140,6 +323,7 @@ return 1; } const double init_ms = now_ms() - t_init0;+ 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()) {@@ -175,12 +359,38 @@ // 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 (n_pre > 0) {+ 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 = cells == n_pre && (int64_t) ids.size() == n_pre;+ 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) { if (!strata::metal_engine::prompt_run_rows(prompt_ids.data(), n_pre, 0, nullptr, 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;