From 258345efa640eb099eb1af3c9ee8f8e6e8e7b0d3 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 31 Aug 2026 01:33:52 +0000 Subject: [PATCH] llama: map each contiguous run of a context's tensors, not one span over all of them get_mapping_range returns a single [first, last) taken as the min and max offset of the tensors a context holds in one file, so it cannot express a gap. When a large tensor belonging to a different buffer type sits between them, that tensor is inside the mapped range. On a backend that copies per tensor this costs nothing, which is why it has gone unnoticed. On a backend that maps the host pointer, buffer_from_host_ptr is called on the whole range and the gap becomes resident for the run. Qwen3.8-Flash-Next puts a 26.8 GiB per_layer_token_embd, which is CPU resident, at tensor index 4 of 373 in its second shard. The GPU context's range then covers the entire 49.84 GiB shard for 20.4 GiB of its own weights, and on Metal all of it is wired. Map each contiguous run separately. Runs closer together than 32 MiB are merged, so a model whose tensors are laid out per buffer type still gets one buffer and one mapping, and only a genuinely large foreign tensor splits the range. If a layout somehow yields more than 64 runs, fall back to the single span rather than create a buffer per tensor. ggml_backend_tensor_alloc already asserts that the address lies inside the buffer it is given, so a tensor resolved against the wrong run aborts rather than mapping silently to the wrong place. Assisted-by: Claude --- src/llama-model-loader.cpp | 57 +++++++++++++++++++++++++++++++++++--- src/llama-model-loader.h | 6 +++- src/llama-model.cpp | 25 +++++++++++------ 3 files changed, 74 insertions(+), 14 deletions(-) diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 7663797ba00..b4d979c9311 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1456,6 +1456,48 @@ void llama_model_loader::get_mapping_range(size_t * first, size_t * last, void * } } +void llama_model_loader::get_mapping_ranges(std::vector> & ranges, void ** addr, int idx, ggml_context * ctx) const { + GGML_ASSERT(!mappings.empty()); + const auto & mapping = mappings.at(idx); + + ranges.clear(); + *addr = mapping->addr(); + + for (ggml_tensor * tensor = ggml_get_first_tensor(ctx); tensor; tensor = ggml_get_next_tensor(ctx, tensor)) { + const auto * weight = get_weight(ggml_get_name(tensor)); + if (!weight || weight->idx != idx) { + continue; + } + ranges.emplace_back(weight->offs, weight->offs + ggml_nbytes(tensor)); + } + + if (ranges.empty()) { + return; + } + + std::sort(ranges.begin(), ranges.end()); + + // below this a gap is padding, and splitting on it would cost a buffer for nothing + constexpr size_t min_gap = 32ull*1024*1024; + + // never trade one oversized buffer for very many small ones + constexpr size_t max_ranges = 64; + + size_t n = 0; + for (size_t i = 1; i < ranges.size(); ++i) { + if (ranges[i].first > ranges[n].second + min_gap) { + ranges[++n] = ranges[i]; + } else { + ranges[n].second = std::max(ranges[n].second, ranges[i].second); + } + } + ranges.resize(n + 1); + + if (ranges.size() > max_ranges) { + ranges = { { ranges.front().first, ranges.back().second } }; + } +} + void llama_model_loader::unmap_weight(const llama_tensor_weight & w) const { if (!use_mmap) { return; } mappings.at(w.idx)->unmap_fragment(w.offs, w.offs + ggml_nbytes(w.tensor)); @@ -1523,7 +1565,7 @@ bool llama_model_loader::load_all_data( } // When not using mmaped io use async uploads from pinned memory to GPU memory. // First determine if the backend supports the necessary features for async uploads. - auto * buf = bufs.count(0) ? bufs.at(0) : nullptr; + auto * buf = bufs.count(0) && !bufs.at(0).empty() ? bufs.at(0).front() : nullptr; if (!buf) { LLAMA_LOG_DEBUG("%s: no buffer found for async uploads\n", func); return nullptr; @@ -1594,7 +1636,7 @@ bool llama_model_loader::load_all_data( if (upload_backend) { LLAMA_LOG_DEBUG("%s: using async uploads for device %s, buffer type %s, backend %s\n", __func__, ggml_backend_dev_name(ggml_backend_get_device(upload_backend)), - ggml_backend_buft_name(ggml_backend_buffer_get_type(bufs.at(0))), + ggml_backend_buft_name(ggml_backend_buffer_get_type(bufs.at(0).front())), ggml_backend_name(upload_backend)); } @@ -1618,10 +1660,17 @@ bool llama_model_loader::load_all_data( if (from_mapping) { const auto & mapping = mappings.at(weight->idx); ggml_backend_buffer_t buf_mmap = nullptr; + uint8_t * data = (uint8_t *) mapping->addr() + weight->offs; + if (bufs.count(weight->idx)) { - buf_mmap = bufs.at(weight->idx); + for (ggml_backend_buffer_t b : bufs.at(weight->idx)) { + uint8_t * base = (uint8_t *) ggml_backend_buffer_get_base(b); + if (data >= base && data + n_size <= base + ggml_backend_buffer_get_size(b)) { + buf_mmap = b; + break; + } + } } - uint8_t * data = (uint8_t *) mapping->addr() + weight->offs; if (check_tensors) { validation_result.emplace_back(std::async(std::launch::async, [cur, data, n_size] { diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index 9e51d0ce750..6204a227457 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -16,7 +16,7 @@ #include #include -using llama_buf_map = std::unordered_map; +using llama_buf_map = std::unordered_map>; // lists of buffer types used for each layer using buft_list_t = std::vector>; @@ -244,6 +244,10 @@ struct llama_model_loader { void get_mapping_range(size_t * first, size_t * last, void ** addr, int idx, ggml_context * ctx) const; + // near-adjacent runs are merged, so a normal layout yields one range and only a large + // foreign tensor splits it + void get_mapping_ranges(std::vector> & ranges, void ** addr, int idx, ggml_context * ctx) const; + // release a weight's mmap pages void unmap_weight(const llama_tensor_weight & w) const; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index e679b24e87f..6e9192512d7 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1727,19 +1727,26 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { // this is important for metal with apple silicon: if the entire model could be mapped to a metal buffer, // then we could just use metal for all layers // this allows using partial offloading when the model size exceeds the metal buffer size, but not the RAM size + // a tensor of another buffer type can sit between this context's, and one span + // over them would map it too, so map each contiguous run separately void * addr = nullptr; - size_t first, last; // NOLINT - ml.get_mapping_range(&first, &last, &addr, idx, ctx); - if (first >= last) { + std::vector> ranges; + ml.get_mapping_ranges(ranges, &addr, idx, ctx); + if (ranges.empty()) { continue; } const size_t max_size = ggml_get_max_tensor_size(ctx); - ggml_backend_buffer_t buf = ggml_backend_dev_buffer_from_host_ptr(dev, (char *) addr + first, last - first, max_size); - if (buf == nullptr) { - throw std::runtime_error(format("unable to allocate %s buffer", ggml_backend_buft_name(buft))); + for (const auto & [first, last] : ranges) { + if (first >= last) { + continue; + } + ggml_backend_buffer_t buf = ggml_backend_dev_buffer_from_host_ptr(dev, (char *) addr + first, last - first, max_size); + if (buf == nullptr) { + throw std::runtime_error(format("unable to allocate %s buffer", ggml_backend_buft_name(buft))); + } + bufs.emplace_back(buf); + buf_map[idx].push_back(buf); } - bufs.emplace_back(buf); - buf_map.emplace(idx, buf); } } else { ggml_backend_buffer_t buf; @@ -1762,7 +1769,7 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { } bufs.emplace_back(buf); for (uint32_t idx = 0; idx < ml.files.size(); idx++) { - buf_map.emplace(idx, buf); + buf_map[idx].push_back(buf); } }