From 78d0e6212eb4fa6e37a5a2da4670ee6d20007a8f Mon Sep 17 00:00:00 2001 From: Gopi Date: Sat, 17 Oct 2026 07:55:51 +0530 Subject: [PATCH] prefill: the prompt path's MMQ scratch rides in the loan (opt-in, STRATA_MMQ_SLAB) llama.cpp's mmq.cuh allocates its stream-k fixup tiles through ggml_cuda_pool, which this build backs with a cudaMalloc 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 instead of failing a request. The Maya checkout (`project-maya`) hit exactly that at a 16,384-token chunk (3.4 MB free of 7,707.7, 1.2 MB wanted) and fixed it by lending the pool the loan's unused tail (local commit `bf844d5`). A layout does not use the whole loan: the carve is a bump, so everything from its stop to the loan's end is provably the prompt's own memory (a loan is whole cache slots and the top one is only partly needed). The pool now takes the tiles from there. It is published per layout, keyed on the tail's ADDRESS, republished at the start of every run and withdrawn when the run ends - the caller refills the borrowed slots as soon as the prompt returns, and a cached buffer in a refilled slot would be silent corruption rather than a crash. A layout that uses its whole loan, or that spilled into owned buffers, publishes no tail and the pool allocates as before. The pool also rounds a request up to a size class (power of two, 128 KiB floor). Keyed on the exact size it held one buffer per J of MMQ's J search - 32 buffers, 31.5 MiB on a 579-token chunk (the Maya checkout's measurement on its rig) - where a class bounds it to a few MB, and the cudaMalloc that fails now names what the pool held, how much was free and how much the request wanted. STRATA_MMQ_SLAB=1 turns it on; unset, empty or 0 leaves the pool on free VRAM exactly as before, so the default path is unchanged and a new feature is opt-in. The size classes stay either way - they are what bounds the pool. The tokens are unchanged in both cases: the tiles are scratch and the kernels touch only the first nsm*J*I floats of whatever they are handed. Checked: `src/prefill/prefill.cpp` compiles clean with the build's own flags (`-fsyntax-only`); `src/prefill/ggml_cuda_host.cu` and `src/prefill/moe_mmq.cu` compile clean with the build's nvcc flags (CUDA 13.2, compute_86 and compute_120); the development branch that also carries the change built fully with `ninja -C build strata`. `engine/strata` was not touched. No HIP toolchain exists on this box, so the HIP compile of the three touched sources was not run - the new code uses only cudaMalloc, cudaMemGetInfo, cudaFree, cudaGetErrorString, cudaError_t and std::remove_if, all of which the HIP compat header maps or the host compiler provides. The check to run here (a prompt reads to completion, and no `strata prefill mmq: CUDA%d the scratch pool takes a new N MiB buffer` line at the start) is the owner's to run. Co-Authored-By: Claude Code --- docs/DETAILS.md | 18 ++++++ include/strata/prefill/moe_mmq.hpp | 7 +++ src/prefill/ggml_cuda_host.cu | 97 +++++++++++++++++++++++++++--- src/prefill/moe_mmq.cu | 5 ++ src/prefill/prefill.cpp | 53 ++++++++++++++++ 5 files changed, 171 insertions(+), 9 deletions(-) diff --git a/docs/DETAILS.md b/docs/DETAILS.md index 07c9b3e68a2..4350cf0aff6 100644 --- a/docs/DETAILS.md +++ b/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)**. diff --git a/include/strata/prefill/moe_mmq.hpp b/include/strata/prefill/moe_mmq.hpp index fee662fcc76..13108d73c31 100644 --- a/include/strata/prefill/moe_mmq.hpp +++ b/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, diff --git a/src/prefill/ggml_cuda_host.cu b/src/prefill/ggml_cuda_host.cu index 1245c2b0b22..730c1eb81cf 100644 --- a/src/prefill/ggml_cuda_host.cu +++ b/src/prefill/ggml_cuda_host.cu @@ -3,6 +3,7 @@ #include "strata/core/emulate.hpp" #include "common.cuh" +#include #include #include #include @@ -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 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 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 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_backend_cuda_context::new_pool_for_device(int, int) { - return std::unique_ptr(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 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_backend_cuda_context::new_pool_for_device(int device, int) { + return std::unique_ptr(new CachingPool(device)); } ggml_backend_cuda_context::~ggml_backend_cuda_context() {} diff --git a/src/prefill/moe_mmq.cu b/src/prefill/moe_mmq.cu index e7ff5d45bc1..a64436ca620 100644 --- a/src/prefill/moe_mmq.cu +++ b/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 diff --git a/src/prefill/prefill.cpp b/src/prefill/prefill.cpp index f00f1cc1347..cc8ec6fd6c7 100644 --- a/src/prefill/prefill.cpp +++ b/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 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;