From 0ac1a7fd35c127a18c0cbc6c57d12e256a6a385b Mon Sep 17 00:00:00 2001 From: Aman Gupta Date: Thu, 29 Oct 2026 19:22:43 +0800 Subject: [PATCH] llama : reuse the tensor staging buffer and skip the direct-io bounce buffer (#31321) --- src/llama-model-loader.cpp | 28 ++++++++++++++++++++++------ 1 file changed, 22 insertions(+), 6 deletions(-) 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> 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> 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>(); + 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))); } }