diff --git a/hosts/focalor/modules/strata.nix b/hosts/focalor/modules/strata.nix index 065dace..292f716 100644 --- a/hosts/focalor/modules/strata.nix +++ b/hosts/focalor/modules/strata.nix @@ -119,6 +119,14 @@ let "4096" "--prompt-cache" "16" + # parked conversations: up to 6 GiB in RAM (kept 2.5 GiB above empty), then 200 GiB on /data; a switch between + # conversations reads the other one back instead of its whole prompt, and what is parked at exit outlasts a restart + "--conversation-cache-mib" + "6144" + "--conversation-disk" + "${data}/conversations" + "--conversation-disk-gib" + "200" ]; cwd = data; tokenizer = "${data}/packs/iq3_xxs/tokenizer"; @@ -218,6 +226,8 @@ in RestartSec = "10s"; TimeoutStartSec = "infinity"; TimeoutStopSec = "5min"; + # SIGTERM to the server only: it sends the engine QUIT, which writes its parked conversations to disk + KillMode = "mixed"; LimitNOFILE = 1048576; LimitMEMLOCK = "infinity"; }; diff --git a/hosts/focalor/modules/strata.patch b/hosts/focalor/modules/strata.patch index 3fac744..462f738 100644 --- a/hosts/focalor/modules/strata.patch +++ b/hosts/focalor/modules/strata.patch @@ -1,8 +1,8 @@ diff --git a/CMakeLists.txt b/CMakeLists.txt -index e28abd9b..2622929a 100644 +index e28abd9b..5a6c0238 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt -@@ -117,6 +117,9 @@ if(STRATA_BUILD_CONVERSATION_TESTS) +@@ -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) @@ -12,7 +12,13 @@ index e28abd9b..2622929a 100644 add_executable(conversation_file_test src/core/conversation_file_test.cpp src/core/conversation_file.cpp) target_include_directories(conversation_file_test PRIVATE include) add_test(NAME conversation_file_test COMMAND conversation_file_test) -@@ -378,6 +381,7 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP) ++ 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 @@ -20,7 +26,7 @@ index e28abd9b..2622929a 100644 src/kernels/cuda/kv_stream.cu src/kernels/cuda/kv_q4.cu src/kernels/cuda/fused_gr.cu -@@ -1791,3 +1795,11 @@ if(STRATA_HIP_GFX906) +@@ -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() @@ -32,6 +38,326 @@ index e28abd9b..2622929a 100644 + 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 @@ -344,6 +670,179 @@ index 00000000..5191305e +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 @@ -1350,10 +1849,18 @@ index 00000000..ee44a125 + return failures == 0 ? 0 : 1; +} diff --git a/src/program/generate.cpp b/src/program/generate.cpp -index ba00f94d..3f1c8eab 100644 +index ba00f94d..d73e3415 100644 --- a/src/program/generate.cpp +++ b/src/program/generate.cpp -@@ -81,6 +81,8 @@ +@@ -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" @@ -1362,7 +1869,38 @@ index ba00f94d..3f1c8eab 100644 #include "strata/core/device.hpp" #include "strata/core/emulate.hpp" #ifndef NOMINMAX -@@ -8127,6 +8129,24 @@ int main(int argc, char** argv) { +@@ -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; @@ -1387,7 +1925,57 @@ index ba00f94d..3f1c8eab 100644 // #879: a window met non-finite logits: the next request drops every cached state first and reads from token 0 bool nan_flush = false; std::vector checks; -@@ -10441,7 +10461,7 @@ int main(int argc, char** argv) { +@@ -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; } @@ -1396,7 +1984,7 @@ index ba00f94d..3f1c8eab 100644 if (o.peer_device >= 1) { refuse("session files do not support --peer-device"); continue; } if (o.batch > 0) { refuse("session files do not support --batch (parallel requests)"); continue; } if (o.prompt_cache <= 0) { refuse("session files need --prompt-cache > 0"); continue; } -@@ -10518,26 +10538,71 @@ int main(int argc, char** argv) { +@@ -10518,26 +10592,71 @@ int main(int argc, char** argv) { refuse(why, strata::core::SessionError::memory); continue; } @@ -1482,7 +2070,7 @@ index ba00f94d..3f1c8eab 100644 } if (st.dir_flush_unsupported) std::fprintf(stderr, "strata serve: session saved to %s; this filesystem cannot flush a " -@@ -10546,6 +10611,23 @@ int main(int argc, char** argv) { +@@ -10546,6 +10665,23 @@ int main(int argc, char** argv) { refuse("not enough RAM to save the session", strata::core::SessionError::memory); continue; } @@ -1506,7 +2094,7 @@ index ba00f94d..3f1c8eab 100644 std::fprintf(stderr, "strata serve: session saved %zu tokens, %zu of %zu checkpoints, %zu bytes to %s " "in %.1f ms (state %.1f ms)\n", live.size(), kept, checks.size(), bytes, path.c_str(), ms(), capture_ms); -@@ -10575,6 +10657,18 @@ int main(int argc, char** argv) { +@@ -10575,6 +10711,18 @@ int main(int argc, char** argv) { refuse(err, strata::core::SessionError::io); continue; } @@ -1525,7 +2113,7 @@ index ba00f94d..3f1c8eab 100644 strata::core::SessionStatus st; if (!strata::core::session_file_read(path, id, image, bytes, err, limits, &st)) { refuse(err, st.error); -@@ -10587,7 +10681,21 @@ int main(int argc, char** argv) { +@@ -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); @@ -1548,7 +2136,7 @@ index ba00f94d..3f1c8eab 100644 refuse(err); continue; } -@@ -10600,8 +10708,16 @@ int main(int argc, char** argv) { +@@ -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); @@ -1567,7 +2155,7 @@ index ba00f94d..3f1c8eab 100644 // validated above: a failure here is a transfer failure, after device writes began - never // decode from a partial state; the server starts the engine again std::fprintf(stderr, "strata serve: session restore %s: transfer failed: %s\n", path.c_str(), -@@ -10613,7 +10729,32 @@ int main(int argc, char** argv) { +@@ -10613,13 +10783,38 @@ int main(int argc, char** argv) { } live = std::move(image.live.ids); live_imgs = std::move(image.live.imgs); @@ -1601,7 +2189,14 @@ index ba00f94d..3f1c8eab 100644 for (const ConvCheckpoint& c : checks) check_clock = std::max(check_clock, c.used); cvec_cached = image.cvec; live_ok = true; -@@ -10730,7 +10871,13 @@ int main(int argc, char** argv) { + 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; } @@ -1616,7 +2211,46 @@ index ba00f94d..3f1c8eab 100644 req_imgs.clear(); if (geni && !o.vision) { std::printf("ERR this engine was started without --vision\n"); continue; } if (geni || !mrope_identity) { -@@ -11009,6 +11156,104 @@ int main(int argc, char** argv) { +@@ -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(); } @@ -1721,7 +2355,7 @@ index ba00f94d..3f1c8eab 100644 // Preserve the outgoing branch before any checkpoint rewind, reset, // or incoming restore overwrites the positional state it requires. // pin=N: this request resumes exactly at the pinned shared prefix of the conversation the session holds -@@ -11020,7 +11265,8 @@ int main(int argc, char** argv) { +@@ -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; @@ -1731,7 +2365,7 @@ index ba00f94d..3f1c8eab 100644 std::printf("ERR %s\n", err.c_str()); return 1; } -@@ -11053,6 +11299,7 @@ int main(int argc, char** argv) { +@@ -11053,6 +11374,7 @@ int main(int argc, char** argv) { live = bs[(size_t) slot_source].ids; } live_imgs.clear(); @@ -1739,7 +2373,7 @@ index ba00f94d..3f1c8eab 100644 resume = slot_tokens; from_live = true; std::fprintf(stderr, "strata batch: slot %d gave back %lld tokens of this conversation (%s) in " -@@ -11103,6 +11350,7 @@ int main(int argc, char** argv) { +@@ -11103,6 +11425,7 @@ int main(int argc, char** argv) { } live = std::move(incoming->live.ids); live_imgs = std::move(incoming->live.imgs); @@ -1747,7 +2381,18 @@ index ba00f94d..3f1c8eab 100644 if (stages.empty()) { checks = std::move(incoming->checkpoints); } else { // give each checkpoint back its stage parts (moved, not copied); none if they do not line up -@@ -11151,11 +11399,12 @@ int main(int argc, char** argv) { +@@ -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) { @@ -1762,7 +2407,7 @@ index ba00f94d..3f1c8eab 100644 cudaStreamSynchronize(st->stream); } checks.clear(); -@@ -11195,6 +11444,83 @@ int main(int argc, char** argv) { +@@ -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); @@ -1846,7 +2491,7 @@ index ba00f94d..3f1c8eab 100644 tr("request", n, geni ? 1 : 0); if (use_mtp) mtp.set_prompt_len(n); const int64_t read_from = reread_to > 0 ? 0 : resume; -@@ -11761,12 +12087,65 @@ int main(int argc, char** argv) { +@@ -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}; @@ -1916,7 +2561,7 @@ index ba00f94d..3f1c8eab 100644 if (to <= at) continue; err.clear(); const bool win = windows_ok(at, to); -@@ -12792,6 +13171,7 @@ int main(int argc, char** argv) { +@@ -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); @@ -1924,7 +2569,7 @@ index ba00f94d..3f1c8eab 100644 live_imgs = imgs_below(req_imgs, (int64_t) live.size()); live_ok = o.prompt_cache > 0 && req_ckpt; // ckpt=0: nothing to continue or park (#830) if (nf_hit) live_ok = false; // #879 (the next request's start drops the rest) -@@ -12801,6 +13181,8 @@ int main(int argc, char** argv) { +@@ -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); @@ -1933,7 +2578,7 @@ index ba00f94d..3f1c8eab 100644 live_imgs = imgs_below(req_imgs, pp_reached); live_ok = true; std::fprintf(stderr, "strata serve: cancelled request keeps its conversation: %lld tokens live\n", -@@ -12935,8 +13317,8 @@ int main(int argc, char** argv) { +@@ -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, @@ -1944,3 +2589,18 @@ index ba00f94d..3f1c8eab 100644 (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