Qwen parallel: FP16 weight cache sized by free VRAM (3 GiB kept free), no loan of cache slots when every expert fits, replica profile keeps 2500 MiB free

include/strata/prefill/gemm.hpp

@@ -100,6 +100,9 @@ private:
     std::vector<W16Entry> w16_;
     int64_t w16_budget_ = -1;   // bytes; -1: env not read yet, 0: off
     int64_t w16_used_ = 0;
+    int64_t w16_reserve_ = 2ll << 30;   // VRAM kept free after a cached weight
+    bool w16_logged_ = false;
+    static constexpr int64_t kW16AutoReserve = 3ll << 30;
     void* hipblaslt_state_ = nullptr;
     void* rocblas_state_ = nullptr;   // HIP: the rocBLAS solution table of f16_inplace (STRATA_ROCBLAS_TUNING)
     bool f16_io_ = false;

src/prefill/gemm.cu

@@ -1394,10 +1394,26 @@ void Gemm::native(const uint16_t* X, int ggml_type, const void* W_blocks, float*
     if (beta == 0.0f && (ldx == 0 || ldx == K) && native_mmq(X, ggml_type, W_blocks, Y, T, N, K, ldy)) return;
 #endif
 #ifndef STRATA_USE_HIP
-    // STRATA_PF_W16_CACHE_MIB=N: the FP16 weight kept on the device after its first dequant (same values, same product)
+    // The FP16 weight kept on the device after its first dequant (same values, same product: bit-exact).  0.2.9.1: on by
+    // default (Linux CUDA) and sized by the card, never by the user.  It takes only VRAM nobody else holds: a weight is
+    // cached while the card keeps kW16AutoReserve free afterwards (the MTP prompt capture, the batch windows and a long
+    // prompt's buffers are allocated later, out of that room; a replica with 309 MiB free ran out of memory at 14K tokens),
+    // so a card whose expert cache fills it caches nothing and a card with room (all experts resident) caches what the
+    // prompt path uses (~6 GB for Qwen3.8).  STRATA_PF_W16_CACHE_MIB=N is an explicit budget (2 GiB kept free, as before),
+    // =0 turns it off.
     if (w16_budget_ < 0) {
         const char* v = std::getenv("STRATA_PF_W16_CACHE_MIB");
-        w16_budget_ = v ? (int64_t) std::atoll(v) * (1ll << 20) : 0;
+        if (v != nullptr) {
+            w16_budget_ = (int64_t) std::atoll(v) * (1ll << 20);
+            w16_reserve_ = 2ll << 30;
+        } else {
+#if defined(_WIN32)
+            w16_budget_ = 0;   // never run on Windows NVIDIA: off until a box has measured it
+#else
+            w16_budget_ = 16ll << 30;
+            w16_reserve_ = kW16AutoReserve;
+#endif
+        }
     }
     if (w16_budget_ > 0 && !side_scratch_) {
         const int64_t bytes = N * K * 2;
@@ -1406,7 +1422,7 @@ void Gemm::native(const uint16_t* X, int ggml_type, const void* W_blocks, float*
             if (e.p == W_blocks && e.type == ggml_type && e.n == N && e.k == K) { cw = e.d; break; }
         if (cw == nullptr && w16_used_ + bytes <= w16_budget_) {
             size_t fb = 0, tb = 0;
-            if (cudaMemGetInfo(&fb, &tb) == cudaSuccess && fb > (size_t) bytes + (2ull << 30) &&
+            if (cudaMemGetInfo(&fb, &tb) == cudaSuccess && fb > (size_t) bytes + (size_t) w16_reserve_ &&
                 cudaMalloc((void**) &cw, (size_t) bytes) == cudaSuccess) {
                 strata::kernels::dequant_f16(ggml_type, W_blocks, 0, N, K, cw, stream_);
                 w16_.push_back({W_blocks, ggml_type, N, K, cw});
@@ -1414,6 +1430,11 @@ void Gemm::native(const uint16_t* X, int ggml_type, const void* W_blocks, float*
             } else {
                 (void) cudaGetLastError();
                 cw = nullptr;
+                if (w16_budget_ > 0 && !w16_logged_) {
+                    w16_logged_ = true;
+                    std::fprintf(stderr, "strata prefill: FP16 weight cache %lld MiB (the card keeps %lld MiB free; STRATA_PF_W16_CACHE_MIB=0 turns it off)\n",
+                                 (long long) (w16_used_ >> 20), (long long) (w16_reserve_ >> 20));
+                }
                 w16_budget_ = w16_used_ > 0 ? w16_used_ : 0;   // out of room: stop trying
             }
         }

src/program/generate.cpp

@@ -5225,7 +5225,7 @@ int main(int argc, char** argv) {
     // a cache from, or --no-prefill-borrow - do the buffers take a reserve, and then this estimate stands in
     // for buffers that cannot be priced exactly yet because the sessions do not exist.  `plan_lend` uses the
     // exact `Prefill::bytes_needed` as soon as it can.
-    const bool pf_borrow = !o.no_prefill_borrow && !o.expert_profile.empty();
+    bool pf_borrow = !o.no_prefill_borrow && !o.expert_profile.empty();   // (0.2.9.1: the auto cache sizing below can clear it, see AUTO NO-BORROW)
     // (#340: the estimate predates the streamed ring: from 1024-token chunks the prompt path also holds a ring of
     // whole expert blobs, which a split without borrowing sizes at 96 (Prefill::set_ring_override below) and books
     // here - without it a `--no-prefill-borrow` split filled the cards and the draft head no longer fit)
@@ -7116,6 +7116,40 @@ int main(int argc, char** argv) {
                                  "(layer, expert) pair lifts it - tools/make_profile.py)\n", o.expert_cache,
                          profile.size(), o.expert_profile.c_str(), profile.size());
         o.expert_cache = (int) sized_slots.size();
+#if !defined(STRATA_USE_HIP) && !defined(STRATA_GPU_METAL) && !defined(_WIN32)
+        // ---- AUTO NO-BORROW (0.2.9.1).  With batch slots on one NVIDIA card the prompt path lent cache slots to its own
+        // buffers, and every lent slot's expert was read again from the host during each chunk.  When the card holds EVERY
+        // expert of the profile and still has room for the path's own buffers (the chunk the loan would have used, 8192)
+        // plus kNoBorrowKeepMib, the loan only costs time: the buffers are allocated instead, the chunk stays what it was, and
+        // the arithmetic is the same (QUAL-PAR2: 12/12 identical, KL 0).  A card that cannot hold all experts keeps the loan -
+        // its cache is worth more than the buffers (a 16 GB A4000 lost 2.5x on a long prompt and ran out of memory with them).
+        // Not with an explicit --prefill N, a split, helper caches, STRATA_PREFILL_BORROW=1 or STRATA_AUTO_NOBORROW=0.
+        if (auto_cache && pf_borrow && o.prefill_auto && o.batch >= 2 && !multi_gpu && stages.empty() && !profile.empty() &&
+            sized_slots.size() == profile.size() && std::none_of(o.expert_cache_remote.begin(), o.expert_cache_remote.end(), [](int n) { return n > 0; })) {
+            const char* e1 = std::getenv("STRATA_AUTO_NOBORROW");
+            const char* e2 = std::getenv("STRATA_PREFILL_BORROW");
+            if (!(e1 != nullptr && e1[0] == '0') && !(e2 != nullptr && e2[0] == '1')) {
+                constexpr int64_t kNoBorrowChunk = 8192, kNoBorrowKeepMib = 3072;
+                uint64_t used_b = 0;
+                for (const int64_t sb : sized_slots) used_b += (uint64_t) strata::core::ExpertCache::slot_stride(sb, slot_guard);
+                const int64_t chunk = std::min<int64_t>(kNoBorrowChunk, std::max<int64_t>(256, o.prefill_auto_max));
+                const int64_t own_mib = ((int64_t) strata::prefill::Prefill::bytes_needed_owned(g, ss, chunk, srcp != nullptr) >> 20) + 64;
+                const size_t fb = strata::core::device_free_bytes();
+                const int64_t room_mib = fb > used_b ? (int64_t) ((fb - used_b) >> 20) - o.vram_reserve_mib - (int64_t) (pipe_first >> 20) : 0;
+                const bool fits = room_mib >= own_mib + kNoBorrowKeepMib;
+                std::fprintf(stderr, "strata generate: every expert fits the cache; the prompt path's own buffers (%lld MiB for a %lld-token chunk) %s "
+                                     "(%lld MiB left beside the %d MiB reserve, %lld MiB kept free)%s\n",
+                             (long long) own_mib, (long long) chunk, fits ? "fit" : "do not fit", (long long) room_mib, o.vram_reserve_mib,
+                             (long long) kNoBorrowKeepMib, fits ? ": no loan of cache slots" : ": the loan stays");
+                if (fits) {
+                    o.no_prefill_borrow = true;
+                    pf_borrow = false;
+                    o.prefill_auto = false;
+                    o.prefill_chunk = chunk;
+                }
+            }
+        }
+#endif
     }
     // #533: --vram-elastic: the cache in physical segments (the VRAM command resizes it between requests).  One GPU,
     // no helper caches, serve mode: anything else keeps the one cudaMalloc, said once.

strata_control/replica_profile.py

@@ -33,6 +33,7 @@ PER_REPLICA_GB = 8.0        # measured: about 7 GB of private RAM per engine bes
 HEADROOM_GB = 12.0          # what the OS, the server and the page cache keep
 MIN_CARD_GB = 24.0          # measured on 32 GB cards; a smaller card is not claimed (an A4000 16 GB is measured separately)
 MIN_REPLICAS = 2
+REPLICA_RESERVE_MIB = 2500   # VRAM each one-card engine keeps free (the draft head's prompt capture needs ~2 GB at 14K tokens)
 SWITCHES = {"STRATA_QSA_SLOTS": "1", "STRATA_BATCH_SELF_COMMIT": "1", "STRATA_SLOT_COPY_DIRECT": "1",
             "STRATA_CKPT_POOL": "1", "STRATA_CKPT_POOL_READY": "4", "STRATA_SLOT_CKPT_MOVE": "1",
             "STRATA_SLOT_COPY_BATCHSYNC": "1"}
@@ -119,6 +120,11 @@ def apply(cards: list[dict], cfg: dict, ram_gb: float, *, parallel_given: bool =
         if key not in environ:                              # an environment variable of the person's own wins
             env.setdefault(key, val)
     cfg["env"] = env
+    # 0.2.9.1 (QUAL-PAR): with the expert cache on auto each engine fills its card and keeps ~300 MiB; a 14K-token read with the draft head
+    # then ran out of memory in the prompt capture and the engine died.  2.5 GB kept free (2,111 MiB left on a 48 GB card) ran
+    # the 61-task suite and 48 minutes of soak with no error.  A reserve the person gave wins.
+    if "--vram-reserve-mib" not in {str(x) for x in cfg.get("args") or []}:
+        cfg["args"] = list(cfg.get("args") or []) + ["--vram-reserve-mib", str(REPLICA_RESERVE_MIB)]
     notes = [f"{k} engines of one card each instead of one layer split of {len(cards)}: requests spread over them "
              f"(about 1.8 times the throughput at 8 to 16 requests at once; one request alone is as fast as before)"]
     if k < len(cards):

tests/terminal/test_replica_profile.py

@@ -66,6 +66,15 @@ class Apply(unittest.TestCase):
         self.assertEqual(c["parallel"], 8)
         self.assertEqual(len(c["replicas"]), 2)
 
+    def test_each_engine_keeps_vram_free_for_the_prompt_capture(self):
+        c = cfg(gpu=[0, 1])
+        P.apply(cards(2), c, 128.0, platform="linux", environ={})
+        self.assertEqual(c["args"][c["args"].index("--vram-reserve-mib") + 1], str(P.REPLICA_RESERVE_MIB))
+        c = cfg(gpu=[0, 1], args=["--pack", "/p/iq3_s", "--vram-reserve-mib", "900"])
+        P.apply(cards(2), c, 128.0, platform="linux", environ={})
+        self.assertEqual(c["args"].count("--vram-reserve-mib"), 1)
+        self.assertEqual(c["args"][c["args"].index("--vram-reserve-mib") + 1], "900")
+
     def test_less_ram_drops_cards_and_says_so(self):
         c = cfg()
         notes = P.apply(cards(4), c, 80.0, platform="linux", environ={})