qwen4exp: direct reads for the lazy PLE table (>2x prefill performance improvement on GB10)

common/arg.cpp

@@ -2737,10 +2737,12 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         {"-lzm", "--lazy-mode"}, "MODE",
         "on-demand reading of certain tensors, for example per-layer embeddings (default: auto)\n"
         "- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)\n"
+        "- on-direct: like on, but the arch reads the rows with explicit pread()s instead of demand paging the mmap (currently only PLE tables)\n"
         "- auto: on, but only for tensors larger than 4 GiB\n"
         "- off: always keep them resident",
         [](common_params & params, const std::string & value) {
             /**/ if (value == "on")   { params.lazy_mode = LLAMA_LAZY_MODE_ON;   }
+            else if (value == "on-direct") { params.lazy_mode = LLAMA_LAZY_MODE_DIRECT; }
             else if (value == "auto") { params.lazy_mode = LLAMA_LAZY_MODE_AUTO; }
             else if (value == "off")  { params.lazy_mode = LLAMA_LAZY_MODE_OFF;  }
             else { throw std::invalid_argument("invalid value"); }

include/llama.h

@@ -218,6 +218,7 @@ extern "C" {
         LLAMA_LAZY_MODE_OFF  = 0, // always read the whole tensor up front
         LLAMA_LAZY_MODE_AUTO = 1, // lazy only for marked tensors larger than 4 GiB (requires mmap)
         LLAMA_LAZY_MODE_ON   = 2, // read the rows of tensors marked by the arch on demand (requires mmap)
+        LLAMA_LAZY_MODE_DIRECT = 3, // like ON, but the arch reads the rows with explicit pread()s instead of demand paging the mmap
     };
 
     enum llama_context_type {

src/llama-lazy-reader.h

@@ -0,0 +1,147 @@
+#pragma once
+
+// Serves rows of a lazy tensor with explicit pread()s instead of demand paging.
+// This works because the row indices of a whole ubatch are known host-side
+// before the graph runs. Hands out F32 rows, like ggml_get_rows does.
+
+#include "ggml.h"
+#include "llama-impl.h"
+
+#include <algorithm>
+#include <cerrno>
+#include <cstdio>
+#include <cstring>
+#include <stdexcept>
+#include <thread>
+#include <utility>
+#include <vector>
+
+#ifndef _WIN32
+#include <fcntl.h>
+#include <unistd.h>
+#endif
+
+struct llama_lazy_reader {
+#ifdef _WIN32
+    // pread()/open() are unavailable on Windows; --lazy-mode on-direct falls
+    // back to the lazy mmap reads there (see llama_model_base::load_lazy_reader)
+    const int64_t head_dim = 0;
+
+    void gather(const int32_t *, int64_t, float *) const {
+        GGML_ABORT("lazy direct reads are not supported on this platform");
+    }
+#else
+    llama_lazy_reader(int fd, size_t base, size_t row_size, int64_t n_rows, int n_threads,
+                      enum ggml_type type, int64_t head_dim)
+        : fd(fd), base(base), row_size(row_size), n_rows(n_rows), n_threads(n_threads),
+          head_dim(head_dim), to_float(type == GGML_TYPE_F32 ? nullptr : ggml_get_type_traits(type)->to_float) {
+        // F32 rows have no dequantizer; they are staged as-is, like ggml_get_rows
+        GGML_ASSERT((type == GGML_TYPE_F32 || to_float != nullptr) && head_dim > 0);
+    }
+
+    llama_lazy_reader(const llama_lazy_reader &) = delete;
+    llama_lazy_reader & operator=(const llama_lazy_reader &) = delete;
+
+    ~llama_lazy_reader() {
+        if (fd >= 0) {
+            ::close(fd);
+        }
+    }
+
+    const int      fd;
+    const size_t   base;       // file offset of row 0
+    const size_t   row_size;   // bytes per quantized row
+    const int64_t  n_rows;
+    const int      n_threads;  // in-flight read workers
+    const int64_t  head_dim;   // F32 elements per staged row
+    ggml_to_float_t to_float;  // same dequantizer the ggml_get_rows CPU kernel uses
+
+    // fill dst with the n gathered rows, dequantized to F32:
+    // dst[slot * head_dim, ...) = to_float(table[rows[slot]])
+    // thread-safe; never lets an exception escape a worker thread
+    void gather(const int32_t * rows, int64_t n, float * dst) const {
+        std::vector<std::pair<int32_t, int32_t>> pairs; // (row, dst slot)
+        pairs.reserve(n);
+        for (int64_t i = 0; i < n; ++i) {
+            GGML_ASSERT(rows[i] >= 0 && (int64_t) rows[i] < n_rows);
+            pairs.emplace_back(rows[i], (int32_t) i);
+        }
+
+        std::sort(pairs.begin(), pairs.end()); // equal rows adjacent, file order
+
+        // small gathers are not worth a thread per row
+        const int n_workers = (int) std::min<int64_t>(n_threads, std::max<int64_t>(1, n / 32));
+
+        // worker w reads rows pairs[n*w/n_workers, n*(w+1)/n_workers)
+        auto run_chunk = [&](int w, std::exception_ptr & err) {
+            try {
+                run_range(pairs, n * w / n_workers, n * (w + 1) / n_workers, dst);
+            } catch (...) {
+                err = std::current_exception();
+            }
+        };
+
+        // an exception leaving a joinable std::thread, or destroying one,
+        // terminates the process; keep worker creation failure-safe
+        std::vector<std::exception_ptr> errs(n_workers);
+        std::vector<std::thread> workers;
+        try {
+            for (int w = 1; w < n_workers; ++w) {
+                workers.emplace_back([&run_chunk, &errs, w]() {
+                    run_chunk(w, errs[w]);
+                });
+            }
+        } catch (...) {
+            for (auto & t : workers) {
+                t.join();
+            }
+            throw;
+        }
+
+        run_chunk(0, errs[0]); // this thread takes the first chunk
+        for (auto & t : workers) {
+            t.join();
+        }
+
+        for (const auto & err : errs) {
+            if (err) {
+                std::rethrow_exception(err);
+            }
+        }
+    }
+
+private:
+    void run_range(const std::vector<std::pair<int32_t, int32_t>> & pairs,
+                   int64_t begin, int64_t end, float * dst) const {
+        std::vector<uint8_t> bounce(row_size);
+        for (int64_t i = begin; i < end; ) {
+            int64_t j = i;
+            while (j + 1 < end && pairs[j + 1].first == pairs[i].first) {
+                ++j; // dedup: one read serves the whole run
+            }
+            const size_t off = base + (size_t) pairs[i].first * row_size;
+            for (size_t done = 0; done < row_size; ) {
+                const ssize_t n_read = ::pread(fd, bounce.data() + done, row_size - done, off + done);
+                if (n_read < 0 && errno == EINTR) {
+                    continue; // interrupted by a signal without SA_RESTART
+                }
+                if (n_read <= 0) {
+                    throw std::runtime_error(format("lazy direct read of %zu bytes at file offset %zu failed: %s",
+                            row_size, off, n_read == 0 ? "unexpected EOF" : strerror(errno)));
+                }
+                done += n_read;
+            }
+            float * first = dst + (size_t) pairs[i].second * head_dim;
+            if (to_float) {
+                to_float(bounce.data(), first, head_dim);
+            } else {
+                memcpy(first, bounce.data(), (size_t) head_dim * sizeof(float));
+            }
+            for (int64_t k = i + 1; k <= j; ++k) {
+                memcpy(dst + (size_t) pairs[k].second * head_dim, first, (size_t) head_dim * sizeof(float));
+            }
+            i = j + 1;
+        }
+    }
+#endif
+};

src/llama-mmap.cpp

@@ -83,7 +83,7 @@ struct llama_file::impl {
         return ret;
     }
 
-    impl(const char * fname, const char * mode, [[maybe_unused]] const bool use_direct_io = false) {
+    impl(const char * fname, const char * mode, [[maybe_unused]] const bool use_direct_io = false) : fname(fname) {
         fp = ggml_fopen(fname, mode);
         if (fp == NULL) {
             throw std::runtime_error(format("failed to open %s: %s", fname, strerror(errno)));
@@ -93,8 +93,7 @@ struct llama_file::impl {
         size = tell();
         seek(0, SEEK_SET);
     }
-
-    impl(FILE * file) : owns_fp(false) {
+    impl(FILE * file) : fname("(file*)"), owns_fp(false) {
         fp = file;
         fp_win32 = (HANDLE) _get_osfhandle(_fileno(fp));
         seek(0, SEEK_END);
@@ -381,9 +380,10 @@ struct llama_file::impl {
         }
     }
     int fd = -1;
-    std::string fname;
 #endif
 
+    std::string fname;
+
     size_t read_alignment() const {
         return alignment;
     }
@@ -408,6 +408,10 @@ size_t llama_file::size() const { return pimpl->size; }
 size_t llama_file::read_alignment() const { return pimpl->read_alignment(); }
 bool llama_file::has_direct_io() const { return pimpl->has_direct_io(); }
 
+const std::string & llama_file::name() const {
+    return pimpl->fname;
+}
+
 int llama_file::file_id() const {
 #ifdef _WIN32
     return _fileno(pimpl->fp);

src/llama-mmap.h

@@ -2,6 +2,7 @@
 
 #include <cstdint>
 #include <memory>
+#include <string>
 #include <utility>
 #include <vector>
 #include <cstdio>
@@ -22,6 +23,8 @@ struct llama_file {
     size_t tell() const;
     size_t size() const;
 
+    const std::string & name() const; // path this file was opened from
+
     int file_id() const; // fileno overload
 
     void seek(size_t offset, int whence) const;

src/llama-model-loader.cpp

@@ -1085,7 +1085,7 @@ bool llama_model_loader::lazy_read::add(const std::string & name, const ggml_ten
 
     // do not lazy-read small tensors, it has significant overhead and is not worth it
     constexpr size_t auto_min_size = 4ull * 1024 * 1024 * 1024;
-    if (mode != LLAMA_LAZY_MODE_ON && ggml_nbytes(t) <= auto_min_size) {
+    if (mode == LLAMA_LAZY_MODE_AUTO && ggml_nbytes(t) <= auto_min_size) {
         return false;
     }
 

src/llama-model.cpp

@@ -28,7 +28,15 @@
 #include <algorithm>
 #include <cassert>
 #include <cfloat>
+#include <cinttypes>
 #include <cstdint>
+#include <thread>
+
+#ifndef _WIN32
+#include <fcntl.h>
+#include <unistd.h>
+#endif
+
 #include <cstring>
 #include <cmath>
 #include <functional>
@@ -1823,6 +1831,65 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
     return true;
 }
 
+#ifndef _WIN32
+const llama_lazy_reader * llama_model_base::load_lazy_reader(llama_model_loader & ml, const char * tensor_name, const ggml_tensor * t) {
+    if (ml.lazy.mode != LLAMA_LAZY_MODE_DIRECT) {
+        return nullptr;
+    }
+
+    if (!t) {
+        return nullptr;
+    }
+
+    if (const auto it = lazy_readers.find(tensor_name); it != lazy_readers.end()) {
+        return it->second.get();
+    }
+
+    const auto * w = ml.get_weight(tensor_name);
+    if (!w) {
+        // e.g. synthesised from metadata, no file rows to read
+        return nullptr;
+    }
+
+    // an independently opened buffered descriptor: dup() would share the
+    // loader's open file description, whose readahead advice and O_DIRECT
+    // flag would fight the small scattered row reads
+    const int fd = ::open(ml.files[w->idx]->name().c_str(), O_RDONLY | O_CLOEXEC);
+    if (fd < 0) {
+        // e.g. a FILE*-backed model has no reopenable path; the tensor is
+        // still lazy, so keep serving it through the mmap reads
+        LLAMA_LOG_WARN("%s: could not open %s for direct reads (%s), using lazy mmap reads\n",
+                __func__, ml.files[w->idx]->name().c_str(), strerror(errno));
+        return nullptr;
+    }
+
+#ifdef __linux__
+    // rows are tiny and scattered, so sequential readahead would be pure waste
+    ::posix_fadvise(fd, 0, 0, POSIX_FADV_RANDOM);
+#endif
+
+    // in-flight reads are IO queue depth, not compute; 2x cores worked well
+    // on NVMe and stays sane on smaller machines
+    const int n_threads = 2 * (int) std::max(1u, std::thread::hardware_concurrency());
+
+    auto reader = std::make_unique<llama_lazy_reader>(fd, w->offs,
+            ggml_row_size(t->type, t->ne[0]), t->ne[1], n_threads, t->type, t->ne[0]);
+
+    LLAMA_LOG_INFO("%s: direct reads enabled for %s: %" PRId64 " rows of %zu bytes at file offset %zu, %d threads\n",
+            __func__, tensor_name, reader->n_rows, reader->row_size, w->offs, n_threads);
+
+    lazy_readers[tensor_name] = std::move(reader);
+    return lazy_readers.at(tensor_name).get();
+}
+#else
+const llama_lazy_reader * llama_model_base::load_lazy_reader(llama_model_loader & ml, const char *, const ggml_tensor *) {
+    if (ml.lazy.mode == LLAMA_LAZY_MODE_DIRECT) {
+        LLAMA_LOG_WARN("%s: --lazy-mode on-direct is not supported on this platform, using lazy mmap reads\n", __func__);
+    }
+    return nullptr;
+}
+#endif
+
 ggml_tensor * llama_model_base::create_tensor(llama_model_loader & ml, const LLM_TN_IMPL & tn, const std::initializer_list<int64_t> & ne, int flags) {
     const buft_list_t * buft_list_layer = tn.bid == -1 ? nullptr : pimpl->dev_layer.at(tn.bid).buft_list;
     return ml.create_tensor(

src/llama-model.h

@@ -5,6 +5,7 @@
 #include "llama-graph.h"
 #include "llama-hparams.h"
 #include "llama-memory.h"
+#include "llama-lazy-reader.h"
 #include "llama-vocab.h"
 
 #include <map>
@@ -816,6 +817,14 @@ struct llama_model_base : public llama_model {
     void load_arch_hparams(llama_model_loader & ml) override = 0;
     void load_arch_tensors(llama_model_loader & ml) override = 0;
     std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override = 0;
+
+    // --lazy-mode on-direct: read the rows of a lazy tensor with explicit
+    // pread()s instead of demand-faulting them in through the mmap. Call once
+    // per lazy tensor after creating it; returns null if the platform cannot
+    // serve direct reads (the tensor then stays on the lazy mmap path).
+    const llama_lazy_reader * load_lazy_reader(llama_model_loader & ml, const char * tensor_name, const ggml_tensor * t);
+
+    std::map<std::string, std::unique_ptr<llama_lazy_reader>> lazy_readers;
 };
 
 const char * llm_type_name(llm_type type);