--- 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 +#include +#include +#include +#include + +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 #include @@ -31,7 +32,10 @@ #include #include #include +#include +#include #include +#include #include #include #include @@ -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& 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& 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& secs, bool write, std::string& err) { + struct Piece { uint8_t* p; size_t n; off_t off; }; + std::vector 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 next{0}; + std::atomic 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 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& secs, const std::vector& gst, const std::vector& ghist, + const mtl::Buffer& ple, const std::vector* 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* 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 ds; + nb_ = false; + fold(ds, pend_fold_, 1 - hp_); + cx_->run_graph(ds); + pend_fold_ = 0; + } + const std::vector* q[14] = {&kq_, &vq_, &ksc_, &vsc_, &tail_, &dead_, &pooled_}; + if (mtp_on_) { + const std::vector* m[7] = {&mkq_, &mvq_, &mksc_, &mvsc_, &mtail_, &mdead_, &mpooled_}; + std::copy(m, m + 7, q + 7); + } + std::vector 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& 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* q[14] = {&kq_, &vq_, &ksc_, &vsc_, &tail_, &dead_, &pooled_}; + if (h.flags & 2u) { + const std::vector* m[7] = {&mkq_, &mvq_, &mksc_, &mvsc_, &mtail_, &mdead_, &mpooled_}; + std::copy(m, m + 7, q + 7); } + std::vector 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 L[5]; + auto add = [&](const std::vector& k, const std::vector& v, const std::vector& ks, + const std::vector& vs, const std::vector& 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 ok{true}; + std::vector 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(), L[1][j]->as(), L[2][j]->as(), L[3][j]->as(), + L[4][j]->as(), 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& 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 #include @@ -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& 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 read_ids(const std::string& path) { + std::vector 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 lp; // [n][n_vocab] + double nll = 0; + double ms = 0; +}; +static bool tf_run(void* w, const std::vector& inputs, const std::vector& 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 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 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 A = read_ids(pa), B = read_ids(pb), D = read_ids(pd); + std::vector 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 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& ids, int64_t pos) { + std::vector 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 a, const std::vector& 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 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 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 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 [--ple-gguf ] " "[--n-expert 256] [--ctx 4096] [--n-prompt 22] [--n-gen 64] [--tokens-file ] " - "[--logits-out ] [--logits-ref ]\n"); + "[--logits-out ] [--logits-ref ] [--save-session | --load-session ] " + "[--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(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 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 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;