Skip to content

llama: map each contiguous run of a context's tensors, not one span over all of them #152

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
danielhanchen wants to merge 1 commit into
base: base/upstream-9723942ad
Choose a base branch
from
+74 −14
Open

llama: map each contiguous run of a context's tensors, not one span over all of them #152

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 53 additions & 4 deletions src/llama-model-loader.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<std::pair<size_t, size_t>> & 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));
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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));
}

Expand All @@ -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] {
Expand Down
6 changes: 5 additions & 1 deletion src/llama-model-loader.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
#include <stdexcept>
#include <unordered_map>

using llama_buf_map = std::unordered_map<uint32_t, ggml_backend_buffer_t>;
using llama_buf_map = std::unordered_map<uint32_t, std::vector<ggml_backend_buffer_t>>;

// lists of buffer types used for each layer
using buft_list_t = std::vector<std::pair<ggml_backend_dev_t, ggml_backend_buffer_type_t>>;
Expand Down Expand Up @@ -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
Comment on lines +247 to +248

@chatgpt-codex-connector chatgpt-codex-connector Bot Aug 31, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Keep each comment sentence on one line

Join this declaration comment into one physical line; it splits a single sentence between lines 247-248, and the new explanation in src/llama-model.cpp does the same at lines 1730-1731, contrary to the repository's explicit comment-formatting rule.

AGENTS.md reference: AGENTS.md:L81-L81

Useful? React with 👍 / 👎.

void get_mapping_ranges(std::vector<std::pair<size_t, size_t>> & ranges, void ** addr, int idx, ggml_context * ctx) const;

// release a weight's mmap pages
void unmap_weight(const llama_tensor_weight & w) const;

Expand Down
25 changes: 16 additions & 9 deletions src/llama-model.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<std::pair<size_t, size_t>> 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;
Expand All @@ -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);
}
}

Expand Down
Loading