From da42e08213a33b9fd8042fedbede4acea97ca944 Mon Sep 17 00:00:00 2001 From: "nekomimi.pet" Date: Fri, 16 Oct 2026 19:46:35 -0400 Subject: [PATCH] focalor: Strata serves niri (Qwen3.8-Flash-Next, SCR, session files) on the 5070; nemotron-bolt off --- hosts/focalor/default.nix | 5 +- hosts/focalor/modules/strata.nix | 213 +++ hosts/focalor/modules/strata.patch | 1946 ++++++++++++++++++++++++++++ 3 files changed, 2162 insertions(+), 2 deletions(-) create mode 100644 hosts/focalor/modules/strata.nix create mode 100644 hosts/focalor/modules/strata.patch diff --git a/hosts/focalor/default.nix b/hosts/focalor/default.nix index ef4dfef..be6fe77 100644 --- a/hosts/focalor/default.nix +++ b/hosts/focalor/default.nix @@ -32,9 +32,10 @@ in ]; }; - # off: it would take the RTX 5070 from nemotron-bolt + # Strata (niri's Qwen3.8-Flash-Next) has the RTX 5070; llama-cpp and nemotron-bolt would take it modules.llama-cpp.enable = false; - modules.nemotron-bolt.enable = true; + modules.nemotron-bolt.enable = false; + modules.strata.enable = true; system.stateVersion = "25.05"; catppuccin = { diff --git a/hosts/focalor/modules/strata.nix b/hosts/focalor/modules/strata.nix new file mode 100644 index 0000000..6fc30ed --- /dev/null +++ b/hosts/focalor/modules/strata.nix @@ -0,0 +1,213 @@ +{ + config, + lib, + pkgs, + ... +}: + +# Strata (Niko1221/Strata) serving Qwen3.8-Flash-Next 125B-A6B, GSQ-RCO IQ3_XXS, for niri. Upstream 296b59d plus +# branch ana/split-sessions of okami.mom/strata (strata.patch): session files for a layer split and for the whole +# checkpoint chain, suffix cache reuse (an edited prompt moves the cached cells instead of reading them again), pads +# kept beside session files, and unbuffered prompt-path reads of an experts.bin pack. One RTX 5070: the layer split +# with the RTX 4060 clamped every prompt chunk to the 4060's small cache and decoded no faster. 16 GiB of the hottest +# experts are pinned in RAM, the rest are read from /data. Measured 2026-10-17: niri's 62.5K-token system prompt +# reads cold in ~50 s (1,400 t/s), decode 27-28 t/s, a memory edit ~15 s, a compaction ~12 s. +let + cfg = config.modules.strata; + port = 8092; + data = "/data/strata-data"; + model = "${data}/models/IQ3_XXS/Qwen3.8-Flash-Next-GSQ-RCO-IQ3_XXS"; + cuda = pkgs.cudaPackages_13_0; # Strata pins CUDA 13.0: 13.2.0/13.2.1 miscompile for sm_120 + llamaCpp = pkgs.fetchFromGitHub { + # the ggml and mtmd sources setup.py pins (LLAMA_CPP_COMMIT) + owner = "ggml-org"; + repo = "llama.cpp"; + rev = "3cf03257f219afbe7334045ff7c6a06ac68c627d"; + hash = "sha256-SRGoXa+4ACBCB3eaG9XFYhMN1i0FyPEy9Rrer+dFGYI="; + }; + strata = cuda.backendStdenv.mkDerivation { + pname = "strata"; + version = "0.1.44.1-ana"; + src = pkgs.fetchFromGitHub { + owner = "Niko1221"; + repo = "Strata"; + rev = "296b59d2f73c8f04a805df29c2fa97655adace78"; + hash = "sha256-Jfo2uldtyXW8Y/wQxA7WcYwtmzCGAOaXgnlMSo7K9LE="; + }; + patches = [ ./strata.patch ]; + nativeBuildInputs = [ + pkgs.cmake + pkgs.ninja + cuda.cudatoolkit + pkgs.autoAddDriverRunpath + ]; + buildInputs = [ cuda.cudatoolkit ]; + dontUseCmakeConfigure = true; + # STRATA_PORTABLE: ggml's CPU backend for an AVX2 baseline instead of the build host (Strata's own CPU kernels + # are chosen at run time either way); the image encoder runs on the CPU + buildPhase = '' + runHook preBuild + export CUDA_HOME=${cuda.cudatoolkit} CUDAHOSTCXX=${cuda.backendStdenv.cc}/bin/c++ + cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release -DSTRATA_ENABLE_CUDA=ON -DSTRATA_BUILD_TESTS=OFF \ + -DSTRATA_PORTABLE=ON -DCMAKE_CUDA_ARCHITECTURES=120 -DSTRATA_GGML_DIR=${llamaCpp} + cmake --build build --target strata + cmake -S tools/vision -B build-vision -G Ninja -DCMAKE_BUILD_TYPE=Release -DLLAMA_DIR=${llamaCpp} \ + -DSTRATA_VISION_CUDA=OFF -DSTRATA_PORTABLE=ON + cmake --build build-vision --target strata-vision + runHook postBuild + ''; + installPhase = '' + runHook preInstall + mkdir -p $out/share/strata/engine + cp -r serve tools data $out/share/strata/ + cp build/strata build-vision/bin/strata-vision $out/share/strata/engine/ + if [ -e build/libstrata_kq.so ]; then cp build/libstrata_kq.so $out/share/strata/engine/; fi + runHook postInstall + ''; + }; + python = pkgs.python3.withPackages (ps: [ + ps.numpy + ps.jinja2 + ps.regex + ps.pyyaml + ps.tqdm + ps.requests + ps.pillow + ps.psutil + ]); + root = "${strata}/share/strata"; + strataConfig = pkgs.writeText "strata.json" ( + builtins.toJSON { + exe = "${root}/engine/strata"; + args = [ + "--pack" + "${data}/packs/iq3_xxs" + "--native" + "${model}-00001-of-00002.gguf" + "--ple-gguf" + "${model}-00002-of-00002.gguf" + "--expert-profile" + "${root}/data/expert-profile.bin" + "--expert-cache" + "auto" + # 8192-token prompt chunks: every chunk streams most of the experts, so a bigger chunk reads them half as + # often; 700 MiB of reserve is the most that keeps it + "--prefill" + "8192" + "--vram-reserve-mib" + "700" + "--spec" + "4" + "--spec-min-p" + "0.5" + "--mtp" + "${data}/mtp/rt" + "--max-context" + "131072" + "--kv" + "int8" + "--mmap-experts" + "--resident-budget-gib" + "16" + "--vision" + # checkpoints every 4096 read tokens, 16 kept: where a suffix splice restarts its read + "--prompt-cache-every" + "4096" + "--prompt-cache" + "16" + ]; + cwd = data; + tokenizer = "${data}/packs/iq3_xxs/tokenizer"; + model_name = "qwen3.8-flash-next-iq3_xxs"; + log = "${data}/strata.log"; + inherit port; + gpu = 1; # nvidia-smi's numbering: the RTX 5070 + gpus_asked = true; + host = "0.0.0.0"; + draft_vocab = "en"; + open_browser = false; + vision = { + exe = "${root}/engine/strata-vision"; + mmproj = "${data}/models/mmproj-Qwen3.8-Flash-Next-BF16.gguf"; + model = "${model}-00001-of-00002.gguf"; + gpu = false; + max_tokens = 1024; + threads = 8; + }; + slot_save_path = "${data}/sessions"; + env = { + STRATA_SCR = "1"; + STRATA_SESSION_CHAIN = "1"; + STRATA_PROMPT_DIRECT = "1"; + }; + } + ); + api = "http://127.0.0.1:${toString port}"; + # the session niri left at the last stop, else her saved system prompt; a file that does not load is a cold read + restore = pkgs.writeShellScript "strata-restore" '' + for _ in $(seq 600); do + ${pkgs.curl}/bin/curl -fsS ${api}/health | ${pkgs.gnugrep}/bin/grep -q '"loaded": true' && break + sleep 1 + done + for f in niri-live.bin niri-prefix.bin; do + [ -e ${data}/sessions/$f ] || continue + ${pkgs.curl}/bin/curl -fsS --max-time 300 -X POST "${api}/slots/0?action=restore" \ + -H 'Content-Type: application/json' -d "{\"filename\": \"$f\"}" && exit 0 + done + exit 0 + ''; + save = pkgs.writeShellScript "strata-save" '' + ${pkgs.curl}/bin/curl -fsS --max-time 120 -X POST "${api}/slots/0?action=save" \ + -H 'Content-Type: application/json' -d '{"filename": "niri-live.bin"}' || true + ''; +in +{ + options.modules.strata.enable = lib.mkEnableOption "the Strata service (Qwen3.8-Flash-Next for niri)"; + + config = lib.mkIf cfg.enable { + networking.firewall.interfaces.br0.allowedTCPPorts = [ port ]; + + systemd.tmpfiles.rules = [ "d ${data}/sessions 0700 regent users -" ]; + + systemd.services.strata = { + description = "Strata: Qwen3.8-Flash-Next for niri"; + unitConfig.ConditionPathExists = "${data}/packs/iq3_xxs/experts.bin"; + after = [ "network-online.target" ]; + wants = [ "network-online.target" ]; + wantedBy = [ "multi-user.target" ]; + + environment = { + CUDA_DEVICE_ORDER = "PCI_BUS_ID"; + # the driver's libcuda ahead of the toolkit's link stub + LD_LIBRARY_PATH = "/run/opengl-driver/lib"; + PYTHONPATH = "${root}:${root}/tools"; + PYTHONDONTWRITEBYTECODE = "1"; + }; + + serviceConfig = { + Type = "simple"; + User = "regent"; + Group = "users"; + WorkingDirectory = data; + ExecStart = lib.escapeShellArgs [ + "${python}/bin/python" + "${root}/serve/server.py" + "--engine" + "strata" + "--config" + "${strataConfig}" + "--port" + (toString port) + ]; + ExecStartPost = "${restore}"; + ExecStop = "${save}"; + Restart = "on-failure"; + RestartSec = "10s"; + TimeoutStartSec = "infinity"; + TimeoutStopSec = "5min"; + LimitNOFILE = 1048576; + LimitMEMLOCK = "infinity"; + }; + }; + }; +} diff --git a/hosts/focalor/modules/strata.patch b/hosts/focalor/modules/strata.patch new file mode 100644 index 0000000..3fac744 --- /dev/null +++ b/hosts/focalor/modules/strata.patch @@ -0,0 +1,1946 @@ +diff --git a/CMakeLists.txt b/CMakeLists.txt +index e28abd9b..2622929a 100644 +--- a/CMakeLists.txt ++++ b/CMakeLists.txt +@@ -117,6 +117,9 @@ if(STRATA_BUILD_CONVERSATION_TESTS) + add_executable(conversation_memory_test src/core/conversation_memory_test.cpp src/core/conversation_memory.cpp) + target_include_directories(conversation_memory_test PRIVATE include) + add_test(NAME conversation_memory_test COMMAND conversation_memory_test) ++ add_executable(scr_test src/core/scr_test.cpp) ++ target_include_directories(scr_test PRIVATE include) ++ add_test(NAME scr_test COMMAND scr_test) + add_executable(conversation_file_test src/core/conversation_file_test.cpp src/core/conversation_file.cpp) + target_include_directories(conversation_file_test PRIVATE include) + add_test(NAME conversation_file_test COMMAND conversation_file_test) +@@ -378,6 +381,7 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) + src/kernels/cuda/native_qsa_score.cu + src/kernels/cuda/ple.cu + src/kernels/cuda/kv_q8.cu ++ src/kernels/cuda/scr_move.cu + src/kernels/cuda/kv_stream.cu + src/kernels/cuda/kv_q4.cu + src/kernels/cuda/fused_gr.cu +@@ -1791,3 +1795,11 @@ if(STRATA_HIP_GFX906) + add_executable(gfx906_qsa_query_swizzle_probe bench/gfx906_qsa_query_swizzle_probe.cpp) + target_link_libraries(gfx906_qsa_query_swizzle_probe PRIVATE strata_kernels ${_strata_gpu_runtime_target}) + endif() ++ ++# the SCR move kernels' parity test (needs a CUDA device) ++option(STRATA_BUILD_SCR_TESTS "Build the SCR kernel parity test" OFF) ++if(STRATA_BUILD_SCR_TESTS AND TARGET strata_kernels) ++ add_executable(scr_move_parity src/kernels/scr_move_parity.cu) ++ target_include_directories(scr_move_parity PRIVATE include) ++ target_link_libraries(scr_move_parity PRIVATE strata_kernels) ++endif() +diff --git a/include/strata/core/conversation_file.hpp b/include/strata/core/conversation_file.hpp +index 069ca242..d27b167d 100644 +--- a/include/strata/core/conversation_file.hpp ++++ b/include/strata/core/conversation_file.hpp +@@ -218,6 +218,9 @@ struct SessionReadLimits { + // per K/V layer (the QSA layers, then the draft): the most bytes of each part (k, v, k_scale, v_scale, pooled). + // A layer past the end of the list is bounded by the payload only. + std::vector> max_kv_bytes; ++ // a layer split: one entry per later stage, its own carve's limits; empty on one GPU. A file saved with another ++ // number of stages is refused. ++ std::vector stage_limits; + // the RAM PREFLIGHT: asked once the header and identity match and before any payload allocation, with what the + // parse holds at its peak (session_read_peak_bytes); false refuses the file. Not a reservation. + std::function admit; +diff --git a/include/strata/core/scr.hpp b/include/strata/core/scr.hpp +new file mode 100644 +index 00000000..2a48f751 +--- /dev/null ++++ b/include/strata/core/scr.hpp +@@ -0,0 +1,227 @@ ++// Suffix cache reuse (SCR): after an edit in the middle of a prompt, the K/V of the text that survived the edit is ++// moved to its new positions instead of being read again. Ported from llama-cpp-bonsai's common/scr.cpp. ++// ++// The moved cells were computed over the text before the edit, so the result approximates a fresh read. The GDN ++// layers have no cells to move: their state continues from a conversation checkpoint at the end of the last moved ++// run (the plan's `adopt`), and everything after that run is read from it. Strata's indexer pools 4 cells into one ++// key and keeps only the pooled keys, so a run starts and ends on a block boundary of the OLD prompt and moves by a ++// multiple of 4: up to 3 pad tokens are read before a run to make its shift one (DeepSeek V4.1's recipe pads to an ++// even shift the same way). The moved blocks then keep their cells, and their pooled keys only turn by the shift. ++#pragma once ++ ++#include ++#include ++#include ++#include ++#include ++ ++namespace strata::core { ++ ++struct ScrParams { ++ int64_t min_span = 256; // shortest run worth moving ++ int64_t max_runs = 64; // most runs moved per request ++ int64_t seed = 32; // n-gram length that seeds a match ++ int64_t block = 4; // the indexer's cells per pooled key: read parts start and end on these ++}; ++ ++// the cells old[old_start, old_start + n) reappear at [new_start, new_start + n) of the prompt, after `pad` pad tokens ++// that the read inserts right before the run; `at` is where the run lands in the padded (physical) prompt ++struct ScrRun { ++ int64_t new_start = 0, old_start = 0, n = 0; ++ int64_t pad = 0, at = 0; ++ int64_t shift() const { return at - old_start; } ++}; ++ ++struct ScrPlan { ++ int64_t n_prefix = 0; // leading tokens equal in both prompts ++ int64_t restore = 0; // the checkpoint the read restarts from (at or before n_prefix) ++ int64_t adopt = -1; // old length of the checkpoint the GDN state continues from after the last run ++ std::vector runs; // in new order, at or after restore ++ ++ bool empty() const { return runs.empty(); } ++ // the prompt position the final read starts at, and the same in the padded prompt ++ int64_t resume() const { return runs.empty() ? restore : runs.back().new_start + runs.back().n; } ++ int64_t resume_at() const { return runs.empty() ? restore : runs.back().at + runs.back().n; } ++ int64_t pads() const { ++ int64_t n = 0; ++ for (const ScrRun& r : runs) n += r.pad; ++ return n; ++ } ++ int64_t reused() const { ++ int64_t n = 0; ++ for (const ScrRun& r : runs) n += r.n; ++ return n; ++ } ++ std::string str() const { ++ std::string s = "prefix " + std::to_string(n_prefix) + ", restore " + std::to_string(restore) + ", adopt " + ++ std::to_string(adopt); ++ for (const ScrRun& r : runs) ++ s += ", R " + std::to_string(r.new_start) + " <- " + std::to_string(r.old_start) + " x " + std::to_string(r.n) + ++ (r.pad ? " pad " + std::to_string(r.pad) : ""); ++ return s; ++ } ++}; ++ ++namespace scr_detail { ++// greedy maximal spans of cur[from, ...) found in old[from, ...), seeded by n-gram hashes, each old token used once ++template ++std::vector find_spans(const std::vector& old, const std::vector& cur, int64_t from, const ScrParams& p) { ++ std::vector spans; ++ const int64_t n_old = (int64_t) old.size(), n_cur = (int64_t) cur.size(), k = p.seed; ++ if (n_old - from < p.min_span || n_cur - from < p.min_span) return spans; ++ auto gram = [k](const auto& v, int64_t i) { ++ uint64_t h = 1469598103934665603ULL; ++ for (int64_t j = 0; j < k; ++j) h = (h ^ (uint32_t) v[(size_t) (i + j)]) * 1099511628211ULL; ++ return h; ++ }; ++ constexpr size_t max_bucket = 16; // the first few old positions of every seed n-gram ++ std::unordered_map> index; ++ index.reserve((size_t) (n_old - from)); ++ for (int64_t o = from; o + k <= n_old; ++o) { ++ auto& b = index[gram(old, o)]; ++ if (b.size() < max_bucket) b.push_back(o); ++ } ++ std::vector used((size_t) n_old, false); ++ for (int64_t q = from; q + k <= n_cur;) { ++ const auto it = index.find(gram(cur, q)); ++ int64_t best_o = -1, best_n = 0; ++ if (it != index.end()) ++ for (const int64_t o : it->second) { ++ int64_t n = 0; ++ while (q + n < n_cur && o + n < n_old && !used[(size_t) (o + n)] && ++ (int64_t) cur[(size_t) (q + n)] == (int64_t) old[(size_t) (o + n)]) ++ ++n; ++ if (n > best_n) { best_n = n; best_o = o; } ++ } ++ if (best_n >= p.min_span) { ++ spans.push_back({q, best_o, best_n}); ++ std::fill(used.begin() + best_o, used.begin() + best_o + best_n, true); ++ q += best_n; ++ } else { ++ ++q; ++ } ++ } ++ return spans; ++} ++inline int64_t round_up(int64_t x, int64_t b) { return (x + b - 1) / b * b; } ++inline int64_t round_down(int64_t x, int64_t b) { return x / b * b; } ++} // namespace scr_detail ++ ++// Plans the reuse of a session that holds old[0, n_old) for the prompt `cur`. `checkpoints`: the old lengths whose ++// GDN state can be restored (the live state counts as old.size()). A run starts on a block boundary of the new ++// prompt; every run but the last also ends on one, and the last one ends at a checkpoint. The plan always leaves at ++// least one token of `cur` to read. ++template ++ScrPlan scr_plan(const std::vector& old, const std::vector& cur, std::vector checkpoints, ++ const ScrParams& p) { ++ using namespace scr_detail; ++ ScrPlan plan; ++ const int64_t n_cur = (int64_t) cur.size(), bl = std::max(p.block, 1); ++ std::sort(checkpoints.begin(), checkpoints.end()); ++ int64_t a = 0; ++ while (a < (int64_t) old.size() && a < n_cur && (int64_t) old[(size_t) a] == (int64_t) cur[(size_t) a]) ++a; ++ plan.n_prefix = a; ++ for (const int64_t q : checkpoints) ++ if (q <= a) plan.restore = std::max(plan.restore, q); ++ auto spans = find_spans(old, cur, a, p); ++ // a run starts on a block boundary of the old prompt and (but the last) ends on one ++ for (ScrRun& s : spans) { ++ const int64_t cut = round_up(s.old_start, bl) - s.old_start; ++ s.new_start += cut; ++ s.old_start += cut; ++ s.n -= cut; ++ s.n = round_down(s.old_start + s.n, bl) - s.old_start; ++ } ++ spans.erase(std::remove_if(spans.begin(), spans.end(), [&](const ScrRun& s) { return s.n < p.min_span; }), ++ spans.end()); ++ // the largest runs, back in prompt order ++ std::sort(spans.begin(), spans.end(), [](const ScrRun& x, const ScrRun& y) { return x.n > y.n; }); ++ if ((int64_t) spans.size() > p.max_runs) spans.resize((size_t) p.max_runs); ++ std::sort(spans.begin(), spans.end(), [](const ScrRun& x, const ScrRun& y) { return x.new_start < y.new_start; }); ++ // the last run ends at a checkpoint, which also ends the reuse: everything after it is read from that state. ++ // It was cut back to a block boundary, so a checkpoint up to block - 1 cells past it may still match. ++ while (!spans.empty()) { ++ ScrRun& last = spans.back(); ++ int64_t end = -1; ++ for (const int64_t q : checkpoints) { ++ const int64_t n = q - last.old_start; ++ if (n < p.min_span || n >= last.n + bl || last.new_start + n >= n_cur || q <= end) continue; ++ if (n <= last.n || std::equal(old.begin() + last.old_start + last.n, old.begin() + q, ++ cur.begin() + last.new_start + last.n, ++ [](const A& x, const B& y) { return (int64_t) x == (int64_t) y; })) ++ end = q; ++ } ++ if (end >= 0) { ++ last.n = end - last.old_start; ++ plan.adopt = end; ++ break; ++ } ++ spans.pop_back(); ++ } ++ // the pads: walking the runs in prompt order, each gets what makes its padded position a multiple of the block ++ // away from its old one ++ int64_t pads = 0; ++ for (ScrRun& s : spans) { ++ s.pad = ((s.old_start - (s.new_start + pads)) % bl + bl) % bl; ++ pads += s.pad; ++ s.at = s.new_start + pads; ++ } ++ plan.runs = std::move(spans); ++ return plan; ++} ++ ++// The pads an earlier splice read (`pads`: their positions in `live`, sorted) that the prompt `cur` still shares with ++// the session: the ones before the last token of `cur` that `live` holds at the same place once the pads are ++// skipped. `out` is `cur` with those pads put back, which makes it start with `live` as far as they agree. ++template ++std::vector scr_carry(const std::vector& cur, const std::vector& live, const std::vector& pads, ++ int64_t pad_token, std::vector& out) { ++ size_t p = 0, i = 0, k = 0, kept = 0; ++ while (p < live.size()) { ++ if (k < pads.size() && pads[k] == (int64_t) p) { ++ if ((int64_t) live[p] != pad_token) break; ++ ++k; ++ ++p; ++ continue; ++ } ++ if (i + 1 >= cur.size() || (int64_t) live[p] != (int64_t) cur[i]) break; ++ ++i; ++ ++p; ++ kept = k; ++ } ++ std::vector carried(pads.begin(), pads.begin() + (long) kept); ++ out.clear(); ++ out.reserve(cur.size() + carried.size()); ++ size_t c = 0; ++ for (const int64_t at : carried) { ++ while ((int64_t) out.size() < at) out.push_back(cur[c++]); ++ out.push_back((B) pad_token); ++ } ++ out.insert(out.end(), cur.begin() + (long) c, cur.end()); ++ return carried; ++} ++ ++// The positions of the pads `plan` inserts, in the padded prompt. ++inline std::vector scr_pad_positions(const ScrPlan& plan) { ++ std::vector at; ++ for (const ScrRun& r : plan.runs) ++ for (int64_t j = r.pad; j > 0; --j) at.push_back(r.at - j); ++ return at; ++} ++ ++// The prompt the engine reads under `plan`: `cur` with each run's pads inserted right before it. ++template ++std::vector scr_padded(const std::vector& cur, const ScrPlan& plan, B pad_token) { ++ std::vector out; ++ out.reserve(cur.size() + (size_t) plan.pads()); ++ int64_t from = 0; ++ for (const ScrRun& r : plan.runs) { ++ out.insert(out.end(), cur.begin() + from, cur.begin() + r.new_start); ++ out.insert(out.end(), (size_t) r.pad, pad_token); ++ from = r.new_start; ++ } ++ out.insert(out.end(), cur.begin() + from, cur.end()); ++ return out; ++} ++ ++} // namespace strata::core +diff --git a/include/strata/core/session.hpp b/include/strata/core/session.hpp +index 4d633205..4c97ac60 100644 +--- a/include/strata/core/session.hpp ++++ b/include/strata/core/session.hpp +@@ -106,6 +106,9 @@ void session_release(SessionState& s); + /// Zeroes every layer's state - the residual to `R_init`, everything else to zero, so a fresh sequence starts + /// from the reference's own `zeros()`. + void session_zero(SessionState& s, const ModelGeometry& g, const float* R_init, void* stream); ++/// The same, keeping the QSA layers' cells and indexer when `keep_kv` (SCR moves cached cells into a session that ++/// restarts its running state at token 0). ++void session_zero_running(SessionState& s, const ModelGeometry& g, const float* R_init, void* stream, bool keep_kv); + + /// One token: layers 0..47 in order, each a `block_layer`, and the residual is updated in place. + /// +diff --git a/include/strata/kernels/scr_move.hpp b/include/strata/kernels/scr_move.hpp +new file mode 100644 +index 00000000..5191305e +--- /dev/null ++++ b/include/strata/kernels/scr_move.hpp +@@ -0,0 +1,45 @@ ++// include/strata/kernels/scr_move.hpp - suffix cache reuse: moving a QSA layer's cells to new positions. ++// ++// Identity-layout INT8 K/V only (no KV streaming, no rotation, not K8V4): the codes of a cell live at ++// `((pos / page_size) * kv_heads + h) * page_size + pos % page_size`. V and K's unrotated dims move by plain copy. ++// K's rotated group (dims [0, n_rot), one INT8 scale group) is turned back to position zero, turned to the target ++// position with the same angle the fresh write uses, and quantized as the fresh write quantizes it. `master` ++// (optional, device-accessible, FP16 [cell][kv_head][n_rot]) keeps that unrotated group per cell: a moved cell with a ++// valid master starts from it, so a cell moved many times is quantized from the same values every time instead of ++// from the previous move's codes. ++#pragma once ++ ++#include "strata/kernels/rope_scaling.hpp" ++ ++#include ++ ++namespace strata::kernels { ++ ++struct ScrKvLayer { ++ int8_t* k_q = nullptr; ++ int8_t* v_q = nullptr; ++ uint16_t* k_scale = nullptr; ++ uint16_t* v_scale = nullptr; ++ uint16_t* master = nullptr; ///< FP16 [max_cells][kv_heads][n_rot], or null ++ uint8_t* master_valid = nullptr; ///< [max_cells], 1 where `master` holds the cell's unrotated group ++ int kv_heads = 0, head_dim = 0, page_size = 0, n_rot = 0; ++ int64_t max_cells = 0; ++}; ++ ++/// Bytes of scratch scr_move_kv needs to move `chunk` cells at a time. ++uint64_t scr_move_kv_scratch_bytes(const ScrKvLayer& l, int64_t chunk); ++ ++/// Moves cells [src, src + n) to [src + shift, src + shift + n), overlapping ranges included. `scratch` is ++/// scr_move_kv_scratch_bytes(l, chunk) of device memory. Returns false (and moves nothing) on bad arguments. ++bool scr_move_kv(const ScrKvLayer& l, int64_t src, int64_t n, int64_t shift, const RopeScaling& rope, ++ void* scratch, int64_t chunk, void* stream); ++ ++/// Moves the indexer's pooled keys of blocks [b0, b0 + n) by `shift_blocks` blocks (F32 [block][dim], the first ++/// `n_rot` dims rotated at position block_size * block), each turned to its new block's angle. ++bool scr_move_pooled(float* pooled, int64_t rows, int dim, int n_rot, int block_size, int64_t b0, int64_t n, ++ int64_t shift_blocks, const RopeScaling& rope, float* scratch, int64_t chunk, void* stream); ++ ++/// Marks cells [from, to) as freshly written: their master no longer describes them. ++void scr_master_invalidate(const ScrKvLayer& l, int64_t from, int64_t to, void* stream); ++ ++} // namespace strata::kernels +diff --git a/src/core/conversation_file.cpp b/src/core/conversation_file.cpp +index 24291be5..38835d0f 100644 +--- a/src/core/conversation_file.cpp ++++ b/src/core/conversation_file.cpp +@@ -46,6 +46,9 @@ namespace { + constexpr char kMagic[8] = {'S', 'T', 'R', 'S', 'E', 'S', 'S', '\x01'}; + constexpr char kEnd[8] = {'S', 'T', 'R', 'S', 'E', 'N', 'D', '\x01'}; + constexpr uint32_t kVersion = 1; ++// v2: a layer split's file - the v1 payload of the first stage, then the count of later stages and one v1 payload ++// for each (its own geometry, layer range, running state and K/V). A single-GPU file stays v1. ++constexpr uint32_t kVersionStages = 2; + constexpr size_t kHeader = 64, kTrailer = 16; + + // Progress for a transfer: one report after every block that was handed to (or read from) the OS, whatever its size +@@ -707,13 +710,21 @@ void put_payload(Out& o, const SavedConversation& s, const std::vector* sources) { ++ put_payload(o, s, sources); ++ if (s.stage_images.empty()) return; ++ o.u64(s.stage_images.size()); ++ for (const auto& stage : s.stage_images) put_payload(o, stage); ++} ++ + // ---- parsing: every count is checked against the bytes left and the caller's limits before anything is allocated + struct In { + FileSource* f; + SessionHasher hash; + uint64_t left; + std::string& error; +- const SessionReadLimits& limits; ++ const SessionReadLimits* lim; // the stage being parsed: the file's limits, then each later stage's + SessionError kind = SessionError::invalid; + size_t kv_layer = 0; + bool fail(const std::string& m) { if (error.empty()) error = "session file: " + m; return false; } +@@ -743,7 +754,7 @@ struct In { + } + bool buffer(ConversationBuffer& b, size_t part) { + uint64_t n = 0; +- const uint64_t limit = kv_layer < limits.max_kv_bytes.size() ? limits.max_kv_bytes[kv_layer][part] : UINT64_MAX; ++ const uint64_t limit = kv_layer < lim->max_kv_bytes.size() ? lim->max_kv_bytes[kv_layer][part] : UINT64_MAX; + if (!count(n, 1, limit)) return false; + b = {}; + b.resize((size_t) n); +@@ -752,12 +763,12 @@ struct In { + }; + + bool get_checkpoint(In& in, ConversationCheckpoint& c) { +- if (!in.vec(c.ids, in.limits.max_tokens)) return false; ++ if (!in.vec(c.ids, in.lim->max_tokens)) return false; + uint64_t n = 0; +- if (!in.count(n, 16, in.limits.max_tokens)) return false; ++ if (!in.count(n, 16, in.lim->max_tokens)) return false; + c.imgs.resize((size_t) n); + for (auto& k : c.imgs) if (!in.i64(k.start) || !in.u64(k.hash)) return false; +- const auto& m = in.limits.max_state_bytes; // byte arrays: element and byte counts are the same ++ const auto& m = in.lim->max_state_bytes; // byte arrays: element and byte counts are the same + return in.vec(c.gdn, m[0]) && in.vec(c.ple, m[1]) && in.vec(c.tails, m[2]) && in.vec(c.dead, m[3]) && + in.vec(c.block_pos, m[4]) && in.u64(c.used); + } +@@ -765,22 +776,22 @@ bool get_checkpoint(In& in, ConversationCheckpoint& c) { + bool get_payload(In& in, SavedConversation& s) { + for (auto& g : s.geometry) if (!in.i64(g)) return false; + // checked before any state array is read: a file of another geometry allocates nothing +- if (in.limits.geometry && *in.limits.geometry != s.geometry) ++ if (in.lim->geometry && *in.lim->geometry != s.geometry) + return in.fail("saved with another model geometry than this runtime's"); + uint64_t cvec = 0; + if (!in.i64(s.layer_lo) || !in.i64(s.layer_hi) || !in.u64(cvec)) return false; +- if (in.limits.layer_range && +- (in.limits.layer_range->first != s.layer_lo || in.limits.layer_range->second != s.layer_hi)) ++ if (in.lim->layer_range && ++ (in.lim->layer_range->first != s.layer_lo || in.lim->layer_range->second != s.layer_hi)) + return in.fail("saved with another layer range than this runtime's"); + if (cvec > 1) return in.fail("invalid cvec flag"); + s.cvec = cvec == 1; + if (!get_checkpoint(in, s.live)) return false; + uint64_t n = 0; + // a checkpoint is at least 8 counts + used = 72 bytes; a K/V layer at least 7 + 5 = 96 +- if (!in.count(n, 72, in.limits.max_checkpoints)) return false; ++ if (!in.count(n, 72, in.lim->max_checkpoints)) return false; + s.checkpoints.resize((size_t) n); + for (auto& c : s.checkpoints) if (!get_checkpoint(in, c)) return false; +- if (!in.count(n, 96, in.limits.max_kv_layers)) return false; ++ if (!in.count(n, 96, in.lim->max_kv_layers)) return false; + s.kv.resize((size_t) n); + for (auto& k : s.kv) { + in.kv_layer = (size_t) (&k - s.kv.data()); +@@ -795,18 +806,36 @@ bool get_payload(In& in, SavedConversation& s) { + return true; + } + ++// the later stages of a v2 file, each parsed against that stage's own limits ++bool get_stages(In& in, SavedConversation& s, const SessionReadLimits& file_limits) { ++ uint64_t n = 0; ++ if (!in.count(n, 1, file_limits.stage_limits.size())) return false; ++ if (n != file_limits.stage_limits.size()) ++ return in.fail("saved with " + std::to_string(n + 1) + " GPUs, this runtime splits the layers over " + ++ std::to_string(file_limits.stage_limits.size() + 1)); ++ s.stage_images.resize((size_t) n); ++ for (size_t i = 0; i < s.stage_images.size(); ++i) { ++ in.lim = &file_limits.stage_limits[i]; ++ in.kv_layer = 0; ++ if (!get_payload(in, s.stage_images[i])) return false; ++ } ++ in.lim = &file_limits; ++ return true; ++} ++ + bool has_stage_parts(const SavedConversation& s) { + if (!s.live.stage_parts.empty()) return true; + for (const auto& c : s.checkpoints) if (!c.stage_parts.empty()) return true; ++ for (const auto& stage : s.stage_images) if (has_stage_parts(stage) || !stage.stage_images.empty()) return true; + return false; + } + + // header v1, little-endian: magic[8] version:u32 header_size:u32 model:u64 config:u64 payload:u64 reserved:u64[2] + // header_hash:u64 (session_hash64 of bytes 0..55, seed 0) +-void header_bytes(uint8_t* h, const SessionFileIdentity& id, uint64_t payload) { ++void header_bytes(uint8_t* h, const SessionFileIdentity& id, uint64_t payload, uint32_t version) { + std::memset(h, 0, kHeader); + std::memcpy(h, kMagic, 8); +- const uint32_t version = kVersion, size = kHeader; ++ const uint32_t size = kHeader; + std::memcpy(h + 8, &version, 4); + std::memcpy(h + 12, &size, 4); + std::memcpy(h + 16, &id.model, 8); +@@ -1014,6 +1043,14 @@ uint64_t session_read_max_file_bytes(const SessionReadLimits& l) { + n = sat_add(n, 7 * 8); + for (uint64_t b : l.max_kv_bytes[(size_t) i]) n = sat_add(n, sat_add(8, b)); + } ++ if (!l.stage_limits.empty()) { // v2: the stage count, then each later stage's payload (without its own framing) ++ n = sat_add(n, 8); ++ for (const auto& s : l.stage_limits) { ++ const uint64_t m = session_read_max_file_bytes(s); ++ if (m == UINT64_MAX) return l.max_file_bytes; ++ n = sat_add(n, m - kHeader - kTrailer); ++ } ++ } + return std::min(n, l.max_file_bytes); + } + +@@ -1029,10 +1066,17 @@ bool write_impl(const std::string& path, const SavedConversation& image, const s + SessionStatus& st) { + st = {}; + st.error = SessionError::invalid; // refusals before any disk access +- if (has_stage_parts(image)) { error = "session file: layer-split state cannot be saved"; return false; } ++ if (has_stage_parts(image)) { ++ error = "session file: a checkpoint's stage parts belong in the stage images"; ++ return false; ++ } ++ if (sources && !image.stage_images.empty()) { ++ error = "session file: a layer split's state is written from the captured image, not from K/V sources"; ++ return false; ++ } + if (sources && !image.kv.empty()) { error = "session file: K/V given both in the image and as sources"; return false; } + Out sizing; +- put_payload(sizing, image, sources); ++ put_image(sizing, image, sources); + if (sizing.overflow || sizing.n > UINT64_MAX - kHeader - kTrailer || + sizing.n + kHeader + kTrailer > (uint64_t) std::numeric_limits::max()) { + error = "session file: the state is too large for this build"; +@@ -1069,13 +1113,13 @@ bool write_impl(const std::string& path, const SavedConversation& image, const s + f.write_fault = injected(fault, "write"); + f.flush_fault = injected(fault, "file_flush"); + uint8_t h[kHeader]; +- header_bytes(h, id, payload); ++ header_bytes(h, id, payload, image.stage_images.empty() ? kVersion : kVersionStages); + bool ok = f.write(h, kHeader); + SessionHasher hash(0); + Out out; + out.f = &f; + out.hash = &hash; +- if (ok) put_payload(out, image, sources); ++ if (ok) put_image(out, image, sources); + ok = ok && out.ok && out.n == payload; + if (ok) { + const uint64_t ph = hash.digest(); +@@ -1204,7 +1248,15 @@ bool read_impl(const std::string& path, const SessionFileIdentity& id, SavedConv + std::memcpy(&r0, h + 40, 8); std::memcpy(&r1, h + 48, 8); std::memcpy(&hh, h + 56, 8); + if (std::memcmp(h, kMagic, 8) != 0) { error = "session file: not a Strata session file (magic)"; return false; } + if (hh != session_hash64(h, 56, 0)) { error = "session file: header checksum mismatch"; return false; } +- if (version != kVersion) { error = "session file: unsupported version " + std::to_string(version); return false; } ++ if (version != kVersion && version != kVersionStages) { ++ error = "session file: unsupported version " + std::to_string(version); ++ return false; ++ } ++ if ((version == kVersionStages) != !limits.stage_limits.empty()) { ++ error = version == kVersionStages ? "session file: saved with a layer split, this runtime runs on one GPU" ++ : "session file: saved on one GPU, this runtime splits the layers"; ++ return false; ++ } + if (hsize != kHeader || r0 || r1) { error = "session file: invalid header fields"; return false; } + if (model != id.model) { error = "session file: saved with another model (model fingerprint differs)"; return false; } + if (config != id.config) { error = "session file: saved with another engine configuration (config fingerprint differs)"; return false; } +@@ -1222,8 +1274,11 @@ bool read_impl(const std::string& path, const SessionFileIdentity& id, SavedConv + } + } + SavedConversation parsed; +- In in{&f, SessionHasher(0), payload, error, limits}; +- if (!get_payload(in, parsed)) { st.error = in.kind; return false; } ++ In in{&f, SessionHasher(0), payload, error, &limits}; ++ if (!get_payload(in, parsed) || (version == kVersionStages && !get_stages(in, parsed, limits))) { ++ st.error = in.kind; ++ return false; ++ } + if (in.left) { error = "session file: payload has " + std::to_string(in.left) + " unparsed bytes"; return false; } + uint8_t t[kTrailer]; + if (!f.read(t, kTrailer)) { st.error = f.kind(); error = "session file: trailer read error: " + f.error(); return false; } +diff --git a/src/core/conversation_file_test.cpp b/src/core/conversation_file_test.cpp +index fd85f2ca..0cb2de3b 100644 +--- a/src/core/conversation_file_test.cpp ++++ b/src/core/conversation_file_test.cpp +@@ -108,6 +108,9 @@ bool same(const SavedConversation& a, const SavedConversation& b) { + return false; + for (size_t i = 0; i < a.checkpoints.size(); ++i) + if (!same_checkpoint(a.checkpoints[i], b.checkpoints[i])) return false; ++ if (a.stage_images.size() != b.stage_images.size()) return false; ++ for (size_t i = 0; i < a.stage_images.size(); ++i) ++ if (!same(a.stage_images[i], b.stage_images[i])) return false; + for (size_t i = 0; i < a.kv.size(); ++i) { + const auto& x = a.kv[i]; const auto& y = b.kv[i]; + if (x.format != y.format || x.cells != y.cells || x.heads != y.heads || x.head_dim != y.head_dim || +@@ -244,7 +247,7 @@ int main() { + // wrong version: header field patched and header hash recomputed so only the version differs + { + auto d = image; +- uint32_t v = 2; std::memcpy(d.data() + 8, &v, 4); ++ uint32_t v = 99; std::memcpy(d.data() + 8, &v, 4); + const uint64_t h = session_hash64(d.data(), 56, 0); std::memcpy(d.data() + 56, &h, 8); + const fs::path p = dir / "ver.bin"; spit(p, d); + check(rejects(p, id, "version"), "unknown version rejected"); +@@ -270,14 +273,45 @@ int main() { + check(!session_file_read((dir / "absent.bin").string(), id, keep, b, error), "missing read fails"); + check(same(keep, original), "failed read does not modify the destination"); + } +- // layer-split images are not written; no file appears ++ // stage parts left inside a checkpoint are refused (they belong in the stage images); no file appears + { + SavedConversation split = sample(); + split.live.stage_parts.push_back(checkpoint(1, 9)); + const fs::path p = dir / "split.bin"; +- check(!session_file_write(p.string(), split, id, written, error), "layer-split image refused"); ++ check(!session_file_write(p.string(), split, id, written, error), "stage parts in a checkpoint refused"); + check(!fs::exists(p) && no_temp(dir), "refused write leaves no file"); + } ++ // a layer split's image (format v2): the first stage's payload, then one per later stage ++ { ++ SavedConversation split = sample(); ++ SavedConversation later = sample(); ++ later.layer_lo = 3; later.layer_hi = 48; ++ later.live = checkpoint(1000, 7); ++ later.kv.pop_back(); // a later stage without the drafter holds no draft K/V ++ split.layer_hi = 3; ++ split.stage_images.push_back(later); ++ const fs::path p = dir / "split2.bin"; ++ error.clear(); ++ check(session_file_write(p.string(), split, id, written, error), "layer-split image written"); ++ const std::vector d = slurp(p); ++ uint32_t version = 0; std::memcpy(&version, d.data() + 8, 4); ++ check(version == 2, "layer-split file is format v2"); ++ SessionReadLimits two; ++ two.stage_limits.emplace_back(); ++ SavedConversation back; ++ size_t read = 0; ++ error.clear(); ++ check(session_file_read(p.string(), id, back, read, error, two), "layer-split file reads back"); ++ check(same(split, back), "layer-split round trip is identical"); ++ check(rejects(p, id, "layer split"), "a one-GPU runtime refuses a layer-split file"); ++ SessionReadLimits three; ++ three.stage_limits.resize(2); ++ check(rejects(p, id, "GPUs", three), "another number of stages refused"); ++ SessionReadLimits range = two; ++ range.stage_limits[0].layer_range = std::make_pair(int64_t(24), int64_t(48)); ++ check(rejects(p, id, "layer range", range), "a later stage's layer range is checked"); ++ check(rejects(good, id, "one GPU", two), "a layer-split runtime refuses a one-GPU file"); ++ } + // a write into a missing directory fails cleanly + check(!session_file_write((dir / "nope" / "x.bin").string(), original, id, written, error), + "write into a missing directory fails"); +diff --git a/src/core/expert_source.cpp b/src/core/expert_source.cpp +index 4878c111..3f1a5bed 100644 +--- a/src/core/expert_source.cpp ++++ b/src/core/expert_source.cpp +@@ -2470,8 +2470,10 @@ void FileExpertSource::begin_layer(int64_t layer, const int32_t* ids, int64_t k) + + bool FileExpertSource::transient(int64_t layer, int64_t expert) const { + // the io-prefetch staging of the experts.bin tier is the decode's: the prompt path keeps reading the mapping +- // (its own read-ahead hints), so this answers as it did without io prefetch +- if ((role_ptr_.empty() && direct_.empty()) || layer < 0 || expert < 0 || layer >= n_layers_ || expert >= n_expert_) return false; ++ // (its own read-ahead hints), so this answers as it did without io prefetch - unless STRATA_PROMPT_DIRECT=1 ++ // opened the prompt path's own unbuffered descriptors, which only copy_blobs (the stager) reads through ++ if ((role_ptr_.empty() && direct_.empty() && pdirect_.empty()) || layer < 0 || expert < 0 || layer >= n_layers_ || ++ expert >= n_expert_) return false; + const size_t index = (size_t) layer * (size_t) n_expert_ + (size_t) expert; + if (complement_ready_ && resident_blob(index) != nullptr) + return false; +diff --git a/src/core/scr_test.cpp b/src/core/scr_test.cpp +new file mode 100644 +index 00000000..2895394d +--- /dev/null ++++ b/src/core/scr_test.cpp +@@ -0,0 +1,209 @@ ++// src/core/scr_test.cpp - the SCR planner (include/strata/core/scr.hpp), CPU only. ++#include "strata/core/scr.hpp" ++ ++#include ++#include ++#include ++ ++using strata::core::ScrParams; ++using strata::core::ScrPlan; ++using strata::core::scr_carry; ++using strata::core::scr_pad_positions; ++using strata::core::scr_padded; ++using strata::core::scr_plan; ++ ++namespace { ++int failures = 0, checks = 0; ++void check(bool ok, const char* label) { ++ ++checks; ++ if (!ok) { ++failures; std::fprintf(stderr, "FAIL: %s\n", label); } ++} ++ ++std::vector random_tokens(std::mt19937& rng, size_t n) { ++ std::uniform_int_distribution d(0, 200000); ++ std::vector v(n); ++ for (auto& t : v) t = d(rng); ++ return v; ++} ++ ++// every invariant execution relies on ++bool sound(const ScrPlan& p, const std::vector& old, const std::vector& cur, ++ const std::vector& cks, const ScrParams& prm, const char** why) { ++ const int64_t n_cur = (int64_t) cur.size(), bl = prm.block; ++ if (p.restore > p.n_prefix) { *why = "restore past the prefix"; return false; } ++ if (p.restore > 0 && std::find(cks.begin(), cks.end(), p.restore) == cks.end()) { *why = "restore not a checkpoint"; return false; } ++ for (int64_t i = 0; i < p.n_prefix; ++i) if (old[i] != cur[i]) { *why = "prefix differs"; return false; } ++ int64_t prev_end = p.n_prefix, pads = 0; ++ std::vector used(old.size(), false); ++ for (size_t r = 0; r < p.runs.size(); ++r) { ++ const auto& run = p.runs[r]; ++ if (run.n < prm.min_span) { *why = "run shorter than min_span"; return false; } ++ if (run.new_start < prev_end) { *why = "runs overlap or are out of order"; return false; } ++ if (run.old_start % bl != 0) { *why = "run start not on an old block boundary"; return false; } ++ const bool last = r + 1 == p.runs.size(); ++ if (!last && (run.old_start + run.n) % bl != 0) { *why = "run end not on an old block boundary"; return false; } ++ if (last && run.old_start + run.n != p.adopt) { *why = "last run does not end at adopt"; return false; } ++ if (run.pad < 0 || run.pad >= bl) { *why = "pad count out of range"; return false; } ++ pads += run.pad; ++ if (run.at != run.new_start + pads) { *why = "padded position wrong"; return false; } ++ if (run.shift() % bl != 0) { *why = "shift not a multiple of the block"; return false; } ++ if (run.old_start < 0 || run.old_start + run.n > (int64_t) old.size()) { *why = "run outside old"; return false; } ++ for (int64_t i = 0; i < run.n; ++i) { ++ if (old[run.old_start + i] != cur[run.new_start + i]) { *why = "run tokens differ"; return false; } ++ if (used[run.old_start + i]) { *why = "old cell moved twice"; return false; } ++ used[run.old_start + i] = true; ++ } ++ prev_end = run.new_start + run.n; ++ } ++ // the padded prompt holds every run at its padded position and is the prompt with the pads added ++ const auto padded = scr_padded(cur, p, int32_t(198)); ++ if ((int64_t) padded.size() != n_cur + p.pads()) { *why = "padded length"; return false; } ++ for (const auto& run : p.runs) ++ for (int64_t i = 0; i < run.n; ++i) ++ if (padded[run.at + i] != old[run.old_start + i]) { *why = "padded prompt misplaces a run"; return false; } ++ if (p.resume_at() - p.resume() != p.pads()) { *why = "resume_at"; return false; } ++ if (!p.runs.empty() && std::find(cks.begin(), cks.end(), p.adopt) == cks.end()) { *why = "adopt not a checkpoint"; return false; } ++ if (p.resume() >= n_cur) { *why = "nothing left to read"; return false; } ++ if ((int64_t) p.runs.size() > prm.max_runs) { *why = "too many runs"; return false; } ++ return true; ++} ++ ++std::vector periodic(int64_t n, int64_t every) { ++ std::vector v; ++ for (int64_t q = every; q < n; q += every) v.push_back(q); ++ v.push_back(n); // the live state ++ return v; ++} ++} // namespace ++ ++int main() { ++ std::mt19937 rng(438); ++ const ScrParams prm; ++ const char* why = ""; ++ ++ // no edit: everything is prefix, nothing to move ++ { ++ auto old = random_tokens(rng, 20000); ++ auto cur = old; ++ cur.push_back(7); ++ const auto cks = periodic((int64_t) old.size(), 4096); ++ const ScrPlan p = scr_plan(old, cur, cks, prm); ++ check(p.n_prefix == 20000 && p.restore == 20000 && p.empty(), "no edit: prefix only"); ++ } ++ // a 10-token replacement at 2000 in a 20000-token prompt: the rest comes back up to the last checkpoint ++ { ++ auto old = random_tokens(rng, 20000); ++ auto cur = old; ++ for (int i = 0; i < 10; ++i) cur[2000 + i] = 300000 + i; ++ cur.push_back(1); ++ const auto cks = periodic((int64_t) old.size(), 4096); ++ const ScrPlan p = scr_plan(old, cur, cks, prm); ++ check(sound(p, old, cur, cks, prm, &why), why); ++ check(p.n_prefix == 2000 && p.restore == 0, "replacement: prefix and restore"); ++ check(p.adopt == 20000 && p.resume() == 20000, "replacement: the live state is adopted"); ++ check(p.reused() >= 20000 - 2000 - 10 - 2 * prm.block, "replacement: almost everything after the edit reused"); ++ } ++ // an insertion of 5 tokens (a shift that is not a multiple of the block) ++ { ++ auto old = random_tokens(rng, 30000); ++ std::vector cur(old.begin(), old.begin() + 9000); ++ for (int i = 0; i < 5; ++i) cur.push_back(400000 + i); ++ cur.insert(cur.end(), old.begin() + 9000, old.end()); ++ cur.push_back(2); ++ const auto cks = periodic((int64_t) old.size(), 4096); ++ const ScrPlan p = scr_plan(old, cur, cks, prm); ++ check(sound(p, old, cur, cks, prm, &why), why); ++ check(p.runs.size() == 1 && p.runs[0].pad == 3 && p.runs[0].shift() == 8, "insertion: 3 pads make the shift 8"); ++ check(p.adopt == 30000, "insertion: ends at the live state"); ++ } ++ // no checkpoint inside the moved text: nothing can be adopted, so nothing moves ++ { ++ auto old = random_tokens(rng, 10000); ++ auto cur = old; ++ cur[100] = 500000; ++ cur.push_back(3); ++ const std::vector cks = {50}; // only one before the edit ++ const ScrPlan p = scr_plan(old, cur, cks, prm); ++ check(p.empty() && p.restore == 50, "no checkpoint after the edit: no runs, restore at 50"); ++ } ++ // the cap on runs: 100 edits keeps the 64 largest runs ++ { ++ auto old = random_tokens(rng, 120000); ++ auto cur = old; ++ for (int e = 0; e < 100; ++e) cur[1000 + e * 1100] = 600000 + e; ++ cur.push_back(4); ++ const auto cks = periodic((int64_t) old.size(), 4096); ++ const ScrPlan p = scr_plan(old, cur, cks, prm); ++ check(sound(p, old, cur, cks, prm, &why), why); ++ check((int64_t) p.runs.size() <= prm.max_runs && p.runs.size() >= 60, "cap: up to 64 runs kept"); ++ } ++ // a moved block of text (the same paragraphs in another order) ++ { ++ auto old = random_tokens(rng, 24000); ++ std::vector cur(old.begin(), old.begin() + 4000); ++ cur.insert(cur.end(), old.begin() + 12000, old.begin() + 20000); ++ cur.insert(cur.end(), old.begin() + 4000, old.begin() + 12000); ++ cur.insert(cur.end(), old.begin() + 20000, old.end()); ++ cur.push_back(5); ++ const auto cks = periodic((int64_t) old.size(), 4096); ++ const ScrPlan p = scr_plan(old, cur, cks, prm); ++ check(sound(p, old, cur, cks, prm, &why), why); ++ check(p.reused() > 12000, "reorder: both moved halves reused"); ++ } ++ // random edits: every plan is sound ++ { ++ int plans = 0; ++ for (int t = 0; t < 400; ++t) { ++ auto old = random_tokens(rng, 2000 + rng() % 40000); ++ auto cur = old; ++ const int edits = 1 + (int) (rng() % 6); ++ for (int e = 0; e < edits; ++e) { ++ const size_t at = rng() % cur.size(); ++ const int kind = (int) (rng() % 3), len = 1 + (int) (rng() % 300); ++ if (kind == 0) cur.erase(cur.begin() + (long) at, cur.begin() + (long) std::min(cur.size(), at + len)); ++ else if (kind == 1) { auto ins = random_tokens(rng, (size_t) len); cur.insert(cur.begin() + (long) at, ins.begin(), ins.end()); } ++ else for (int i = 0; i < len && at + i < cur.size(); ++i) cur[at + i] = 700000 + i; ++ } ++ cur.push_back(6); ++ std::vector cks = periodic((int64_t) old.size(), 1024 + (int64_t) (rng() % 8192)); ++ const int extra = (int) (rng() % 3); // turn boundaries at arbitrary lengths ++ for (int i = 0; i < extra; ++i) cks.push_back(1 + (int64_t) (rng() % old.size())); ++ const ScrPlan p = scr_plan(old, cur, cks, prm); ++ const bool ok = sound(p, old, cur, cks, prm, &why); ++ if (!ok) std::fprintf(stderr, " plan %d: %s (%s)\n", t, why, p.str().c_str()); ++ check(ok, "random edits: plan is sound"); ++ plans += !p.empty(); ++ } ++ check(plans > 200, "random edits: most plans reuse something"); ++ } ++ // pad carry: the session read a spliced prompt with pads; the next turn (the same text, more after it) gets ++ // them back, and an edit before a pad drops that pad and every later one ++ { ++ auto old = random_tokens(rng, 30000); ++ std::vector cur(old.begin(), old.begin() + 9000); ++ for (int i = 0; i < 5; ++i) cur.push_back(400000 + i); ++ cur.insert(cur.end(), old.begin() + 9000, old.end()); ++ cur.push_back(2); ++ const auto cks = periodic((int64_t) old.size(), 4096); ++ const ScrPlan p = scr_plan(old, cur, cks, prm); ++ const auto live = scr_padded(cur, p, int32_t(198)); ++ const auto pads = scr_pad_positions(p); ++ bool marked = !pads.empty(); ++ for (const int64_t at : pads) marked = marked && live[(size_t) at] == 198; ++ check(marked && (int64_t) pads.size() == p.pads(), "pad positions point at the pads"); ++ std::vector next = cur; // the next turn: the server's prompt has no pads ++ next.push_back(11); next.push_back(12); ++ std::vector out; ++ const auto carried = scr_carry(next, live, pads, 198, out); ++ check(carried == pads, "carry: every pad comes back"); ++ check(std::equal(live.begin(), live.end(), out.begin()), "carry: the prompt starts with the session again"); ++ std::vector edited = cur; ++ edited[100] = 900000; ++ out.clear(); ++ check(scr_carry(edited, live, pads, 198, out).empty() && out == edited, "carry: an edit before the pads drops them"); ++ std::vector none; ++ check(scr_carry(next, live, std::vector{}, 198, none).empty() && none == next, "carry: no pads, no change"); ++ } ++ std::printf("scr_test: %d checks, %d failed\n", checks, failures); ++ return failures == 0 ? 0 : 1; ++} +diff --git a/src/core/session.cpp b/src/core/session.cpp +index 157c4cc4..8e493189 100644 +--- a/src/core/session.cpp ++++ b/src/core/session.cpp +@@ -146,6 +146,10 @@ namespace strata::core { + #endif + + void session_zero(SessionState& s, const ModelGeometry& g, const float* R_init, void* stream) { ++ session_zero_running(s, g, R_init, stream, false); ++} ++ ++void session_zero_running(SessionState& s, const ModelGeometry& g, const float* R_init, void* stream, bool keep_kv) { + #if defined(STRATA_GPU_METAL) + // the running state lives in the Metal decode window; the session's own device buffers are not read there + (void) g; (void) R_init; (void) stream; +@@ -168,7 +172,8 @@ void session_zero(SessionState& s, const ModelGeometry& g, const float* R_init, + // every owned GDN layer's recurrence and conv history + cudaMemsetAsync(s.gdn_state, 0, (size_t) s.gdn_alloc * gdn_state_floats(g) * 4, cs); + // and every owned QSA layer's cache and indexer +- for (int64_t j = 0; j < s.qsa_alloc; ++j) qsa_state_zero(s.qsa_states[s.qsa_ord0 + j], g, stream); ++ if (!keep_kv) ++ for (int64_t j = 0; j < s.qsa_alloc; ++j) qsa_state_zero(s.qsa_states[s.qsa_ord0 + j], g, stream); + // **AND THE PLE'S CONV HISTORY AND TOKEN WINDOW.** A sequence that started with a warm history would + // convolve over rows belonging to a different sequence - the conv reads NG_HIST previous NORMALIZED rows, + // so a stale one is a real contribution and not a zero. The token window resets to `NG_HIST`-many nulls +diff --git a/src/kernels/cuda/scr_move.cu b/src/kernels/cuda/scr_move.cu +new file mode 100644 +index 00000000..37e398b6 +--- /dev/null ++++ b/src/kernels/cuda/scr_move.cu +@@ -0,0 +1,218 @@ ++// src/kernels/cuda/scr_move.cu - see include/strata/kernels/scr_move.hpp. ++#include "strata/kernels/scr_move.hpp" ++#include "strata/kernels/f16_bits.hpp" ++ ++#include ++ ++#include ++#include ++ ++namespace strata::kernels { ++namespace { ++ ++constexpr int GROUP = 64; // KV_Q8_GROUP: one FP16 scale per 64 values ++ ++__device__ __forceinline__ long long row_of(long long pos, int h, int kv_heads, int page_size) { ++ return ((pos / page_size) * kv_heads + h) * page_size + pos % page_size; ++} ++ ++// cos/sin at `pos` for rotated pair `pair`, the fresh K write's angle (native_rope.cu `apply`) ++__device__ __forceinline__ void angle(long long pos, int pair, float theta_scale, const RopeKernelArgs& k, float& c, ++ float& s) { ++ const float theta_extrap = (float) pos * powf(theta_scale, float(pair)); ++ rope_scaled_angle(theta_extrap, k.freq_scale, k.corr_low, k.corr_high, k.ext_factor, k.attn_factor, pair, c, s); ++} ++ ++// scratch per cell and head: K codes, V codes (head_dim each), the scales (head_dim / 64 each), the unrotated K group ++// (n_rot FP16) ++struct Scratch { ++ int8_t* k; int8_t* v; uint16_t* ks; uint16_t* vs; uint16_t* pre; ++}; ++__host__ __device__ inline Scratch scratch_of(void* base, int64_t chunk, int kv_heads, int head_dim, int n_rot) { ++ const int64_t cells = chunk * kv_heads, groups = head_dim / GROUP; ++ auto* p = static_cast(base); ++ Scratch s; ++ s.k = reinterpret_cast(p); p += cells * head_dim; ++ s.v = reinterpret_cast(p); p += cells * head_dim; ++ s.ks = reinterpret_cast(p); p += cells * groups * 2; ++ s.vs = reinterpret_cast(p); p += cells * groups * 2; ++ s.pre = reinterpret_cast(p); ++ return s; ++} ++ ++// one block per (cell, head): the cell's codes and scales into scratch, and its K group back at position zero ++__global__ void gather(ScrKvLayer l, long long src, Scratch s, float theta_scale, RopeKernelArgs k) { ++ const long long i = blockIdx.x; // cell within the chunk ++ const int h = blockIdx.y, t = threadIdx.x; ++ const long long pos = src + i; ++ const long long row = row_of(pos, h, l.kv_heads, l.page_size); ++ const long long j = i * l.kv_heads + h; ++ const int groups = l.head_dim / GROUP; ++ for (int d = t; d < l.head_dim; d += blockDim.x) { ++ s.k[j * l.head_dim + d] = l.k_q[row * l.head_dim + d]; ++ s.v[j * l.head_dim + d] = l.v_q[row * l.head_dim + d]; ++ } ++ if (t < groups) { ++ s.ks[j * groups + t] = l.k_scale[row * groups + t]; ++ s.vs[j * groups + t] = l.v_scale[row * groups + t]; ++ } ++ const int half = l.n_rot / 2; ++ if (t >= half) return; ++ uint16_t* pre = s.pre + j * l.n_rot; ++ if (l.master != nullptr && l.master_valid[pos]) { ++ const uint16_t* m = l.master + (pos * l.kv_heads + h) * l.n_rot; ++ pre[t] = m[t]; ++ pre[t + half] = m[t + half]; ++ return; ++ } ++ const float sc = f32_from_f16(l.k_scale[row * groups]); ++ const float a = (float) l.k_q[row * l.head_dim + t] * sc, b = (float) l.k_q[row * l.head_dim + t + half] * sc; ++ float c, sn; ++ angle(pos, t, theta_scale, k, c, sn); ++ // the inverse of (a, b) -> (a c - b s, a s + b c), whose matrix is mscale times a rotation ++ const float m2 = c * c + sn * sn; ++ pre[t] = f16_from_f32((a * c + b * sn) / m2); ++ pre[t + half] = f16_from_f32((b * c - a * sn) / m2); ++} ++ ++// one block per (cell, head), GROUP threads: the codes back from scratch at the target, the K group turned to the ++// target's angle and quantized as kv_append_q8 quantizes it ++__global__ void scatter(ScrKvLayer l, long long dst, Scratch s, float theta_scale, RopeKernelArgs k) { ++ const long long i = blockIdx.x; ++ const int h = blockIdx.y, t = threadIdx.x; ++ const long long pos = dst + i; ++ const long long row = row_of(pos, h, l.kv_heads, l.page_size); ++ const long long j = i * l.kv_heads + h; ++ const int groups = l.head_dim / GROUP; ++ for (int d = t; d < l.head_dim; d += blockDim.x) { ++ if (d >= GROUP) l.k_q[row * l.head_dim + d] = s.k[j * l.head_dim + d]; ++ l.v_q[row * l.head_dim + d] = s.v[j * l.head_dim + d]; ++ } ++ if (t < groups) { ++ if (t > 0) l.k_scale[row * groups + t] = s.ks[j * groups + t]; ++ l.v_scale[row * groups + t] = s.vs[j * groups + t]; ++ } ++ // the rotated group: n_rot == GROUP, one value a thread ++ const int half = l.n_rot / 2, pair = t < half ? t : t - half; ++ const uint16_t* pre = s.pre + j * l.n_rot; ++ const float a = f32_from_f16(pre[pair]), b = f32_from_f16(pre[pair + half]); ++ float c, sn; ++ angle(pos, pair, theta_scale, k, c, sn); ++ const float x = t < half ? a * c - b * sn : a * sn + b * c; ++ float m = fabsf(x); ++ for (int o = 16; o > 0; o >>= 1) m = fmaxf(m, __shfl_xor_sync(0xffffffffu, m, o)); ++ __shared__ float warp_max[2]; ++ if ((t & 31) == 0) warp_max[t >> 5] = m; ++ __syncthreads(); ++ const float amax = fmaxf(warp_max[0], warp_max[1]); ++ const uint16_t sbits = f16_from_f32(amax / 127.0f); ++ const float sf = f32_from_f16(sbits); ++ int q = 0; ++ if (sf > 0.0f) { ++ q = __float2int_rn(x / sf); ++ q = q < -127 ? -127 : (q > 127 ? 127 : q); ++ } ++ l.k_q[row * l.head_dim + t] = (int8_t) q; ++ if (t == 0) l.k_scale[row * groups] = sbits; ++ if (l.master != nullptr) { ++ l.master[(pos * l.kv_heads + h) * l.n_rot + t] = pre[t]; ++ if (t == 0 && h == 0) l.master_valid[pos] = 1; ++ } ++} ++ ++__global__ void pooled_gather(const float* pooled, int dim, long long b0, float* scratch) { ++ const long long i = blockIdx.x; ++ for (int d = threadIdx.x; d < dim; d += blockDim.x) scratch[i * dim + d] = pooled[(b0 + i) * dim + d]; ++} ++ ++__global__ void pooled_scatter(float* pooled, int dim, int n_rot, int block_size, long long b0, long long shift, ++ const float* scratch, float theta_scale, RopeKernelArgs k) { ++ const long long i = blockIdx.x; ++ const int half = n_rot / 2; ++ const long long from = block_size * (b0 + i), to = block_size * (b0 + i + shift); ++ for (int d = threadIdx.x; d < dim; d += blockDim.x) { ++ const float* row = scratch + i * dim; ++ float y = row[d]; ++ if (d < n_rot) { ++ const int pair = d % half; ++ float c0, s0, c1, s1; ++ angle(from, pair, theta_scale, k, c0, s0); ++ angle(to, pair, theta_scale, k, c1, s1); ++ const float m2 = c0 * c0 + s0 * s0; ++ const float a = (row[pair] * c0 + row[pair + half] * s0) / m2; ++ const float z = (row[pair + half] * c0 - row[pair] * s0) / m2; ++ y = d < half ? a * c1 - z * s1 : a * s1 + z * c1; ++ } ++ pooled[(b0 + i + shift) * dim + d] = y; ++ } ++} ++ ++__global__ void invalidate(uint8_t* valid, long long from, long long n) { ++ const long long i = blockIdx.x * (long long) blockDim.x + threadIdx.x; ++ if (i < n) valid[from + i] = 0; ++} ++ ++// the chunks of a move in an order that never overwrites a source still to be read ++template ++void chunks(int64_t n, int64_t shift, int64_t chunk, F f) { ++ const int64_t count = (n + chunk - 1) / chunk; ++ for (int64_t c = 0; c < count; ++c) { ++ const int64_t k = shift > 0 ? count - 1 - c : c; ++ const int64_t a = k * chunk; ++ f(a, std::min(chunk, n - a)); ++ } ++} ++ ++} // namespace ++ ++uint64_t scr_move_kv_scratch_bytes(const ScrKvLayer& l, int64_t chunk) { ++ const uint64_t cells = (uint64_t) chunk * l.kv_heads, groups = (uint64_t) l.head_dim / GROUP; ++ return cells * (2ull * l.head_dim + 4ull * groups + 2ull * l.n_rot); ++} ++ ++bool scr_move_kv(const ScrKvLayer& l, int64_t src, int64_t n, int64_t shift, const RopeScaling& rope, void* scratch, ++ int64_t chunk, void* stream) { ++ if (!l.k_q || !l.v_q || !l.k_scale || !l.v_scale || !scratch || chunk < 1 || n < 0 || l.n_rot != GROUP || ++ l.head_dim % GROUP != 0 || l.kv_heads < 1 || l.page_size < 1 || (l.master && !l.master_valid) || src < 0 || ++ src + n > l.max_cells || src + shift < 0 || src + shift + n > l.max_cells || chunk > 2147483647ll || ++ rope_scaling_invalid(rope) != nullptr) ++ return false; ++ if (n == 0 || shift == 0) return true; ++ auto st = static_cast(stream); ++ const float theta_scale = powf((float) rope.freq_base, -2.0f / l.n_rot); ++ const RopeKernelArgs k = rope.kernel_args(l.n_rot); ++ const Scratch s = scratch_of(scratch, chunk, l.kv_heads, l.head_dim, l.n_rot); ++ chunks(n, shift, chunk, [&](int64_t a, int64_t m) { ++ const dim3 grid((unsigned) m, (unsigned) l.kv_heads); ++ gather<<>>(l, src + a, s, theta_scale, k); ++ scatter<<>>(l, src + a + shift, s, theta_scale, k); ++ }); ++ return cudaGetLastError() == cudaSuccess; ++} ++ ++bool scr_move_pooled(float* pooled, int64_t rows, int dim, int n_rot, int block_size, int64_t b0, int64_t n, ++ int64_t shift_blocks, const RopeScaling& rope, float* scratch, int64_t chunk, void* stream) { ++ if (!pooled || !scratch || n < 0 || chunk < 1 || n_rot % 2 || n_rot > dim || b0 < 0 || b0 + n > rows || ++ b0 + shift_blocks < 0 || b0 + shift_blocks + n > rows || rope_scaling_invalid(rope) != nullptr) ++ return false; ++ if (n == 0 || shift_blocks == 0) return true; ++ auto st = static_cast(stream); ++ const float theta_scale = powf((float) rope.freq_base, -2.0f / n_rot); ++ const RopeKernelArgs k = rope.kernel_args(n_rot); ++ chunks(n, shift_blocks, chunk, [&](int64_t a, int64_t m) { ++ pooled_gather<<<(unsigned) m, 128, 0, st>>>(pooled, dim, b0 + a, scratch); ++ pooled_scatter<<<(unsigned) m, 128, 0, st>>>(pooled, dim, n_rot, block_size, b0 + a, shift_blocks, scratch, ++ theta_scale, k); ++ }); ++ return cudaGetLastError() == cudaSuccess; ++} ++ ++void scr_master_invalidate(const ScrKvLayer& l, int64_t from, int64_t to, void* stream) { ++ if (!l.master_valid || to <= from) return; ++ from = std::max(from, 0); ++ to = std::min(to, l.max_cells); ++ const long long n = to - from; ++ if (n > 0) invalidate<<<(unsigned) ((n + 255) / 256), 256, 0, static_cast(stream)>>>(l.master_valid, from, n); ++} ++ ++} // namespace strata::kernels +diff --git a/src/kernels/scr_move_parity.cu b/src/kernels/scr_move_parity.cu +new file mode 100644 +index 00000000..ee44a125 +--- /dev/null ++++ b/src/kernels/scr_move_parity.cu +@@ -0,0 +1,236 @@ ++// src/kernels/scr_move_parity.cu - scr_move.cu against a CPU reference. Needs a CUDA device (a few MB). ++#include "strata/kernels/scr_move.hpp" ++#include "strata/kernels/f16_bits.hpp" ++ ++#include ++ ++#include ++#include ++#include ++#include ++#include ++ ++using namespace strata::kernels; ++ ++namespace { ++int checks = 0, failures = 0; ++void check(bool ok, const char* what) { ++ ++checks; ++ if (!ok) { ++failures; std::fprintf(stderr, "FAIL: %s\n", what); } ++} ++ ++constexpr int H = 2, D = 256, PS = 64, R = 64, G = 64; ++constexpr int64_t CELLS = 4096; ++const RopeScaling rope{}; ++ ++void cpu_angle(int64_t pos, int pair, float& c, float& s) { ++ const float ts = powf((float) rope.freq_base, -2.0f / R); ++ const RopeKernelArgs k = rope.kernel_args(R); ++ rope_scaled_angle((float) pos * powf(ts, float(pair)), k.freq_scale, k.corr_low, k.corr_high, k.ext_factor, ++ k.attn_factor, pair, c, s); ++} ++int64_t row_of(int64_t pos, int h) { return ((pos / PS) * H + h) * PS + pos % PS; } ++ ++struct Host { ++ std::vector k, v; ++ std::vector ks, vs; ++ std::vector pre; // the unrotated K group the "fresh write" started from, per cell and head ++}; ++ ++// quantize 64 values as kv_append_q8 does ++void quant(const float* x, int8_t* q, uint16_t* sc) { ++ float amax = 0; ++ for (int i = 0; i < G; ++i) amax = std::max(amax, std::fabs(x[i])); ++ *sc = f16_from_f32(amax / 127.0f); ++ const float sf = f32_from_f16(*sc); ++ for (int i = 0; i < G; ++i) { ++ int c = sf > 0 ? (int) std::nearbyint(x[i] / sf) : 0; ++ q[i] = (int8_t) std::max(-127, std::min(127, c)); ++ } ++} ++ ++Host fresh(std::mt19937& rng) { ++ Host h; ++ h.k.resize(CELLS * H * D); h.v.resize(CELLS * H * D); ++ h.ks.resize(CELLS * H * (D / G)); h.vs.resize(CELLS * H * (D / G)); ++ h.pre.resize(CELLS * H * R); ++ std::normal_distribution n(0.0f, 1.0f); ++ std::vector x(D); ++ for (int64_t p = 0; p < CELLS; ++p) ++ for (int hh = 0; hh < H; ++hh) { ++ const int64_t row = row_of(p, hh); ++ float* pre = &h.pre[(p * H + hh) * R]; ++ for (int i = 0; i < R; ++i) pre[i] = n(rng); ++ for (int i = 0; i < R / 2; ++i) { ++ float c, s; ++ cpu_angle(p, i, c, s); ++ x[i] = pre[i] * c - pre[i + R / 2] * s; ++ x[i + R / 2] = pre[i] * s + pre[i + R / 2] * c; ++ } ++ for (int i = R; i < D; ++i) x[i] = n(rng); ++ for (int g = 0; g < D / G; ++g) quant(&x[g * G], &h.k[row * D + g * G], &h.ks[row * (D / G) + g]); ++ for (int i = 0; i < D; ++i) x[i] = n(rng); ++ for (int g = 0; g < D / G; ++g) quant(&x[g * G], &h.v[row * D + g * G], &h.vs[row * (D / G) + g]); ++ } ++ return h; ++} ++ ++struct Dev { ++ ScrKvLayer l; ++ void* scratch = nullptr; ++ int64_t chunk = 300; ++ Dev(const Host& h, bool master) { ++ l.kv_heads = H; l.head_dim = D; l.page_size = PS; l.n_rot = R; l.max_cells = CELLS; ++ cudaMalloc(&l.k_q, h.k.size()); cudaMalloc(&l.v_q, h.v.size()); ++ cudaMalloc(&l.k_scale, h.ks.size() * 2); cudaMalloc(&l.v_scale, h.vs.size() * 2); ++ cudaMemcpy(l.k_q, h.k.data(), h.k.size(), cudaMemcpyHostToDevice); ++ cudaMemcpy(l.v_q, h.v.data(), h.v.size(), cudaMemcpyHostToDevice); ++ cudaMemcpy(l.k_scale, h.ks.data(), h.ks.size() * 2, cudaMemcpyHostToDevice); ++ cudaMemcpy(l.v_scale, h.vs.data(), h.vs.size() * 2, cudaMemcpyHostToDevice); ++ if (master) { ++ cudaMalloc(&l.master, CELLS * H * R * 2); ++ cudaMalloc(&l.master_valid, CELLS); ++ cudaMemset(l.master_valid, 0, CELLS); ++ } ++ cudaMalloc(&scratch, scr_move_kv_scratch_bytes(l, chunk)); ++ } ++ ~Dev() { ++ cudaFree(l.k_q); cudaFree(l.v_q); cudaFree(l.k_scale); cudaFree(l.v_scale); ++ if (l.master) { cudaFree(l.master); cudaFree(l.master_valid); } ++ cudaFree(scratch); ++ } ++ bool move(int64_t src, int64_t n, int64_t shift) { ++ const bool ok = scr_move_kv(l, src, n, shift, rope, scratch, chunk, nullptr); ++ return cudaDeviceSynchronize() == cudaSuccess && ok; ++ } ++ Host get() const { ++ Host h; ++ h.k.resize(CELLS * H * D); h.v.resize(CELLS * H * D); ++ h.ks.resize(CELLS * H * (D / G)); h.vs.resize(CELLS * H * (D / G)); ++ cudaMemcpy(h.k.data(), l.k_q, h.k.size(), cudaMemcpyDeviceToHost); ++ cudaMemcpy(h.v.data(), l.v_q, h.v.size(), cudaMemcpyDeviceToHost); ++ cudaMemcpy(h.ks.data(), l.k_scale, h.ks.size() * 2, cudaMemcpyDeviceToHost); ++ cudaMemcpy(h.vs.data(), l.v_scale, h.vs.size() * 2, cudaMemcpyDeviceToHost); ++ return h; ++ } ++}; ++ ++// after moving [src, src+n) by `shift` from `before`: the moved cells' bytes and their rotated group ++void check_move(const Host& before, const Host& after, int64_t src, int64_t n, int64_t shift, const char* label) { ++ bool bytes = true, angle_ok = true; ++ double worst = 0; ++ for (int64_t i = 0; i < n; ++i) ++ for (int hh = 0; hh < H; ++hh) { ++ const int64_t a = row_of(src + i, hh), b = row_of(src + i + shift, hh); ++ if (std::memcmp(&before.v[a * D], &after.v[b * D], D) != 0 || ++ std::memcmp(&before.k[a * D + G], &after.k[b * D + G], D - G) != 0 || ++ std::memcmp(&before.vs[a * (D / G)], &after.vs[b * (D / G)], (D / G) * 2) != 0 || ++ std::memcmp(&before.ks[a * (D / G) + 1], &after.ks[b * (D / G) + 1], (D / G - 1) * 2) != 0) ++ bytes = false; ++ // the moved group against the exact rotation of the original values at the target ++ const float* pre = &before.pre[((src + i) * H + hh) * R]; ++ const float sf = f32_from_f16(after.ks[b * (D / G)]); ++ for (int j = 0; j < R / 2; ++j) { ++ float c, s; ++ cpu_angle(src + i + shift, j, c, s); ++ const float want0 = pre[j] * c - pre[j + R / 2] * s, want1 = pre[j] * s + pre[j + R / 2] * c; ++ const double e0 = std::fabs(after.k[b * D + j] * sf - want0), e1 = std::fabs(after.k[b * D + j + R / 2] * sf - want1); ++ worst = std::max(worst, std::max(e0, e1) / (sf > 0 ? sf : 1)); ++ } ++ } ++ // two quantizations (the stored codes, then the move): at most about 1.5 steps from the exact value ++ if (worst > 1.6) angle_ok = false; ++ char buf[256]; ++ std::snprintf(buf, sizeof buf, "%s: unrotated bytes moved exactly", label); check(bytes, buf); ++ std::snprintf(buf, sizeof buf, "%s: rotated group within 1.6 steps (worst %.3f)", label, worst); check(angle_ok, buf); ++} ++ ++void pooled_test(std::mt19937& rng) { ++ constexpr int DIM = 128, BR = 64, BS = 4; ++ constexpr int64_t ROWS = 1024; ++ std::normal_distribution n(0.0f, 1.0f); ++ std::vector pre(ROWS * DIM), rows(ROWS * DIM); ++ for (int64_t b = 0; b < ROWS; ++b) ++ for (int d = 0; d < DIM; ++d) pre[b * DIM + d] = n(rng); ++ auto rot = [&](const float* p, int64_t pos, float* out) { ++ for (int d = 0; d < DIM; ++d) out[d] = p[d]; ++ for (int j = 0; j < BR / 2; ++j) { ++ float c, s; ++ cpu_angle(pos, j, c, s); ++ out[j] = p[j] * c - p[j + BR / 2] * s; ++ out[j + BR / 2] = p[j] * s + p[j + BR / 2] * c; ++ } ++ }; ++ for (int64_t b = 0; b < ROWS; ++b) rot(&pre[b * DIM], BS * b, &rows[b * DIM]); ++ float *d_rows = nullptr, *d_scratch = nullptr; ++ cudaMalloc(&d_rows, rows.size() * 4); ++ cudaMalloc(&d_scratch, 100 * DIM * 4); ++ for (int64_t shift : {37, -29, 3, -500}) { ++ cudaMemcpy(d_rows, rows.data(), rows.size() * 4, cudaMemcpyHostToDevice); ++ const int64_t b0 = shift > 0 ? 100 : 600, cnt = 300; ++ const bool ok = scr_move_pooled(d_rows, ROWS, DIM, BR, BS, b0, cnt, shift, rope, d_scratch, 100, nullptr); ++ cudaDeviceSynchronize(); ++ std::vector got(rows.size()); ++ cudaMemcpy(got.data(), d_rows, got.size() * 4, cudaMemcpyDeviceToHost); ++ double worst = 0; ++ std::vector want(DIM); ++ for (int64_t i = 0; i < cnt; ++i) { ++ rot(&pre[(b0 + i) * DIM], BS * (b0 + i + shift), want.data()); ++ for (int d = 0; d < DIM; ++d) worst = std::max(worst, (double) std::fabs(got[(b0 + i + shift) * DIM + d] - want[d])); ++ } ++ char buf[128]; ++ std::snprintf(buf, sizeof buf, "pooled shift %lld: rows at their new angle (worst %.2e)", (long long) shift, worst); ++ check(ok && worst < 2e-4, buf); ++ } ++ cudaFree(d_rows); ++ cudaFree(d_scratch); ++} ++} // namespace ++ ++int main() { ++ std::mt19937 rng(7); ++ const Host h = fresh(rng); ++ // overlapping moves, both directions, several chunks ++ for (int64_t shift : {36, -40, 4, -1000, 2000}) { ++ Dev d(h, false); ++ const int64_t src = shift > 0 ? 100 : 1500, n = 1000; ++ char label[64]; ++ std::snprintf(label, sizeof label, "move %lld", (long long) shift); ++ check(d.move(src, n, shift), label); ++ check_move(h, d.get(), src, n, shift, label); ++ } ++ // the master: two moves equal one move by their sum, byte for byte ++ { ++ Dev once(h, true), twice(h, true); ++ check(once.move(200, 1000, 36), "master: one move"); ++ check(twice.move(200, 1000, 8) && twice.move(208, 1000, 28), "master: two moves"); ++ const Host a = once.get(), b = twice.get(); ++ bool same = true; ++ for (int64_t i = 0; i < 1000; ++i) ++ for (int hh = 0; hh < H; ++hh) { ++ const int64_t r = row_of(236 + i, hh); ++ if (std::memcmp(&a.k[r * D], &b.k[r * D], D) || a.ks[r * (D / G)] != b.ks[r * (D / G)]) same = false; ++ } ++ check(same, "master: two moves are byte-identical to one move by the sum"); ++ // without the master the second move starts from the first one's codes, so they differ ++ Dev plain(h, false); ++ plain.move(200, 1000, 8); ++ plain.move(208, 1000, 28); ++ const Host c = plain.get(); ++ int64_t differ = 0; ++ for (int64_t i = 0; i < 1000; ++i) { ++ const int64_t r = row_of(236 + i, 0); ++ differ += std::memcmp(&a.k[r * D], &c.k[r * D], G) != 0; ++ } ++ std::printf(" without the master %lld of 1000 cells differ after two moves\n", (long long) differ); ++ } ++ // invalid arguments move nothing ++ { ++ Dev d(h, false); ++ check(!scr_move_kv(d.l, 4000, 200, 0, rope, d.scratch, d.chunk, nullptr), "past the end refused"); ++ check(!scr_move_kv(d.l, 10, 100, -20, rope, d.scratch, d.chunk, nullptr), "before zero refused"); ++ } ++ pooled_test(rng); ++ std::printf("scr_move_parity: %d checks, %d failed\n", checks, failures); ++ return failures == 0 ? 0 : 1; ++} +diff --git a/src/program/generate.cpp b/src/program/generate.cpp +index ba00f94d..3f1c8eab 100644 +--- a/src/program/generate.cpp ++++ b/src/program/generate.cpp +@@ -81,6 +81,8 @@ + #include "strata/program/batch_groups.hpp" + #include "strata/kernels/cvec.hpp" + #include "strata/core/progress.hpp" ++#include "strata/core/scr.hpp" ++#include "strata/kernels/scr_move.hpp" + #include "strata/core/device.hpp" + #include "strata/core/emulate.hpp" + #ifndef NOMINMAX +@@ -8127,6 +8129,24 @@ int main(int argc, char** argv) { + std::vector live; + std::vector live_imgs, req_imgs; + bool live_ok = false; ++ // SCR (STRATA_SCR=1, core/scr.hpp): after an edit in the middle of a prompt the cells of the text that survived ++ // are moved instead of read. `live_pads`: where `live` holds the pad tokens a splice read (the server's prompt ++ // has none, so the next request gets them back); `scr_master_ok`: the cell ranges whose K master is current. ++ static const bool scr_on = [] { const char* e = std::getenv("STRATA_SCR"); return e != nullptr && e[0] == '1'; }(); ++ static const int64_t scr_min_gain = [] { // a splice must save at least this many reads ++ const char* e = std::getenv("STRATA_SCR_MIN_GAIN"); ++ return e != nullptr ? (int64_t) std::max(0LL, std::atoll(e)) : (int64_t) 1024; ++ }(); ++ static const int64_t scr_pad = [] { // the pad token ("\n" in Qwen3.8's vocabulary) ++ const char* e = std::getenv("STRATA_SCR_PAD"); ++ return e != nullptr ? (int64_t) std::atoll(e) : (int64_t) 198; ++ }(); ++ std::vector live_pads; ++ std::vector> scr_master_ok; ++ struct ScrLayerMem { uint16_t* master = nullptr; uint8_t* valid = nullptr; }; ++ std::map scr_mem; ++ std::map> scr_scratch; // per device: the K/V move scratch, the pooled scratch ++ constexpr int64_t kScrChunk = 2048, kScrPooledChunk = 512; + // #879: a window met non-finite logits: the next request drops every cached state first and reads from token 0 + bool nan_flush = false; + std::vector checks; +@@ -10441,7 +10461,7 @@ int main(int argc, char** argv) { + }; + strata::core::progress_at(save ? "saving a session" : "restoring a session"); + if (path.empty()) { refuse("missing path"); continue; } +- if (!stages.empty() || multi_gpu) { refuse("session files do not support --layer-split"); continue; } ++ if (stages.empty() && multi_gpu) { refuse("session files do not support --layer-split"); continue; } + if (o.peer_device >= 1) { refuse("session files do not support --peer-device"); continue; } + if (o.batch > 0) { refuse("session files do not support --batch (parallel requests)"); continue; } + if (o.prompt_cache <= 0) { refuse("session files need --prompt-cache > 0"); continue; } +@@ -10518,26 +10538,71 @@ int main(int argc, char** argv) { + refuse(why, strata::core::SessionError::memory); + continue; + } +- kept = disk_checks.size(); +- const strata::core::ConversationView view{live, live_imgs, disk_checks, cvec_cached}; +- strata::core::SavedConversation meta; +- std::vector sources; +- // the live running state comes off the device in one synchronous copy +- blocking("capture", sl.state_bytes == UINT64_MAX ? 0 : sl.state_bytes); +- if (!strata::core::conversation_snapshot_sources(meta, sources, view, ss, g, mtp.kv_state(), +- err)) { +- refuse(err, strata::core::SessionError::io); +- continue; ++ // STRATA_SESSION_CHAIN=1: the whole chain, so a restored session keeps its root (the end of the ++ // system prompt) and the periodic checkpoints SCR restarts from, not only the deepest one ++ static const bool save_chain = [] { const char* e = std::getenv("STRATA_SESSION_CHAIN"); return e != nullptr && e[0] == '1'; }(); ++ if (save_chain && checks.size() > disk_checks.size()) { ++ uint64_t chain_bytes = 0; ++ for (const ConvCheckpoint& c : checks) chain_bytes += c.bytes(); ++ std::string low; ++ if (admit(chain_bytes, low)) disk_checks = checks; ++ else std::fprintf(stderr, "strata serve: session save keeps only the deepest checkpoint: %s\n", low.c_str()); + } +- capture_ms = ms(); ++ kept = disk_checks.size(); + strata::core::SessionWriteOptions wo; + wo.min_free_bytes = (uint64_t) o.session_min_free_mib << 20; + wo.progress = moving; + wo.phase = blocking; + strata::core::SessionStatus st; +- if (!strata::core::session_file_write(path, meta, sources, id, bytes, err, wo, &st)) { +- refuse(err, st.error, st.published); +- continue; ++ if (!stages.empty()) { ++ // a layer split: one image per stage, captured as parking captures them (the draft ++ // layer's K/V with the stage that holds the drafter), from copies of the checkpoints ++ // moved apart into their stage parts; a checkpoint without every part is not saved ++ const size_t n_st = stages.size(); ++ const strata::core::QsaState* draft0 = (n_st > 0 && !mtp_first) || !use_mtp ? nullptr : &mtp.kv_state(); ++ auto draft_of = [&](size_t k) -> const strata::core::QsaState* { ++ return use_mtp && !mtp_first && k + 1 == n_st ? &mtp.kv_state() : nullptr; ++ }; ++ auto cs = strata::core::conversation_checkpoints_split(std::move(disk_checks), n_st); ++ kept = cs.stage0.size(); ++ strata::core::SavedConversation image; ++ blocking("capture", sl.state_bytes == UINT64_MAX ? 0 : sl.state_bytes * (n_st + 1)); ++ const strata::core::ConversationView view{live, live_imgs, cs.stage0, cvec_cached}; ++ bool captured = strata::core::conversation_snapshot_save(image, view, ss, g, draft0, err); ++ for (size_t k = 0; captured && k < n_st; ++k) { ++ const strata::core::OnDevice on(stages[k]->dev); ++ if (cudaDeviceSynchronize() != cudaSuccess) { err = "stage sync failed"; captured = false; break; } ++ const strata::core::ConversationView view_k{live, live_imgs, cs.parts[k], cvec_cached}; ++ strata::core::SavedConversation part; ++ captured = strata::core::conversation_snapshot_save(part, view_k, stages[k]->ss, g, ++ draft_of(k), err); ++ if (captured) image.stage_images.push_back(std::move(part)); ++ } ++ if (!captured) { ++ refuse(err, strata::core::SessionError::io); ++ continue; ++ } ++ capture_ms = ms(); ++ if (!strata::core::session_file_write(path, image, id, bytes, err, wo, &st)) { ++ refuse(err, st.error, st.published); ++ continue; ++ } ++ } else { ++ const strata::core::ConversationView view{live, live_imgs, disk_checks, cvec_cached}; ++ strata::core::SavedConversation meta; ++ std::vector sources; ++ // the live running state comes off the device in one synchronous copy ++ blocking("capture", sl.state_bytes == UINT64_MAX ? 0 : sl.state_bytes); ++ if (!strata::core::conversation_snapshot_sources(meta, sources, view, ss, g, mtp.kv_state(), ++ err)) { ++ refuse(err, strata::core::SessionError::io); ++ continue; ++ } ++ capture_ms = ms(); ++ if (!strata::core::session_file_write(path, meta, sources, id, bytes, err, wo, &st)) { ++ refuse(err, st.error, st.published); ++ continue; ++ } + } + if (st.dir_flush_unsupported) + std::fprintf(stderr, "strata serve: session saved to %s; this filesystem cannot flush a " +@@ -10546,6 +10611,23 @@ int main(int argc, char** argv) { + refuse("not enough RAM to save the session", strata::core::SessionError::memory); + continue; + } ++ // SCR: the pads the session read go beside the file (its format has no place for them); a restore ++ // takes them back only when the file's tokens are the ones they were saved with ++ { ++ const std::string pads_path = path + ".pads"; ++ if (live_pads.empty()) { ++ std::remove(pads_path.c_str()); ++ } else if (std::FILE* pf = std::fopen(pads_path.c_str(), "wb")) { ++ const uint64_t head[3] = {(uint64_t) live.size(), fnv1a(live.data(), live.size() * sizeof(int32_t)), ++ (uint64_t) live_pads.size()}; ++ const bool ok = std::fwrite(head, sizeof head, 1, pf) == 1 && ++ std::fwrite(live_pads.data(), sizeof(int64_t), live_pads.size(), pf) == live_pads.size(); ++ if (std::fclose(pf) != 0 || !ok) { ++ std::remove(pads_path.c_str()); ++ std::fprintf(stderr, "strata serve: session save: the pads were not written beside %s\n", path.c_str()); ++ } ++ } ++ } + std::fprintf(stderr, "strata serve: session saved %zu tokens, %zu of %zu checkpoints, %zu bytes to %s " + "in %.1f ms (state %.1f ms)\n", live.size(), kept, checks.size(), bytes, + path.c_str(), ms(), capture_ms); +@@ -10575,6 +10657,18 @@ int main(int argc, char** argv) { + refuse(err, strata::core::SessionError::io); + continue; + } ++ bool stage_limits_ok = true; ++ for (const auto& st_k : stages) { // a layer split: every later stage's own carve ++ strata::core::SessionReadLimits lk; ++ if (!strata::core::conversation_session_read_limits( ++ lk, st_k->ss, g, mtp.kv_state(), (uint64_t) o.max_context, ++ (uint64_t) std::max(o.prompt_cache, 1), err)) { stage_limits_ok = false; break; } ++ limits.stage_limits.push_back(std::move(lk)); ++ } ++ if (!stage_limits_ok) { ++ refuse(err, strata::core::SessionError::io); ++ continue; ++ } + strata::core::SessionStatus st; + if (!strata::core::session_file_read(path, id, image, bytes, err, limits, &st)) { + refuse(err, st.error); +@@ -10587,7 +10681,21 @@ int main(int argc, char** argv) { + const double read_ms = ms(); + // the whole image against this engine, still without any device write + blocking("validate", bytes); +- if (!strata::core::conversation_snapshot_validate(image, ss, g, mtp.kv_state(), err)) { ++ // with a layer split the draft layer's K/V is in the image of the stage that holds the drafter ++ const size_t n_st = stages.size(); ++ const strata::core::QsaState* draft_main = ++ n_st == 0 ? &mtp.kv_state() : use_mtp && mtp_first ? &mtp.kv_state() : nullptr; ++ auto draft_of = [&](size_t k) -> const strata::core::QsaState* { ++ return use_mtp && !mtp_first && k + 1 == n_st ? &mtp.kv_state() : nullptr; ++ }; ++ bool valid = image.stage_images.size() == n_st; ++ if (!valid) err = "the file holds another number of GPU stages than this runtime"; ++ for (size_t k = 0; valid && k < n_st; ++k) { ++ const strata::core::OnDevice on(stages[k]->dev); ++ valid = strata::core::conversation_snapshot_validate(image.stage_images[k], stages[k]->ss, g, ++ draft_of(k), err); ++ } ++ if (!valid || !strata::core::conversation_snapshot_validate(image, ss, g, draft_main, err)) { + refuse(err); + continue; + } +@@ -10600,8 +10708,16 @@ int main(int argc, char** argv) { + live_ok = false; + // host -> device in synchronous copies of the whole state: one bounded allowance + blocking("transfer", bytes); +- if (strata::core::conversation_snapshot_restore(image, ss, g, mtp.kv_state(), err) != +- strata::core::ConversationRestore::restored) { ++ bool transferred = strata::core::conversation_snapshot_restore(image, ss, g, draft_main, err) == ++ strata::core::ConversationRestore::restored; ++ for (size_t k = 0; transferred && k < n_st; ++k) { ++ const strata::core::OnDevice on(stages[k]->dev); ++ transferred = strata::core::conversation_snapshot_restore(image.stage_images[k], stages[k]->ss, g, ++ draft_of(k), err) == ++ strata::core::ConversationRestore::restored; ++ if (transferred && cudaDeviceSynchronize() != cudaSuccess) { err = "stage sync failed"; transferred = false; } ++ } ++ if (!transferred) { + // validated above: a failure here is a transfer failure, after device writes began - never + // decode from a partial state; the server starts the engine again + std::fprintf(stderr, "strata serve: session restore %s: transfer failed: %s\n", path.c_str(), +@@ -10613,7 +10729,32 @@ int main(int argc, char** argv) { + } + live = std::move(image.live.ids); + live_imgs = std::move(image.live.imgs); +- checks = std::move(image.checkpoints); ++ live_pads.clear(); ++ scr_master_ok.clear(); ++ if (std::FILE* pf = std::fopen((path + ".pads").c_str(), "rb")) { // the pads saved beside it ++ uint64_t head[3] = {}; ++ std::vector pads; ++ bool ok = std::fread(head, sizeof head, 1, pf) == 1 && head[0] == (uint64_t) live.size() && ++ head[1] == fnv1a(live.data(), live.size() * sizeof(int32_t)) && head[2] <= live.size(); ++ if (ok) { ++ pads.resize((size_t) head[2]); ++ ok = std::fread(pads.data(), sizeof(int64_t), pads.size(), pf) == pads.size(); ++ } ++ std::fclose(pf); ++ for (size_t i = 0; ok && i < pads.size(); ++i) ++ ok = pads[i] >= 0 && pads[i] < (int64_t) live.size() && live[(size_t) pads[i]] == (int32_t) scr_pad && ++ (i == 0 || pads[i] > pads[i - 1]); ++ if (ok) live_pads = std::move(pads); ++ else std::fprintf(stderr, "strata serve: session restore: the pads beside %s do not match it\n", path.c_str()); ++ } ++ if (n_st == 0) { ++ checks = std::move(image.checkpoints); ++ } else { // give each checkpoint back its stage parts; none if they do not line up ++ strata::core::ConversationCheckpointSplit cs; ++ cs.stage0 = std::move(image.checkpoints); ++ for (auto& si : image.stage_images) cs.parts.push_back(std::move(si.checkpoints)); ++ if (!strata::core::conversation_checkpoints_merge(std::move(cs), checks)) checks.clear(); ++ } + for (const ConvCheckpoint& c : checks) check_clock = std::max(check_clock, c.used); + cvec_cached = image.cvec; + live_ok = true; +@@ -10730,7 +10871,13 @@ int main(int argc, char** argv) { + std::printf("ERR bad request: %s\n", pe.empty() ? "max_new" : pe.c_str()); + continue; + } +- const int64_t n = (int64_t) ids.size(); ++ std::vector req_pads; // SCR: the positions in `ids` of the pads this request reads ++ if (scr_on && !geni && live_ok && !live_pads.empty()) { ++ std::vector carried; ++ req_pads = strata::core::scr_carry(ids, live, live_pads, scr_pad, carried); ++ if (!req_pads.empty()) ids = std::move(carried); ++ } ++ int64_t n = (int64_t) ids.size(); + req_imgs.clear(); + if (geni && !o.vision) { std::printf("ERR this engine was started without --vision\n"); continue; } + if (geni || !mrope_identity) { +@@ -11009,6 +11156,104 @@ int main(int argc, char** argv) { + incoming.reset(); + err.clear(); + } ++ // SCR: a prompt the live session holds most of, edited in the middle. The cells of the text that survived ++ // are moved to their new places (scr_move.hpp), the parts between them read, and the GDN state taken from ++ // the checkpoint at the end of the last moved run (scr.hpp); everything after it is read as usual. ++ const bool scr_cells_kept = !incoming && slot_source < 0; ++ strata::core::ScrPlan scr; ++ std::vector scr_order; // the runs in an order that never overwrites a source still to move ++ ConvCheckpoint scr_adopt; ++ // the QSA states SCR moves: each stage's, on its device and stream ++ struct ScrState { strata::core::QsaState* st; int dev; void* stream; }; ++ auto scr_states = [&] { ++ std::vector v; ++ for (int64_t j = 0; j < ss.qsa_alloc; ++j) v.push_back({&ss.qsa_states[ss.qsa_ord0 + j], 0, main_cs}); ++ for (auto& sg : stages) ++ for (int64_t j = 0; j < sg->ss.qsa_alloc; ++j) ++ v.push_back({&sg->ss.qsa_states[sg->ss.qsa_ord0 + j], sg->dev, (void*) sg->stream}); ++ return v; ++ }; ++ auto scr_supported = [&] { ++ for (const ScrState& s : scr_states()) { ++ const strata::core::QsaState& q = *s.st; ++ if (!q.kv_int8 || q.kv_hybrid || q.kv_rot || q.kv_q4 || q.kv_mode != 0 || q.kv_elastic >= 0 || ++ !q.k_q || !q.idx_pooled || !q.idx_dead) return false; ++ } ++ return true; ++ }; ++ if (scr_on && scr_cells_kept && !from_live && req_imgs.empty() && live_imgs.empty() && live_ok && ++ want_cvec == cvec_cached && !kvg.on && o.prompt_cache > 0 && scr_supported()) { ++ const auto t0 = Clock::now(); ++ std::vector cks; // the GDN states the plan may continue from: the chain and the live state ++ for (const ConvCheckpoint& c : checks) if (c.imgs.empty()) cks.push_back((int64_t) c.ids.size()); ++ cks.push_back((int64_t) live.size()); ++ strata::core::ScrPlan plan = strata::core::scr_plan(live, ids, cks, strata::core::ScrParams{}); ++ const int64_t spliced = n + plan.pads() - plan.restore - plan.reused(); ++ // a run moves once its target holds no source of a run still to move (llama-cpp-bonsai's rule) ++ std::vector pending = plan.runs; ++ while (!pending.empty()) { ++ const auto next = std::find_if(pending.begin(), pending.end(), [&](const strata::core::ScrRun& r) { ++ return std::none_of(pending.begin(), pending.end(), [&](const strata::core::ScrRun& q) { ++ return &q != &r && q.old_start < r.at + r.n && r.at < q.old_start + q.n; ++ }); ++ }); ++ if (next == pending.end()) break; ++ scr_order.push_back(*next); ++ pending.erase(next); ++ } ++ std::string why; ++ if (plan.empty()) why = "nothing to move"; ++ else if (n + plan.pads() + max_new + 8 > o.max_context) why = "the pads do not fit the context"; ++ else if (spliced + scr_min_gain > n - resume) why = "not worth it"; ++ else if (!pending.empty()) why = "the runs cannot be moved in any order"; ++ if (why.empty()) { ++ // the GDN state after the last run: the live state, or the checkpoint there; its indexer's last ++ // block moves with the run ++ bool taken = false; ++ if (plan.adopt == (int64_t) live.size()) { ++ scr_adopt.ids = live; ++ taken = gpu_sync_ok() && checkpoint_save(scr_adopt, ss, g); ++ for (auto& sg : stages) { ++ const strata::core::OnDevice on(sg->dev); ++ ConvCheckpoint part; ++ part.ids = live; ++ taken = taken && gpu_sync_ok() && checkpoint_save(part, sg->ss, g); ++ scr_adopt.stage_parts.push_back(std::move(part)); ++ } ++ } else { ++ for (const ConvCheckpoint& c : checks) ++ if ((int64_t) c.ids.size() == plan.adopt) { scr_adopt = c; taken = c.stage_parts.size() == stages.size(); } ++ } ++ const int32_t shift = (int32_t) plan.runs.back().shift(); ++ auto move_block_pos = [shift](ConvCheckpoint& c) { ++ for (size_t b = 0; b + sizeof(int32_t) <= c.block_pos.size(); b += sizeof(int32_t)) { ++ int32_t v; ++ std::memcpy(&v, c.block_pos.data() + b, sizeof v); ++ v += shift; ++ std::memcpy(c.block_pos.data() + b, &v, sizeof v); ++ } ++ }; ++ move_block_pos(scr_adopt); ++ for (ConvCheckpoint& part : scr_adopt.stage_parts) move_block_pos(part); ++ if (!taken) why = "the state after the last run could not be taken"; ++ } ++ if (!why.empty()) { ++ std::fprintf(stderr, "strata serve: scr: %s (%s)\n", why.c_str(), plan.str().c_str()); ++ scr_order.clear(); ++ } else { ++ for (const int64_t p : strata::core::scr_pad_positions(plan)) req_pads.push_back(p); ++ std::sort(req_pads.begin(), req_pads.end()); ++ ids = strata::core::scr_padded(ids, plan, scr_pad); ++ n = (int64_t) ids.size(); ++ cur = ids; ++ resume = plan.restore; ++ scr = std::move(plan); ++ std::fprintf(stderr, "strata serve: scr: %s; moves %lld, reads %lld instead of %lld (planned in %.1f ms)\n", ++ scr.str().c_str(), (long long) scr.reused(), (long long) spliced, ++ (long long) (n - scr.pads() - resume), ++ std::chrono::duration(Clock::now() - t0).count()); ++ } ++ } + // Preserve the outgoing branch before any checkpoint rewind, reset, + // or incoming restore overwrites the positional state it requires. + // pin=N: this request resumes exactly at the pinned shared prefix of the conversation the session holds +@@ -11020,7 +11265,8 @@ int main(int argc, char** argv) { + resume == req_pin && live_ok) + for (const ConvCheckpoint& c : checks) + if ((int64_t) c.ids.size() == resume && c.pinned) pin_sibling = true; +- if ((!from_live || incoming || slot_source >= 0) && !pin_sibling && !park_current(incoming ? incoming->bytes() : 0)) { ++ if (scr.empty() && (!from_live || incoming || slot_source >= 0) && !pin_sibling && ++ !park_current(incoming ? incoming->bytes() : 0)) { + std::printf("ERR %s\n", err.c_str()); + return 1; + } +@@ -11053,6 +11299,7 @@ int main(int argc, char** argv) { + live = bs[(size_t) slot_source].ids; + } + live_imgs.clear(); ++ live_pads.clear(); + resume = slot_tokens; + from_live = true; + std::fprintf(stderr, "strata batch: slot %d gave back %lld tokens of this conversation (%s) in " +@@ -11103,6 +11350,7 @@ int main(int argc, char** argv) { + } + live = std::move(incoming->live.ids); + live_imgs = std::move(incoming->live.imgs); ++ live_pads.clear(); + if (stages.empty()) { + checks = std::move(incoming->checkpoints); + } else { // give each checkpoint back its stage parts (moved, not copied); none if they do not line up +@@ -11151,11 +11399,12 @@ int main(int argc, char** argv) { + } + int64_t reread_to = -1; // STRATA_CKPT_REREAD only: read [0, reread_to) again instead of restoring + if (resume == 0) { +- strata::core::session_zero(ss, g, nullptr, main_cs); ++ // a splice from token 0 keeps the cells it is about to move ++ strata::core::session_zero_running(ss, g, nullptr, main_cs, !scr.empty()); + cudaStreamSynchronize(main_stream); + for (auto& st : stages) { + const strata::core::OnDevice on(st->dev); +- strata::core::session_zero(st->ss, g, nullptr, (void*) st->stream); ++ strata::core::session_zero_running(st->ss, g, nullptr, (void*) st->stream, !scr.empty()); + cudaStreamSynchronize(st->stream); + } + checks.clear(); +@@ -11195,6 +11444,83 @@ int main(int argc, char** argv) { + // KV streaming: the drafter's ring may hold cells past `resume` from a longer turn; the main layers' + // host copies and slots are always current (every writer writes both), so they need nothing + if (use_mtp && resume > 0 && reread_to <= 0) mtp.kv_restore(resume); ++ // SCR: the cells move before any part is read (a part may write where a later run's source still lies). ++ // A K master is current only where no read has written since it was taken: up to where this request ++ // starts reading, and in the cells a splice moves. ++ auto scr_trim = [&](int64_t upto) { ++ std::vector> kept; ++ for (const auto& [a, b] : scr_master_ok) ++ if (a < upto) kept.push_back({a, std::min(b, upto)}); ++ scr_master_ok = std::move(kept); ++ }; ++ if (!scr_cells_kept) scr_master_ok.clear(); ++ else scr_trim(resume); ++ std::vector> scr_first; // per QSA state, per run: the pooled key of its first block ++ if (!scr.empty()) { ++ const auto tm = Clock::now(); ++ const auto& rope = strata::kernels::rope_scaling(); ++ const int64_t page = strata::kernels::qsa_real_shapes().page_size, D = g.idx_key_dim, ++ bl = strata::kernels::qsa_real_shapes().idx_block; ++ bool moved = true; ++ for (const ScrState& s : scr_states()) { ++ const strata::core::OnDevice on(s.dev); ++ strata::core::QsaState& q = *s.st; ++ ScrLayerMem& mem = scr_mem[&q]; ++ if (mem.master == nullptr) { // pinned host memory every device reads (FP16 [cell][head][n_rot]) ++ const size_t mb = (size_t) q.max_cells * g.n_head_kv * g.n_rot * 2; ++ moved = moved && cudaHostAlloc((void**) &mem.master, mb, cudaHostAllocMapped | cudaHostAllocPortable) == cudaSuccess && ++ cudaHostAlloc((void**) &mem.valid, (size_t) q.max_cells, cudaHostAllocMapped | cudaHostAllocPortable) == cudaSuccess; ++ if (!moved) break; ++ } ++ std::memset(mem.valid, 0, (size_t) q.max_cells); ++ for (const auto& [a, b] : scr_master_ok) std::memset(mem.valid + a, 1, (size_t) (b - a)); ++ auto& sc = scr_scratch[s.dev]; ++ strata::kernels::ScrKvLayer l; ++ l.k_q = q.k_q; l.v_q = q.v_q; l.k_scale = q.k_scale; l.v_scale = q.v_scale; ++ l.master = mem.master; l.master_valid = mem.valid; ++ l.kv_heads = (int) g.n_head_kv; l.head_dim = (int) g.head_dim; l.page_size = (int) page; ++ l.n_rot = (int) g.n_rot; l.max_cells = q.max_cells; ++ if (sc.first == nullptr) { ++ moved = moved && cudaMalloc(&sc.first, strata::kernels::scr_move_kv_scratch_bytes(l, kScrChunk)) == cudaSuccess && ++ cudaMalloc(&sc.second, (size_t) kScrPooledChunk * D * sizeof(float)) == cudaSuccess; ++ if (!moved) break; ++ } ++ for (const strata::core::ScrRun& r : scr_order) { ++ moved = moved && strata::kernels::scr_move_kv(l, r.old_start, r.n, r.shift(), rope, sc.first, kScrChunk, s.stream) && ++ strata::kernels::scr_move_pooled(q.idx_pooled, q.idx_pooled_rows, (int) D, (int) g.n_rot, (int) bl, ++ r.old_start / bl, r.n / bl, r.shift() / bl, rope, sc.second, ++ kScrPooledChunk, s.stream); ++ } ++ // each run's first pooled key, which the part read before it overwrites with the spare; the spare ++ // after each run, where the next part (or the tail) goes on ++ std::vector first; ++ for (const strata::core::ScrRun& r : scr.runs) { ++ float* row = nullptr; ++ moved = moved && cudaMalloc(&row, (size_t) D * sizeof(float)) == cudaSuccess; ++ if (!moved) break; ++ cudaMemcpyAsync(row, q.idx_pooled + (r.at / bl) * D, (size_t) D * sizeof(float), ++ cudaMemcpyDeviceToDevice, (cudaStream_t) s.stream); ++ first.push_back(row); ++ } ++ for (const strata::core::ScrRun& r : scr.runs) ++ cudaMemcpyAsync(q.idx_pooled + ((r.at + r.n) / bl) * D, q.idx_dead, (size_t) D * sizeof(float), ++ cudaMemcpyDeviceToDevice, (cudaStream_t) s.stream); ++ scr_first.push_back(std::move(first)); ++ moved = moved && cudaStreamSynchronize((cudaStream_t) s.stream) == cudaSuccess; ++ if (!moved) break; ++ } ++ if (!moved) { ++ // the cells may be half moved: the engine stops, as for any CUDA failure in the session ++ std::printf("ERR scr: moving the cached cells failed\n"); ++ std::fflush(stdout); ++ return 1; ++ } ++ scr_trim(scr.restore); ++ for (const strata::core::ScrRun& r : scr.runs) scr_master_ok.push_back({r.at, r.at + r.n}); ++ std::fprintf(stderr, "strata serve: scr: moved %lld cells in %zu runs over %zu QSA states in %.1f ms\n", ++ (long long) scr.reused(), scr.runs.size(), scr_first.size(), ++ std::chrono::duration(Clock::now() - tm).count()); ++ } + tr("request", n, geni ? 1 : 0); + if (use_mtp) mtp.set_prompt_len(n); + const int64_t read_from = reread_to > 0 ? 0 : resume; +@@ -11761,12 +12087,65 @@ int main(int argc, char** argv) { + "%lld)\n", (long long) req_pin, (long long) read_from); + } + std::vector cuts = {reread_to, root_at, message_at, turn_at, n - 1}; +- if (pin_at >= 0) { +- cuts.push_back(pin_at); +- std::sort(cuts.begin(), cuts.end()); // the skipped -1s first, n - 1 still last +- } ++ if (pin_at >= 0) cuts.push_back(pin_at); ++ for (const strata::core::ScrRun& r : scr.runs) cuts.push_back(r.at); // SCR: each part ends at a run ++ if (pin_at >= 0 || !scr.empty()) std::sort(cuts.begin(), cuts.end()); // the skipped -1s first, n - 1 still last + int64_t at = read_from; ++ // SCR: the read reached a run. Its first pooled key comes back (the part's spare overwrote it), the read ++ // goes on after it, and after the last run the GDN state and the indexer's tail are the checkpoint's there. ++ size_t scr_next = 0; ++ auto scr_take = [&]() -> bool { ++ const strata::core::ScrRun& r = scr.runs[scr_next]; ++ const int64_t D = g.idx_key_dim, bl = strata::kernels::qsa_real_shapes().idx_block; ++ const auto states = scr_states(); ++ bool ok = true; ++ for (size_t i = 0; i < states.size(); ++i) { ++ const strata::core::OnDevice on(states[i].dev); ++ float*& row = scr_first[i][scr_next]; ++ ok = ok && cudaMemcpyAsync(states[i].st->idx_pooled + (r.at / bl) * D, row, (size_t) D * sizeof(float), ++ cudaMemcpyDeviceToDevice, (cudaStream_t) states[i].stream) == cudaSuccess && ++ cudaStreamSynchronize((cudaStream_t) states[i].stream) == cudaSuccess; ++ cudaFree(row); ++ row = nullptr; ++ } ++ at = r.at + r.n; ++ pp_reached = at; ++ if (++scr_next == scr.runs.size()) { ++ // STRATA_SCR_ADOPT_GDN=0: the indexer's tail from the checkpoint, but the GDN and PLE state as the ++ // parts read it (the checkpoint's state also holds the text the edit removed) ++ static const bool adopt_gdn = [] { const char* e = std::getenv("STRATA_SCR_ADOPT_GDN"); return e == nullptr || e[0] != '0'; }(); ++ if (!adopt_gdn) { ++ ConvCheckpoint now; ++ ok = ok && gpu_sync_ok() && checkpoint_save(now, ss, g); ++ scr_adopt.gdn = std::move(now.gdn); ++ scr_adopt.ple = std::move(now.ple); ++ for (size_t i = 0; ok && i < stages.size(); ++i) { ++ const strata::core::OnDevice on(stages[i]->dev); ++ ConvCheckpoint part; ++ ok = gpu_sync_ok() && checkpoint_save(part, stages[i]->ss, g); ++ scr_adopt.stage_parts[i].gdn = std::move(part.gdn); ++ scr_adopt.stage_parts[i].ple = std::move(part.ple); ++ } ++ } ++ ok = ok && gpu_sync_ok() && checkpoint_restore(scr_adopt, ss, g); ++ for (size_t i = 0; ok && i < stages.size(); ++i) { ++ const strata::core::OnDevice on(stages[i]->dev); ++ ok = gpu_sync_ok() && checkpoint_restore(scr_adopt.stage_parts[i], stages[i]->ss, g); ++ } ++ } ++ // the n-gram hash's predecessors are the run's last tokens ++ const int32_t p0 = at >= 2 ? (int32_t) ids[(size_t) (at - 2)] : -1, p1 = (int32_t) ids[(size_t) (at - 1)]; ++ ss.ple_prev[0] = p0; ss.ple_prev[1] = p1; ++ for (auto& sg : stages) { sg->ss.ple_prev[0] = p0; sg->ss.ple_prev[1] = p1; } ++ return ok; ++ }; + for (const int64_t to : cuts) { ++ while (scr_next < scr.runs.size() && scr.runs[scr_next].at == at) ++ if (!scr_take()) { ++ std::printf("ERR scr: restoring the state after a moved run failed\n"); ++ std::fflush(stdout); ++ return 1; ++ } + if (to <= at) continue; + err.clear(); + const bool win = windows_ok(at, to); +@@ -12792,6 +13171,7 @@ int main(int argc, char** argv) { + // a prompt stopped halfway leaves the session somewhere between two chunks: nothing to continue from + // (the checkpoints taken while reading it are still good) + live.swap(consumed); ++ live_pads = req_pads; + live_imgs = imgs_below(req_imgs, (int64_t) live.size()); + live_ok = o.prompt_cache > 0 && req_ckpt; // ckpt=0: nothing to continue or park (#830) + if (nf_hit) live_ok = false; // #879 (the next request's start drops the rest) +@@ -12801,6 +13181,8 @@ int main(int argc, char** argv) { + // pp_reached, so it stays the live conversation - the next request continues from it or parks it. + // Without this a cancelled request that had restored a parked conversation dropped it. + live.assign(cur.begin(), cur.begin() + pp_reached); ++ live_pads.clear(); ++ for (const int64_t p : req_pads) if (p < pp_reached) live_pads.push_back(p); + live_imgs = imgs_below(req_imgs, pp_reached); + live_ok = true; + std::fprintf(stderr, "strata serve: cancelled request keeps its conversation: %lld tokens live\n", +@@ -12935,8 +13317,8 @@ int main(int argc, char** argv) { + if (yielded_at >= 0) std::printf("YIELDED %d %lld\n", yielded_slot, (long long) yielded_at); + std::printf("DONE %lld %lld %.1f %.1f %s %lld %lld %lld %lld %lld %lld %lld %.1f %lld %lld%s\n", + (long long) produced_n, +- (long long) n, prompt_ms, decode_ms, finish, (long long) draft_accepted, (long long) draft_offered, +- (long long) resume, (long long) req_hits, (long long) req_look, ++ (long long) (n - (int64_t) req_pads.size()), prompt_ms, decode_ms, finish, (long long) draft_accepted, (long long) draft_offered, ++ (long long) (resume + scr.reused()), (long long) req_hits, (long long) req_look, + (long long) (src.ram_reads() - ram0), (long long) (src.file_reads() - files0), + (double) (src.file_read_bytes() - file_bytes0) / 1e6, (long long) read_n, + (long long) req_offload, chain_txt); -- 2.54.0