llama: batched readahead for lazily read gather tables

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<size_t> 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<WIN32_MEMORY_RANGE_ENTRY> 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

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:

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;
 }

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;
 

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<llm_graph_input_embd>(n_embd);
+    auto inp = std::make_unique<llm_graph_input_embd_per_layer>(n_embd, model);
 
     ggml_tensor * inp_per_layer;
     float tok_embd_scale = sqrtf((float) n_embd_per_layer);

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));
 }