UPSTREAM PR #21821: llama : add --hugepages for HugeTLB-backed weight loading (Linux)

common/arg.cpp

@@ -574,6 +574,21 @@ static bool common_params_parse_ex(int argc, char ** argv, common_params_context
         throw std::invalid_argument("error: --prompt-cache-all not supported in interactive mode yet\n");
     }
 
+    // --hugepages compatibility checks. PR #1 covers weights on the --mmap path only;
+    // --no-mmap and --direct-io both route around the hugetlb branch, so surface the
+    // conflict at the user rather than silently ignoring the request. Linux-only.
+    if (params.use_hugepages) {
+#ifndef __linux__
+        throw std::invalid_argument("error: --hugepages is Linux only.\n");
+#endif
+        if (!params.use_mmap) {
+            throw std::invalid_argument("error: --hugepages requires --mmap. Coverage for --no-mmap is added in a follow-up PR. Drop --no-mmap to use --hugepages.\n");
+        }
+        if (params.use_direct_io) {
+            throw std::invalid_argument("error: --hugepages and --direct-io are incompatible. Drop one.\n");
+        }
+    }
+
     // handle model and download
     if (!skip_model_download) {
         auto res = common_params_handle_model(params.model, params.hf_token, params.offline);
@@ -2225,6 +2240,15 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.use_mmap = value;
         }
     ).set_env("LLAMA_ARG_MMAP"));
