From de5ccf357a086cfca37eeb1270c72e6efd4c33d5 Mon Sep 17 00:00:00 2001 From: Aman Gupta Date: Tue, 29 Sep 2026 22:31:04 +0800 Subject: [PATCH] avoid exposing llama-mmap in model code, route via llama-impl --- src/llama-impl.cpp | 18 ++++++++++++++++++ src/llama-impl.h | 3 +++ src/llama-mmap.cpp | 32 +++++++++++++++----------------- src/llama-mmap.h | 11 +++++++---- src/models/gemma4.cpp | 2 +- src/models/qwen4exp.cpp | 1 - 6 files changed, 44 insertions(+), 23 deletions(-) diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp index b3a94b946d..d3a62970a5 100644 --- a/src/llama-impl.cpp +++ b/src/llama-impl.cpp @@ -1,5 +1,7 @@ #include "llama-impl.h" +#include "llama-mmap.h" +#include "ggml-backend.h" #include "gguf.h" #include "llama.h" @@ -17,6 +19,22 @@ struct llama_logger_state { static llama_logger_state g_logger_state; +void llama_prefetch_rows(const ggml_tensor * tensor, const int32_t * rows, size_t n_rows) { + if (!tensor || !tensor->data || !tensor->buffer || !ggml_backend_buffer_is_host(tensor->buffer) || n_rows == 0) { + return; + } + + const size_t row_bytes = ggml_row_size(tensor->type, tensor->ne[0]); + const auto * base = (const char *) tensor->data; + std::vector ranges; + ranges.reserve(n_rows); + for (size_t i = 0; i < n_rows; ++i) { + GGML_ASSERT(rows[i] >= 0 && rows[i] < tensor->ne[1]); + ranges.push_back({ base + (size_t) rows[i] * tensor->nb[1], row_bytes }); + } + llama_prefetch(std::move(ranges)); +} + time_meas::time_meas(int64_t & t_acc, bool disable) : t_start_us(disable ? -1 : ggml_time_us()), t_acc(t_acc) {} time_meas::~time_meas() { diff --git a/src/llama-impl.h b/src/llama-impl.h index 4988b06d2c..aa2473d134 100644 --- a/src/llama-impl.h +++ b/src/llama-impl.h @@ -74,6 +74,9 @@ static inline ggml_tensor * llama_mul_mat_hadamard( return res; } +// Prefetch the host pages needed to gather these rows. +void llama_prefetch_rows(const ggml_tensor * tensor, const int32_t * rows, size_t n_rows); + struct time_meas { time_meas(int64_t & t_acc, bool disable = false); ~time_meas(); diff --git a/src/llama-mmap.cpp b/src/llama-mmap.cpp index 19fe482e8d..5b17a633c8 100644 --- a/src/llama-mmap.cpp +++ b/src/llama-mmap.cpp @@ -3,7 +3,6 @@ #include "llama-impl.h" #include "ggml.h" -#include "ggml-backend.h" #include #include @@ -679,9 +678,9 @@ const bool llama_mmap::SUPPORTED = true; const bool llama_mmap::SUPPORTED = false; #endif -void llama_prefetch_rows(const ggml_tensor * tensor, const int32_t * rows, size_t n_rows) { +void llama_prefetch(std::vector ranges) { #if defined(__linux__) || (defined(_WIN32) && _WIN32_WINNT >= 0x602) - if (!tensor || !tensor->data || !tensor->buffer || !ggml_backend_buffer_is_host(tensor->buffer) || n_rows == 0) { + if (ranges.empty()) { return; } @@ -699,10 +698,9 @@ void llama_prefetch_rows(const ggml_tensor * tensor, const int32_t * rows, size_ } const size_t page = (size_t) page_size; - const size_t row_bytes = ggml_row_size(tensor->type, tensor->ne[0]); - const uintptr_t base = (uintptr_t) tensor->data; - std::vector sorted(rows, rows + n_rows); - std::sort(sorted.begin(), sorted.end()); + std::sort(ranges.begin(), ranges.end(), [](const llama_memory_range & a, const llama_memory_range & b) { + return (uintptr_t) a.addr < (uintptr_t) b.addr; + }); uintptr_t begin = 0, end = 0; #if defined(_WIN32) @@ -715,17 +713,19 @@ void llama_prefetch_rows(const ggml_tensor * tensor, const int32_t * rows, size_ #else auto prefetch = [&]() { if (madvise((void *) begin, end - begin, MADV_WILLNEED) != 0) { - LLAMA_LOG_WARN("llama_prefetch_rows: madvise(MADV_WILLNEED) failed: %s\n", strerror(errno)); + LLAMA_LOG_WARN("llama_prefetch: madvise(MADV_WILLNEED) failed: %s\n", strerror(errno)); return false; } return true; }; #endif - for (const int32_t row : sorted) { - GGML_ASSERT(row >= 0 && row < tensor->ne[1]); - const uintptr_t pointer = base + (size_t) row * tensor->nb[1]; + for (const auto & range : ranges) { + if (!range.addr || range.size == 0) { + continue; + } + const uintptr_t pointer = (uintptr_t) range.addr; const uintptr_t first = pointer / page * page; - const uintptr_t last = (pointer + row_bytes + page - 1) / page * page; + const uintptr_t last = (pointer + range.size + page - 1) / page * page; if (end && first > end) { if (!prefetch()) { return; @@ -741,15 +741,13 @@ void llama_prefetch_rows(const ggml_tensor * tensor, const int32_t * rows, size_ prefetch(); } #if defined(_WIN32) - if (!PrefetchVirtualMemory(GetCurrentProcess(), (ULONG_PTR) entries.size(), entries.data(), 0)) { - LLAMA_LOG_WARN("llama_prefetch_rows: PrefetchVirtualMemory failed: %s\n", + if (!entries.empty() && !PrefetchVirtualMemory(GetCurrentProcess(), (ULONG_PTR) entries.size(), entries.data(), 0)) { + LLAMA_LOG_WARN("llama_prefetch: PrefetchVirtualMemory failed: %s\n", llama_format_win_err(GetLastError()).c_str()); } #endif #else - GGML_UNUSED(tensor); - GGML_UNUSED(rows); - GGML_UNUSED(n_rows); + GGML_UNUSED(ranges); #endif } diff --git a/src/llama-mmap.h b/src/llama-mmap.h index 7a9da61e5d..a34daed303 100644 --- a/src/llama-mmap.h +++ b/src/llama-mmap.h @@ -6,8 +6,6 @@ #include #include -struct ggml_tensor; - struct llama_file; struct llama_mmap; struct llama_mlock; @@ -78,7 +76,12 @@ private: std::unique_ptr pimpl; }; -// Prefetch the host pages needed to gather these rows. -void llama_prefetch_rows(const ggml_tensor * tensor, const int32_t * rows, size_t n_rows); +struct llama_memory_range { + const void * addr; + size_t size; +}; + +// Prefetch the host pages covering these memory ranges. +void llama_prefetch(std::vector ranges); size_t llama_path_max(); diff --git a/src/models/gemma4.cpp b/src/models/gemma4.cpp index ed45553c70..73b6ee661e 100644 --- a/src/models/gemma4.cpp +++ b/src/models/gemma4.cpp @@ -1,5 +1,5 @@ #include "models.h" -#include "llama-mmap.h" +#include "llama-impl.h" void llama_model_gemma4::load_arch_hparams(llama_model_loader & ml) { hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index b0df4aca1a..b6f670225c 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -1,6 +1,5 @@ #include "models.h" #include "llama-impl.h" -#include "llama-mmap.h" #include "llama-memory-hybrid-idx.h" #include "llama-memory-recurrent.h"