--- 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;