+    add_opt(common_arg(
+        {"--hugepages"},
+        "back model weights with anonymous hugetlb 2 MiB pages (Linux only).\n"
+        "reserve the pool first with e.g. `sysctl -w vm.nr_hugepages=N` — no reboot required.\n"
+        "pinned in RAM, bypasses swap and page cache; primary win is vmemmap reclamation via HVO",
+        [](common_params & params) {
+            params.use_hugepages = true;
+        }
+    ).set_env("LLAMA_ARG_HUGEPAGES"));
     add_opt(common_arg(
         {"-dio", "--direct-io"},
         {"-ndio", "--no-direct-io"},

common/common.cpp

@@ -1421,6 +1421,7 @@ struct llama_model_params common_model_params_to_llama(common_params & params) {
     mparams.tensor_split    = params.tensor_split;
     mparams.use_mmap        = params.use_mmap;
     mparams.use_direct_io   = params.use_direct_io;
+    mparams.use_hugepages   = params.use_hugepages;
     mparams.use_mlock       = params.use_mlock;
     mparams.check_tensors   = params.check_tensors;
     mparams.use_extra_bufts = !params.no_extra_bufts;

common/common.h

@@ -534,6 +534,7 @@ struct common_params {
     bool input_prefix_bos  = false; // prefix BOS to user inputs, preceding input_prefix
     bool use_mmap          = true;  // enable mmap to use filesystem cache
     bool use_direct_io     = false; // read from disk without buffering
+    bool use_hugepages     = false; // back weight mappings with anonymous hugetlb pages (Linux only)
     bool use_mlock         = false; // use mlock to keep model in memory
     bool verbose_prompt    = false; // print prompt tokens before generation
     bool display_prompt    = true;  // print prompt before generation

include/llama.h

@@ -314,6 +314,7 @@ extern "C" {
         bool vocab_only;      // only load the vocabulary, no weights
         bool use_mmap;        // use mmap if possible
         bool use_direct_io;   // use direct io, takes precedence over use_mmap when supported
+        bool use_hugepages;   // back model memory with anonymous hugetlb pages (Linux only)
         bool use_mlock;       // force system to keep model in RAM
         bool check_tensors;   // validate model tensor data
         bool use_extra_bufts; // use extra buffer types (used for weight repacking)

src/llama-mmap.cpp

@@ -40,6 +40,20 @@
 #include <TargetConditionals.h>
 #endif
 
+#ifdef __linux__
+// Older glibc headers may miss these; upstream kernel has had them for years.
+#ifndef MAP_HUGETLB
+#define MAP_HUGETLB 0x40000
+#endif
+#ifndef MAP_HUGE_2MB
+#define MAP_HUGE_2MB (21 << MAP_HUGE_SHIFT)
+#endif
+#endif
+
+// 2 MiB hugepage size, used for both the hugetlb mmap length and the
+// unmap_fragment alignment granularity when the mapping is hugetlb-backed.
+static constexpr size_t LLAMA_HUGE_PAGE_SIZE = 2ull * 1024 * 1024;
+
 // TODO: consider moving to llama-impl.h if needed in more places
 #if defined(_WIN32)
 static std::string llama_format_win_err(DWORD err) {
@@ -434,9 +448,38 @@ struct llama_mmap::impl {
 #ifdef _POSIX_MAPPED_FILES
     std::vector<std::pair<size_t, size_t>> mapped_fragments;
 
-    impl(struct llama_file * file, size_t prefetch, bool numa) {
+    impl(struct llama_file * file, size_t prefetch, bool numa, bool hugetlb) {
         size = file->size();
         int fd = file->file_id();
+#ifdef __linux__
+        if (hugetlb) {
+            // Anonymous hugetlb mapping rounded up to 2 MiB. PROT_WRITE lets
+            // load_all_data pread the file in (downgraded to PROT_READ after);
+            // MAP_POPULATE surfaces pool-exhaustion here, not mid-load.
+            mmap_size = (size + LLAMA_HUGE_PAGE_SIZE - 1) & ~(LLAMA_HUGE_PAGE_SIZE - 1);
+            addr = mmap(NULL, mmap_size, PROT_READ | PROT_WRITE,
+                        MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB | MAP_HUGE_2MB | MAP_POPULATE, -1, 0);
+            if (addr == MAP_FAILED) {
+                int saved = errno;
+                if (saved == ENOMEM) {
+                    size_t need = mmap_size / LLAMA_HUGE_PAGE_SIZE;
+                    long   have = -1;
+                    if (FILE * f = std::fopen("/sys/kernel/mm/hugepages/hugepages-2048kB/free_hugepages", "r")) {
+                        if (std::fscanf(f, "%ld", &have) != 1) { have = -1; }
+                        std::fclose(f);
+                    }
+                    throw std::runtime_error(format("hugetlb mmap failed: need %zu free 2 MiB pages, pool has %ld. "
+                                                    "Try: sudo sysctl -w vm.nr_hugepages=%zu", need, have, need));
+                }
+                throw std::runtime_error(format("hugetlb mmap failed: %s", strerror(saved)));
+            }
+            is_hugetlb_ = true;
+            mapped_fragments.emplace_back(0, mmap_size);
+            return;
+        }
+#else
+        (void) hugetlb;
+#endif
         int flags = MAP_SHARED;
         if (numa) { prefetch = 0; }
 #ifdef __linux__
@@ -480,7 +523,9 @@ struct llama_mmap::impl {
     }
 
     void unmap_fragment(size_t first, size_t last) {
-        int page_size = sysconf(_SC_PAGESIZE);
+        // Hugetlb munmaps must be 2 MiB-aligned; the file-backed path uses
+        // the kernel base page size as before.
+        size_t page_size = is_hugetlb_ ? LLAMA_HUGE_PAGE_SIZE : (size_t) sysconf(_SC_PAGESIZE);
         align_range(&first, &last, page_size);
         size_t len = last - first;
 
@@ -525,8 +570,9 @@ struct llama_mmap::impl {
 #elif defined(_WIN32)
     HANDLE hMapping = nullptr;
 
-    impl(struct llama_file * file, size_t prefetch, bool numa) {
+    impl(struct llama_file * file, size_t prefetch, bool numa, bool hugetlb) {
         GGML_UNUSED(numa);
+        GGML_UNUSED(hugetlb);
 
         size = file->size();
 
@@ -589,10 +635,11 @@ struct llama_mmap::impl {
         }
     }
 #else
-    impl(struct llama_file * file, size_t prefetch, bool numa) {
+    impl(struct llama_file * file, size_t prefetch, bool numa, bool hugetlb) {
         GGML_UNUSED(file);
         GGML_UNUSED(prefetch);
         GGML_UNUSED(numa);
+        GGML_UNUSED(hugetlb);
 
         throw std::runtime_error("mmap not supported");
     }
@@ -605,15 +652,21 @@ struct llama_mmap::impl {
     }
 #endif
 
-    void * addr;
-    size_t size;
+    void * addr = nullptr;
+    size_t size = 0;
+    // Hugetlb: physical mapping length (file size rounded up to 2 MiB) used
+    // by munmap; `size` continues to report the underlying file length.
+    size_t mmap_size = 0;
+    bool   is_hugetlb_ = false;
 };
 
-llama_mmap::llama_mmap(struct llama_file * file, size_t prefetch, bool numa) : pimpl(std::make_unique<impl>(file, prefetch, numa)) {}
+llama_mmap::llama_mmap(struct llama_file * file, size_t prefetch, bool numa, bool hugetlb) : pimpl(std::make_unique<impl>(file, prefetch, numa, hugetlb)) {}
 llama_mmap::~llama_mmap() = default;
 
 size_t llama_mmap::size() const { return pimpl->size; }
+size_t llama_mmap::mmap_size() const { return pimpl->mmap_size ? pimpl->mmap_size : pimpl->size; }
 void * llama_mmap::addr() const { return pimpl->addr; }
+bool   llama_mmap::is_hugetlb() const { return pimpl->is_hugetlb_; }
 
 void llama_mmap::unmap_fragment(size_t first, size_t last) { pimpl->unmap_fragment(first, last); }
 

src/llama-mmap.h

@@ -42,11 +42,13 @@ struct llama_file {
 
 struct llama_mmap {
     llama_mmap(const llama_mmap &) = delete;
-    llama_mmap(struct llama_file * file, size_t prefetch = (size_t) -1, bool numa = false);
+    llama_mmap(struct llama_file * file, size_t prefetch = (size_t) -1, bool numa = false, bool hugetlb = false);
     ~llama_mmap();
 
     size_t size() const;
+    size_t mmap_size() const; // physical mapping length (>= size() when hugetlb-rounded)
     void * addr() const;
+    bool   is_hugetlb() const;
 
     void unmap_fragment(size_t first, size_t last);
 

src/llama-model-loader.cpp

@@ -7,12 +7,17 @@
 
 #include <algorithm>
 #include <array>
+#include <cerrno>
 #include <cinttypes>
 #include <cstdint>
 #include <cstring>
 #include <future>
 #include <regex>
 
+#ifdef __linux__
+#include <sys/mman.h>
+#endif
+
 static const size_t kiB = 1024;
 static const size_t MiB = 1024*kiB;
 static const size_t GiB = 1024*MiB;
@@ -516,6 +521,7 @@ llama_model_loader::llama_model_loader(
         FILE * file,
         bool use_mmap,
         bool use_direct_io,
+        bool use_hugepages,
         bool check_tensors,
         bool no_alloc,
         const llama_model_kv_override * param_overrides_p,
@@ -814,6 +820,7 @@ llama_model_loader::llama_model_loader(
 
     this->use_mmap = use_mmap;
     this->use_direct_io = use_direct_io;
+    this->use_hugepages = use_hugepages;
     this->check_tensors = check_tensors;
     this->no_alloc = no_alloc;
 }
@@ -1339,7 +1346,7 @@ void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps
                 }
             }
 
-            std::unique_ptr<llama_mmap> mapping = std::make_unique<llama_mmap>(file.get(), prefetch ? -1 : 0, is_numa);
+            std::unique_ptr<llama_mmap> mapping = std::make_unique<llama_mmap>(file.get(), prefetch ? -1 : 0, is_numa, use_hugepages);
             mmaps_used.emplace_back(mapping->size(), 0);
             if (mlock_mmaps) {
                 std::unique_ptr<llama_mlock> mlock_mmap(new llama_mlock());
@@ -1534,6 +1541,14 @@ bool llama_model_loader::load_all_data(
             }
             uint8_t * data = (uint8_t *) mapping->addr() + weight->offs;
 
+            // Hugetlb mappings are anonymous and zero-filled; populate the
+            // region per-tensor before check_tensors / view alloc consume it.
+            if (mapping->is_hugetlb()) {
+                const auto & file = files.at(weight->idx);
+                file->seek(weight->offs, SEEK_SET);
+                file->read_raw(data, n_size);
+            }
+
             if (check_tensors) {
                 validation_result.emplace_back(std::async(std::launch::async, [cur, data, n_size] {
                     return std::make_pair(cur, ggml_validate_row_data(cur->type, data, n_size));
@@ -1664,6 +1679,16 @@ bool llama_model_loader::load_all_data(
             for (uint32_t idx = 0; idx < mappings.size(); idx++) {
                 const auto & mmap_used = mmaps_used.at(idx);
                 auto & mapping = mappings.at(idx);
+#ifdef __linux__
+                // Downgrade the hugetlb region to PROT_READ while it is still
+                // fully mapped — unmap_fragment below may drop 2 MiB chunks
+                // from the head/tail, so mprotect must come first.
+                if (mapping->is_hugetlb()) {
+                    if (mprotect(mapping->addr(), mapping->mmap_size(), PROT_READ)) {
+                        LLAMA_LOG_WARN("warning: mprotect(PROT_READ) failed: %s\n", strerror(errno));
+                    }
+                }
+#endif
                 mapping->unmap_fragment(0, mmap_used.first);
                 if (mmap_used.second != 0) {
                     mapping->unmap_fragment(mmap_used.second, mapping->size());

src/llama-model-loader.h

@@ -77,6 +77,7 @@ struct llama_model_loader {
 
     bool use_mmap = false;
     bool use_direct_io = false;
+    bool use_hugepages = false;
     bool check_tensors;
     bool no_alloc;
 
@@ -128,6 +129,7 @@ struct llama_model_loader {
         FILE * file,
         bool use_mmap,
         bool use_direct_io,
+        bool use_hugepages,
         bool check_tensors,
         bool no_alloc,
         const llama_model_kv_override * param_overrides_p,

src/llama-model.cpp

@@ -9302,6 +9302,7 @@ llama_model_params llama_model_default_params() {
         /*.vocab_only                  =*/ false,
         /*.use_mmap                    =*/ true,
         /*.use_direct_io               =*/ false,
+        /*.use_hugepages               =*/ false,
         /*.use_mlock                   =*/ false,
         /*.check_tensors               =*/ false,
         /*.use_extra_bufts             =*/ true,

src/llama-quant.cpp

@@ -879,7 +879,7 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::
     const llama_model_kv_override * kv_overrides = params->kv_overrides;
     std::vector<std::string> splits = {};
     llama_model_loader ml(/*metadata*/ nullptr, /*set_tensor_data*/ nullptr, /*set_tensor_data_ud*/ nullptr,
-        fname_inp, splits, /*file*/ nullptr, use_mmap, /*use_direct_io*/ false, /*check_tensors*/ true, /*no_alloc*/ false, kv_overrides, nullptr);
+        fname_inp, splits, /*file*/ nullptr, use_mmap, /*use_direct_io*/ false, /*use_hugepages*/ false, /*check_tensors*/ true, /*no_alloc*/ false, kv_overrides, nullptr);
     ml.init_mappings(false); // no prefetching
 
     llama_model model(llama_model_default_params());

src/llama.cpp

@@ -845,7 +845,7 @@ static int llama_model_load(struct gguf_context * metadata, llama_model_set_tens
 
     try {
         llama_model_loader ml(metadata, set_tensor_data, set_tensor_data_ud, fname, splits, file, params.use_mmap, params.use_direct_io,
-            params.check_tensors, params.no_alloc, params.kv_overrides, params.tensor_buft_overrides);
+            params.use_hugepages, params.check_tensors, params.no_alloc, params.kv_overrides, params.tensor_buft_overrides);
 
         ml.print_info();