diff --git a/CMakeLists.txt b/CMakeLists.txt index e28abd9b..5a6c0238 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -117,9 +117,15 @@ 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) + add_executable(conversation_disk_test src/core/conversation_disk_test.cpp src/core/conversation_file.cpp) + target_include_directories(conversation_disk_test PRIVATE include) + add_test(NAME conversation_disk_test COMMAND conversation_disk_test) add_executable(session_save_reclaim_test src/core/session_save_reclaim_test.cpp src/core/conversation_file.cpp) target_include_directories(session_save_reclaim_test PRIVATE include) add_test(NAME session_save_reclaim_test COMMAND session_save_reclaim_test) @@ -378,6 +384,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 +1798,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_cache.hpp b/include/strata/core/conversation_cache.hpp index 08e9cb80..f8848f85 100644 --- a/include/strata/core/conversation_cache.hpp +++ b/include/strata/core/conversation_cache.hpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include @@ -284,6 +285,7 @@ public: auto victim = std::find_if(entries_.begin(), entries_.end(), [](const SavedConversation& e) { return !e.pinned(); }); if (victim == entries_.end()) return false; bytes_ -= victim->bytes(); + if (spill_) spill_(*victim); // the NVMe tier (ConversationDisk) keeps it entries_.erase(victim); ++evictions_; return true; @@ -299,6 +301,20 @@ public: return n; } + /// Hands every parked conversation to the spill (pinned ones too) and empties the cache. Returns how many. + size_t spill_all() { + const size_t n = entries_.size(); + if (spill_) + for (const auto& e : entries_) spill_(e); + entries_.clear(); + bytes_ = 0; + reuse_ = {}; + return n; + } + + /// Called with each conversation evict_oldest() or spill_all() drops (not with superseded or cleared ones). + void set_spill(std::function spill) { spill_ = std::move(spill); } + // The slot count, so a caller that evicts in a loop has a bound it did not invent. size_t slots() const { return slots_; } @@ -351,6 +367,7 @@ private: int64_t min_tokens_ = 0; // --conversation-cache-min-tokens: 0 parks every conversation, whatever its size std::deque entries_; // least recently active first ConversationKvReuse reuse_; + std::function spill_; }; } // namespace strata::core diff --git a/include/strata/core/conversation_disk.hpp b/include/strata/core/conversation_disk.hpp new file mode 100644 index 00000000..9960f3f6 --- /dev/null +++ b/include/strata/core/conversation_disk.hpp @@ -0,0 +1,265 @@ +#pragma once +// The NVMe tier of the conversation cache: parked conversations the RAM tier (ConversationCache) evicts are written +// here as session files instead of being dropped, under a byte budget, least recently used first out. Each file has +// a small index beside it (.idx: the token ids and pictures of its live state and of each checkpoint), so a +// lookup and a restart never read a state. A request whose prompt one of them covers deeper than the live state, an +// idle slot or the RAM tier reads it back (session_file_read) and removes it, as ConversationCache::take does. + +#include "strata/core/conversation_cache.hpp" +#include "strata/core/conversation_file.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace strata::core { + +class ConversationDisk { +public: + using Match = ConversationCache::Match; + + /// One parked conversation on disk: its file and the keys (ids and pictures, no state) of its live state and its + /// checkpoints, so conversation_prefix() matches it like a RAM entry. + struct Entry { + std::string path; + ConversationCheckpoint live; + std::vector checkpoints; + bool cvec = true; + uint64_t bytes = 0; + }; + + ConversationDisk(std::string dir, uint64_t budget) : dir_(std::move(dir)), budget_(budget) {} + bool enabled() const { return budget_ > 0 && !dir_.empty(); } + size_t size() const { return entries_.size(); } + uint64_t bytes() const { return bytes_; } + const std::deque& entries() const { return entries_; } + + /// Read the index files left by an earlier run, oldest first; a file without its index (or the reverse, or a + /// half-written .tmp) is removed. Over budget, the oldest go. Returns how many were found. + size_t load() { + if (!enabled()) return 0; + namespace fs = std::filesystem; + std::error_code ec; + fs::create_directories(dir_, ec); + struct Found { Entry e; fs::file_time_type t; }; + std::vector found; + for (const auto& de : fs::directory_iterator(dir_, ec)) { + const fs::path p = de.path(); + const std::string ext = p.extension().string(); + if (ext == ".tmp") { fs::remove(p, ec); continue; } + if (ext == ".idx") { + if (!fs::exists(fs::path(p).replace_extension(".kvs"), ec)) fs::remove(p, ec); + continue; + } + if (ext != ".kvs") continue; + Entry e; + e.path = p.string(); + if (!read_index(index_of(e.path), e)) { fs::remove(p, ec); fs::remove(index_of(e.path), ec); continue; } + e.bytes = (uint64_t) fs::file_size(p, ec); + if (ec) { ec.clear(); continue; } + found.push_back({std::move(e), fs::last_write_time(p, ec)}); + } + std::sort(found.begin(), found.end(), [](const Found& a, const Found& b) { return a.t < b.t; }); + for (auto& f : found) { + bytes_ += f.e.bytes; + seq_ = std::max(seq_, sequence_of(f.e.path) + 1); + entries_.push_back(std::move(f.e)); + } + trim(0); + return entries_.size(); + } + + /// Write `image` as a session file and index it; the entries it supersedes (the same conversation a turn back) + /// and, over budget, the least recently used ones are removed first. False (with `error`) when it cannot be + /// written - the caller loses only this copy. + bool put(const SavedConversation& image, const SessionFileIdentity& id, std::string& error, + const SessionWriteOptions& options = {}) { + if (!enabled()) return false; + drop_superseded(image); + std::error_code ec; + std::filesystem::create_directories(dir_, ec); + char name[48]; + std::snprintf(name, sizeof name, "conv-%012llu.kvs", (unsigned long long) seq_++); + Entry e; + e.path = (std::filesystem::path(dir_) / name).string(); + e.live = key_of(image.live); + for (const auto& c : image.checkpoints) e.checkpoints.push_back(key_of(c)); + e.cvec = image.cvec; + size_t written = 0; + if (!session_file_write(e.path, image, id, written, error, options)) return false; + if (!write_index(index_of(e.path), e)) { + std::filesystem::remove(e.path, ec); + error = "cannot write the index of " + e.path; + return false; + } + e.bytes = written; + trim(e.bytes); + bytes_ += e.bytes; + entries_.push_back(std::move(e)); + return true; + } + + /// The entry and checkpoint that hold the longest prefix of `prompt` (ConversationCache::best's rule; ties prefer + /// the most recently written). + template + Match best(const std::vector& prompt, const std::vector& images, bool cvec) const { + Match best; + for (size_t i = entries_.size(); i-- > 0;) { + const auto& e = entries_[i]; + if (e.cvec != cvec) continue; + auto consider = [&](const ConversationCheckpoint& c, bool live) { + const int64_t n = conversation_prefix(c, prompt, images); + if (n > best.tokens) best = {i, n, live}; + }; + consider(e.live, true); + for (const auto& c : e.checkpoints) consider(c, false); + } + return best; + } + + /// Read entry `index` back and remove it from disk either way (a file that does not read is not worth keeping). + bool take(size_t index, const SessionFileIdentity& id, SavedConversation& out, std::string& error, + const SessionReadLimits& limits = {}) { + Entry e = std::move(entries_.at(index)); + entries_.erase(entries_.begin() + (std::ptrdiff_t) index); + bytes_ -= e.bytes; + size_t read = 0; + const bool ok = session_file_read(e.path, id, out, read, error, limits); + remove_files(e.path); + return ok; + } + + /// Every entry goes (#879: after non-finite logits any of them may carry the poison). Returns how many. + size_t clear() { + const size_t n = entries_.size(); + for (const auto& e : entries_) remove_files(e.path); + entries_.clear(); + bytes_ = 0; + return n; + } + +private: + static constexpr uint32_t kMagic = 0x58444353; // "SCDX" + static constexpr uint32_t kVersion = 1; + + static ConversationCheckpoint key_of(const ConversationCheckpoint& c) { + ConversationCheckpoint k; + k.ids = c.ids; + k.imgs = c.imgs; + return k; + } + static std::string index_of(const std::string& path) { + return std::filesystem::path(path).replace_extension(".idx").string(); + } + static uint64_t sequence_of(const std::string& path) { + unsigned long long n = 0; + return std::sscanf(std::filesystem::path(path).filename().string().c_str(), "conv-%llu.kvs", &n) == 1 ? n : 0; + } + static void remove_files(const std::string& path) { + std::error_code ec; + std::filesystem::remove(path, ec); + std::filesystem::remove(index_of(path), ec); + } + + // The same rule as ConversationCache::drop_superseded: an entry whose deepest checkpoint the incoming image + // still holds is this conversation a turn back. + void drop_superseded(const SavedConversation& image) { + auto held = [&](const ConversationCheckpoint& c) { + if (c.ids == image.live.ids && c.imgs == image.live.imgs) return true; + for (const auto& k : image.checkpoints) + if (k.ids == c.ids && k.imgs == c.imgs) return true; + return false; + }; + for (size_t i = 0; i < entries_.size();) { + const auto& e = entries_[i]; + const ConversationCheckpoint* deepest = nullptr; + for (const auto& c : e.checkpoints) + if (!deepest || c.ids.size() > deepest->ids.size()) deepest = &c; + if (e.cvec == image.cvec && deepest && !deepest->ids.empty() && held(*deepest)) { + bytes_ -= e.bytes; + remove_files(e.path); + entries_.erase(entries_.begin() + (std::ptrdiff_t) i); + continue; + } + ++i; + } + } + + // the oldest entries leave until `incoming` more bytes fit the budget + void trim(uint64_t incoming) { + while (!entries_.empty() && bytes_ + incoming > budget_) { + bytes_ -= entries_.front().bytes; + remove_files(entries_.front().path); + entries_.pop_front(); + } + } + + static bool write_key(std::ofstream& f, const ConversationCheckpoint& c) { + const uint64_t n = c.ids.size(), m = c.imgs.size(); + f.write((const char*) &n, sizeof n); + f.write((const char*) c.ids.data(), (std::streamsize) (n * sizeof(int32_t))); + f.write((const char*) &m, sizeof m); + for (const auto& im : c.imgs) { + f.write((const char*) &im.start, sizeof im.start); + f.write((const char*) &im.hash, sizeof im.hash); + } + return (bool) f; + } + static bool read_key(std::ifstream& f, ConversationCheckpoint& c) { + uint64_t n = 0, m = 0; + if (!f.read((char*) &n, sizeof n) || n > (1ull << 26)) return false; + c.ids.resize(n); + if (!f.read((char*) c.ids.data(), (std::streamsize) (n * sizeof(int32_t)))) return false; + if (!f.read((char*) &m, sizeof m) || m > (1ull << 20)) return false; + c.imgs.resize(m); + for (auto& im : c.imgs) + if (!f.read((char*) &im.start, sizeof im.start) || !f.read((char*) &im.hash, sizeof im.hash)) return false; + return true; + } + static bool write_index(const std::string& path, const Entry& e) { + const std::string tmp = path + ".tmp"; + { + std::ofstream f(tmp, std::ios::binary | std::ios::trunc); + const uint32_t head[2] = {kMagic, kVersion}; + const uint8_t cvec = e.cvec; + const uint64_t nc = e.checkpoints.size(); + f.write((const char*) head, sizeof head); + f.write((const char*) &cvec, sizeof cvec); + f.write((const char*) &nc, sizeof nc); + if (!write_key(f, e.live)) return false; + for (const auto& c : e.checkpoints) + if (!write_key(f, c)) return false; + f.flush(); + if (!f) return false; + } + std::error_code ec; + std::filesystem::rename(tmp, path, ec); + return !ec; + } + static bool read_index(const std::string& path, Entry& e) { + std::ifstream f(path, std::ios::binary); + uint32_t head[2] = {0, 0}; + uint8_t cvec = 1; + uint64_t nc = 0; + if (!f.read((char*) head, sizeof head) || head[0] != kMagic || head[1] != kVersion) return false; + if (!f.read((char*) &cvec, sizeof cvec) || !f.read((char*) &nc, sizeof nc) || nc > 4096) return false; + e.cvec = cvec != 0; + if (!read_key(f, e.live)) return false; + e.checkpoints.resize(nc); + for (auto& c : e.checkpoints) + if (!read_key(f, c)) return false; + return true; + } + + std::string dir_; + uint64_t budget_ = 0, bytes_ = 0, seq_ = 0; + std::deque entries_; // least recently written first +}; + +} // namespace strata::core 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_disk_test.cpp b/src/core/conversation_disk_test.cpp new file mode 100644 index 00000000..de500ede --- /dev/null +++ b/src/core/conversation_disk_test.cpp @@ -0,0 +1,167 @@ +// CPU-only tests for the NVMe tier of the conversation cache (include/strata/core/conversation_disk.hpp): what it +// keeps, finds, reads back and drops, and that the RAM tier hands it what it evicts. +// Built with -DSTRATA_BUILD_CONVERSATION_TESTS=ON; no CUDA, no model. +#include "strata/core/conversation_disk.hpp" + +#include +#include +#include +#include +#include + +using namespace strata::core; +namespace fs = std::filesystem; + +namespace { + +int checks = 0; +void check(bool ok, const char* label) { + ++checks; + if (!ok) { std::fprintf(stderr, "FAIL: %s\n", label); std::exit(1); } +} + +std::vector bytes_of(size_t n, uint8_t seed) { + std::vector v(n); + for (size_t i = 0; i < n; ++i) v[i] = uint8_t(seed + i * 7); + return v; +} + +ConversationBuffer pattern(size_t n, uint8_t seed) { + ConversationBuffer b; + b.resize(n); + b.visit(0, n, [&](uint8_t* p, size_t c, size_t at) { + for (size_t i = 0; i < c; ++i) p[i] = uint8_t(seed + (at + i) * 13); + return true; + }); + return b; +} + +// a checkpoint over tokens base, base+1, ... (count of them) +ConversationCheckpoint checkpoint(int32_t base, size_t count, uint8_t seed) { + ConversationCheckpoint c; + for (size_t i = 0; i < count; ++i) c.ids.push_back(base + int32_t(i)); + c.gdn = bytes_of(4099, seed); + c.ple = bytes_of(77, seed + 1); + c.tails = bytes_of(301, seed + 2); + c.dead = bytes_of(64, seed + 3); + c.block_pos = bytes_of(8, seed + 4); + c.used = 42 + seed; + return c; +} + +// one conversation: tokens base.. of `live` length, checkpoints at each of `at` +SavedConversation conversation(int32_t base, size_t live, std::vector at, uint8_t seed) { + SavedConversation s; + for (size_t i = 0; i < s.geometry.size(); ++i) s.geometry[i] = int64_t(100 + i); + s.layer_lo = 0; s.layer_hi = 48; + s.live = checkpoint(base, live, seed); + for (size_t n : at) s.checkpoints.push_back(checkpoint(base, n, uint8_t(seed + n))); + ConversationKv kv; + kv.format = 2; kv.cells = 1024; kv.heads = 2; kv.head_dim = 256; kv.page_size = 64; kv.pooled_rows = 9; kv.idx_dim = 128; + kv.k = pattern(65536 + seed, seed); + kv.v = pattern(4096, uint8_t(seed + 10)); + kv.k_scale = pattern(512, uint8_t(seed + 20)); + kv.v_scale = pattern(512, uint8_t(seed + 30)); + kv.pooled = pattern(9 * 128 * 4, uint8_t(seed + 40)); + s.kv.push_back(std::move(kv)); + return s; +} + +std::vector prompt(int32_t base, size_t count) { + std::vector p; + for (size_t i = 0; i < count; ++i) p.push_back(base + int32_t(i)); + return p; +} + +size_t files_in(const fs::path& dir, const char* ext) { + size_t n = 0; + for (const auto& e : fs::directory_iterator(dir)) n += e.path().extension() == ext; + return n; +} + +} // namespace + +int main() { + const fs::path dir = fs::temp_directory_path() / ("strata-conv-disk-test-" + std::to_string( + std::chrono::steady_clock::now().time_since_epoch().count())); + const SessionFileIdentity id{0x1111222233334444ull, 0x5555666677778888ull}; + std::string error; + + // two conversations are kept, each found by the prompt that continues it, and read back whole + { + ConversationDisk disk(dir.string(), 1ull << 30); + check(disk.load() == 0, "an empty folder holds nothing"); + const SavedConversation a = conversation(1000, 900, {100, 600}, 1); + const SavedConversation b = conversation(5000, 400, {100}, 2); + check(disk.put(a, id, error) && disk.put(b, id, error), "both are written"); + check(disk.size() == 2 && files_in(dir, ".kvs") == 2 && files_in(dir, ".idx") == 2, "a file and an index each"); + auto m = disk.best(prompt(1000, 950), {}, true); + check(m.tokens == 900 && m.live, "the prompt that continues a is matched to a's live state"); + m = disk.best(prompt(1000, 650), {}, true); + check(m.tokens == 600 && !m.live, "a shorter prompt is matched to a's checkpoint"); + check(disk.best(prompt(9000, 500), {}, true).tokens == 0, "an unrelated prompt matches nothing"); + check(disk.best(prompt(1000, 950), {}, false).tokens == 0, "another control vector matches nothing"); + SavedConversation back; + m = disk.best(prompt(5000, 450), {}, true); + check(disk.take(m.index, id, back, error), "b is read back"); + check(back.live.ids == b.live.ids && back.live.gdn == b.live.gdn && back.checkpoints.size() == 1 && + back.kv.size() == 1 && back.kv[0].k.size() == b.kv[0].k.size(), "b comes back as it was written"); + check(disk.size() == 1 && files_in(dir, ".kvs") == 1 && files_in(dir, ".idx") == 1, "a taken entry leaves the disk"); + } + + // a restart finds what an earlier run wrote, without reading the states; stray files go + { + std::ofstream(dir / "conv-000000000099.kvs.tmp") << "half"; + std::ofstream(dir / "conv-000000000098.kvs") << "no index"; + ConversationDisk disk(dir.string(), 1ull << 30); + check(disk.load() == 1, "the remaining conversation is found again"); + check(disk.best(prompt(1000, 950), {}, true).tokens == 900, "and matched as before"); + check(!fs::exists(dir / "conv-000000000099.kvs.tmp") && !fs::exists(dir / "conv-000000000098.kvs"), + "a half-written file and a file without its index are removed"); + SessionFileIdentity other = id; + other.model ^= 1; + SavedConversation back; + check(!disk.take(0, other, back, error) && !error.empty(), "another model's engine cannot read it"); + check(disk.size() == 0 && files_in(dir, ".kvs") == 0, "and a file that does not read is not kept"); + } + + // the same conversation a turn later replaces its earlier copy; over budget the oldest goes + { + ConversationDisk disk(dir.string(), 1ull << 30); + check(disk.put(conversation(1000, 900, {100, 600}, 1), id, error), "turn 1 is written"); + check(disk.put(conversation(1000, 1400, {100, 600, 1200}, 1), id, error), "turn 2 is written"); + check(disk.size() == 1 && disk.best(prompt(1000, 1500), {}, true).tokens == 1400, "turn 2 supersedes turn 1"); + disk.clear(); + check(disk.put(conversation(7000, 900, {100}, 4), id, error), "one of the size below is written"); + const uint64_t one = disk.bytes(); + disk.clear(); + ConversationDisk small(dir.string(), one * 2 + one / 2); + check(small.put(conversation(1000, 900, {100}, 1), id, error) && + small.put(conversation(5000, 900, {100}, 2), id, error) && + small.put(conversation(9000, 900, {100}, 3), id, error), "three are written"); + check(small.size() == 2 && small.bytes() <= one * 2 + one / 2, "the budget keeps two"); + check(small.best(prompt(1000, 950), {}, true).tokens == 0 && small.best(prompt(9000, 950), {}, true).tokens == 900, + "the least recently written left"); + check(small.clear() == 2 && files_in(dir, ".kvs") == 0 && files_in(dir, ".idx") == 0, "clear removes every file"); + } + + // the RAM tier hands each conversation it evicts to the spill, and nothing else + { + ConversationCache ram(1ull << 30, 1); + std::vector> spilled; + ram.set_spill([&](const SavedConversation& s) { spilled.push_back(s.live.ids); }); + check(ram.put(conversation(1000, 900, {100}, 1)), "a is parked"); + check(ram.put(conversation(5000, 900, {100}, 2)), "b is parked, evicting a"); + check(spilled.size() == 1 && spilled[0] == prompt(1000, 900), "a was spilled"); + check(ram.put(conversation(5000, 1300, {100, 900}, 2)) && spilled.size() == 1, + "b's next turn supersedes b without a spill"); + ram.clear(); + check(spilled.size() == 1, "clear spills nothing"); + check(ram.put(conversation(1000, 900, {100}, 1)) && ram.spill_all() == 1 && spilled.size() == 2 && ram.size() == 0, + "spill_all hands over every parked conversation"); + } + + fs::remove_all(dir); + std::printf("conversation_disk_test: %d checks passed\n", checks); + return 0; +} 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..d73e3415 100644 --- a/src/program/generate.cpp +++ b/src/program/generate.cpp @@ -23,6 +23,7 @@ #include "strata/core/remote_expert_opt.hpp" #include "strata/core/expert_cache.hpp" #include "strata/core/conversation_snapshot.hpp" +#include "strata/core/conversation_disk.hpp" #include "strata/core/conversation_file.hpp" #include "strata/core/session_save_reclaim.hpp" #include "strata/core/conversation_memory.hpp" @@ -81,6 +82,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 @@ -637,6 +640,10 @@ struct Options { /// conversation_cache_slots whatever its length, so without it a client that interleaves short side requests /// with one long conversation evicts the long one (oldest first). 0 = park every conversation. int64_t conversation_cache_min_tokens = 0; + /// --serve: the NVMe tier. Parked conversations the RAM budget evicts are written to this folder as session files + /// (with --conversation-disk-gib above 0) and read back when a request continues one of them. + std::string conversation_disk; + int64_t conversation_disk_gib = 0; /// --memory-limit-mib N / STRATA_MEMORY_LIMIT_MIB: the total RAM this engine's container or cgroup may use, for a /// container that cannot see its real limit (0 or unset = none; the flag wins over the variable) int64_t memory_limit_mib = -1; @@ -792,6 +799,10 @@ void usage() { " wins; 0 = none). Parking and session save/restore keep their floor under it\n" " --conversation-cache-min-tokens N --serve: park only conversations of at least N tokens, so\n" " short side requests do not take the slots a long one needs (default 0 = park any)\n" + " --conversation-disk DIR --serve: write the parked conversations the RAM budget evicts to DIR and\n" + " read them back when a request continues one (needs --conversation-disk-gib)\n" + " --conversation-disk-gib N --serve: disk budget of --conversation-disk, least recently written out\n" + " first (default 0 = off)\n" " --session-min-free-mib N --serve: disk space a SAVE must leave free (default 4096; 0 = no check)\n" " --session-save-reclaim --serve: reclaim unpinned host caches if SAVE needs RAM (opt-in)\n" " --vram-elastic --serve (#533, opt-in, NVIDIA): the expert cache in segments (--vram-segment-mib,\n" @@ -2123,6 +2134,8 @@ int main(int argc, char** argv) { else if (a == "--memory-limit-mib") o.memory_limit_mib = number; else o.conversation_cache_slots = (int) number; } + else if (a == "--conversation-disk") o.conversation_disk = next("--conversation-disk"); + else if (a == "--conversation-disk-gib") o.conversation_disk_gib = std::max(0LL, std::atoll(next("--conversation-disk-gib"))); else if (a == "--prompt-cache-tail") o.prompt_cache_tail = true; else if (a == "--prompt-cache-every") o.prompt_cache_every = std::max(0LL, std::atoll(next("--prompt-cache-every"))); else if (a == "--prompt-cache-root") o.prompt_cache_root = std::max(0LL, std::atoll(next("--prompt-cache-root"))); @@ -8127,6 +8140,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; @@ -8212,6 +8243,49 @@ int main(int argc, char** argv) { id.config = *config_fp; return true; }; + // The NVMe tier: what the RAM tier evicts is written here, and what a request continues is read back. Its + // reads have the bounds and the RAM floor a RESTORE has. + strata::core::ConversationDisk conv_disk(o.prompt_cache > 0 && conversations.enabled() ? o.conversation_disk : "", + (uint64_t) o.conversation_disk_gib << 30); + auto conv_disk_limits = [&](strata::core::SessionReadLimits& limits, std::string& e) -> bool { + const uint64_t floor = (uint64_t) o.conversation_cache_min_free_mib << 20; + limits.admit = [floor](uint64_t need, std::string& why) { + const auto mem = strata::core::sample_host_memory(); + if (strata::core::conversation_memory_admit(mem ? std::optional(mem->available) : std::nullopt, + need, floor)) return true; + why = "not enough RAM to read it (" + std::to_string(need >> 20) + " MiB plus the floor)"; + return false; + }; + if (!strata::core::conversation_session_read_limits(limits, ss, g, mtp.kv_state(), (uint64_t) o.max_context, + (uint64_t) std::max(o.prompt_cache, 1), e)) return false; + for (const auto& st_k : stages) { + 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), e)) return false; + limits.stage_limits.push_back(std::move(lk)); + } + return true; + }; + if (conv_disk.enabled()) { + const size_t found = conv_disk.load(); + std::fprintf(stderr, "strata serve: conversation disk: %s, %lld GiB; %zu parked conversation%s from before " + "(%.1f GiB)\n", o.conversation_disk.c_str(), (long long) o.conversation_disk_gib, found, + found == 1 ? "" : "s", conv_disk.bytes() / 1073741824.0); + conversations.set_spill([&](const strata::core::SavedConversation& image) { + const auto t0 = Clock::now(); + strata::core::SessionFileIdentity id; + std::string e; + strata::core::SessionWriteOptions wo; + wo.min_free_bytes = (uint64_t) o.session_min_free_mib << 20; + wo.durable = false; // a cache: a copy lost to a crash is read again from the prompt + if (session_identity(id, e, {}) && conv_disk.put(image, id, e, wo)) + std::fprintf(stderr, "strata serve: conversation disk: wrote %zu tokens in %.0f ms; on disk %zu (%.1f GiB)\n", + image.live.ids.size(), std::chrono::duration(Clock::now() - t0).count(), + conv_disk.size(), conv_disk.bytes() / 1073741824.0); + else + std::fprintf(stderr, "strata serve: conversation disk: not written (%s)\n", e.c_str()); + }); + } // Save only on a switch/rewind, not on each continuing request. No graph // addresses change: all parked images live in ordinary host vectors. auto park_current_body = [&](size_t held) -> bool { @@ -10441,7 +10515,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 +10592,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 +10665,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 +10711,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 +10735,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 +10762,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,13 +10783,38 @@ 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; if (nan_flush) { // #879: the file replaced the poisoned session; the parked ones go (no slots here) nan_flush = false; - const size_t n_parked = conversations.clear(); + const size_t n_parked = conversations.clear() + conv_disk.clear(); std::fprintf(stderr, "strata serve: after non-finite logits (#879): the restored session " "replaces the live one; dropped %zu parked conversation%s\n", n_parked, n_parked == 1 ? "" : "s"); @@ -10730,7 +10925,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) { @@ -10898,7 +11099,7 @@ int main(int argc, char** argv) { live_ok = false; checks.clear(); tail_ckpt_len = -1; - const size_t n_parked = conversations.clear(); + const size_t n_parked = conversations.clear() + conv_disk.clear(); size_t n_slots = 0, n_running = 0; for (BSlot& sl : bs) if (sl.active) { // a running slot finishes its reply, but is not kept as a cache after it @@ -10984,7 +11185,28 @@ int main(int argc, char** argv) { } const auto parked = conversations.best(ids, req_imgs, want_cvec); std::optional incoming; - if (parked.tokens > std::max(resume, slot_tokens)) incoming.emplace(conversations.take(parked.index)); + auto incoming_match = parked; // where in the incoming image this prompt resumes + // the NVMe tier, when it holds this prompt deeper than everything in memory + const auto on_disk = conv_disk.enabled() ? conv_disk.best(ids, req_imgs, want_cvec) + : strata::core::ConversationDisk::Match{}; + if (on_disk.tokens > std::max({resume, slot_tokens, parked.tokens})) { + const auto t0 = Clock::now(); + strata::core::SessionFileIdentity id; + strata::core::SessionReadLimits limits; + strata::core::SavedConversation image; + std::string e; + if (session_identity(id, e, {}) && conv_disk_limits(limits, e) && + conv_disk.take(on_disk.index, id, image, e, limits)) { + std::fprintf(stderr, "strata serve: conversation disk: read %lld of this prompt's tokens back in %.0f ms\n", + (long long) on_disk.tokens, + std::chrono::duration(Clock::now() - t0).count()); + incoming.emplace(std::move(image)); + incoming_match = on_disk; + } else { + std::fprintf(stderr, "strata serve: conversation disk: cannot read it back (%s)\n", e.c_str()); + } + } + if (!incoming && parked.tokens > std::max(resume, slot_tokens)) incoming.emplace(conversations.take(parked.index)); if (incoming) slot_source = -1; // Reject the entire image before parking/overwriting the outgoing // state. Invalid entries can safely fall back to its existing prefix. @@ -11009,6 +11231,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 +11340,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 +11374,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 +11425,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 @@ -11112,8 +11435,8 @@ int main(int argc, char** argv) { if (!strata::core::conversation_checkpoints_merge(std::move(cs), checks)) checks.clear(); } cvec_cached = incoming->cvec; - resume = parked.tokens; - from_live = parked.live; + resume = incoming_match.tokens; + from_live = incoming_match.live; if (std::getenv("STRATA_SNAPSHOT_FULL_CAPTURE") == nullptr) { std::vector> stage_kv; // every stage's, for its next park for (auto& si : incoming->stage_images) stage_kv.push_back(std::move(si.kv)); @@ -11151,11 +11474,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 +11519,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 +12162,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 +13246,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 +13256,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 +13392,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); @@ -13137,6 +13594,14 @@ int main(int argc, char** argv) { } } save_profile("exit"); // #477: QUIT, or the server closed stdin + // the NVMe tier keeps every conversation this engine holds, the live one too, for the next start + if (conv_disk.enabled()) { + const auto t0 = Clock::now(); + if (!park_current(0)) err.clear(); + const size_t n = conversations.spill_all(); + std::fprintf(stderr, "strata serve: conversation disk: wrote %zu parked conversation%s at exit in %.1f s\n", n, + n == 1 ? "" : "s", std::chrono::duration(Clock::now() - t0).count()); + } #if defined(STRATA_GPU_METAL) // the stdin reader (detached above) may still take in_mu when the server closes the pipe after QUIT; on macOS // locking that mutex once this scope has destroyed it throws (the engine ended with an abort): end here