CPU prefill share: skip the default when host memory is short, and never fail a chat on its buffers

0.1.41 report (RTX 5090, tight RAM): chats failed and the engine restarted until STRATA_PREFILL_CPU_SHARE=0.
- share_rules: default_share_memory_verdict. The default is not armed when the RAM this process can get
  (cgroup-aware) is below the share's buffers + 3 GiB, when the Windows commit left is below buffers + 4 GiB,
  or when the experts it would read are file pages the free RAM cannot keep. One log line says why and that
  STRATA_PREFILL_CPU_SHARE=auto turns it on anyway. An explicit setting is kept and only warns.
- prefill: the share's two pinned buffers are taken once at their largest (an expert takes at most MAXT
  tokens' rows); a failed allocation ends the share for the run instead of failing the request.
- unit test, docs/DETAILS.md.

docs/DETAILS.md

@@ -1503,6 +1503,8 @@ experts it reads from RAM as they are: the arena, the page-locked copy of the re
 `experts.bin`'s pages in the file cache (`--mmap-experts`), not only page-locked ones. Serve without `--batch`; on a
 layer split every stage has the pool and one stage at a time takes it for a chunk.
 
+**0.1.42: not armed by default when memory is short.** A 0.1.41 report (RTX 5090, RAM already tight) had chats failing and the engine restarting until `STRATA_PREFILL_CPU_SHARE=0` was set. The default is now skipped, with one line in the log saying why, when the RAM this process can get (the cgroup limit counts on Linux) is below the share's own buffers plus 3 GiB, when on Windows the commit limit (RAM + page file) has less than 4 GiB left beyond them, or when the experts it would read are file pages (the mapped `experts.bin` outside the GPU cache and the RAM copy) that the free RAM plus 3 GiB cannot keep. An explicit `STRATA_PREFILL_CPU_SHARE=auto` (or a number) is kept and only warns. The share's two page-locked buffers (about 50 MB) are now taken once, at their largest, and a failed allocation ends the share for that run instead of failing the request.
+
 When on, the CPU's rows are computed in the CPU's own activation format, so the output changes in the last bits (first
 token KL against off: mean 0.006, max 0.026 nats over 22 prompts; about half of the 32-token greedy answers on 500 and
 1,000-token prompts are identical, the rest part at a near tie after about 23 tokens). Chunks of 3,072 tokens and more

include/strata/prefill/share_rules.hpp

@@ -7,6 +7,8 @@
 #pragma once
 
 #include <cstdint>
+#include <cstdio>
+#include <string>
 
 namespace strata::prefill::share_rules {
 
@@ -15,4 +17,46 @@ constexpr bool default_share_has_work(int64_t resident_slots, int64_t total_pair
     return total_pairs <= 0 || resident_slots < total_pairs;
 }
 
+
+// ---- memory (0.1.41 report, RTX 5090, tight RAM: chats failed and the engine restarted until STRATA_PREFILL_CPU_SHARE=0) ----
+// The default share reads experts with the CPU and keeps two pinned buffers on the host.  When the host has little
+// memory left, or the experts it would read are file pages the RAM cannot keep (each is a page-in on a worker thread),
+// the default is not armed.  An explicit STRATA_PREFILL_CPU_SHARE is the user's: it is armed and the engine only warns.
+constexpr uint64_t kShareRamMargin = 3ull << 30;     // RAM kept free beside the share's own buffers and the file pages
+constexpr uint64_t kShareCommitFloor = 4ull << 30;   // Windows commit kept free (the #1607 floor)
+constexpr uint64_t kUnknown = ~uint64_t{0};
+
+struct ShareMemory {
+    uint64_t avail_ram = kUnknown;       ///< RAM this process can get (cgroup-aware); kUnknown: not known
+    uint64_t avail_commit = kUnknown;    ///< Windows commit left (RAM + page file); kUnknown elsewhere
+    uint64_t buffer_bytes = 0;           ///< host bytes the share itself allocates (pinned buffers, scratch)
+    uint64_t inplace_bytes = 0;          ///< expert bytes the CPU would read as file pages (not in the GPU cache, not in RAM)
+};
+
+struct ShareVerdict {
+    bool arm = true;
+    std::string why;   ///< when !arm: what the engine says, in one line
+};
+
+inline ShareVerdict default_share_memory_verdict(const ShareMemory& m) {
+    ShareVerdict v;
+    const double gib = 1073741824.0;
+    char buf[300];
+    if (m.avail_ram != kUnknown && m.avail_ram < m.buffer_bytes + kShareRamMargin) {
+        std::snprintf(buf, sizeof buf, "only %.1f GiB of RAM is free (the share needs %.2f GiB and %.1f GiB left over)",
+                      (double) m.avail_ram / gib, (double) m.buffer_bytes / gib, (double) kShareRamMargin / gib);
+        v.arm = false;
+    } else if (m.avail_commit != kUnknown && m.avail_commit < m.buffer_bytes + kShareCommitFloor) {
+        std::snprintf(buf, sizeof buf, "only %.1f GiB of the Windows commit limit is left (the share needs %.2f GiB and %.1f GiB left over)",
+                      (double) m.avail_commit / gib, (double) m.buffer_bytes / gib, (double) kShareCommitFloor / gib);
+        v.arm = false;
+    } else if (m.inplace_bytes > 0 && m.avail_ram != kUnknown && m.inplace_bytes + kShareRamMargin > m.avail_ram) {
+        std::snprintf(buf, sizeof buf, "the experts it would read are %.1f GiB of file pages and only %.1f GiB of RAM is free to keep them",
+                      (double) m.inplace_bytes / gib, (double) m.avail_ram / gib);
+        v.arm = false;
+    }
+    if (!v.arm) v.why = buf;
+    return v;
+}
+
 }  // namespace strata::prefill::share_rules

