From 58325573c9c553d97d30a06b291b7b49bef7b3d4 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Thu, 27 Aug 2026 06:06:07 +0000 Subject: [PATCH] llama: narrow the random-access mmap advice to the gather table The advice was applied per mapping: every mapping the model kept got POSIX_MADV_RANDOM plus a whole-file POSIX_FADV_RANDOM, and the eager pull-in was skipped for every file. On qwen4exp that also hit token_embd.weight, which sits 0.33 GiB past the PLE table in the same shard and is read densely, not by sparse gathers. Measured over -c 512 --chunks 60 on IQ1_S it fell to 8.45% resident, against 100% with the feature off. A model now nominates its gather tables (qwen4exp: per_layer_tok_embd) and only those byte ranges are advised. The range is rounded out to whole pages, which on this model takes in 832 bytes before and 192 after. token_embd goes back to 86.55% resident and the PLE table still drops to 4.44%; smaps shows one VM_RAND_READ VMA of exactly the table instead of one over all 27.16 GiB that stays mapped. posix_fadvise is dropped from the narrowed path. POSIX_FADV_RANDOM ignores its offset and length and marks the whole open file, and the FMODE_RANDOM it sets is only read by page_cache_sync_ra() on the read() path, which a fault on a MADV_RANDOM vma never reaches. POSIX_FADV_ DONTNEED does take a range, so the drop mode keeps it. The eager pull-in is now skipped only for the files holding a nominated table, and re-issued as WILLNEED over the rest of such a file, so other shards load exactly as before. prefetch_rows() keys off the tensor being nominated rather than off a mapping-level flag, so the batched readahead lands only where the advice did. -c 512 --chunks 60, cold, IQ1_S, mean of 3, total wall: default 32.50 s whole mapping 30.05 s narrowed 30.35 s PPL 4.2061 in all three. IQ1_S KLD is bit-identical with the feature on and off, including Mean KLD 0.396070 +/- 0.001931 and Same top p 77.325%. tg128 73.65 +/- 0.33 narrowed against 73.49 +/- 0.34 whole. Assisted-by: Claude --- src/llama-mmap.cpp | 141 ++++++++++++++++++++++++++++--------- src/llama-mmap.h | 19 ++--- src/llama-model-loader.cpp | 15 +++- src/llama-model-loader.h | 4 ++ src/llama-model.cpp | 64 ++++++++++++----- src/llama-model.h | 14 +++- src/models/models.h | 8 +++ 7 files changed, 199 insertions(+), 66 deletions(-) diff --git a/src/llama-mmap.cpp b/src/llama-mmap.cpp index 33167c7280..a615eb6ef9 100644 --- a/src/llama-mmap.cpp +++ b/src/llama-mmap.cpp @@ -531,10 +531,6 @@ struct llama_mmap::impl { fd_advise = fd; int flags = MAP_SHARED; if (numa) { prefetch = 0; } - // with the random hints the eager pull-in is pure waste: it reads the whole file to - // populate pages that the gathers will then hit at most a few percent of. the loader - // still gets sequential readahead for the tensors it actually copies out. - if (llama_mmap_random_mode_get() != LLAMA_MMAP_RANDOM_OFF) { prefetch = 0; } #ifdef __linux__ if (posix_fadvise(fd, 0, 0, POSIX_FADV_SEQUENTIAL)) { LLAMA_LOG_WARN("warning: posix_fadvise(.., POSIX_FADV_SEQUENTIAL) failed: %s\n", @@ -565,30 +561,63 @@ struct llama_mmap::impl { // the load path asks for POSIX_FADV_SEQUENTIAL, which is right while the file is being // streamed once into buffers and wrong for whatever stays host-resident afterwards: those - // tensors are read by sparse gathers, where readahead and fault-around buy nothing and cost - // page cache. flipping the advice only after load keeps both halves happy. - void advise_random(bool drop) { + // tensors are read by sparse gathers, where readahead buys nothing and costs page cache. + // flipping the advice only after load, and only over the tensor, keeps everything else on + // the loader's behaviour. + void advise_random_range(size_t offset, size_t len, bool drop) { + if (offset >= size || len == 0) { + return; + } + len = std::min(len, size - offset); + + // madvise rejects an unaligned start and rounds the length up, so round both out. that + // can take in the tail of the tensor before and the head of the one after, one page each + const size_t page = llama_mmap_page_size(); + const size_t first = offset & ~(page - 1); + const size_t last = std::min(size, (offset + len + page - 1) & ~(page - 1)); + #if defined(__linux__) if (drop) { - // hand back what the load pulled in; the hot pages fault back on demand - if (madvise(addr, size, MADV_DONTNEED)) { + // on a shared file map this only tears down our page tables + if (madvise((char *) addr + first, last - first, MADV_DONTNEED)) { LLAMA_LOG_WARN("warning: madvise(.., MADV_DONTNEED) failed: %s\n", strerror(errno)); } - if (fd_advise >= 0 && posix_fadvise(fd_advise, 0, 0, POSIX_FADV_DONTNEED)) { + // and this frees the page cache. it takes the range and spares partial pages, so a + // tensor sharing the first or last page keeps its cache + if (fd_advise >= 0 && posix_fadvise(fd_advise, (off_t) offset, (off_t) len, POSIX_FADV_DONTNEED)) { LLAMA_LOG_WARN("warning: posix_fadvise(.., POSIX_FADV_DONTNEED) failed: %s\n", strerror(errno)); } } - if (fd_advise >= 0 && posix_fadvise(fd_advise, 0, 0, POSIX_FADV_RANDOM)) { - LLAMA_LOG_WARN("warning: posix_fadvise(.., POSIX_FADV_RANDOM) failed: %s\n", strerror(errno)); - } #else GGML_UNUSED(drop); #endif -#if defined(_POSIX_MAPPED_FILES) - if (posix_madvise(addr, size, POSIX_MADV_RANDOM)) { + // no POSIX_FADV_RANDOM to go with this: it ignores the range and marks the whole open + // file, and the FMODE_RANDOM it sets is only read by the read() path, never by a fault + if (posix_madvise((char *) addr + first, last - first, POSIX_MADV_RANDOM)) { LLAMA_LOG_WARN("warning: posix_madvise(.., POSIX_MADV_RANDOM) failed: %s\n", strerror(errno)); } -#endif + } + + void prefetch_except(const std::vector> & skip) { + const size_t page = llama_mmap_page_size(); + + size_t pos = 0; + for (const auto & [off, len] : skip) { + const size_t first = off & ~(page - 1); + if (first > pos) { + prefetch_range(pos, first - pos); + } + pos = std::max(pos, std::min(size, (off + len + page - 1) & ~(page - 1))); + } + if (pos < size) { + prefetch_range(pos, size - pos); + } + } + + void prefetch_range(size_t offset, size_t len) const { + if (posix_madvise((char *) addr + offset, len, POSIX_MADV_WILLNEED)) { + LLAMA_LOG_WARN("warning: posix_madvise(.., POSIX_MADV_WILLNEED) failed: %s\n", strerror(errno)); + } } void prefetch_rows(const void * base, size_t stride, size_t row_size, @@ -683,9 +712,6 @@ struct llama_mmap::impl { throw std::runtime_error(format("CreateFileMappingA failed: %s", llama_format_win_err(error).c_str())); } - // see the POSIX branch: opting in means the eager pull-in is waste - if (llama_mmap_random_mode_get() != LLAMA_MMAP_RANDOM_OFF) { prefetch = 0; } - addr = MapViewOfFile(hMapping, FILE_MAP_READ, 0, 0, 0); DWORD error = GetLastError(); @@ -721,13 +747,55 @@ struct llama_mmap::impl { GGML_UNUSED(last); } - // Windows has no per-mapping "read this randomly" hint. skipping the eager - // PrefetchVirtualMemory in the constructor is what keeps the pages out; there is nothing - // further to say here, and nothing to drop back. - void advise_random(bool drop) { + // Windows has no "read this range randomly" hint. not pulling the range in is what keeps the + // pages out; there is nothing further to say here, and nothing to drop back. + void advise_random_range(size_t offset, size_t len, bool drop) { + GGML_UNUSED(offset); + GGML_UNUSED(len); GGML_UNUSED(drop); } + void prefetch_except(const std::vector> & skip) { +#if _WIN32_WINNT >= 0x602 + BOOL (WINAPI *pPrefetchVirtualMemory) (HANDLE, ULONG_PTR, PWIN32_MEMORY_RANGE_ENTRY, ULONG); + HMODULE hKernel32 = GetModuleHandleW(L"kernel32.dll"); + + pPrefetchVirtualMemory = (decltype(pPrefetchVirtualMemory))(void *) GetProcAddress(hKernel32, "PrefetchVirtualMemory"); + if (!pPrefetchVirtualMemory) { + return; + } + + const size_t page = llama_mmap_page_size(); + + std::vector entries; + size_t pos = 0; + for (const auto & [off, len] : skip) { + const size_t first = off & ~(page - 1); + if (first > pos) { + WIN32_MEMORY_RANGE_ENTRY e; + e.VirtualAddress = (char *) addr + pos; + e.NumberOfBytes = (SIZE_T) (first - pos); + entries.push_back(e); + } + pos = std::max(pos, std::min(size, (off + len + page - 1) & ~(page - 1))); + } + if (pos < size) { + WIN32_MEMORY_RANGE_ENTRY e; + e.VirtualAddress = (char *) addr + pos; + e.NumberOfBytes = (SIZE_T) (size - pos); + entries.push_back(e); + } + + 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()); + } +#else + GGML_UNUSED(skip); + LLAMA_LOG_DEBUG("skipping PrefetchVirtualMemory because _WIN32_WINNT < 0x602\n"); +#endif + } + // PrefetchVirtualMemory takes the whole set of ranges in one call, which is exactly the // batching this wants: the reads are issued together instead of one fault at a time. void prefetch_rows(const void * base, size_t stride, size_t row_size, @@ -795,12 +863,20 @@ struct llama_mmap::impl { throw std::runtime_error("mmap not supported"); } - void advise_random(bool drop) { + void advise_random_range(size_t offset, size_t len, bool drop) { + GGML_UNUSED(offset); + GGML_UNUSED(len); GGML_UNUSED(drop); throw std::runtime_error("mmap not supported"); } + void prefetch_except(const std::vector> & skip) { + GGML_UNUSED(skip); + + throw std::runtime_error("mmap not supported"); + } + void prefetch_rows(const void * base, size_t stride, size_t row_size, const int32_t * rows, size_t n_rows) const { GGML_UNUSED(base); @@ -824,8 +900,7 @@ struct llama_mmap::impl { size_t size; // the fd is kept only to re-advise the file; the mapping owns no reference to it - int fd_advise = -1; - bool random = false; + int fd_advise = -1; }; llama_mmap::llama_mmap(struct llama_file * file, size_t prefetch, bool numa) : pimpl(std::make_unique(file, prefetch, numa)) {} @@ -836,22 +911,18 @@ void * llama_mmap::addr() const { return pimpl->addr; } void llama_mmap::unmap_fragment(size_t first, size_t last) { pimpl->unmap_fragment(first, last); } -void llama_mmap::advise_random(bool drop) { - pimpl->advise_random(drop); - // set here rather than per platform: the flag means "the caller opted this mapping in", and - // the batched prefetch is worth doing even where the advice itself is a no-op - pimpl->random = true; +void llama_mmap::advise_random_range(size_t offset, size_t len, bool drop) { + pimpl->advise_random_range(offset, len, drop); } -bool llama_mmap::is_random() const { return pimpl->random; } +void llama_mmap::prefetch_except(const std::vector> & skip) { + pimpl->prefetch_except(skip); +} bool llama_mmap::contains(const void * ptr, size_t len) const { return pimpl->contains(ptr, len); } void llama_mmap::prefetch_rows(const void * base, size_t stride, size_t row_size, const int32_t * rows, size_t n_rows) const { - if (!pimpl->random) { - return; - } pimpl->prefetch_rows(base, stride, row_size, rows, n_rows); } diff --git a/src/llama-mmap.h b/src/llama-mmap.h index ab4ef60408..07a1f6f685 100644 --- a/src/llama-mmap.h +++ b/src/llama-mmap.h @@ -3,6 +3,7 @@ #include #include #include +#include #include struct llama_file; @@ -50,19 +51,21 @@ struct llama_mmap { void unmap_fragment(size_t first, size_t last); - // opt-in, see llama_mmap_random_mode(). marks the whole mapping as randomly accessed, which - // is only correct once loading is done: until then the loader streams the file sequentially. - void advise_random(bool drop); + // opt-in, see llama_mmap_random_mode(). marks one byte range as randomly accessed, which is + // only correct once loading is done: until then the loader streams the file sequentially. + // offsets are into the file, which is also the offset into the mapping - the whole file is + // always mapped from zero. the range is rounded out to whole pages, since madvise needs that. + void advise_random_range(size_t offset, size_t len, bool drop); - // true once advise_random() has been applied. everything that keys off "this mapping is read - // randomly" tests this, so nothing changes for a mapping the user did not opt in for. - bool is_random() const; + // eager pull-in for everything outside the given ranges, used in place of the constructor's + // whole-file one when part of the file must not be read ahead. ranges must be sorted. + void prefetch_except(const std::vector> & skip); // true if [ptr, ptr + len) lies inside this mapping bool contains(const void * ptr, size_t len) const; // ask the kernel to start reading the given rows. issued as one batch so the faults overlap - // instead of serializing; a no-op unless the mapping is marked random. + // instead of serializing. void prefetch_rows(const void * base, size_t stride, size_t row_size, const int32_t * rows, size_t n_rows) const; @@ -78,7 +81,7 @@ private: // host-resident tensors are read sequentially, so this cannot be a default. enum llama_mmap_random_mode { LLAMA_MMAP_RANDOM_OFF = 0, // upstream behaviour - LLAMA_MMAP_RANDOM_ON = 1, // skip MAP_POPULATE/WILLNEED, advise random after load + LLAMA_MMAP_RANDOM_ON = 1, // advise the gather tables random after load, do not pull them in LLAMA_MMAP_RANDOM_DROP = 2, // additionally drop what the load pulled in }; diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 5071556a1f..0654895283 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1354,7 +1354,8 @@ 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 (size_t i = 0; i < files.size(); ++i) { + const auto & file = files[i]; bool is_numa = false; auto * dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU); @@ -1366,7 +1367,17 @@ void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps } } - std::unique_ptr mapping = std::make_unique(file.get(), prefetch ? -1 : 0, is_numa); + const auto no_prefetch = mmap_no_prefetch.find((uint16_t) i); + + // the eager pull-in would read a gather table in full to populate pages the gathers + // hit a few percent of. skip it for this file and ask for everything else instead, + // so the tensors that really are streamed once keep the readahead they had. + const bool split_prefetch = prefetch && !is_numa && no_prefetch != mmap_no_prefetch.end(); + + std::unique_ptr mapping = std::make_unique(file.get(), prefetch && !split_prefetch ? -1 : 0, is_numa); + if (split_prefetch) { + mapping->prefetch_except(no_prefetch->second); + } mmaps_used.emplace_back(mapping->size(), 0); if (mlock_mmaps) { std::unique_ptr mlock_mmap(new llama_mlock()); diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index e9fe3592d4..96c9d338b1 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -88,6 +88,10 @@ struct llama_model_loader { llama_mmaps mappings; + // byte ranges, per source file, that init_mappings() must not pull in eagerly: gather tables + // the model reads a few percent of. set under LLAMA_MMAP_RANDOM only, sorted by offset. + std::map>> mmap_no_prefetch; + std::map weights_map; std::unordered_map kv_overrides; const llama_model_tensor_buft_override * tensor_buft_overrides; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 925e2aca61..b4a26830b9 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1149,9 +1149,16 @@ struct llama_model::impl { // model memory mapped files llama_mmaps mappings; - // set once loading is done, if any mapping was advised random. lets prefetch_rows() bail out - // without walking the mappings, which is the only cost the feature has when it is off. - bool mappings_random = false; + // gather tables that really came out of a mapping, resolved from gather_tables() during load. + // empty unless the user opted in, which is the only cost the feature has when it is off. + struct gather_range { + const ggml_tensor * tensor; + uint16_t idx; // source file, and so the mapping + size_t offs; // byte offset into that file + size_t len; + }; + + std::vector gather_ranges; // objects representing data potentially being locked in memory llama_mlocks mlock_bufs; @@ -1683,6 +1690,22 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { } } + // kept local until the mappings exist: pimpl->gather_ranges must only ever hold ranges that + // were checked against a live mapping, since everything downstream indexes one + std::vector nominated; + if (llama_mmap_random_mode_get() != LLAMA_MMAP_RANDOM_OFF) { + for (const ggml_tensor * t : gather_tables()) { + const auto * w = t ? ml.get_weight(ggml_get_name(t)) : nullptr; + if (w) { + nominated.push_back({ t, w->idx, w->offs, ggml_nbytes(w->tensor) }); + ml.mmap_no_prefetch[w->idx].emplace_back(w->offs, ggml_nbytes(w->tensor)); + } + } + for (auto & [_, ranges] : ml.mmap_no_prefetch) { + std::sort(ranges.begin(), ranges.end()); + } + } + ml.init_mappings(true, use_mlock ? &pimpl->mlock_mmaps : nullptr); pimpl->mappings.reserve(ml.mappings.size()); @@ -1817,17 +1840,23 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { pimpl->mappings.emplace_back(std::move(mapping)); } - // only now that every tensor has been read is it safe to say the file is read randomly: + // only now that every tensor has been read is it safe to say a range is read randomly: // the load itself is a sequential pass and wants the readahead it has been getting. const llama_mmap_random_mode random_mode = llama_mmap_random_mode_get(); - if (random_mode != LLAMA_MMAP_RANDOM_OFF) { - for (auto & mapping : pimpl->mappings) { - mapping->advise_random(random_mode == LLAMA_MMAP_RANDOM_DROP); - } - pimpl->mappings_random = !pimpl->mappings.empty(); - LLAMA_LOG_INFO("%s: LLAMA_MMAP_RANDOM: advised %zu mapping(s) for random access%s\n", - __func__, pimpl->mappings.size(), + // a nominated tensor that did not end up served from its mapping was offloaded or copied + // into a buffer, and nothing will gather out of the file. drop it rather than advise it + for (const auto & r : nominated) { + if (r.idx < pimpl->mappings.size() && pimpl->mappings[r.idx]->contains(r.tensor->data, r.len)) { + pimpl->gather_ranges.push_back(r); + } + } + + for (const auto & r : pimpl->gather_ranges) { + pimpl->mappings[r.idx]->advise_random_range(r.offs, r.len, random_mode == LLAMA_MMAP_RANDOM_DROP); + + LLAMA_LOG_INFO("%s: LLAMA_MMAP_RANDOM: %s advised for random access, %.2f MiB%s\n", + __func__, ggml_get_name(r.tensor), r.len / 1024.0 / 1024.0, random_mode == LLAMA_MMAP_RANDOM_DROP ? ", dropped cached pages" : ""); } } @@ -1836,19 +1865,18 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { } void llama_model::prefetch_rows(const struct ggml_tensor * t, const int32_t * rows, size_t n_rows) const { - if (!pimpl->mappings_random || t == nullptr || t->data == nullptr || n_rows == 0) { + if (pimpl->gather_ranges.empty() || t == nullptr || t->data == nullptr || n_rows == 0) { return; } if (!llama_mmap_random_prefetch_enabled()) { return; } - // rows are addressed off the tensor, so the whole tensor has to sit in the mapping we find - const size_t nbytes = ggml_nbytes(t); - - for (const auto & mapping : pimpl->mappings) { - if (mapping->is_random() && mapping->contains(t->data, nbytes)) { - mapping->prefetch_rows(t->data, t->nb[1], ggml_row_size(t->type, t->ne[0]), rows, n_rows); + // keyed off the tensor, not off its mapping: the readahead must land where the advice did, + // and the mapping now holds ranges that still want the kernel's own readahead + for (const auto & r : pimpl->gather_ranges) { + if (r.tensor == t) { + pimpl->mappings[r.idx]->prefetch_rows(t->data, t->nb[1], ggml_row_size(t->type, t->ne[0]), rows, n_rows); return; } } diff --git a/src/llama-model.h b/src/llama-model.h index cf7c75a476..e101de2994 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -737,11 +737,19 @@ struct llama_model { // ask the kernel to start reading the rows a gather is about to take out of a host-mapped // tensor, so the faults overlap instead of serializing one NVMe latency at a time. // - // does nothing unless the tensor lives in a mapping that was advised random, which only - // happens under LLAMA_MMAP_RANDOM. off, and for anything not read out of a mapping (offloaded - // tensors, --load-mode none, non-POSIX hosts), this is one bool test. + // does nothing unless the tensor was nominated by gather_tables() and really is read out of + // a mapping. off, and for anything else (offloaded tensors, --load-mode none, non-POSIX + // hosts), this is one empty-vector test. void prefetch_rows(const struct ggml_tensor * t, const int32_t * rows, size_t n_rows) const; + // tensors that stay host-resident and are read by sparse row gathers rather than streamed + // once. under LLAMA_MMAP_RANDOM these get the random-access advice and the batched readahead + // of prefetch_rows(); every other tensor keeps the loader's sequential behaviour. + // + // nominated by the model, not guessed from size: a big host-resident tensor read in full, + // such as token_embd on a CPU-only run, wants the readahead this takes away. + virtual std::vector gather_tables() const { return {}; } + float get_rope_freq_base (const llama_cparams & cparams, int il) const; float get_rope_freq_scale(const llama_cparams & cparams, int il) const; diff --git a/src/models/models.h b/src/models/models.h index dacca61600..dc8a279c12 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2281,6 +2281,14 @@ struct llama_model_qwen4exp : public llama_model_base { void load_arch_hparams(llama_model_loader & ml) override; void load_arch_tensors(llama_model_loader & ml) override; + // the PLE n-gram table is far too big to offload and is read by 16 tiny gathers per token + std::vector gather_tables() const override { + if (per_layer_tok_embd == nullptr) { + return {}; + } + return { per_layer_tok_embd }; + } + struct graph : public llm_build_delta_net_base { graph(const llama_model & model, const llm_graph_params & params); private: