[Engine] avoid CUDA OOM during conversation-cache restore

include/strata/program/conv_cache.hpp

@@ -77,6 +77,18 @@ inline size_t eviction_victim(const uint64_t* stamps, size_t n, int64_t cap, con
     return v;
 }
 
+/// A parked snapshot is restored on every GPU in a layer split.  If any one of
+/// them is already below the configured reserve, rereading the prompt is safer
+/// than turning a recoverable cache hit into a process-fatal CUDA allocation.
+/// A single-GPU restore has no split-stage residency to admit, so it keeps the
+/// existing restore path even when the sizing reserve exceeds current free VRAM.
+inline bool restore_has_headroom(bool layer_split, const uint64_t* free_bytes, size_t n, uint64_t floor) {
+    if (!layer_split || free_bytes == nullptr || n == 0 || floor == 0) return true;
+    for (size_t i = 0; i < n; ++i)
+        if (free_bytes[i] < floor) return false;
+    return true;
+}
+
 /// --prompt-cache-spread: the item whose removal leaves the smallest gap for its distance from the newest, so the kept
 /// points thin out geometrically towards the root and an edit d tokens before the end reads about d again.  `lens`:
 /// the chain's token counts.  Never the root, the newest or a pinned item; the tail rule first, as above.

src/program/conv_cache_test.cpp

@@ -20,6 +20,7 @@
 #include <vector>
 
 using strata::program::conv_cache::eviction_victim;
+using strata::program::conv_cache::restore_has_headroom;
 using strata::program::conv_cache::spread_victim;
 
 namespace {
@@ -132,6 +133,14 @@ int main() {
         check(eviction_victim(s4.data(), s4.size(), 4, nullptr, all) == 2, "everything pinned: the oldest pin leaves, never the root");
         check(eviction_victim(s4.data(), s4.size(), 1, nullptr, all) == 0, "a one-slot budget keeps no pin: the oldest leaves");
     }
+    {   // A layer-split restore must leave headroom on every device, not only CUDA0.
+        const uint64_t free_bytes[] = {900, 1200};
+        check(restore_has_headroom(true, free_bytes, 2, 700), "both GPUs above the restore floor: admit the cache hit");
+        const uint64_t low_later_stage[] = {1200, 699};
+        check(!restore_has_headroom(true, low_later_stage, 2, 700), "a low later-stage GPU: reread instead of restoring");
+        check(restore_has_headroom(true, nullptr, 0, 700), "no sampled devices: do not reject on a missing sample");
+        check(restore_has_headroom(false, low_later_stage, 2, 700), "single-GPU restore keeps its existing path");
+    }
     std::printf(g_fail ? "FAIL\n" : "PASS\n");
     return g_fail ? 1 : 0;
 }

src/program/generate.cpp

@@ -12044,6 +12044,45 @@ int main(int argc, char** argv) {
                 incoming.reset();
                 err.clear();
             }
+            // A parked snapshot is authoritative host state, but restoring it also
+            // repopulates each device's streaming/ring residency.  On a layer split
+            // the primary and later stages can have very different free VRAM.  Do
+            // not turn a cache hit into a process-fatal CUDA OOM on the first stage
+            // switch: return the snapshot to the cache and reread this prompt.
+            auto conversation_restore_has_headroom = [&]() -> bool {
+                if (!incoming || o.vram_reserve_mib <= 0) return true;
+                const uint64_t floor = (uint64_t) o.vram_reserve_mib << 20;
+                std::vector<uint64_t> free_bytes;
+                free_bytes.reserve(stages.size() + 1);
+                auto sample = [&](int dev) -> bool {
+                    const strata::core::OnDevice on(dev);
+                    size_t free_b = 0, total_b = 0;
+                    if (cudaMemGetInfo(&free_b, &total_b) != cudaSuccess) {
+                        (void) cudaGetLastError();
+                        return false;
+                    }
+                    free_bytes.push_back((uint64_t) free_b);
+                    return true;
+                };
+                if (!sample(-1)) return true;
+                for (const auto& st : stages)
+                    if (!sample(st->dev)) return true;
+                if (strata::program::conv_cache::restore_has_headroom(true, free_bytes.data(), free_bytes.size(), floor))
+                    return true;
+                uint64_t least = free_bytes.front();
+                for (const uint64_t bytes : free_bytes) least = std::min(least, bytes);
+                std::fprintf(stderr, "strata serve: conversation cache: skip restore (%lld MiB free on the most constrained "
+                                     "GPU, %lld MiB reserve); rereading the prompt\n",
+                             (long long) (least >> 20), (long long) o.vram_reserve_mib);
+                return false;
+            };
+            if (!stages.empty() && incoming && !conversation_restore_has_headroom()) {
+                auto retry = std::move(*incoming);
+                incoming.reset();
+                if (!conversations.put(std::move(retry)))
+                    std::fprintf(stderr, "strata serve: conversation cache: could not return skipped snapshot\n");
+                err.clear();
+            }
             if (borrow) {
                 const strata::core::SavedConversation* donor = conversations.find(parked.serial);
                 if (donor == nullptr ||