src/prefill/prefill.cpp

@@ -746,6 +746,7 @@ struct Prefill::Impl {
     // writes (Dm's tail, in Dm's row order), both pinned; and the pool's per-token activations and jobs.
     float *cpu_x = nullptr, *cpu_rows = nullptr;
     size_t cpu_x_n = 0, cpu_rows_n = 0;
+    bool cpu_dead = false;   // a pinned buffer of the share could not be had (host memory is short): the GPU does it all
     std::vector<uint8_t> cpu_nact;
     std::vector<kernels::cpu::ActQ> cpu_actq;   // a Q2_0 layer's activations (the pool's Q2_0 kernels read ActQ)
     std::vector<kernels::cpu::ExpertJobMulti> cpu_jobs;
@@ -3005,8 +3006,18 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
                         int32_t* src_h = grp_mapped ? m.grp_host + 2 * m.grp_tk : m.src_host.data();
                         if (grp_mapped) copy_i32(m.grp_dev, m.ids, T * K, m.cs);
                         else cudaMemcpyAsync(m.ids_host.data(), m.ids, (size_t) T * K * 4, cudaMemcpyDeviceToHost, m.cs);
-                        const bool cpu_maybe = pool_hold.held && !stream_all && m.src != nullptr &&
-                                               lay.native && !m.pp && !lay.fmt.empty();
+                        bool cpu_maybe = pool_hold.held && !stream_all && m.src != nullptr && !m.cpu_dead &&
+                                         lay.native && !m.pp && !lay.fmt.empty();
+                        // The share's two pinned buffers are small (~50 MB) but are had mid-request, when the host may have
+                        // nothing left (0.1.41 report: tight RAM, chats failing).  A buffer that cannot be had ends the share
+                        // for this engine - the GPU takes every expert, as with STRATA_PREFILL_CPU_SHARE=0 - instead of
+                        // failing the request.
+                        auto cpu_share_off = [&](const char* what) {
+                            m.cpu_dead = true;
+                            cpu_maybe = false;
+                            std::fprintf(stderr, "prefill: the CPU share is off for this run: the host could not give %s "
+                                                 "(RAM is short); the GPU streams every expert\n", what);
+                        };
                         // auto shares only while the layers that share cost less per non-resident expert than their
                         // neighbours that do not.  #1282 (RX 7900 GRE + 5700X3D): every share lost, auto's 0.47 by
                         // 4.5%, the host's issue time per streamed expert rising 107 -> 200 us with the CPU busy -
@@ -3034,12 +3045,14 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
                                 m.cpu_x = nullptr;
                                 m.cpu_x_n = 0;
                                 if (cudaHostAlloc((void**) &m.cpu_x, want * sizeof(float), cudaHostAllocDefault) != cudaSuccess) {
-                                    err = "prefill: cannot allocate the CPU experts' activations";
-                                    return false;
+                                    (void) cudaGetLastError();
+                                    m.cpu_x = nullptr;
+                                    cpu_share_off("the activation buffer");
+                                } else {
+                                    m.cpu_x_n = want;
                                 }
-                                m.cpu_x_n = want;
                             }
