Something went wrong. Try again.
strata but optimized for m5 (and maybe more later)
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144// metal/src/engine/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,728constexpr size_t GHIST_BYTES = (size_t) 4 * CH * 3; // 122,880constexpr size_t PLE_BYTES = (size_t) 4 * NG_HIST * CH; // 368,640constexpr 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 layerconstexpr 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