Public/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-19 17:24:57 +02:00
Code Issues Packages Projects Releases Wiki Activity

llama: model_loader: add TENSOR_GET_ROW_LAZY

Browse Source
This commit is contained in:
Xuan Son Nguyen
2026-08-27 12:54:09 +02:00
parent f29551215b
commit e072597f9f
7 changed files with 83 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-mmap.cpp
+59 -13
View File
@@ -438,11 +438,34 @@ void llama_file::write_u32(uint32_t val) const { pimpl->write_u32(val); }
// llama_mmap
#if defined(_POSIX_MAPPED_FILES) || defined(_WIN32)
// merge `ranges` and return their complement within [0, limit)
static std::vector<std::pair<size_t, size_t>> ranges_complement(std::vector<std::pair<size_t, size_t>> ranges, size_t limit) {
std::vector<std::pair<size_t, size_t>> res;
std::sort(ranges.begin(), ranges.end());
size_t pos = 0;
for (const auto & range : ranges) {
const size_t beg = std::min(range.first, limit);
const size_t end = std::min(range.second, limit);
if (beg > pos) {
res.emplace_back(pos, beg);
}
pos = std::max(pos, end);
}
if (pos < limit) {
res.emplace_back(pos, limit);
}
return res;
}
#endif
struct llama_mmap::impl {
#ifdef _POSIX_MAPPED_FILES
std::vector<std::pair<size_t, size_t>> mapped_fragments;
impl(struct llama_file * file, size_t prefetch, bool numa) {
impl(struct llama_file * file, size_t prefetch, bool numa, const std::vector<std::pair<size_t, size_t>> & lazy_ranges) {
size = file->size();
int fd = file->file_id();
int flags = MAP_SHARED;
@@ -452,18 +475,34 @@ struct llama_mmap::impl {
LLAMA_LOG_WARN("warning: posix_fadvise(.., POSIX_FADV_SEQUENTIAL) failed: %s\n",
strerror(errno));
}
if (prefetch) { flags |= MAP_POPULATE; }
// MAP_POPULATE would fault in the lazy ranges too
if (prefetch && lazy_ranges.empty()) { flags |= MAP_POPULATE; }
#endif
addr = mmap(NULL, file->size(), PROT_READ, flags, fd, 0);
if (addr == MAP_FAILED) {
throw std::runtime_error(format("mmap failed: %s", strerror(errno)));
}
if (prefetch > 0) {
if (posix_madvise(addr, std::min(file->size(), prefetch), POSIX_MADV_WILLNEED)) {
LLAMA_LOG_WARN("warning: posix_madvise(.., POSIX_MADV_WILLNEED) failed: %s\n",
strerror(errno));
// page-aligned madvise over [beg, end), clamped to the file
auto advise = [&](size_t beg, size_t end, int advice, const char * name) {
const size_t page_size = sysconf(_SC_PAGESIZE);
beg = beg & ~(page_size - 1);
end = std::min((end + page_size - 1) & ~(page_size - 1), file->size());
if (beg >= end) {
return;
}
if (posix_madvise((char *) addr + beg, end - beg, advice)) {
LLAMA_LOG_WARN("warning: posix_madvise(.., %s) failed: %s\n", name, strerror(errno));
}
};
if (prefetch > 0) {
for (const auto & range : ranges_complement(lazy_ranges, std::min(file->size(), prefetch))) {
advise(range.first, range.second, POSIX_MADV_WILLNEED, "POSIX_MADV_WILLNEED");
}
}
for (const auto & range : lazy_ranges) {
advise(range.first, range.second, POSIX_MADV_RANDOM, "POSIX_MADV_RANDOM");
}
if (numa) {
if (posix_madvise(addr, file->size(), POSIX_MADV_RANDOM)) {
@@ -533,7 +572,7 @@ struct llama_mmap::impl {
#elif defined(_WIN32)
HANDLE hMapping = nullptr;
impl(struct llama_file * file, size_t prefetch, bool numa) {
impl(struct llama_file * file, size_t prefetch, bool numa, const std::vector<std::pair<size_t, size_t>> & lazy_ranges) {
GGML_UNUSED(numa);
size = file->size();
@@ -563,10 +602,15 @@ struct llama_mmap::impl {
pPrefetchVirtualMemory = (decltype(pPrefetchVirtualMemory))(void *) GetProcAddress(hKernel32, "PrefetchVirtualMemory");
if (pPrefetchVirtualMemory) {
WIN32_MEMORY_RANGE_ENTRY range;
range.VirtualAddress = addr;
range.NumberOfBytes = (SIZE_T) std::min(size, prefetch);
if (!pPrefetchVirtualMemory(GetCurrentProcess(), 1, &range, 0)) {
std::vector<WIN32_MEMORY_RANGE_ENTRY> entries;
for (const auto & range : ranges_complement(lazy_ranges, std::min(size, prefetch))) {
WIN32_MEMORY_RANGE_ENTRY entry;
entry.VirtualAddress = (char *) addr + range.first;
entry.NumberOfBytes = (SIZE_T) (range.second - range.first);
entries.push_back(entry);
}
if (!entries.empty() &&
!pPrefetchVirtualMemory(GetCurrentProcess(), (ULONG_PTR) entries.size(), entries.data(), 0)) {
LLAMA_LOG_WARN("warning: PrefetchVirtualMemory failed: %s\n",
llama_format_win_err(GetLastError()).c_str());
}
@@ -597,10 +641,11 @@ struct llama_mmap::impl {
}
}
#else
impl(struct llama_file * file, size_t prefetch, bool numa) {
impl(struct llama_file * file, size_t prefetch, bool numa, const std::vector<std::pair<size_t, size_t>> & lazy_ranges) {
GGML_UNUSED(file);
GGML_UNUSED(prefetch);
GGML_UNUSED(numa);
GGML_UNUSED(lazy_ranges);
throw std::runtime_error("mmap not supported");
}
@@ -617,7 +662,8 @@ struct llama_mmap::impl {
size_t size;
};
llama_mmap::llama_mmap(struct llama_file * file, size_t prefetch, bool numa) : pimpl(std::make_unique<impl>(file, prefetch, numa)) {}
llama_mmap::llama_mmap(struct llama_file * file, size_t prefetch, bool numa,
const std::vector<std::pair<size_t, size_t>> & lazy_ranges) : pimpl(std::make_unique<impl>(file, prefetch, numa, lazy_ranges)) {}
llama_mmap::~llama_mmap() = default;
size_t llama_mmap::size() const { return pimpl->size; }
src/llama-mmap.h
+3 -1
View File
@@ -2,6 +2,7 @@
#include <cstdint>
#include <memory>
#include <utility>
#include <vector>
#include <cstdio>
@@ -42,7 +43,8 @@ private:
struct llama_mmap {
llama_mmap(const llama_mmap &) = delete;
llama_mmap(struct llama_file * file, size_t prefetch = (size_t) -1, bool numa = false);
llama_mmap(struct llama_file * file, size_t prefetch = (size_t) -1, bool numa = false,
const std::vector<std::pair<size_t, size_t>> & lazy_ranges = {});
~llama_mmap();
size_t size() const;
src/llama-model-loader.cpp
+13 -2
View File
@@ -1282,6 +1282,11 @@ struct ggml_tensor * llama_model_loader::create_tensor(
return NULL;
}
if ((flags & TENSOR_GET_ROW_LAZY) && use_mmap) {
const auto & w = require_weight(tn.str().c_str());
lazy_tensor_ranges[w.idx].emplace_back(w.offs, w.offs + ggml_nbytes(cur));
}
ggml_tensor t_meta = *cur;
if (flags & TENSOR_ALLOW_RESHAPE) {
for (size_t dim = 0; dim < GGML_MAX_DIMS; dim++) {
@@ -1349,7 +1354,9 @@ void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps
if (use_mmap) {
mappings.reserve(files.size());
mmaps_used.reserve(files.size());
for (const auto & file : files) {
for (uint32_t idx = 0; idx < files.size(); idx++) {
const auto & file = files[idx];
bool is_numa = false;
auto * dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
@@ -1361,7 +1368,11 @@ void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps
}
}
std::unique_ptr<llama_mmap> mapping = std::make_unique<llama_mmap>(file.get(), prefetch ? -1 : 0, is_numa);
const auto it_lazy = lazy_tensor_ranges.find(idx);
static const std::vector<std::pair<size_t, size_t>> no_lazy_ranges;
std::unique_ptr<llama_mmap> mapping = std::make_unique<llama_mmap>(file.get(), prefetch ? -1 : 0, is_numa,
it_lazy != lazy_tensor_ranges.end() ? it_lazy->second : no_lazy_ranges);
mmaps_used.emplace_back(mapping->size(), 0);
if (mlock_mmaps) {
std::unique_ptr<llama_mlock> mlock_mmap(new llama_mlock());
src/llama-model-loader.h
+4
View File
@@ -68,6 +68,7 @@ struct llama_model_loader {
static const int TENSOR_SKIP = 1 << 2;
static const int TENSOR_SKIP_IF_VIRTUAL = 1 << 3;
static const int TENSOR_ALLOW_RESHAPE = 1 << 4;
static const int TENSOR_GET_ROW_LAZY = 1 << 5; // read rows on demand instead of loading whole tensor; requires mmap for now
int n_kv = 0;
int n_tensors = 0;
@@ -88,6 +89,9 @@ struct llama_model_loader {
llama_mmaps mappings;
// byte ranges of TENSOR_GET_ROW_LAZY tensors, per file index
std::map<uint32_t, std::vector<std::pair<size_t, size_t>>> lazy_tensor_ranges;
std::map<std::string, llama_tensor_weight, weight_name_comparer> weights_map;
std::unordered_map<std::string, llama_model_kv_override> kv_overrides;
const llama_model_tensor_buft_override * tensor_buft_overrides;
src/llama-model.cpp
+2 -1
View File
@@ -3067,7 +3067,8 @@ llama_model_base::llama_model_base(const struct llama_model_params & params) : l
TENSOR_NOT_REQUIRED (llama_model_loader::TENSOR_NOT_REQUIRED),
TENSOR_SKIP (llama_model_loader::TENSOR_SKIP),
TENSOR_SKIP_IF_VIRTUAL(llama_model_loader::TENSOR_SKIP_IF_VIRTUAL),
TENSOR_ALLOW_RESHAPE (llama_model_loader::TENSOR_ALLOW_RESHAPE) {}
TENSOR_ALLOW_RESHAPE (llama_model_loader::TENSOR_ALLOW_RESHAPE),
TENSOR_GET_ROW_LAZY (llama_model_loader::TENSOR_GET_ROW_LAZY) {}
ggml_tensor * llama_model_base::create_tensor(const LLM_TN_IMPL & tn, const std::initializer_list<int64_t> & ne, int flags) {
GGML_ASSERT(ml != nullptr);
src/llama-model.h
+1
View File
@@ -756,6 +756,7 @@ struct llama_model_base : public llama_model {
const int TENSOR_SKIP;
const int TENSOR_SKIP_IF_VIRTUAL;
const int TENSOR_ALLOW_RESHAPE;
const int TENSOR_GET_ROW_LAZY;
explicit llama_model_base(const llama_model_params & params);
virtual ~llama_model_base() = default;
src/models/gemma4.cpp
+1 -1
View File
@@ -50,7 +50,7 @@ void llama_model_gemma4::load_arch_tensors(llama_model_loader &) {
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
if (n_embd_per_layer > 0) {
per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), {n_embd_per_layer * n_layer, n_vocab}, 0);
per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), {n_embd_per_layer * n_layer, n_vocab}, TENSOR_GET_ROW_LAZY);
per_layer_model_proj = create_tensor(tn(LLM_TENSOR_PER_LAYER_MODEL_PROJ, "weight", 0), {n_embd, n_embd_per_layer * n_layer}, 0);
per_layer_proj_norm = create_tensor(tn(LLM_TENSOR_PER_LAYER_PROJ_NORM, "weight", 0), {n_embd_per_layer}, 0);
}