-                            cudaMemcpyAsync(m.cpu_x, m.mixed, want * sizeof(float), cudaMemcpyDeviceToHost, m.cs);
+                            if (cpu_maybe) cudaMemcpyAsync(m.cpu_x, m.mixed, want * sizeof(float), cudaMemcpyDeviceToHost, m.cs);
                         }
                         // #579: a stall here is the GPU (this layer's attention and router, or the previous layer's
                         // work), not the host: the watchdog's report says so (only its text changes)
@@ -3047,6 +3060,21 @@ bool Prefill::run_impl(const int64_t* tokens, int64_t n, int64_t pos0, std::stri
                                           l, p0);
                         cudaStreamSynchronize(m.cs);
                         core::progress_at("reading the prompt (batched): layer", l, p0);
+                        if (cpu_maybe) {   // the rows buffer, once, at its largest (an expert takes at most MAXT tokens' rows)
+                            const size_t cap = (size_t) std::min<int64_t>(T * K, (int64_t) strata::kernels::cpu::MAXT * m.g->n_expert) * N;
+                            if (m.cpu_rows_n < cap) {
+                                if (m.cpu_rows) cudaFreeHost(m.cpu_rows);
+                                m.cpu_rows = nullptr;
+                                m.cpu_rows_n = 0;
+                                if (cudaHostAlloc((void**) &m.cpu_rows, cap * sizeof(float), cudaHostAllocDefault) != cudaSuccess) {
+                                    (void) cudaGetLastError();
+                                    m.cpu_rows = nullptr;
+                                    cpu_share_off("the rows buffer");
+                                } else {
+                                    m.cpu_rows_n = cap;
+                                }
+                            }
+                        }
                         if (m.cpu_pend) {   // the measured share: the last CPU-sharing layer's GPU time is final now
                             m.cpu_pend = false;
                             float g_ms = 0;

src/program/generate.cpp

@@ -6521,6 +6521,35 @@ int main(int argc, char** argv) {
         std::fprintf(stderr, "strata generate: the mapped view of the experts %s (%s)\n",
                      dropped ? "is closed" : "stays open", why.c_str());
     }
+    // 0.1.42 (H11): the default CPU share is armed only when the host has the memory for it.  One line says why not;
+    // an explicit STRATA_PREFILL_CPU_SHARE (other than 0) is armed anyway, with a warning when memory is short.
+    auto share_default_ok = [&]() -> bool {
+        strata::prefill::share_rules::ShareMemory sm;
+        strata::core::detail::HostMemory hm;
+        if (strata::core::detail::host_available_memory(hm)) {
+            sm.avail_ram = hm.available;
+            sm.avail_commit = hm.commit == ~uint64_t{0} ? strata::prefill::share_rules::kUnknown : hm.commit;
+        }
+        sm.buffer_bytes = (uint64_t) (1024 + 8 * g.n_expert) * (uint64_t) g.n_embd * 4 + (256ull << 20);
+        if (srcp == &src) {
+            const auto& lay = strata::kernels::cpu::expert_layout();
+            for (int64_t l = 0; l < g.n_layers; ++l)
+                for (int64_t e = 0; e < g.n_expert; ++e)
+                    if (xcache.slot_of(l, e) < 0 && src.resident_blob(l, e) == nullptr && !src.transient(l, e))
+                        sm.inplace_bytes += (uint64_t) lay.blob_bytes(l);
+        }
+        const auto v = strata::prefill::share_rules::default_share_memory_verdict(sm);
+        if (v.arm) return true;
+        const char* ev = std::getenv("STRATA_PREFILL_CPU_SHARE");
+        if (ev == nullptr)
+            std::fprintf(stderr, "strata: the CPU prefill share is not armed by default: %s. STRATA_PREFILL_CPU_SHARE=auto "
+                                 "turns it on anyway\n", v.why.c_str());
+        else if (std::strcmp(ev, "0") != 0)
+            std::fprintf(stderr, "strata: WARNING STRATA_PREFILL_CPU_SHARE=%s is set and kept, but %s: the share can slow "
+                                 "prompts or fail on a machine this short of memory (STRATA_PREFILL_CPU_SHARE=0 turns it off)\n",
+                         ev, v.why.c_str());
+        return false;
+    };
     if (o.serve) {
         if (o.spec < 2 || o.prefill_chunk <= 0 ||
             (graph_hits && (thits.d_res == nullptr || host_res.empty()))) {
@@ -6555,12 +6584,13 @@ int main(int argc, char** argv) {
         }
 #endif
         const bool share_pool = o.batch <= 0 && !o.no_pool;
-        const bool share_has_work = strata::prefill::share_rules::default_share_has_work(
+        bool share_has_work = strata::prefill::share_rules::default_share_has_work(
             (int64_t) xcache.slots(), (int64_t) (g.n_layers * g.n_expert));   // #1595
         if (!share_has_work)
             std::fprintf(stderr, "strata serve: every expert is in the GPU cache (%lld of %lld): the CPU prefill share has "
                                  "nothing to take, so it is not armed by default (STRATA_PREFILL_CPU_SHARE=auto forces it)\n",
                          (long long) xcache.slots(), (long long) (g.n_layers * g.n_expert));
+        if (share_pool && stages.empty() && !multi_gpu && share_has_work && !share_default_ok()) share_has_work = false;   // H11
         strata::prefill::Prefill::arm_cpu_share(share_pool, share_pool && stages.empty() && !multi_gpu && share_has_work);
         void* borrow = nullptr;
         uint64_t borrow_bytes = 0;
@@ -11595,7 +11625,8 @@ int main(int argc, char** argv) {
     // #1595: not by default when the cache holds every expert (nothing to take)
     strata::prefill::Prefill::arm_cpu_share(!multi_gpu && !o.no_pool,
         !multi_gpu && !o.no_pool && o.batch <= 0 &&
-        strata::prefill::share_rules::default_share_has_work((int64_t) xcache.slots(), (int64_t) (g.n_layers * g.n_expert)));   // before the chunk below sizes the loan
+        strata::prefill::share_rules::default_share_has_work((int64_t) xcache.slots(), (int64_t) (g.n_layers * g.n_expert)) &&
+        share_default_ok());   // H11   // before the chunk below sizes the loan
     bool kvg_started = false;   // the elastic K/V took this run's cells
     double prefill_batched_ms = 0;
     std::FILE* final_r = o.dump_final_r.empty() ? nullptr : std::fopen(o.dump_final_r.c_str(), "wb");

tests/prefill/share_rules_test.cpp

@@ -1,5 +1,6 @@
 // The default CPU share is not armed when every expert is resident (#1595): no GPU, no model.
 #include <cstdio>
+#include <string>
 
 #include "strata/prefill/share_rules.hpp"
 
@@ -17,6 +18,39 @@ int main() {
     CHECK(default_share_has_work(0, pairs));               // empty cache
     CHECK(default_share_has_work(100000, 0));              // total unknown: keep today's behaviour
     static_assert(!default_share_has_work(10, 10), "constexpr");
+
+    // H11: memory verdicts
+    using namespace strata::prefill::share_rules;
+    const uint64_t G = 1ull << 30;
+    ShareMemory m;
+    m.buffer_bytes = G / 3;
+    CHECK(default_share_memory_verdict(m).arm);                                   // nothing known: today's behaviour
+    m.avail_ram = 20 * G; m.avail_commit = 30 * G;
+    CHECK(default_share_memory_verdict(m).arm);                                   // plenty of RAM and commit
+    m.avail_ram = 3 * G;
+    CHECK(!default_share_memory_verdict(m).arm);                                  // RAM below buffers + margin
+    CHECK(default_share_memory_verdict(m).why.find("RAM is free") != std::string::npos);
+    m.avail_ram = kShareRamMargin + m.buffer_bytes;
+    CHECK(default_share_memory_verdict(m).arm);                                   // exactly enough
+    m.avail_ram = kShareRamMargin + m.buffer_bytes - 1;
+    CHECK(!default_share_memory_verdict(m).arm);                                  // one byte short
+    m.avail_ram = 20 * G; m.avail_commit = 3 * G;
+    {
+        const ShareVerdict v = default_share_memory_verdict(m);
+        CHECK(!v.arm);                                                            // Windows commit nearly used up
+        CHECK(v.why.find("commit") != std::string::npos);
+    }
+    m.avail_commit = kUnknown;
+    m.inplace_bytes = 18 * G;                                                     // file pages the RAM cannot keep
+    {
+        const ShareVerdict v = default_share_memory_verdict(m);
+        CHECK(!v.arm);
+        CHECK(v.why.find("file pages") != std::string::npos);
+    }
+    m.inplace_bytes = 10 * G;                                                     // they fit with the margin
+    CHECK(default_share_memory_verdict(m).arm);
+    m.avail_ram = kUnknown; m.inplace_bytes = 500 * G;
+    CHECK(default_share_memory_verdict(m).arm);                                   // RAM unknown: never blocks
     if (failed == 0) std::printf("share_rules_test: ok\n");
     return failed ? 1 : 0;
 }