Commit 0ac1a7fd3 for llama.cpp
commit 0ac1a7fd35c127a18c0cbc6c57d12e256a6a385b
Author: Aman Gupta <amangupta052@gmail.com>
Date: Thu Oct 29 19:22:43 2026 +0800
llama : reuse the tensor staging buffer and skip the direct-io bounce buffer (#31321)
diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp
index a71e3017a..a42992673 100644
--- a/src/llama-model-loader.cpp
+++ b/src/llama-model-loader.cpp
@@ -1691,6 +1691,8 @@ bool llama_model_loader::load_all_data(
});
}
+ // staging buffer for non-host buffers without async uploads, reused because a fresh allocation per tensor page faults on every load
+ std::vector<no_init<uint8_t>> read_buf;
for (struct ggml_tensor * cur : tensors) {
const auto * weight = get_weight(ggml_get_name(cur));
if (weight == nullptr) {
@@ -1810,12 +1812,26 @@ bool llama_model_loader::load_all_data(
buffer_idx %= n_buffers;
}
} else {
- // scoped to one tensor so only one staging buffer is alive at a time
- std::vector<no_init<uint8_t>> read_buf(n_size);
- file->seek(weight->offs, SEEK_SET);
- file->read_raw(read_buf.data(), n_size);
- ggml_backend_tensor_set(cur, read_buf.data(), 0, n_size);
- if (check_tensors && !ggml_validate_row_data(cur->type, read_buf.data(), n_size)) {
+ // with direct IO read the aligned range around the tensor and use the data inside it, so no bounce buffer is needed
+ const size_t align = file->read_alignment();
+ const size_t aligned_offset = weight->offs & ~(align - 1);
+ const size_t pad = weight->offs - aligned_offset;
+ const size_t read_size = (pad + n_size + align - 1) & ~(align - 1);
+ const size_t buf_size = read_size + align - 1;
+
+ // tensors come biggest-first: shrink the buffer when it gets too large, so peak memory stays close to one tensor
+ if (read_buf.size() < buf_size || read_buf.size() > 2*buf_size) {
+ // free the old buffer first, `read_buf = {}` would keep its capacity
+ read_buf = std::vector<no_init<uint8_t>>();
+ read_buf.resize(buf_size);
+ }
+ uint8_t * read_dst = (uint8_t *) GGML_PAD((uintptr_t) read_buf.data(), align);
+ file->seek(aligned_offset, SEEK_SET);
+ file->read_raw_unsafe(read_dst, read_size);
+
+ const uint8_t * data = read_dst + pad;
+ ggml_backend_tensor_set(cur, data, 0, n_size);
+ if (check_tensors && !ggml_validate_row_data(cur->type, data, n_size)) {
throw std::runtime_error(format("tensor '%s' has invalid data", ggml_get_name(cur)));
}
}