// 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,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