diff --git a/src/llama-mmap.cpp b/src/llama-mmap.cpp index 4d183cbc9c4..fe17fe3104b 100644 --- a/src/llama-mmap.cpp +++ b/src/llama-mmap.cpp @@ -439,6 +439,56 @@ void llama_file::write_u32(uint32_t val) const { pimpl->write_u32(val); } // llama_mmap #if defined(_POSIX_MAPPED_FILES) || defined(_WIN32) +static size_t llama_mmap_page_size() { +#if defined(_WIN32) + SYSTEM_INFO si; + GetSystemInfo(&si); + return (size_t) si.dwPageSize; +#else + return (size_t) sysconf(_SC_PAGESIZE); +#endif +} + +// the pages the given rows fall on, merged into runs. rows are smaller than a page and +// repeat within a batch, so this turns a hint per row into a hint per page. +static llama_mmap::ranges llama_mmap_row_pages( + size_t base_off, size_t stride, size_t row_size, size_t map_size, + const int32_t * rows, size_t n_rows, size_t page_size) { + std::vector pages; + pages.reserve(n_rows); + + for (size_t i = 0; i < n_rows; ++i) { + if (rows[i] < 0) { + continue; + } + const size_t first = base_off + (size_t) rows[i] * stride; + const size_t last = first + row_size; + // an unexpected index must not turn into a hint outside the mapping + if (row_size == 0 || last > map_size || last < first) { + continue; + } + for (size_t p = first / page_size; p <= (last - 1) / page_size; ++p) { + pages.push_back(p); + } + } + + std::sort(pages.begin(), pages.end()); + pages.erase(std::unique(pages.begin(), pages.end()), pages.end()); + + llama_mmap::ranges res; + for (size_t i = 0; i < pages.size(); ) { + size_t j = i + 1; + while (j < pages.size() && pages[j] == pages[j - 1] + 1) { + ++j; + } + const size_t off = pages[i] * page_size; + res.emplace_back(off, off + std::min((pages[j - 1] - pages[i] + 1) * page_size, map_size - off)); + i = j; + } + + return res; +} + // merge `ranges` and return their complement within [0, limit) static llama_mmap::ranges ranges_complement(llama_mmap::ranges ranges, size_t limit) { llama_mmap::ranges res; @@ -671,6 +721,60 @@ void * llama_mmap::addr() const { return pimpl->addr; } void llama_mmap::unmap_fragment(size_t first, size_t last) { pimpl->unmap_fragment(first, last); } +bool llama_mmap::contains(const void * ptr, size_t len) const { + const char * addr = (const char *) pimpl->addr; + const char * p = (const char *) ptr; + return addr != nullptr && p >= addr && p + len <= addr + pimpl->size; +} + +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 defined(_POSIX_MAPPED_FILES) || defined(_WIN32) + const size_t base_off = (const char *) base - (const char *) pimpl->addr; + const auto ranges = llama_mmap_row_pages(base_off, stride, row_size, pimpl->size, + rows, n_rows, llama_mmap_page_size()); +#endif + +#if defined(_POSIX_MAPPED_FILES) + for (const auto & range : ranges) { + // unchecked: a failed hint only costs the fault it would have avoided + posix_madvise((char *) pimpl->addr + range.first, range.second - range.first, + POSIX_MADV_WILLNEED); + } +#elif defined(_WIN32) + #if _WIN32_WINNT >= 0x602 + // PrefetchVirtualMemory takes all ranges in one call, which is the batching we want + 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; + } + + std::vector entries; + entries.reserve(ranges.size()); + for (const auto & range : ranges) { + WIN32_MEMORY_RANGE_ENTRY e; + e.VirtualAddress = (char *) pimpl->addr + range.first; + e.NumberOfBytes = (SIZE_T) (range.second - range.first); + entries.push_back(e); + } + + if (!entries.empty()) { + // unchecked, same as the POSIX branch + pPrefetchVirtualMemory(GetCurrentProcess(), (ULONG_PTR) entries.size(), entries.data(), 0); + } + #endif +#else + GGML_UNUSED(base); + GGML_UNUSED(stride); + GGML_UNUSED(row_size); + GGML_UNUSED(rows); + GGML_UNUSED(n_rows); +#endif +} + #if defined(_POSIX_MEMLOCK_RANGE) || defined(_WIN32) const bool llama_mmap::SUPPORTED = true; #else diff --git a/src/llama-mmap.h b/src/llama-mmap.h index cc28c8a73fa..4475dc3b26c 100644 --- a/src/llama-mmap.h +++ b/src/llama-mmap.h @@ -55,6 +55,15 @@ struct llama_mmap { void unmap_fragment(size_t first, size_t last); + // true if [ptr, ptr + len) is inside this mapping + bool contains(const void * ptr, size_t len) const; + + // start reading the given rows of a tensor in this mapping, in one batch. + // lazy ranges are MADV_RANDOM, which turns kernel readahead off, so without this a + // sparse gather costs one synchronous fault per row. hints only, never changes results. + void prefetch_rows(const void * base, size_t stride, size_t row_size, + const int32_t * rows, size_t n_rows) const; + static const bool SUPPORTED; private: diff --git a/src/llama-model.cpp b/src/llama-model.cpp index fc83658dd7f..90c9c87a2d5 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2180,6 +2180,19 @@ const ggml_tensor * llama_model::get_tensor(const char * name) const { return it->second; } +void llama_model::prefetch_rows(const ggml_tensor * t, const int32_t * rows, size_t n_rows) const { + if (t == nullptr || t->data == nullptr || n_rows == 0) { + return; + } + + for (const auto & mapping : pimpl->mappings) { + if (mapping->contains(t->data, ggml_nbytes(t))) { + mapping->prefetch_rows(t->data, t->nb[1], ggml_row_size(t->type, t->ne[0]), rows, n_rows); + return; + } + } +} + float llama_model::get_rope_freq_base (const llama_cparams & cparams, int il) const { return hparams.is_swa(il) ? hparams.rope_freq_base_train_swa : cparams.rope_freq_base; } diff --git a/src/llama-model.h b/src/llama-model.h index 38066538ed1..21aeb9a43be 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -744,6 +744,9 @@ struct llama_model { const struct ggml_tensor * get_tensor(const char * name) const; + // queue readahead for rows a gather is about to read. no-op unless the tensor is mmap'd + void prefetch_rows(const struct ggml_tensor * t, const int32_t * rows, size_t n_rows) const; + 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/gemma4.cpp b/src/models/gemma4.cpp index aa518c6df50..627d43abdc9 100644 --- a/src/models/gemma4.cpp +++ b/src/models/gemma4.cpp @@ -408,10 +408,27 @@ llama_model_gemma4::graph::graph(const llama_model & model, const llm_graph_para ggml_build_forward_expand(gf, cur); } +// same as llm_graph_input_embd, but also hints the per-layer table rows this batch reads. +// the table can stay on disk (TENSOR_READ_LAZY), where a gather is one fault per token +class llm_graph_input_embd_per_layer : public llm_graph_input_embd { +public: + llm_graph_input_embd_per_layer(int64_t n_embd, const llama_model & model) : + llm_graph_input_embd(n_embd), model(model) {} + + void set_input(const llama_ubatch * ubatch) override { + if (ubatch->token) { + model.prefetch_rows(model.per_layer_tok_embd, ubatch->token, ubatch->n_tokens); + } + llm_graph_input_embd::set_input(ubatch); + } + + const llama_model & model; +}; + // equivalent to get_per_layer_inputs() in python code // output shape: [n_embd_per_layer, n_layer, n_tokens] ggml_tensor * llama_model_gemma4::graph::build_inp_per_layer() { - auto inp = std::make_unique(n_embd); + auto inp = std::make_unique(n_embd, model); ggml_tensor * inp_per_layer; float tok_embd_scale = sqrtf((float) n_embd_per_layer); diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index acfdd5b5003..5ed4d4957be 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -1039,6 +1039,10 @@ void llm_graph_input_ple::set_input(const llama_ubatch * ubatch) { } } + // the table stays host side and is read by 16 gathers per token, no two on the same page. + // queued here they are in flight before the graph runs + pmodel.prefetch_rows(pmodel.per_layer_tok_embd, idx.data(), idx.size()); + ggml_backend_tensor_set(rows, idx.data(), 0, idx.size()*ggml_element_size(rows)); }