[Engine] the MMQ scratch now comes out of the prompt loan

docs/DETAILS.md

@@ -2197,6 +2197,24 @@ the document, +0.4% on the chat. Details: `bench/results/2026-09-27-esp/`.
   against the FP16 prompt path (IQ3_XXS, `--prefill auto`): 8K prompt KL 0.042 (0.1.39: 0.054), top-1 agreement
   93.5% (92.2%); 32K prompt KL 0.020 (0.019), top-1 95.3% (95.0%) - the same band as before. `STRATA_RING_BYTES=0`
   restores 0.1.39's ring, loan and chunk choice.
+- **The prompt path's MMQ scratch (its stream-k fixup tiles) can come out of the loan's unused tail, not out of
+  free VRAM** (opt-in, `STRATA_MMQ_SLAB=1`; off when unset). llama.cpp's mmq.cuh allocates those tiles lazily
+  through `ggml_cuda_pool`, which this build backs with a
+  `cudaMalloc` on whatever is free at that moment - a few MB on a card the expert cache has filled, at the one
+  moment a prompt is running, and a malloc that fails there aborts the engine mid-read instead of failing the
+  request. A layout does not use the whole loan: the carve is a bump, so everything from its stop to the loan's end
+  is the prompt's own memory, provably (a loan is whole cache slots and the top one is only partly needed). The
+  pool takes the tiles from there, in size classes (a power of two, 128 KiB floor: MMQ's J search picks a different
+  J per launch, so a cache keyed on the exact size holds one buffer per J - 32 buffers, 31.5 MiB on a 579-token
+  chunk - the Maya checkout's measurement on its rig - where a class bounds the pool to a few MB). It is the Maya
+  checkout's (`project-maya`) local commit `bf844d5`, which hit it at a 16,384-token chunk on a
+  5060 (3.4 MB free of 7,707.7, 1.2 MB wanted) and measured the slab costing 42-47 more lent experts per part and no
+  change to the chunk arithmetic - its numbers on its live split (5060/5060/3060/3060). **The check here** is the start's
+  `strata prefill mmq: CUDA%d the scratch pool takes a new N MiB buffer` line (it prints for a buffer of 4 MiB or
+  more): with `STRATA_MMQ_SLAB=1` and a loan that has a tail it should not appear at all, and a prompt should read
+  to completion as before. Unset, the scratch stays on free VRAM (the size classes stay - they are what bounds the
+  pool), and a layout that uses its whole loan, or that spilled into owned buffers, publishes no tail - the pool
+  then allocates as it always did.
 
 The full story, with measurements, bottlenecks and what comes next: **[docs/paper/Strata-Paper.pdf](paper/Strata-Paper.pdf)**.
 

include/strata/prefill/moe_mmq.hpp

@@ -30,6 +30,13 @@ size_t matrix_bytes(int ggml_type, int64_t rows, int64_t cols);
 /// Bytes of `rows` activation rows of `cols` values quantized for MMQ (the row padded to 512 values).
 size_t q8_bytes(int64_t rows, int64_t cols);
 
+/// Lend the MMQ kernels' device scratch (their stream-k fixup tiles, which mmq.cuh allocates lazily through
+/// `ggml_cuda_pool` and this build's pool backs with a cudaMalloc on the card's free VRAM) the part of the borrowed
+/// expert-cache region a chunk's layout did not use, so it costs VRAM the prompt has already taken.  Call per
+/// layout from the prompt path, with the tail its carve left (a null `base` - the loan is gone, or the layout used
+/// all of it - drops the buffers cached from the previous layout).
+void pool_slab(int dev, void* base, size_t bytes);
+
 /// q8_1 activations for MMQ against weights of `ggml_type`: row i of the output is row ids[i] of x (or row i when
 /// ids is null); `x` has `ld` floats per row.
 void quantize(const float* x, const int32_t* ids, void* xq, int ggml_type, int64_t cols, int64_t ld, int64_t rows,

src/prefill/ggml_cuda_host.cu

@@ -3,6 +3,7 @@
 #include "strata/core/emulate.hpp"
 #include "common.cuh"
 
+#include <algorithm>
 #include <cstdio>
 #include <cstdlib>
 #include <mutex>
@@ -114,19 +115,86 @@ const ggml_cuda_device_info & ggml_cuda_info() {
 }
 
 namespace {
+// The prompt path's MMQ scratch, where the pool has no business taking VRAM: the part of the borrowed expert-cache
+// region (a request's loan, Prefill::init/relayout) that this chunk's layout did not use.  It is the prompt's memory
+// for as long as the prompt holds the loan, so the fixup tiles cost the card nothing - and on a card whose free VRAM
+// is a handful of MB at prompt time the cudaMalloc below is what used to abort the engine mid-prompt.  Set per
+// layout (mmq::pool_slab, from Prefill::init and Prefill::relayout on the device that will run the chunk), so the
+// buffers cached from the previous one are dropped when it moves.
+struct Slab {
+    uint8_t * base = nullptr;
+    size_t bytes = 0, used = 0;
+};
+Slab g_slab[GGML_CUDA_MAX_DEVICES];
+std::mutex g_slab_mu;
+
+// MMQ's stream-k fixup asks for nsm*J*I floats, and its J search (8..128 by 8) picks a different J per launch, so a
+// cache keyed on the exact size holds one buffer per J: 32 buffers, 31.5 MiB on a 579-token chunk (and the same at
+// 16k, since J does not grow with the chunk - the Maya checkout measured it on its rig).  Rounding a request up to
+// a size class (a power of two, 128 KiB floor - 128 KiB..2 MiB covers the whole J range on a 30-SM card) bounds the
+// pool to a few MiB whatever J wins.  Handing back more than was asked is what every caching allocator does; the
+// kernels only touch the first nsm*J*I floats.
+size_t size_class(size_t n) {
+    size_t c = 128u << 10;
+    while (c < n) c <<= 1;
+    return c;
+}
+
 // Buffers are kept and reused: MMQ asks for the same few sizes every launch (its stream-k fixup tiles).
 struct CachingPool : ggml_cuda_pool {
-    struct Buf { void * p; size_t size; bool used; };
+    struct Buf { void * p; size_t size; bool used; bool slab; };
     std::vector<Buf> bufs;
     std::mutex mu;
+    int dev;
+    void * slab_seen = nullptr;   // the slab base the cached buffers came from; a move makes them stale
+    explicit CachingPool(int device) : dev(device) {}
+    size_t held() const {   // the bytes this pool holds on the card (never given back)
+        size_t n = 0;
+        for (const auto & b : bufs) n += b.size;
+        return n;
+    }
     void * alloc(size_t size, size_t * actual_size) override {
         std::lock_guard<std::mutex> lk(mu);
-        for (auto & b : bufs)
-            if (!b.used && b.size >= size) { b.used = true; *actual_size = b.size; return b.p; }
+        const size_t cls = size_class(size);
+        {
+            std::lock_guard<std::mutex> sl(g_slab_mu);
+            Slab & s = g_slab[dev];
+            if (s.base != slab_seen) {
+                slab_seen = s.base;
+                // the layout moved under the cached buffers: nothing is in flight between chunks, so drop them
+                bufs.erase(std::remove_if(bufs.begin(), bufs.end(),
+                                          [](const Buf & b) { return b.slab; }), bufs.end());
+            }
+            for (auto & b : bufs)
+                if (!b.used && b.size >= cls) { b.used = true; *actual_size = b.size; return b.p; }
+            if (s.base != nullptr && s.bytes >= s.used && s.bytes - s.used >= cls) {
+                void * p = s.base + s.used;
+                s.used += cls;
+                *actual_size = cls;
+                bufs.push_back({p, cls, true, true});
+                return p;
+            }
+        }
+        size_t free_b = 0, total_b = 0;
+        cudaMemGetInfo(&free_b, &total_b);
         void * p = nullptr;
-        CUDA_CHECK(cudaMalloc(&p, size));
-        bufs.push_back({p, size, true});
-        *actual_size = size;
+        const cudaError_t e = cudaMalloc(&p, cls);
+        if (e != cudaSuccess) {
+            // this is the whole reason a prompt used to abort mid-read with "out of memory" and nothing else
+            std::fprintf(stderr,
+                         "ggml-cuda (strata mmq): CUDA%d wanted %.1f MiB of scratch (the prompt path's pool held %.1f "
+                         "MiB in %zu buffers) and only %.1f MiB of %.1f MiB was free: %s\n",
+                         dev, (double) cls / 1048576.0, (double) held() / 1048576.0, bufs.size(),
+                         (double) free_b / 1048576.0, (double) total_b / 1048576.0, cudaGetErrorString(e));
+            CUDA_CHECK(e);
+        }
+        if (cls >= ((size_t) 4 << 20))
+            std::fprintf(stderr, "strata prefill mmq: CUDA%d the scratch pool takes a new %.1f MiB buffer (%.1f MiB free of "
+                                 "%.1f MiB, %.1f MiB held)\n",
+                         dev, (double) cls / 1048576.0, (double) free_b / 1048576.0, (double) total_b / 1048576.0,
+                         (double) held() / 1048576.0);
+        bufs.push_back({p, cls, true, false});
+        *actual_size = cls;
         return p;
     }
     void free(void * ptr, size_t) override {
@@ -135,13 +203,24 @@ struct CachingPool : ggml_cuda_pool {
             if (b.p == ptr) { b.used = false; return; }
     }
     ~CachingPool() override {
-        for (auto & b : bufs) cudaFree(b.p);
+        for (auto & b : bufs)
+            if (!b.slab) cudaFree(b.p);   // slab buffers are the borrowed region's tail, not ours to free
     }
 };
 }  // namespace
 
-std::unique_ptr<ggml_cuda_pool> ggml_backend_cuda_context::new_pool_for_device(int, int) {
-    return std::unique_ptr<ggml_cuda_pool>(new CachingPool());
+// Called by the prompt path once per layout, before that layout's MMQ launches: `base`/`bytes` are the tail of the
+// borrowed expert-cache region this chunk's carve left over (prefill.cpp's carve, which knows both).  The bump
+// restarts there; a tail somewhere else invalidates the buffers cached from the old one.  A null `base` (the loan is
+// gone, or the layout used all of it) drops them and the pool allocates on the device again.
+extern "C" void strata_mmq_pool_slab(int dev, void * base, size_t bytes) {
+    std::lock_guard<std::mutex> lk(g_slab_mu);
+    if (dev < 0 || dev >= GGML_CUDA_MAX_DEVICES) return;
+    g_slab[dev] = {(uint8_t *) base, base != nullptr ? bytes : 0, 0};
+}
+
+std::unique_ptr<ggml_cuda_pool> ggml_backend_cuda_context::new_pool_for_device(int device, int) {
+    return std::unique_ptr<ggml_cuda_pool>(new CachingPool(device));
 }
 
 ggml_backend_cuda_context::~ggml_backend_cuda_context() {}

src/prefill/moe_mmq.cu

@@ -684,4 +684,9 @@ void set_bounds(int32_t* dst, int32_t rows, void* stream) {
 }
 #endif
 
+// src/prefill/ggml_cuda_host.cu's pool: the slab it may hand out instead of a cudaMalloc (see the header)
+extern "C" void strata_mmq_pool_slab(int dev, void* base, size_t bytes);
+
+void pool_slab(int dev, void* base, size_t bytes) { strata_mmq_pool_slab(dev, base, bytes); }
+
 }  // namespace strata::prefill::mmq

src/prefill/prefill.cpp

@@ -70,6 +70,7 @@ bool supported(int) { return false; }
 bool fits(int, int64_t) { return false; }
 size_t matrix_bytes(int, int64_t, int64_t) { return 0; }
 size_t q8_bytes(int64_t, int64_t) { return 0; }
+void pool_slab(int, void*, size_t) {}
 void quantize(const float*, const int32_t*, void*, int, int64_t, int64_t, int64_t, void*, float*) {}
 bool fp4_activations(int) { return false; }
 void quantize_scatter(const float*, const int32_t*, const int32_t*, void*, int, int64_t, int64_t, int64_t, int, void*) {}
@@ -995,6 +996,11 @@ struct Prefill::Impl {
     // layer split: the next stage's GPU as a stream-only peer (set_stage_helper); its buffers are that stage's own
     // prompt buffers, bound per prompt, and it takes the place of `pp` for a run it helps
     std::unique_ptr<PeerPrefill> help_pp;
+    // THE MMQ SCRATCH'S TAIL (lend_mmq_scratch): the part of the current layout's loan the carve did not take, which
+    // the MMQ pool may allocate its stream-k fixup tiles in.  Re-published at the start of every run (a request that
+    // reuses the layout unchanged still lends it) and withdrawn when the run ends (the loan is given back then).
+    uint8_t* mmq_slab = nullptr;
+    size_t mmq_slab_bytes = 0;
 };
 
 namespace {
@@ -1050,6 +1056,9 @@ void Prefill::reset() {
 void Prefill::release() {
     if (!impl_) return;
     impl_->hand_worker.reset();   // idle here (run drains the chain before it returns); joins its thread
+    mmq::pool_slab(impl_->device, nullptr, 0);   // the layout this slab belonged to is going away
+    impl_->mmq_slab = nullptr;                   // and a run before the next carve publishes nothing
+    impl_->mmq_slab_bytes = 0;
     if (impl_->cs) cudaStreamSynchronize(impl_->cs);
     if (impl_->copy) cudaStreamSynchronize(impl_->copy);
     if (impl_->kv_copy) cudaStreamSynchronize(impl_->kv_copy);
@@ -1328,6 +1337,39 @@ void warm_gemm_libraries(int device) {
 #endif
 }
 
+
+// The MMQ kernels' stream-k fixup tiles (llama.cpp's mmq.cuh) come out of this build's ggml_cuda_pool, which
+// cudaMallocs them on whatever VRAM is free at that moment - a few MB on a card the expert cache has filled, at the
+// one moment a prompt is running, and a malloc that fails there aborts the engine mid-read (the Maya checkout saw
+// exactly that at a 16k chunk, local commit bf844d5: 3.4 MB free of 7,707, 1.2 MB wanted).  The part of the loan
+// this layout did not take is provably the prompt's own memory - the carve is a bump, `used` is where it stopped
+// and every byte below it is a named buffer - so the pool may take the tiles from there instead of from the card's
+// leftovers.
+//
+// Three things make it safe, and all three are load-bearing.  (1) It is published per layout, keyed on the tail's
+// ADDRESS: a request whose chunk moves `used` moves the tail, and the pool drops the buffers it cached from the old
+// one.  (2) `Prefill::run` drops it again once the prompt is read - the caller refills the borrowed slots as soon as
+// it returns (a loan is whole slots, so the tail can hold one), and a cached buffer in a refilled slot is silent
+// corruption, not a crash.  (3) The loan is the participant's own expert cache, `cudaMalloc`ed on this device
+// (ExpertCache::open, under core::OnDevice(p.dev)), and this device's kernels already write experts into the carved
+// part of it, so the tail is this card's memory to write - there is nothing to ask the driver.  A tail that is empty
+// (the layout used the whole loan, or spilled past it into owned buffers) publishes nothing, and the pool allocates
+// on the device as it always did.  It is opt-in: STRATA_MMQ_SLAB=1 turns it on, and unset leaves the pool on free
+// VRAM as it always was.
+static void lend_mmq_scratch(int dev, uint8_t** slab, size_t* slab_bytes, const Alloc& o) {
+    static const bool on = [] {
+        const char* v = std::getenv("STRATA_MMQ_SLAB");
+        return v != nullptr && v[0] != '\0' && std::atoi(v) != 0;
+    }();
+    *slab = nullptr;
+    *slab_bytes = 0;
+    if (on && o.base != nullptr && o.used < o.cap) {
+        *slab = o.base + o.used;
+        *slab_bytes = (size_t) (o.cap - o.used);
+    }
+    mmq::pool_slab(dev, *slab, *slab_bytes);
+}
+
 bool Prefill::init(const core::WeightTable& wt, const core::ModelGeometry& g, core::SessionState& ss,
                    core::ExpertSource* src, const core::ExpertCache* cache, const int32_t* host_res, int64_t chunk,
                    void* stream, std::string& err, void* borrow, uint64_t borrow_bytes, bool borrow_spill) {
@@ -1520,6 +1562,7 @@ bool Prefill::init(const core::WeightTable& wt, const core::ModelGeometry& g, co
         return false;
     }
     { int wd = 0; if (cudaGetDevice(&wd) == cudaSuccess) warm_gemm_libraries(wd); else (void) cudaGetLastError(); }
+    lend_mmq_scratch(m.device, &m.mmq_slab, &m.mmq_slab_bytes, o);
     return true;
 }
 
@@ -1662,6 +1705,7 @@ bool Prefill::relayout(int64_t chunk, void* borrow, uint64_t borrow_bytes, std::
               std::to_string(o.failed_bytes >> 20) + " MiB)";
         return false;
     }
+    lend_mmq_scratch(m.device, &m.mmq_slab, &m.mmq_slab_bytes, o);   // the chunk moved the tail: the pool follows it
     return true;
 }
 
@@ -5060,6 +5104,15 @@ bool Prefill::run(const int64_t* tokens, int64_t n, int64_t pos0, std::string& e
         return ok;
     }
 #endif
+    Impl& m = *impl_;
+    // This prompt's MMQ scratch: the tail of its loan's unused part, or nothing (see lend_mmq_scratch).  A request
+    // that reuses the layout unchanged still lends it (this call), and the withdraw below is a guard, so every exit
+    // takes it with it - a read that failed gives the loan back too.
+    mmq::pool_slab(m.device, m.mmq_slab, m.mmq_slab_bytes);
+    struct SlabEnd {
+        int dev;
+        ~SlabEnd() { mmq::pool_slab(dev, nullptr, 0); }   // the loan goes back: the pool must not hand these out
+    } slab_end{m.device};
     const bool body_ok = run_impl(tokens, n, pos0, err);
 
     std::string drain_err;