From 7dacae023a995f92c0e24c34a9abaf98e7b14326 Mon Sep 17 00:00:00 2001 From: piggidragon Date: Fri, 18 Sep 2026 23:30:05 +0200 Subject: [PATCH] ggml-meta : bound the stack use of the split state walk get_split_state recurses into the sources of a tensor. A stale entry clears the whole cache, so a late node derives the full chain again, one large frame per node: a 65 layer graph needs ~66 MB of stack and dies on the default 8 MiB. Fill the sources from the leaves up first, iteratively, so the recursion finds them cached and stays one level deep. Assisted-by: Claude Opus 5 --- ggml/src/ggml-backend-meta.cpp | 78 +++++++++++++++++++++++++++++++++- 1 file changed, 77 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 3ec40fb1af7f..428610ba74b1 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -429,6 +429,9 @@ struct ggml_backend_meta_buffer_context { static constexpr size_t nbtc = GGML_TENSOR_SIZE - sizeof(ggml_tensor::padding); std::map, std::pair> split_state_cache; + // bumped on every clear of the cache, so a caller can tell that the entries it filled are gone + uint64_t split_state_cache_generation = 0; + int debug; ggml_backend_meta_buffer_context( @@ -1110,6 +1113,7 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( auto it = buf_ctx->split_state_cache.find(key); if (it != buf_ctx->split_state_cache.end() && memcmp(it->second.second, (const char *) tensor, sizeof(it->second.second)) != 0) { buf_ctx->split_state_cache.clear(); + buf_ctx->split_state_cache_generation++; it = buf_ctx->split_state_cache.end(); } @@ -1167,9 +1171,80 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( return ret; } +// A stale entry makes ggml_backend_meta_get_split_state clear the whole cache. Clear it here +// instead, before the walk below fills entries that the clear would drop again. +static bool ggml_backend_meta_split_state_cached( + ggml_backend_meta_buffer_context * buf_ctx, const ggml_tensor * tensor, bool assume_sync) { + const auto it = buf_ctx->split_state_cache.find(std::make_pair(tensor, assume_sync)); + if (it == buf_ctx->split_state_cache.end()) { + return false; + } + if (memcmp(it->second.second, (const char *) tensor, sizeof(it->second.second)) != 0) { + buf_ctx->split_state_cache.clear(); + buf_ctx->split_state_cache_generation++; + return false; + } + return true; +} + +// ggml_backend_meta_get_split_state reads the state of the sources of a tensor, so on a cold cache it +// walks the whole graph in one recursion and needs one large stack frame per node. Fill the sources +// from the leaves up first, iteratively, so that the recursion finds them and stays one level deep. +// The walk reaches what the recursion reaches: a leaf holds its own state and has no sources. +static void ggml_backend_meta_warm_split_states( + ggml_backend_meta_simple_tensor_container & stc, const ggml_tensor * tensor, bool assume_sync) { + ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) tensor->buffer->context; + + // a clear drops what the walk filled, but it leaves the cache empty, so a second walk cannot hit one + for (int attempt = 0; attempt < 2; attempt++) { + const uint64_t generation = buf_ctx->split_state_cache_generation; + if (ggml_backend_meta_split_state_cached(buf_ctx, tensor, assume_sync)) { + return; // the caller reads the cache and never recurses + } + + std::vector> stack; // tensor, sources already pushed + std::set pushed; + stack.emplace_back(tensor, false); + pushed.insert(tensor); + + while (!stack.empty()) { + const ggml_tensor * node = stack.back().first; + if (!stack.back().second) { + stack.back().second = true; + if (ggml_nelements(node) == 0) { + continue; // the state does not come from the sources + } + for (size_t i = 0; i < GGML_MAX_SRC; i++) { + const ggml_tensor * src = node->src[i]; + if (src == nullptr || src == node || pushed.count(src) > 0) { + continue; + } + // a cached source covers its own sources as well + if (ggml_backend_meta_split_state_cached(buf_ctx, src, /*assume_sync =*/ true)) { + continue; + } + pushed.insert(src); + stack.emplace_back(src, false); + } + continue; + } + stack.pop_back(); + if (node != tensor) { + ggml_backend_meta_get_split_state(stc, node, /*assume_sync =*/ true); + } + } + + if (buf_ctx->split_state_cache_generation == generation) { + return; + } + } +} + static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync) { ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) tensor->buffer->context; - return ggml_backend_meta_get_split_state(buf_ctx->get_simple_tensor_container(tensor), tensor, assume_sync); + ggml_backend_meta_simple_tensor_container & stc = buf_ctx->get_simple_tensor_container(tensor); + ggml_backend_meta_warm_split_states(stc, tensor, assume_sync); + return ggml_backend_meta_get_split_state(stc, tensor, assume_sync); } static void * ggml_backend_meta_buffer_get_base(ggml_backend_buffer_t buffer) { @@ -1182,6 +1257,7 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor_impl(ggml_backend_m ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) tensor->buffer->context; const size_t n_simple_bufs = ggml_backend_meta_buffer_n_bufs(tensor->buffer); + ggml_backend_meta_warm_split_states(stc, tensor, /*assume_sync =*/ true); const ggml_backend_meta_split_state split_state = ggml_backend_meta_get_split_state(stc, tensor, /*assume_sync =*/ true); GGML_ASSERT(ggml_nelements(tensor) == 0 || split_state.axis != GGML_BACKEND_SPLIT_AXIS_UNKNOWN); GGML_ASSERT(split_state.n_segments <= 16);