From b44aac9b11b9dd8fe66a250f38cffc04949e3e3f Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 08:44:37 -0700 Subject: [PATCH 1/8] llama : add bounded routed expert store Adapt the positional I/O and cache reservation design from ggml-org/llama.cpp#25294, with lease-based in-flight safety informed by ggml-org/llama.cpp#27861. Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/CMakeLists.txt | 1 + src/llama-expert-store.cpp | 632 ++++++++++++++++++++++++++++++++++++ src/llama-expert-store.h | 103 ++++++ src/llama-mmap.cpp | 20 +- src/llama-mmap.h | 2 +- src/llama-model-loader.cpp | 58 +++- src/llama-model-loader.h | 39 +++ src/llama-model.cpp | 25 +- tests/CMakeLists.txt | 1 + tests/test-expert-store.cpp | 503 ++++++++++++++++++++++++++++ 10 files changed, 1365 insertions(+), 19 deletions(-) create mode 100644 src/llama-expert-store.cpp create mode 100644 src/llama-expert-store.h create mode 100644 tests/test-expert-store.cpp diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 255e8fae1efb..10462e314aef 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -17,6 +17,7 @@ add_library(llama llama-chat.cpp llama-context.cpp llama-cparams.cpp + llama-expert-store.cpp llama-grammar.cpp llama-graph.cpp llama-hparams.cpp diff --git a/src/llama-expert-store.cpp b/src/llama-expert-store.cpp new file mode 100644 index 000000000000..7be7b6144ff6 --- /dev/null +++ b/src/llama-expert-store.cpp @@ -0,0 +1,632 @@ +#include "llama-expert-store.h" + +#include "llama-impl.h" +#include "llama-mmap.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if !defined(_WIN32) +#include +#endif + +// The positional I/O and reservation model is adapted from ggml-org/llama.cpp#25294. +// Leases add the in-flight publication safety described in ggml-org/llama.cpp#27861. + +namespace { + +bool checked_add_u64(uint64_t a, uint64_t b, uint64_t * result) { + if (b > std::numeric_limits::max() - a) { + return false; + } + *result = a + b; + return true; +} + +bool checked_mul_u64(uint64_t a, uint64_t b, uint64_t * result) { + if (a != 0 && b > std::numeric_limits::max() / a) { + return false; + } + *result = a * b; + return true; +} + +bool is_power_of_two(size_t value) { + return value != 0 && (value & (value - 1)) == 0; +} + +struct expert_key { + int32_t layer; + llama_expert_projection projection; + int32_t expert_id; + + bool operator<(const expert_key & other) const { + if (layer != other.layer) { + return layer < other.layer; + } + if (projection != other.projection) { + return projection < other.projection; + } + return expert_id < other.expert_id; + } + + bool operator==(const expert_key & other) const { + return layer == other.layer && projection == other.projection && expert_id == other.expert_id; + } +}; + +struct tensor_key { + int32_t layer; + llama_expert_projection projection; + + bool operator<(const tensor_key & other) const { + if (layer != other.layer) { + return layer < other.layer; + } + return projection < other.projection; + } +}; + +struct aligned_buffer { + uint8_t * data = nullptr; + size_t size = 0; + + aligned_buffer() = default; + + aligned_buffer(size_t size, size_t alignment) { + reset(size, alignment); + } + + aligned_buffer(aligned_buffer && other) noexcept : data(other.data), size(other.size) { + other.data = nullptr; + other.size = 0; + } + + aligned_buffer & operator=(aligned_buffer && other) noexcept { + if (this != &other) { + clear(); + data = other.data; + size = other.size; + other.data = nullptr; + other.size = 0; + } + return *this; + } + + ~aligned_buffer() { + clear(); + } + + aligned_buffer(const aligned_buffer &) = delete; + aligned_buffer & operator=(const aligned_buffer &) = delete; + + void reset(size_t new_size, size_t alignment) { + clear(); + if (new_size == 0) { + return; + } + alignment = std::max(alignment, alignof(void *)); +#if defined(_WIN32) + data = static_cast(_aligned_malloc(new_size, alignment)); + if (data == nullptr) { + throw std::bad_alloc(); + } +#else + void * ptr = nullptr; + if (posix_memalign(&ptr, alignment, new_size) != 0) { + throw std::bad_alloc(); + } + data = static_cast(ptr); +#endif + size = new_size; + } + + void clear() { +#if defined(_WIN32) + _aligned_free(data); +#else + free(data); +#endif + data = nullptr; + size = 0; + } +}; + +struct expert_file { + std::string fname; + uint64_t size = 0; + bool direct = false; + std::unique_ptr file; + + expert_file(const std::string & fname, bool direct_io) : fname(fname) { + reopen(direct_io); + } + + expert_file(const expert_file &) = delete; + expert_file & operator=(const expert_file &) = delete; + + void reopen(bool direct_io) { + file = std::make_unique(fname.c_str(), "rb", direct_io); + size = file->size(); + direct = direct_io && file->has_direct_io(); + if (direct_io && !direct) { + LLAMA_LOG_WARN("%s: direct I/O is unavailable for %s; using buffered reads\n", __func__, fname.c_str()); + } + } + + size_t pread_at_least(void * dst, size_t len, uint64_t offset, size_t need) const { + if (need > len) { + throw std::runtime_error("llama_expert_store: invalid read requirement"); + } + size_t total = 0; + while (total < len) { +#if defined(_WIN32) + file->seek(offset + total, SEEK_SET); + file->read_raw(static_cast(dst) + total, len - total); + const size_t n = len - total; +#else + if (offset + total > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("llama_expert_store: file offset exceeds off_t"); + } + const ssize_t result = pread(file->file_id(), static_cast(dst) + total, len - total, + static_cast(offset + total)); + if (result < 0) { + if (errno == EINTR) { + continue; + } + throw std::runtime_error(format("llama_expert_store: pread failed for %s at %llu: %s", + fname.c_str(), (unsigned long long) (offset + total), strerror(errno))); + } + const size_t n = static_cast(result); +#endif + if (n == 0) { + break; + } + total += n; + } + if (total < need) { + throw std::runtime_error(format("llama_expert_store: short read for %s: %zu bytes, need %zu at %llu", + fname.c_str(), total, need, (unsigned long long) offset)); + } + return total; + } +}; + +} + +llama_expert_store_aligned_read llama_expert_store_align_read( + uint64_t offset, size_t size, size_t alignment, uint64_t file_size) { + if (!is_power_of_two(alignment)) { + throw std::runtime_error("llama_expert_store: I/O alignment must be a power of two"); + } + + uint64_t end; + if (!checked_add_u64(offset, size, &end) || end > file_size) { + throw std::runtime_error("llama_expert_store: read is outside the source file"); + } + + const uint64_t aligned_offset = offset & ~static_cast(alignment - 1); + const size_t prefix = static_cast(offset - aligned_offset); + uint64_t needed; + if (!checked_add_u64(prefix, size, &needed)) { + throw std::runtime_error("llama_expert_store: aligned read size overflow"); + } + uint64_t rounded; + if (!checked_add_u64(needed, alignment - 1, &rounded)) { + throw std::runtime_error("llama_expert_store: aligned read size overflow"); + } + rounded &= ~static_cast(alignment - 1); + llama_expert_store_aligned_read result; + result.offset = aligned_offset; + if (rounded > std::numeric_limits::max()) { + throw std::runtime_error("llama_expert_store: aligned read exceeds addressable memory"); + } + result.size = static_cast(rounded); + result.prefix = prefix; + return result; +} + +void llama_expert_store_validate_tensor(const llama_expert_store_tensor & tensor) { + if (tensor.name.empty() || tensor.fname.empty()) { + throw std::runtime_error("llama_expert_store: tensor name and source file are required"); + } + if (tensor.layer < 0) { + throw std::runtime_error(format("llama_expert_store: tensor %s has an invalid layer", tensor.name.c_str())); + } + if (tensor.projection < LLAMA_EXPERT_PROJECTION_GATE || tensor.projection > LLAMA_EXPERT_PROJECTION_DOWN) { + throw std::runtime_error(format("llama_expert_store: tensor %s has an invalid projection", tensor.name.c_str())); + } + const ggml_type expected_type = tensor.projection == LLAMA_EXPERT_PROJECTION_DOWN ? GGML_TYPE_Q2_K : GGML_TYPE_IQ2_XXS; + if (tensor.type != expected_type) { + throw std::runtime_error(format("llama_expert_store: tensor %s must be %s, got %s", + tensor.name.c_str(), ggml_type_name(expected_type), ggml_type_name(tensor.type))); + } + if (tensor.ne[0] <= 0 || tensor.ne[1] <= 0 || tensor.ne[2] <= 0) { + throw std::runtime_error(format("llama_expert_store: tensor %s has invalid dimensions", tensor.name.c_str())); + } + if (tensor.ne[2] > std::numeric_limits::max()) { + throw std::runtime_error(format("llama_expert_store: tensor %s has too many experts", tensor.name.c_str())); + } + if (tensor.ne[0] % ggml_blck_size(tensor.type) != 0) { + throw std::runtime_error(format("llama_expert_store: tensor %s rows are not whole quantization blocks", tensor.name.c_str())); + } + + const size_t row_size = ggml_row_size(tensor.type, tensor.ne[0]); + uint64_t plane_size; + uint64_t tensor_size; + if (!checked_mul_u64(row_size, static_cast(tensor.ne[1]), &plane_size) || + !checked_mul_u64(plane_size, static_cast(tensor.ne[2]), &tensor_size)) { + throw std::runtime_error(format("llama_expert_store: tensor %s size overflows", tensor.name.c_str())); + } + if (tensor.nb[0] != ggml_type_size(tensor.type) || tensor.nb[1] != row_size || tensor.nb[2] != plane_size) { + throw std::runtime_error(format("llama_expert_store: tensor %s is not a contiguous merged-expert tensor", tensor.name.c_str())); + } + uint64_t tensor_end; + if (!checked_add_u64(tensor.file_offset, tensor_size, &tensor_end) || tensor_end > tensor.file_size) { + throw std::runtime_error(format("llama_expert_store: tensor %s is outside the source file", tensor.name.c_str())); + } +} + +struct llama_expert_store::impl { + struct slot { + bool occupied = false; + expert_key key = {}; + const llama_expert_store_tensor * tensor = nullptr; + aligned_buffer bytes; + uint64_t last_use = 0; + uint32_t pins = 0; + }; + + llama_expert_store_params params; + std::map tensors; + std::map> files; + std::vector slots; + size_t bytes_resident = 0; + uint64_t use_clock = 0; + llama_expert_store_stats counters; + mutable std::mutex mutex; + + impl(std::vector tensors, const llama_expert_store_params & params) : params(params) { + if (params.cache_bytes == 0 || params.cache_slots == 0) { + throw std::runtime_error("llama_expert_store: cache byte and slot budgets must be non-zero"); + } + if (params.cache_slots > std::numeric_limits::max()) { + throw std::runtime_error("llama_expert_store: cache slot budget exceeds the slot ID range"); + } + if (!is_power_of_two(params.io_alignment)) { + throw std::runtime_error("llama_expert_store: I/O alignment must be a power of two"); + } + + for (auto & tensor : tensors) { + llama_expert_store_validate_tensor(tensor); + const tensor_key key = { tensor.layer, tensor.projection }; + if (this->tensors.count(key) != 0) { + throw std::runtime_error(format("llama_expert_store: duplicate tensor for layer %d projection %d", + tensor.layer, static_cast(tensor.projection))); + } + if (tensor.nb[2] > params.cache_bytes) { + throw std::runtime_error(format("llama_expert_store: tensor %s expert plane exceeds the cache byte budget", + tensor.name.c_str())); + } + auto file_it = files.find(tensor.fname); + if (file_it == files.end()) { + file_it = files.emplace(tensor.fname, + std::make_unique(tensor.fname, params.direct_io)).first; + } + if (file_it->second->size != tensor.file_size) { + throw std::runtime_error(format("llama_expert_store: source file size changed for %s", tensor.fname.c_str())); + } + this->tensors.emplace(key, std::move(tensor)); + } + + if (this->tensors.empty()) { + throw std::runtime_error("llama_expert_store: no tensors registered"); + } + for (auto it = this->tensors.begin(); it != this->tensors.end();) { + const int32_t layer = it->first.layer; + const auto gate = this->tensors.find({ layer, LLAMA_EXPERT_PROJECTION_GATE }); + const auto up = this->tensors.find({ layer, LLAMA_EXPERT_PROJECTION_UP }); + const auto down = this->tensors.find({ layer, LLAMA_EXPERT_PROJECTION_DOWN }); + if (gate == this->tensors.end() || up == this->tensors.end() || down == this->tensors.end()) { + throw std::runtime_error(format("llama_expert_store: layer %d must register gate, up, and down tensors", layer)); + } + if (gate->second.ne[0] != up->second.ne[0] || + gate->second.ne[1] != up->second.ne[1] || + gate->second.ne[2] != up->second.ne[2] || + down->second.ne[0] != gate->second.ne[1] || + down->second.ne[1] != gate->second.ne[0] || + down->second.ne[2] != gate->second.ne[2]) { + throw std::runtime_error(format("llama_expert_store: layer %d expert tensor dimensions do not match", layer)); + } + it = this->tensors.upper_bound({ layer, LLAMA_EXPERT_PROJECTION_DOWN }); + } + slots.resize(params.cache_slots); + } + + const llama_expert_store_tensor & get_tensor(const expert_key & key) const { + const auto it = tensors.find({ key.layer, key.projection }); + if (it == tensors.end()) { + throw std::runtime_error(format("llama_expert_store: no tensor for layer %d projection %d", + key.layer, static_cast(key.projection))); + } + if (key.expert_id < 0 || key.expert_id >= it->second.ne[2]) { + throw std::runtime_error(format("llama_expert_store: expert ID %d is outside [0, %lld)", + key.expert_id, (long long) it->second.ne[2])); + } + return it->second; + } + + aligned_buffer read_expert(const llama_expert_store_tensor & tensor, int32_t expert_id, uint64_t * bytes_read) const { + uint64_t expert_delta; + uint64_t expert_offset; + if (!checked_mul_u64(static_cast(expert_id), tensor.nb[2], &expert_delta) || + !checked_add_u64(tensor.file_offset, expert_delta, &expert_offset)) { + throw std::runtime_error(format("llama_expert_store: expert offset overflow for %s", tensor.name.c_str())); + } + + aligned_buffer payload(tensor.nb[2], params.io_alignment); + auto & file = *files.at(tensor.fname); + if (file.direct) { + const llama_expert_store_aligned_read read = + llama_expert_store_align_read(expert_offset, tensor.nb[2], params.io_alignment, tensor.file_size); + aligned_buffer bounce(read.size, params.io_alignment); + try { + *bytes_read += file.pread_at_least(bounce.data, read.size, read.offset, read.prefix + tensor.nb[2]); + memcpy(payload.data, bounce.data + read.prefix, tensor.nb[2]); + } catch (const std::runtime_error & e) { + LLAMA_LOG_WARN("%s: direct read failed for %s; retrying with buffered I/O: %s\n", + __func__, tensor.fname.c_str(), e.what()); + file.reopen(false); + *bytes_read += file.pread_at_least(payload.data, tensor.nb[2], expert_offset, tensor.nb[2]); + } + } else { + *bytes_read += file.pread_at_least(payload.data, tensor.nb[2], expert_offset, tensor.nb[2]); + } + + if (!ggml_validate_row_data(tensor.type, payload.data, tensor.nb[2])) { + throw std::runtime_error(format("llama_expert_store: tensor %s expert %d has invalid payload", + tensor.name.c_str(), expert_id)); + } + return payload; + } + + void unpin(const std::vector & slot_ids) { + std::lock_guard lock(mutex); + for (uint32_t slot_id : slot_ids) { + if (slot_id >= slots.size() || slots[slot_id].pins == 0) { + GGML_ABORT("llama_expert_store: invalid lease slot"); + } + slots[slot_id].pins--; + } + } + + std::vector get_payloads(const std::vector & slot_ids) const { + std::lock_guard lock(mutex); + std::vector result; + result.reserve(slot_ids.size()); + for (uint32_t slot_id : slot_ids) { + const slot & entry = slots.at(slot_id); + GGML_ASSERT(entry.occupied && entry.pins > 0); + result.push_back({ + entry.key.layer, + entry.key.projection, + entry.key.expert_id, + slot_id, + entry.tensor->type, + entry.bytes.data, + entry.bytes.size, + }); + } + return result; + } +}; + +struct llama_expert_store::lease::impl { + std::shared_ptr store; + std::vector pinned_slots; + std::vector> remapped_slots; + + ~impl() { + if (store) { + store->unpin(pinned_slots); + } + } +}; + +llama_expert_store::lease::lease() = default; +llama_expert_store::lease::lease(lease && other) noexcept = default; +llama_expert_store::lease & llama_expert_store::lease::operator=(lease && other) noexcept = default; +llama_expert_store::lease::~lease() = default; + +const std::vector> & llama_expert_store::lease::slot_ids() const { + static const std::vector> empty; + return pimpl ? pimpl->remapped_slots : empty; +} + +std::vector llama_expert_store::lease::payloads() const { + return pimpl ? pimpl->store->get_payloads(pimpl->pinned_slots) : std::vector(); +} + +llama_expert_store::llama_expert_store( + std::vector tensors, const llama_expert_store_params & params) + : pimpl(std::make_shared(std::move(tensors), params)) { +} + +llama_expert_store::~llama_expert_store() = default; + +llama_expert_store::lease llama_expert_store::acquire(const std::vector & requests) { + std::lock_guard lock(pimpl->mutex); + + std::vector> request_keys; + std::vector unique_keys; + request_keys.reserve(requests.size()); + for (const auto & request : requests) { + std::vector keys; + keys.reserve(request.expert_ids.size()); + for (int32_t expert_id : request.expert_ids) { + const expert_key key = { request.layer, request.projection, expert_id }; + pimpl->get_tensor(key); + keys.push_back(key); + unique_keys.push_back(key); + } + request_keys.push_back(std::move(keys)); + } + std::sort(unique_keys.begin(), unique_keys.end()); + unique_keys.erase(std::unique(unique_keys.begin(), unique_keys.end()), unique_keys.end()); + + std::map resident; + for (uint32_t i = 0; i < pimpl->slots.size(); ++i) { + if (pimpl->slots[i].occupied) { + resident.emplace(pimpl->slots[i].key, i); + } + } + + std::vector misses; + std::vector hit_slots; + size_t miss_bytes = 0; + for (const expert_key & key : unique_keys) { + const auto hit = resident.find(key); + if (hit != resident.end()) { + hit_slots.push_back(hit->second); + continue; + } + const auto & tensor = pimpl->get_tensor(key); + if (miss_bytes > pimpl->params.cache_bytes || tensor.nb[2] > pimpl->params.cache_bytes - miss_bytes) { + throw std::runtime_error("llama_expert_store: requested expert union exceeds the cache byte budget"); + } + miss_bytes += tensor.nb[2]; + misses.push_back(key); + } + + std::vector empty_slots; + std::vector candidates; + std::sort(hit_slots.begin(), hit_slots.end()); + for (uint32_t i = 0; i < pimpl->slots.size(); ++i) { + const auto & entry = pimpl->slots[i]; + if (!entry.occupied) { + empty_slots.push_back(i); + } else if (entry.pins == 0 && !std::binary_search(hit_slots.begin(), hit_slots.end(), i)) { + candidates.push_back(i); + } + } + std::sort(candidates.begin(), candidates.end(), [&](uint32_t a, uint32_t b) { + const auto & lhs = pimpl->slots[a]; + const auto & rhs = pimpl->slots[b]; + if (lhs.last_use != rhs.last_use) { + return lhs.last_use < rhs.last_use; + } + return a < b; + }); + + const size_t min_victims = misses.size() > empty_slots.size() ? misses.size() - empty_slots.size() : 0; + std::vector victims; + if (miss_bytes > std::numeric_limits::max() - pimpl->bytes_resident) { + throw std::runtime_error("llama_expert_store: cache byte accounting overflow"); + } + size_t bytes_after = pimpl->bytes_resident + miss_bytes; + for (uint32_t candidate : candidates) { + if (victims.size() >= min_victims && bytes_after <= pimpl->params.cache_bytes) { + break; + } + victims.push_back(candidate); + bytes_after -= pimpl->slots[candidate].bytes.size; + } + if (victims.size() < min_victims || bytes_after > pimpl->params.cache_bytes) { + throw std::runtime_error("llama_expert_store: cache capacity is pinned or too small for the requested expert union"); + } + + std::vector target_slots = empty_slots; + target_slots.insert(target_slots.end(), victims.begin(), victims.end()); + std::sort(target_slots.begin(), target_slots.end()); + + for (uint32_t victim : victims) { + auto & entry = pimpl->slots[victim]; + pimpl->bytes_resident -= entry.bytes.size; + entry.bytes.clear(); + entry.occupied = false; + entry.tensor = nullptr; + entry.last_use = 0; + pimpl->counters.evictions++; + } + + uint64_t bytes_read = 0; + size_t loaded = 0; + try { + for (; loaded < misses.size(); ++loaded) { + const expert_key & key = misses[loaded]; + const auto & tensor = pimpl->get_tensor(key); + auto & entry = pimpl->slots[target_slots[loaded]]; + entry.bytes = pimpl->read_expert(tensor, key.expert_id, &bytes_read); + entry.occupied = true; + entry.key = key; + entry.tensor = &tensor; + entry.pins = 0; + pimpl->bytes_resident += entry.bytes.size; + resident[entry.key] = target_slots[loaded]; + } + } catch (...) { + for (size_t i = 0; i < loaded; ++i) { + auto & entry = pimpl->slots[target_slots[i]]; + pimpl->bytes_resident -= entry.bytes.size; + resident.erase(entry.key); + entry.bytes.clear(); + entry.occupied = false; + entry.tensor = nullptr; + entry.last_use = 0; + } + throw; + } + + auto lease_impl = std::make_unique(); + lease_impl->store = pimpl; + lease_impl->remapped_slots.reserve(request_keys.size()); + for (const auto & keys : request_keys) { + std::vector remapped; + remapped.reserve(keys.size()); + for (const expert_key & key : keys) { + remapped.push_back(resident.at(key)); + } + lease_impl->remapped_slots.push_back(std::move(remapped)); + } + for (const expert_key & key : unique_keys) { + const uint32_t slot_id = resident.at(key); + auto & entry = pimpl->slots[slot_id]; + entry.last_use = ++pimpl->use_clock; + entry.pins++; + lease_impl->pinned_slots.push_back(slot_id); + } + + pimpl->counters.hits += unique_keys.size() - misses.size(); + pimpl->counters.misses += misses.size(); + pimpl->counters.bytes_read += bytes_read; + + lease result; + result.pimpl = std::move(lease_impl); + return result; +} + +llama_expert_store_stats llama_expert_store::stats() const { + std::lock_guard lock(pimpl->mutex); + return pimpl->counters; +} + +size_t llama_expert_store::resident_bytes() const { + std::lock_guard lock(pimpl->mutex); + return pimpl->bytes_resident; +} + +size_t llama_expert_store::resident_entries() const { + std::lock_guard lock(pimpl->mutex); + size_t result = 0; + for (const auto & slot : pimpl->slots) { + result += slot.occupied ? 1 : 0; + } + return result; +} diff --git a/src/llama-expert-store.h b/src/llama-expert-store.h new file mode 100644 index 000000000000..8b5082120968 --- /dev/null +++ b/src/llama-expert-store.h @@ -0,0 +1,103 @@ +#pragma once + +#include "ggml.h" + +#include +#include +#include +#include +#include + +enum llama_expert_projection { + LLAMA_EXPERT_PROJECTION_GATE = 0, + LLAMA_EXPERT_PROJECTION_UP, + LLAMA_EXPERT_PROJECTION_DOWN, +}; + +struct llama_expert_store_tensor { + std::string name; + std::string fname; + int32_t layer = -1; + llama_expert_projection projection = LLAMA_EXPERT_PROJECTION_GATE; + ggml_type type = GGML_TYPE_COUNT; + int64_t ne[3] = {}; + size_t nb[3] = {}; + uint64_t file_offset = 0; + uint64_t file_size = 0; +}; + +struct llama_expert_store_params { + size_t cache_bytes = 0; + size_t cache_slots = 0; + size_t io_alignment = 4096; + bool direct_io = true; +}; + +struct llama_expert_store_request { + int32_t layer = -1; + llama_expert_projection projection = LLAMA_EXPERT_PROJECTION_GATE; + std::vector expert_ids; +}; + +struct llama_expert_store_stats { + uint64_t hits = 0; + uint64_t misses = 0; + uint64_t bytes_read = 0; + uint64_t evictions = 0; +}; + +struct llama_expert_store_aligned_read { + uint64_t offset = 0; + size_t size = 0; + size_t prefix = 0; +}; + +llama_expert_store_aligned_read llama_expert_store_align_read( + uint64_t offset, size_t size, size_t alignment, uint64_t file_size); + +void llama_expert_store_validate_tensor(const llama_expert_store_tensor & tensor); + +struct llama_expert_store { + struct payload { + int32_t layer = -1; + llama_expert_projection projection = LLAMA_EXPERT_PROJECTION_GATE; + int32_t expert_id = -1; + uint32_t slot_id = 0; + ggml_type type = GGML_TYPE_COUNT; + const uint8_t * data = nullptr; + size_t size = 0; + }; + + struct lease { + lease(); + lease(lease && other) noexcept; + lease & operator=(lease && other) noexcept; + ~lease(); + + lease(const lease &) = delete; + lease & operator=(const lease &) = delete; + + const std::vector> & slot_ids() const; + std::vector payloads() const; + + private: + friend struct llama_expert_store; + + struct impl; + std::unique_ptr pimpl; + }; + + llama_expert_store(std::vector tensors, const llama_expert_store_params & params); + ~llama_expert_store(); + + // The lease pins every unique returned slot. Keep it until the backend upload completes. + lease acquire(const std::vector & requests); + + llama_expert_store_stats stats() const; + size_t resident_bytes() const; + size_t resident_entries() const; + +private: + struct impl; + std::shared_ptr pimpl; +}; diff --git a/src/llama-mmap.cpp b/src/llama-mmap.cpp index 4d183cbc9c45..9bd9d98f625c 100644 --- a/src/llama-mmap.cpp +++ b/src/llama-mmap.cpp @@ -465,7 +465,7 @@ struct llama_mmap::impl { #ifdef _POSIX_MAPPED_FILES std::vector> mapped_fragments; - impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & lazy_ranges) { + impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & excluded_ranges) { size = file->size(); int fd = file->file_id(); int flags = MAP_SHARED; @@ -475,8 +475,8 @@ struct llama_mmap::impl { LLAMA_LOG_WARN("warning: posix_fadvise(.., POSIX_FADV_SEQUENTIAL) failed: %s\n", strerror(errno)); } - // MAP_POPULATE would fault in the lazy ranges too - if (prefetch && lazy_ranges.empty()) { flags |= MAP_POPULATE; } + // MAP_POPULATE would fault in excluded ranges too + if (prefetch && excluded_ranges.empty()) { flags |= MAP_POPULATE; } #endif addr = mmap(NULL, file->size(), PROT_READ, flags, fd, 0); if (addr == MAP_FAILED) { @@ -497,11 +497,11 @@ struct llama_mmap::impl { }; if (prefetch > 0) { - for (const auto & range : ranges_complement(lazy_ranges, std::min(file->size(), prefetch))) { + for (const auto & range : ranges_complement(excluded_ranges, std::min(file->size(), prefetch))) { advise(range.first, range.second, POSIX_MADV_WILLNEED, "POSIX_MADV_WILLNEED"); } } - for (const auto & range : lazy_ranges) { + for (const auto & range : excluded_ranges) { advise(range.first, range.second, POSIX_MADV_RANDOM, "POSIX_MADV_RANDOM"); } if (numa) { @@ -572,7 +572,7 @@ struct llama_mmap::impl { #elif defined(_WIN32) HANDLE hMapping = nullptr; - impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & lazy_ranges) { + impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & excluded_ranges) { GGML_UNUSED(numa); size = file->size(); @@ -603,7 +603,7 @@ struct llama_mmap::impl { if (pPrefetchVirtualMemory) { std::vector entries; - for (const auto & range : ranges_complement(lazy_ranges, std::min(size, prefetch))) { + for (const auto & range : ranges_complement(excluded_ranges, std::min(size, prefetch))) { WIN32_MEMORY_RANGE_ENTRY entry; entry.VirtualAddress = (char *) addr + range.first; entry.NumberOfBytes = (SIZE_T) (range.second - range.first); @@ -641,11 +641,11 @@ struct llama_mmap::impl { } } #else - impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & lazy_ranges) { + impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & excluded_ranges) { GGML_UNUSED(file); GGML_UNUSED(prefetch); GGML_UNUSED(numa); - GGML_UNUSED(lazy_ranges); + GGML_UNUSED(excluded_ranges); throw std::runtime_error("mmap not supported"); } @@ -663,7 +663,7 @@ struct llama_mmap::impl { }; llama_mmap::llama_mmap(struct llama_file * file, size_t prefetch, bool numa, - const ranges & lazy_ranges) : pimpl(std::make_unique(file, prefetch, numa, lazy_ranges)) {} + const ranges & excluded_ranges) : pimpl(std::make_unique(file, prefetch, numa, excluded_ranges)) {} llama_mmap::~llama_mmap() = default; size_t llama_mmap::size() const { return pimpl->size; } diff --git a/src/llama-mmap.h b/src/llama-mmap.h index cc28c8a73fa5..c5c8d340cbe8 100644 --- a/src/llama-mmap.h +++ b/src/llama-mmap.h @@ -47,7 +47,7 @@ struct llama_mmap { llama_mmap(const llama_mmap &) = delete; llama_mmap(struct llama_file * file, size_t prefetch = (size_t) -1, bool numa = false, - const ranges & lazy_ranges = {}); + const ranges & excluded_ranges = {}); ~llama_mmap(); size_t size() const; diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 2dfcd6eb9074..b9ed24150586 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1131,6 +1131,14 @@ bool llama_model_loader::lazy_read::add(const std::string & name, const ggml_ten return true; } +void llama_model_loader::external_read::add(const llama_tensor_weight & w) { + const std::string name = ggml_get_name(w.tensor); + if (!tensors.insert(name).second) { + throw std::runtime_error(format("external tensor '%s' is already registered", name.c_str())); + } + ranges[w.idx].emplace_back(w.offs, w.offs + ggml_nbytes(w.tensor)); +} + struct ggml_tensor * llama_model_loader::create_tensor( const llama_hparams & hparams, const buft_list_t * buft_list_cpu, const buft_list_t * buft_list_input, const buft_list_t * buft_list_output, const buft_list_t * buft_list_layer, const LLM_TN_IMPL & tn, const std::initializer_list & ne, int flags) { @@ -1407,6 +1415,40 @@ struct ggml_tensor * llama_model_loader::create_tensor( return tensor; } +llama_expert_store_tensor llama_model_loader::register_external_tensor( + const std::string & name, + int32_t layer, + llama_expert_projection projection, + const std::initializer_list & ne) { + const ggml_tensor * tensor = check_tensor_dims(name, ne, true, false); + GGML_ASSERT(tensor != nullptr); + if (tensor->ne[3] != 1) { + throw std::runtime_error(format("external tensor '%s' must have three dimensions", name.c_str())); + } + + const llama_tensor_weight & weight = require_weight(name.c_str()); + if (fnames.at(weight.idx).empty() || fnames.at(weight.idx) == "(file*)") { + throw std::runtime_error(format("external tensor '%s' requires a reopenable source file", name.c_str())); + } + llama_expert_store_tensor result; + result.name = name; + result.fname = fnames.at(weight.idx); + result.layer = layer; + result.projection = projection; + result.type = tensor->type; + for (size_t i = 0; i < 3; ++i) { + result.ne[i] = tensor->ne[i]; + result.nb[i] = tensor->nb[i]; + } + result.file_offset = weight.offs; + result.file_size = files.at(weight.idx)->size(); + + llama_expert_store_validate_tensor(result); + external.add(weight); + n_created++; + return result; +} + void llama_model_loader::done_getting_tensors(bool partial) const { if (n_created > n_tensors) { throw std::runtime_error(format("%s: too many tensors created; expected %d, got %d", __func__, n_tensors, n_created)); @@ -1446,10 +1488,16 @@ void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps const size_t prefetch_size = prefetch && use_mmap ? -1 : 0; - std::unique_ptr mapping = std::make_unique(file.get(), prefetch_size, is_numa, - lazy.for_file(idx)); + llama_mmap::ranges excluded = lazy.for_file(idx); + const auto & external_ranges = external.for_file(idx); + excluded.insert(excluded.end(), external_ranges.begin(), external_ranges.end()); + + std::unique_ptr mapping = std::make_unique(file.get(), prefetch_size, is_numa, excluded); + for (const auto & range : external_ranges) { + mapping->unmap_fragment(range.first, range.second); + } mmaps_used.emplace_back(mapping->size(), 0); - if (mlock_mmaps) { + if (mlock_mmaps && external_ranges.empty()) { std::unique_ptr mlock_mmap(new llama_mlock()); mlock_mmap->init(mapping->addr()); mlock_mmaps->emplace_back(std::move(mlock_mmap)); @@ -1460,7 +1508,9 @@ void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps // compute the total size of all tensors for progress reporting for (const auto & it : weights_map) { - size_data += ggml_nbytes(it.second.tensor); + if (!external.has(it.second.tensor)) { + size_data += ggml_nbytes(it.second.tensor); + } } } diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index 72bbd53e7d03..5edb02555e7b 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -4,6 +4,7 @@ #include "llama-impl.h" #include "llama-arch.h" +#include "llama-expert-store.h" #include "llama-hparams.h" #include "llama-mmap.h" @@ -117,6 +118,38 @@ struct llama_model_loader { std::set tensors; } lazy; + struct external_read { + void add(const llama_tensor_weight & w); + + bool any() const { + return !ranges.empty(); + } + + bool has(const ggml_tensor * t) const { + return tensors.count(ggml_get_name(t)) > 0; + } + + const llama_mmap::ranges & for_file(uint32_t idx) const { + static const llama_mmap::ranges none; + + const auto it = ranges.find(idx); + return it == ranges.end() ? none : it->second; + } + + bool intersects(uint32_t idx, size_t first, size_t last) const { + for (const auto & range : for_file(idx)) { + if (range.first < last && first < range.second) { + return true; + } + } + return false; + } + + private: + std::map ranges; + std::set tensors; + } external; + llama_files files; std::vector fnames; // one per entry of files, for readers that outlive the loader llama_ftype ftype; @@ -239,6 +272,12 @@ struct llama_model_loader { const llama_hparams & hparams, const buft_list_t * buft_list_cpu, const buft_list_t * buft_list_input, const buft_list_t * buft_list_output, const buft_list_t * buft_list_layer, const LLM_TN_IMPL & tn, const std::initializer_list & ne, int flags); + llama_expert_store_tensor register_external_tensor( + const std::string & name, + int32_t layer, + llama_expert_projection projection, + const std::initializer_list & ne); + void done_getting_tensors(bool partial = false) const; void init_mappings(bool prefetch = true, llama_mlocks * mlock_mmaps = nullptr); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 734bd3edce98..98619e1fef32 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1689,12 +1689,14 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { // With the n-gram table left on disk, a populated mapping would pull the table's // third of the file resident for nothing; readahead alone carries the sequential load. - ml.init_mappings(!params.ple_on_disk, use_mlock ? &pimpl->mlock_mmaps : nullptr); + llama_mlocks * mmap_locks = use_mlock && !ml.external.any() ? &pimpl->mlock_mmaps : nullptr; + ml.init_mappings(!params.ple_on_disk, mmap_locks); pimpl->mappings.reserve(ml.mappings.size()); // create the backend buffers std::vector> ctx_buf_maps; ctx_buf_maps.reserve(ml.ctx_map.size()); + bool keep_mappings = false; // Ensure we have enough capacity for the maximum backend buffer we will potentially create const size_t n_max_backend_buffer = ml.ctx_map.size() * ml.files.size(); @@ -1730,9 +1732,24 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { // a lazy context is mapped whatever the load mode, but the memory-fit pass maps nothing const bool is_lazy_mapped = ctx_key.lazy && !ml.no_alloc; + bool context_mmap_safe = true; + if (ml.use_mmap && ml.external.any() && !is_lazy_mapped) { + for (uint32_t idx = 0; idx < ml.files.size(); ++idx) { + void * addr = nullptr; + size_t first; + size_t last; + ml.get_mapping_range(&first, &last, &addr, idx, ctx); + if (first < last && ml.external.intersects(idx, first, last)) { + context_mmap_safe = false; + break; + } + } + } - if ((ml.use_mmap || is_lazy_mapped) && use_mmap_buffer && buffer_from_host_ptr_supported && is_default_buft) { + if ((ml.use_mmap || is_lazy_mapped) && (use_mmap_buffer || is_lazy_mapped) && + context_mmap_safe && buffer_from_host_ptr_supported && is_default_buft) { GGML_ASSERT(!ml.no_alloc); + keep_mappings = true; for (uint32_t idx = 0; idx < ml.files.size(); idx++) { // only the mmap region containing the tensors in the model is mapped to the backend buffer // this is important for metal with apple silicon: if the entire model could be mapped to a metal buffer, @@ -1826,12 +1843,12 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { // load tensor data for (auto & [ctx, buf_map] : ctx_buf_maps) { - if (!ml.load_all_data(ctx, buf_map, use_mlock ? &pimpl->mlock_mmaps : NULL, params.progress_callback, params.progress_callback_user_data)) { + if (!ml.load_all_data(ctx, buf_map, mmap_locks, params.progress_callback, params.progress_callback_user_data)) { return false; } } - if (use_mmap_buffer) { + if (keep_mappings) { for (auto & mapping : ml.mappings) { pimpl->mappings.emplace_back(std::move(mapping)); } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index ea937784c5a2..8719bcf61e13 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -290,6 +290,7 @@ llama_build_and_test(test-thread-safety.cpp ARGS -m "${MODEL_DEST}" -ngl 99 -p " set_tests_properties(test-thread-safety PROPERTIES FIXTURES_REQUIRED test-download-model) llama_build_and_test(test-arg-parser.cpp) +llama_build_and_test(test-expert-store.cpp) llama_build_and_test(test-model-resolution.cpp) # the test serves its repos from an httplib server, and the library links it privately target_link_libraries(test-model-resolution PRIVATE cpp-httplib) diff --git a/tests/test-expert-store.cpp b/tests/test-expert-store.cpp new file mode 100644 index 000000000000..b301c27a22c8 --- /dev/null +++ b/tests/test-expert-store.cpp @@ -0,0 +1,503 @@ +#include "../src/llama-expert-store.h" +#include "../src/llama-model-loader.h" + +#include "ggml.h" +#include "gguf.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define REQUIRE(cond) do { if (!(cond)) { throw std::runtime_error("requirement failed: " #cond); } } while (0) + +namespace { + +struct temp_file { + std::filesystem::path path; + + explicit temp_file(const char * suffix) { + static uint64_t sequence = 0; + const auto stamp = std::chrono::steady_clock::now().time_since_epoch().count(); + path = std::filesystem::temp_directory_path() / + ("llama-expert-store-" + std::to_string(stamp) + "-" + std::to_string(++sequence) + suffix); + } + + ~temp_file() { + std::error_code ec; + std::filesystem::remove(path, ec); + } +}; + +struct fixture { + static constexpr int64_t n_embd = 512; + static constexpr int64_t n_ff = 256; + static constexpr int64_t n_expert = 4; + + temp_file file { ".gguf" }; + std::vector tensors; + + fixture() { + const size_t gate_plane = ggml_row_size(GGML_TYPE_IQ2_XXS, n_embd) * n_ff; + const size_t down_plane = ggml_row_size(GGML_TYPE_Q2_K, n_ff) * n_embd; + const size_t data_size = 2 * gate_plane * n_expert + down_plane * n_expert; + + ggml_init_params ggml_params = { + /*.mem_size =*/ data_size + 8 * ggml_tensor_overhead() + 4096, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ false, + }; + ggml_context * ctx = ggml_init(ggml_params); + REQUIRE(ctx != nullptr); + + ggml_tensor * gate = ggml_new_tensor_3d(ctx, GGML_TYPE_IQ2_XXS, n_embd, n_ff, n_expert); + ggml_tensor * up = ggml_new_tensor_3d(ctx, GGML_TYPE_IQ2_XXS, n_embd, n_ff, n_expert); + ggml_tensor * down = ggml_new_tensor_3d(ctx, GGML_TYPE_Q2_K, n_ff, n_embd, n_expert); + ggml_set_name(gate, "blk.0.ffn_gate_exps.weight"); + ggml_set_name(up, "blk.0.ffn_up_exps.weight"); + ggml_set_name(down, "blk.0.ffn_down_exps.weight"); + + fill_tensor(gate, 0x10); + fill_tensor(up, 0x20); + fill_tensor(down, 0x30); + + gguf_context * gguf = gguf_init_empty(); + REQUIRE(gguf != nullptr); + gguf_set_val_str(gguf, "general.architecture", "deepseek41"); + gguf_add_tensor(gguf, gate); + gguf_add_tensor(gguf, up); + gguf_add_tensor(gguf, down); + REQUIRE(gguf_write_to_file(gguf, file.path.string().c_str(), false)); + gguf_free(gguf); + ggml_free(ctx); + + std::vector splits; + llama_model_loader loader( + nullptr, nullptr, nullptr, file.path.string(), splits, nullptr, + LLAMA_LOAD_MODE_MMAP, false, false, false, nullptr, nullptr); + + REQUIRE(loader.ctx_map.empty()); + tensors.push_back(loader.register_external_tensor( + "blk.0.ffn_gate_exps.weight", 0, LLAMA_EXPERT_PROJECTION_GATE, { n_embd, n_ff, n_expert })); + tensors.push_back(loader.register_external_tensor( + "blk.0.ffn_up_exps.weight", 0, LLAMA_EXPERT_PROJECTION_UP, { n_embd, n_ff, n_expert })); + tensors.push_back(loader.register_external_tensor( + "blk.0.ffn_down_exps.weight", 0, LLAMA_EXPERT_PROJECTION_DOWN, { n_ff, n_embd, n_expert })); + loader.done_getting_tensors(); + + REQUIRE(loader.external.any()); + REQUIRE(loader.ctx_map.empty()); + for (const auto & tensor : tensors) { + REQUIRE(loader.external.has(loader.require_tensor_meta(tensor.name))); + } + loader.init_mappings(true); + REQUIRE(loader.mappings.size() == 1); + REQUIRE(loader.ctx_map.empty()); + const auto & ranges = loader.external.for_file(0); + REQUIRE(ranges.size() == 3); + for (size_t i = 0; i < ranges.size(); ++i) { + REQUIRE(ranges[i].first == tensors[i].file_offset); + REQUIRE(ranges[i].second == tensors[i].file_offset + tensors[i].nb[2] * tensors[i].ne[2]); + } + } + + static void fill_tensor(ggml_tensor * tensor, uint8_t tag) { + memset(tensor->data, 0, ggml_nbytes(tensor)); + for (int64_t expert = 0; expert < tensor->ne[2]; ++expert) { + uint8_t * plane = static_cast(tensor->data) + expert * tensor->nb[2]; + plane[tensor->nb[2] - 1] = tag + expert; + } + } + + llama_expert_store make_store(size_t slots, size_t bytes) const { + llama_expert_store_params params; + params.cache_slots = slots; + params.cache_bytes = bytes; + params.direct_io = false; + return llama_expert_store(tensors, params); + } + + size_t max_plane_size() const { + size_t result = 0; + for (const auto & tensor : tensors) { + result = std::max(result, tensor.nb[2]); + } + return result; + } + + size_t all_projection_bytes() const { + size_t result = 0; + for (const auto & tensor : tensors) { + result += tensor.nb[2]; + } + return result; + } +}; + +template +void require_throws(F && fn) { + bool threw = false; + try { + fn(); + } catch (const std::exception &) { + threw = true; + } + REQUIRE(threw); +} + +const llama_expert_store::payload & find_payload( + const std::vector & payloads, + llama_expert_projection projection, + int32_t expert_id) { + for (const auto & payload : payloads) { + if (payload.projection == projection && payload.expert_id == expert_id) { + return payload; + } + } + throw std::runtime_error("payload not found"); +} + +void test_layout_and_offsets(const fixture & f) { + const auto & gate = f.tensors[0]; + const auto & up = f.tensors[1]; + const auto & down = f.tensors[2]; + + REQUIRE(gate.type == GGML_TYPE_IQ2_XXS); + REQUIRE(up.type == GGML_TYPE_IQ2_XXS); + REQUIRE(down.type == GGML_TYPE_Q2_K); + REQUIRE(gate.nb[2] == ggml_row_size(GGML_TYPE_IQ2_XXS, fixture::n_embd) * fixture::n_ff); + REQUIRE(down.nb[2] == ggml_row_size(GGML_TYPE_Q2_K, fixture::n_ff) * fixture::n_embd); + REQUIRE(gate.file_offset + 3 * gate.nb[2] > gate.file_offset); + + llama_expert_store store = f.make_store(3, f.all_projection_bytes()); + REQUIRE(store.resident_entries() == 0); + REQUIRE(store.resident_bytes() == 0); + auto lease = store.acquire({ + { 0, LLAMA_EXPERT_PROJECTION_GATE, { 2 } }, + { 0, LLAMA_EXPERT_PROJECTION_UP, { 2 } }, + { 0, LLAMA_EXPERT_PROJECTION_DOWN, { 2 } }, + }); + const auto payloads = lease.payloads(); + REQUIRE(payloads.size() == 3); + REQUIRE(find_payload(payloads, LLAMA_EXPERT_PROJECTION_GATE, 2).data[gate.nb[2] - 1] == 0x12); + REQUIRE(find_payload(payloads, LLAMA_EXPERT_PROJECTION_UP, 2).data[up.nb[2] - 1] == 0x22); + REQUIRE(find_payload(payloads, LLAMA_EXPERT_PROJECTION_DOWN, 2).data[down.nb[2] - 1] == 0x32); +} + +void test_alignment_and_large_offsets() { + const auto aligned = llama_expert_store_align_read(4097, 5000, 4096, 20000); + REQUIRE(aligned.offset == 4096); + REQUIRE(aligned.prefix == 1); + REQUIRE(aligned.size == 8192); + + const uint64_t large_offset = (uint64_t(1) << 32) + 123; + const auto large = llama_expert_store_align_read(large_offset, 777, 4096, large_offset + 777); + REQUIRE(large.offset > std::numeric_limits::max()); + REQUIRE(large.prefix == 123); + REQUIRE(large.size == 4096); + + require_throws([] { + llama_expert_store_align_read(0, 1, 3000, 1); + }); + require_throws([] { + llama_expert_store_align_read(UINT64_MAX - 4, 8, 4096, UINT64_MAX); + }); +} + +void test_published_layout_accounting() { + const int64_t n_embd = 7680; + const int64_t n_ff = 1536; + const int64_t n_expert = 384; + const int64_t n_layer = 40; + + const uint64_t gate_plane = ggml_row_size(GGML_TYPE_IQ2_XXS, n_embd) * n_ff; + const uint64_t up_plane = ggml_row_size(GGML_TYPE_IQ2_XXS, n_embd) * n_ff; + const uint64_t down_plane = ggml_row_size(GGML_TYPE_Q2_K, n_ff) * n_embd; + const uint64_t slot_bytes = (gate_plane + up_plane + down_plane) * n_layer; + const uint64_t routed_bytes = slot_bytes * n_expert; + const uint64_t dense_bytes = 10067427328; + const uint64_t engram_bytes = 202758045696; + + REQUIRE(gate_plane == 3041280); + REQUIRE(up_plane == 3041280); + REQUIRE(down_plane == 3870720); + REQUIRE(gate_plane + up_plane + down_plane == 9953280); + REQUIRE(slot_bytes == 398131200); + REQUIRE(routed_bytes == 152882380800); + REQUIRE(224 * slot_bytes == 89181388800); + REQUIRE(256 * slot_bytes == 101921587200); + REQUIRE(dense_bytes + 224 * slot_bytes == 99248816128); + REQUIRE(dense_bytes + 256 * slot_bytes == 111989014528); + REQUIRE(engram_bytes > routed_bytes); +} + +void test_large_offset_read() { + temp_file sparse { ".bin" }; + const int64_t ne0 = 256; + const int64_t ne1 = 256; + const int64_t n_expert = 1; + const size_t gate_plane = ggml_row_size(GGML_TYPE_IQ2_XXS, ne0) * ne1; + const size_t down_plane = ggml_row_size(GGML_TYPE_Q2_K, ne0) * ne1; + const uint64_t base = (uint64_t(1) << 32) + 4096; + const uint64_t file_size = base + 2 * gate_plane + down_plane; + + { + std::ofstream out(sparse.path, std::ios::binary | std::ios::trunc); + REQUIRE(out.good()); + out.seekp(static_cast(file_size - 1)); + out.put('\0'); + const std::vector> markers = { + { base + gate_plane - 1, 0x41 }, + { base + 2 * gate_plane - 1, 0x42 }, + { file_size - 1, 0x43 }, + }; + for (const auto & marker : markers) { + out.seekp(static_cast(marker.first)); + out.put(static_cast(marker.second)); + } + } + + auto make_tensor = [&](const char * name, llama_expert_projection projection, ggml_type type, uint64_t offset) { + llama_expert_store_tensor tensor; + tensor.name = name; + tensor.fname = sparse.path.string(); + tensor.layer = 0; + tensor.projection = projection; + tensor.type = type; + tensor.ne[0] = ne0; + tensor.ne[1] = ne1; + tensor.ne[2] = n_expert; + tensor.nb[0] = ggml_type_size(type); + tensor.nb[1] = ggml_row_size(type, ne0); + tensor.nb[2] = tensor.nb[1] * ne1; + tensor.file_offset = offset; + tensor.file_size = file_size; + return tensor; + }; + + std::vector tensors; + tensors.push_back(make_tensor("gate", LLAMA_EXPERT_PROJECTION_GATE, GGML_TYPE_IQ2_XXS, base)); + tensors.push_back(make_tensor("up", LLAMA_EXPERT_PROJECTION_UP, GGML_TYPE_IQ2_XXS, base + gate_plane)); + tensors.push_back(make_tensor("down", LLAMA_EXPERT_PROJECTION_DOWN, GGML_TYPE_Q2_K, base + 2 * gate_plane)); + + llama_expert_store_params params { 2 * gate_plane + down_plane, 3, 4096, false }; + llama_expert_store store(std::move(tensors), params); + auto lease = store.acquire({ + { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } }, + { 0, LLAMA_EXPERT_PROJECTION_UP, { 0 } }, + { 0, LLAMA_EXPERT_PROJECTION_DOWN, { 0 } }, + }); + const auto payloads = lease.payloads(); + REQUIRE(find_payload(payloads, LLAMA_EXPERT_PROJECTION_GATE, 0).data[gate_plane - 1] == 0x41); + REQUIRE(find_payload(payloads, LLAMA_EXPERT_PROJECTION_UP, 0).data[gate_plane - 1] == 0x42); + REQUIRE(find_payload(payloads, LLAMA_EXPERT_PROJECTION_DOWN, 0).data[down_plane - 1] == 0x43); +} + +void test_cache_and_remapping(const fixture & f) { + const size_t gate_plane = f.tensors[0].nb[2]; + llama_expert_store store = f.make_store(2, 2 * f.max_plane_size()); + + { + auto lease = store.acquire({ + { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0, 0 } }, + { 0, LLAMA_EXPERT_PROJECTION_GATE, { 1, 0 } }, + }); + REQUIRE(lease.slot_ids().size() == 2); + REQUIRE(lease.slot_ids()[0][0] == lease.slot_ids()[0][1]); + REQUIRE(lease.slot_ids()[1][1] == lease.slot_ids()[0][0]); + REQUIRE(lease.slot_ids()[1][0] != lease.slot_ids()[1][1]); + REQUIRE(store.resident_entries() == 2); + REQUIRE(store.resident_bytes() == 2 * gate_plane); + } + + { + auto hit = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + REQUIRE(hit.slot_ids()[0][0] == 0); + } + { + auto miss = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 2 } } }); + REQUIRE(miss.slot_ids()[0][0] == 1); + } + + const llama_expert_store_stats stats = store.stats(); + REQUIRE(stats.hits == 1); + REQUIRE(stats.misses == 3); + REQUIRE(stats.evictions == 1); + REQUIRE(stats.bytes_read == 3 * gate_plane); + + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { -1 } } }); + }); + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { fixture::n_expert } } }); + }); +} + +void test_direct_io(const fixture & f) { + llama_expert_store_params params; + params.cache_slots = 1; + params.cache_bytes = f.max_plane_size(); + params.io_alignment = 4096; + params.direct_io = true; + + llama_expert_store store(f.tensors, params); + auto lease = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 3 } } }); + const auto payloads = lease.payloads(); + REQUIRE(payloads.size() == 1); + REQUIRE(payloads[0].data[payloads[0].size - 1] == 0x13); + REQUIRE(store.stats().bytes_read >= payloads[0].size); + REQUIRE(store.stats().bytes_read <= payloads[0].size + 2 * params.io_alignment); +} + +void test_pins_and_atomic_failure(const fixture & f) { + llama_expert_store store = f.make_store(1, f.max_plane_size()); + auto pinned = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + const llama_expert_store_stats before = store.stats(); + + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 1 } } }); + }); + REQUIRE(store.resident_entries() == 1); + REQUIRE(store.stats().hits == before.hits); + REQUIRE(store.stats().misses == before.misses); + REQUIRE(pinned.payloads()[0].expert_id == 0); + + pinned = {}; + auto replacement = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 1 } } }); + REQUIRE(replacement.payloads()[0].expert_id == 1); + REQUIRE(store.stats().evictions == 1); + + llama_expert_store::lease surviving; + { + llama_expert_store short_lived = f.make_store(1, f.max_plane_size()); + surviving = short_lived.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 2 } } }); + } + REQUIRE(surviving.payloads()[0].expert_id == 2); +} + +void test_limits_and_validation(const fixture & f) { + require_throws([&] { + f.make_store(0, f.max_plane_size()); + }); + require_throws([&] { + f.make_store(1, f.tensors[2].nb[2] - 1); + }); + { + llama_expert_store_params params { f.max_plane_size(), 1, 1, false }; + llama_expert_store store(f.tensors, params); + auto lease = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + REQUIRE(lease.payloads().size() == 1); + } + { + llama_expert_store store = f.make_store(3, f.max_plane_size()); + require_throws([&] { + store.acquire({ + { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } }, + { 0, LLAMA_EXPERT_PROJECTION_UP, { 0 } }, + }); + }); + REQUIRE(store.resident_entries() == 0); + REQUIRE(store.stats().misses == 0); + } + { + llama_expert_store store = f.make_store(1, 2 * f.tensors[0].nb[2]); + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0, 1 } } }); + }); + REQUIRE(store.resident_entries() == 0); + REQUIRE(store.stats().misses == 0); + } + + auto bad_type = f.tensors; + bad_type[0].type = GGML_TYPE_Q2_K; + require_throws([&] { + llama_expert_store_params params { f.all_projection_bytes(), 3, 4096, false }; + llama_expert_store store(std::move(bad_type), params); + }); + + auto bad_stride = f.tensors; + bad_stride[1].nb[2]++; + require_throws([&] { + llama_expert_store_params params { f.all_projection_bytes(), 3, 4096, false }; + llama_expert_store store(std::move(bad_stride), params); + }); + + auto bad_bounds = f.tensors; + bad_bounds[2].file_size = bad_bounds[2].file_offset + bad_bounds[2].nb[2] - 1; + require_throws([&] { + llama_expert_store_params params { f.all_projection_bytes(), 3, 4096, false }; + llama_expert_store store(std::move(bad_bounds), params); + }); +} + +void test_payload_validation(const fixture & f) { + temp_file copy { ".gguf" }; + std::filesystem::copy_file(f.file.path, copy.path); + auto tensors = f.tensors; + for (auto & tensor : tensors) { + tensor.fname = copy.path.string(); + tensor.file_size = std::filesystem::file_size(copy.path); + } + + { + std::fstream io(copy.path, std::ios::binary | std::ios::in | std::ios::out); + REQUIRE(io.good()); + io.seekp(static_cast(tensors[0].file_offset)); + const uint8_t invalid_scale[2] = { 0x00, 0x7c }; + io.write(reinterpret_cast(invalid_scale), sizeof(invalid_scale)); + } + + llama_expert_store_params params { f.all_projection_bytes(), 3, 4096, false }; + llama_expert_store store(std::move(tensors), params); + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + }); + REQUIRE(store.resident_entries() == 0); +} + +void test_truncated_file(const fixture & f) { + temp_file copy { ".gguf" }; + std::filesystem::copy_file(f.file.path, copy.path); + auto tensors = f.tensors; + for (auto & tensor : tensors) { + tensor.fname = copy.path.string(); + tensor.file_size = std::filesystem::file_size(copy.path); + } + + llama_expert_store_params params { f.all_projection_bytes(), 3, 4096, false }; + llama_expert_store store(std::move(tensors), params); + std::filesystem::resize_file(copy.path, f.tensors[0].file_offset + f.tensors[0].nb[2] - 1); + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + }); + REQUIRE(store.resident_entries() == 0); + REQUIRE(store.stats().misses == 0); +} + +} + +int main() { + try { + fixture f; + test_layout_and_offsets(f); + test_alignment_and_large_offsets(); + test_published_layout_accounting(); + test_large_offset_read(); + test_cache_and_remapping(f); + test_direct_io(f); + test_pins_and_atomic_failure(f); + test_limits_and_validation(f); + test_payload_validation(f); + test_truncated_file(f); + } catch (const std::exception & e) { + fprintf(stderr, "test-expert-store: %s\n", e.what()); + return 1; + } + return 0; +} From 1572ea7fa79d64d0120795e41fbed236986b35bc Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 09:02:53 -0700 Subject: [PATCH 2/8] llama : require direct I/O for expert store Stop aligned reads once the complete expert payload is available and make buffered fallback explicit opt-in so page-cache growth cannot bypass the production memory budget. Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/llama-expert-store.cpp | 30 ++++++++++++++++++++++++------ src/llama-expert-store.h | 2 ++ tests/test-expert-store.cpp | 26 ++++++++++++++++++++++++++ 3 files changed, 52 insertions(+), 6 deletions(-) diff --git a/src/llama-expert-store.cpp b/src/llama-expert-store.cpp index 7be7b6144ff6..5ba4784d901f 100644 --- a/src/llama-expert-store.cpp +++ b/src/llama-expert-store.cpp @@ -145,8 +145,15 @@ struct expert_file { bool direct = false; std::unique_ptr file; - expert_file(const std::string & fname, bool direct_io) : fname(fname) { + expert_file(const std::string & fname, bool direct_io, bool allow_buffered_io) : fname(fname) { reopen(direct_io); + if (direct_io && !direct && !allow_buffered_io) { + throw std::runtime_error(format("llama_expert_store: direct I/O is required but unavailable for %s", fname.c_str())); + } + if (direct_io && !direct) { + LLAMA_LOG_WARN("%s: direct I/O is unavailable for %s; using explicitly allowed buffered reads\n", + __func__, fname.c_str()); + } } expert_file(const expert_file &) = delete; @@ -156,9 +163,6 @@ struct expert_file { file = std::make_unique(fname.c_str(), "rb", direct_io); size = file->size(); direct = direct_io && file->has_direct_io(); - if (direct_io && !direct) { - LLAMA_LOG_WARN("%s: direct I/O is unavailable for %s; using buffered reads\n", __func__, fname.c_str()); - } } size_t pread_at_least(void * dst, size_t len, uint64_t offset, size_t need) const { @@ -166,7 +170,7 @@ struct expert_file { throw std::runtime_error("llama_expert_store: invalid read requirement"); } size_t total = 0; - while (total < len) { + while (total < need) { #if defined(_WIN32) file->seek(offset + total, SEEK_SET); file->read_raw(static_cast(dst) + total, len - total); @@ -318,7 +322,7 @@ struct llama_expert_store::impl { auto file_it = files.find(tensor.fname); if (file_it == files.end()) { file_it = files.emplace(tensor.fname, - std::make_unique(tensor.fname, params.direct_io)).first; + std::make_unique(tensor.fname, params.direct_io, params.allow_buffered_io)).first; } if (file_it->second->size != tensor.file_size) { throw std::runtime_error(format("llama_expert_store: source file size changed for %s", tensor.fname.c_str())); @@ -381,6 +385,10 @@ struct llama_expert_store::impl { *bytes_read += file.pread_at_least(bounce.data, read.size, read.offset, read.prefix + tensor.nb[2]); memcpy(payload.data, bounce.data + read.prefix, tensor.nb[2]); } catch (const std::runtime_error & e) { + if (!params.allow_buffered_io) { + throw std::runtime_error(format("llama_expert_store: direct I/O failed for %s and buffered fallback is disabled: %s", + tensor.fname.c_str(), e.what())); + } LLAMA_LOG_WARN("%s: direct read failed for %s; retrying with buffered I/O: %s\n", __func__, tensor.fname.c_str(), e.what()); file.reopen(false); @@ -630,3 +638,13 @@ size_t llama_expert_store::resident_entries() const { } return result; } + +bool llama_expert_store::direct_io_active() const { + std::lock_guard lock(pimpl->mutex); + for (const auto & item : pimpl->files) { + if (!item.second->direct) { + return false; + } + } + return true; +} diff --git a/src/llama-expert-store.h b/src/llama-expert-store.h index 8b5082120968..07a6188aebdb 100644 --- a/src/llama-expert-store.h +++ b/src/llama-expert-store.h @@ -31,6 +31,7 @@ struct llama_expert_store_params { size_t cache_slots = 0; size_t io_alignment = 4096; bool direct_io = true; + bool allow_buffered_io = false; // opt-in only; page-cache bytes are outside cache_bytes }; struct llama_expert_store_request { @@ -96,6 +97,7 @@ struct llama_expert_store { llama_expert_store_stats stats() const; size_t resident_bytes() const; size_t resident_entries() const; + bool direct_io_active() const; private: struct impl; diff --git a/tests/test-expert-store.cpp b/tests/test-expert-store.cpp index b301c27a22c8..59457217ac7c 100644 --- a/tests/test-expert-store.cpp +++ b/tests/test-expert-store.cpp @@ -345,6 +345,7 @@ void test_direct_io(const fixture & f) { params.cache_bytes = f.max_plane_size(); params.io_alignment = 4096; params.direct_io = true; + params.allow_buffered_io = true; llama_expert_store store(f.tensors, params); auto lease = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 3 } } }); @@ -355,6 +356,28 @@ void test_direct_io(const fixture & f) { REQUIRE(store.stats().bytes_read <= payloads[0].size + 2 * params.io_alignment); } +#if defined(__linux__) +void test_direct_io_file_tail(const fixture & f) { + const auto & down = f.tensors[2]; + REQUIRE(down.file_offset + down.nb[2] * down.ne[2] == down.file_size); + + llama_expert_store_params params; + params.cache_slots = 1; + params.cache_bytes = f.max_plane_size(); + params.io_alignment = 4096; + params.direct_io = true; + params.allow_buffered_io = false; + + llama_expert_store store(f.tensors, params); + REQUIRE(store.direct_io_active()); + auto lease = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_DOWN, { fixture::n_expert - 1 } } }); + const auto payloads = lease.payloads(); + REQUIRE(payloads.size() == 1); + REQUIRE(payloads[0].data[payloads[0].size - 1] == 0x33); + REQUIRE(store.direct_io_active()); +} +#endif + void test_pins_and_atomic_failure(const fixture & f) { llama_expert_store store = f.make_store(1, f.max_plane_size()); auto pinned = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); @@ -491,6 +514,9 @@ int main() { test_large_offset_read(); test_cache_and_remapping(f); test_direct_io(f); +#if defined(__linux__) + test_direct_io_file_tail(f); +#endif test_pins_and_atomic_failure(f); test_limits_and_validation(f); test_payload_validation(f); From a8364460999b94ac2ef52cf28bd1d0d4c925eef2 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Sat, 12 Sep 2026 09:09:06 -0700 Subject: [PATCH 3/8] llama : fail closed on Windows expert direct I/O Report direct I/O unavailable until the Windows file path supports FILE_FLAG_NO_BUFFERING, and cover strict failure plus explicit buffered opt-in. Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/llama-mmap.cpp | 3 ++- tests/test-expert-store.cpp | 23 +++++++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/src/llama-mmap.cpp b/src/llama-mmap.cpp index 9bd9d98f625c..fd6caaf44640 100644 --- a/src/llama-mmap.cpp +++ b/src/llama-mmap.cpp @@ -171,7 +171,8 @@ struct llama_file::impl { } bool has_direct_io() const { - return true; + // Windows uses cached CRT I/O until FILE_FLAG_NO_BUFFERING support is added. + return false; } ~impl() { diff --git a/tests/test-expert-store.cpp b/tests/test-expert-store.cpp index 59457217ac7c..4b5847d7f8a6 100644 --- a/tests/test-expert-store.cpp +++ b/tests/test-expert-store.cpp @@ -378,6 +378,26 @@ void test_direct_io_file_tail(const fixture & f) { } #endif +#if defined(_WIN32) +void test_windows_direct_io_policy(const fixture & f) { + llama_expert_store_params params; + params.cache_slots = 1; + params.cache_bytes = f.max_plane_size(); + params.direct_io = true; + params.allow_buffered_io = false; + + require_throws([&] { + llama_expert_store store(f.tensors, params); + }); + + params.allow_buffered_io = true; + llama_expert_store store(f.tensors, params); + REQUIRE(!store.direct_io_active()); + auto lease = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + REQUIRE(lease.payloads().size() == 1); +} +#endif + void test_pins_and_atomic_failure(const fixture & f) { llama_expert_store store = f.make_store(1, f.max_plane_size()); auto pinned = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); @@ -516,6 +536,9 @@ int main() { test_direct_io(f); #if defined(__linux__) test_direct_io_file_tail(f); +#endif +#if defined(_WIN32) + test_windows_direct_io_policy(f); #endif test_pins_and_atomic_failure(f); test_limits_and_validation(f); From 0c1ae86bffbcaf30c4dd519c2206f305f46ea456 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Mon, 14 Sep 2026 00:37:14 -0700 Subject: [PATCH 4/8] llama : close expert store lifetime bounds Load contexts that cross disk-owned holes without mmap, discard copied source pages where supported, preserve mlock for safe mapped spans, and make lease publication allocation-safe. Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/llama-expert-store.cpp | 8 +++++--- src/llama-mmap.cpp | 33 +++++++++++++++++++++++++++++++ src/llama-mmap.h | 3 +++ src/llama-model-loader.cpp | 11 ++++++++--- src/llama-model-loader.h | 2 ++ src/llama-model.cpp | 39 +++++++++++++++++++++++++++---------- tests/test-expert-store.cpp | 13 +++++++++++++ 7 files changed, 93 insertions(+), 16 deletions(-) diff --git a/src/llama-expert-store.cpp b/src/llama-expert-store.cpp index 5ba4784d901f..82f7e2dbc7ae 100644 --- a/src/llama-expert-store.cpp +++ b/src/llama-expert-store.cpp @@ -593,7 +593,6 @@ llama_expert_store::lease llama_expert_store::acquire(const std::vector(); - lease_impl->store = pimpl; lease_impl->remapped_slots.reserve(request_keys.size()); for (const auto & keys : request_keys) { std::vector remapped; @@ -603,13 +602,16 @@ llama_expert_store::lease llama_expert_store::acquire(const std::vectorremapped_slots.push_back(std::move(remapped)); } + lease_impl->pinned_slots.reserve(unique_keys.size()); for (const expert_key & key : unique_keys) { - const uint32_t slot_id = resident.at(key); + lease_impl->pinned_slots.push_back(resident.at(key)); + } + for (uint32_t slot_id : lease_impl->pinned_slots) { auto & entry = pimpl->slots[slot_id]; entry.last_use = ++pimpl->use_clock; entry.pins++; - lease_impl->pinned_slots.push_back(slot_id); } + lease_impl->store = pimpl; pimpl->counters.hits += unique_keys.size() - misses.size(); pimpl->counters.misses += misses.size(); diff --git a/src/llama-mmap.cpp b/src/llama-mmap.cpp index bd3471767204..fc6651043c69 100644 --- a/src/llama-mmap.cpp +++ b/src/llama-mmap.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -176,6 +177,11 @@ struct llama_file::impl { return false; } + void discard_cache(size_t offset, size_t length) const { + GGML_UNUSED(offset); + GGML_UNUSED(length); + } + ~impl() { if (fp && owns_fp) { std::fclose(fp); @@ -375,6 +381,19 @@ struct llama_file::impl { return fd != -1 && alignment > 1; } + void discard_cache(size_t offset, size_t length) const { +#if defined(POSIX_FADV_DONTNEED) + const int file_id = fd == -1 ? fileno(fp) : fd; + const int result = posix_fadvise(file_id, offset, length, POSIX_FADV_DONTNEED); + if (result != 0) { + LLAMA_LOG_WARN("warning: posix_fadvise(.., POSIX_FADV_DONTNEED) failed: %s\n", strerror(result)); + } +#else + GGML_UNUSED(offset); + GGML_UNUSED(length); +#endif + } + ~impl() { if (fd != -1) { close(fd); @@ -409,6 +428,7 @@ 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(); } +void llama_file::discard_cache(size_t offset, size_t length) const { pimpl->discard_cache(offset, length); } int llama_file::file_id() const { #ifdef _WIN32 @@ -783,6 +803,18 @@ struct llama_mlock::impl { impl() : addr(NULL), size(0), failed_already(false) {} + static void align_range(size_t * first, size_t * last) { + const size_t granularity = lock_granularity(); + *first &= ~(granularity - 1); + const size_t remainder = *last & (granularity - 1); + if (remainder != 0) { + if (*last > std::numeric_limits::max() - (granularity - remainder)) { + throw std::runtime_error("mlock range overflow"); + } + *last += granularity - remainder; + } + } + void init(void * ptr) { GGML_ASSERT(addr == NULL && size == 0); addr = ptr; @@ -815,6 +847,7 @@ llama_mlock::~llama_mlock() = default; void llama_mlock::init(void * ptr) { pimpl->init(ptr); } void llama_mlock::grow_to(size_t target_size) { pimpl->grow_to(target_size); } +void llama_mlock::align_range(size_t * first, size_t * last) { impl::align_range(first, last); } #if defined(_POSIX_MEMLOCK_RANGE) || defined(_WIN32) const bool llama_mlock::SUPPORTED = true; diff --git a/src/llama-mmap.h b/src/llama-mmap.h index c5c8d340cbe8..a99bc716c5c0 100644 --- a/src/llama-mmap.h +++ b/src/llama-mmap.h @@ -31,6 +31,8 @@ struct llama_file { void read_aligned_chunk(void * dest, size_t size); uint32_t read_u32(); + void discard_cache(size_t offset, size_t length) const; + void write_raw(const void * ptr, size_t len) const; void write_u32(uint32_t val) const; @@ -69,6 +71,7 @@ struct llama_mlock { void init(void * ptr); void grow_to(size_t target_size); + static void align_range(size_t * first, size_t * last); static const bool SUPPORTED; private: diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index b9ed24150586..3558f623132e 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1561,6 +1561,8 @@ const void * llama_model_loader::load_data_range(const llama_tensor_weight & w, bool llama_model_loader::load_all_data( struct ggml_context * ctx, llama_buf_map & bufs, + bool load_from_mmap, + bool discard_file_cache, llama_mlocks * lmlocks, llama_progress_callback progress_callback, void * progress_callback_user_data) { @@ -1592,7 +1594,7 @@ bool llama_model_loader::load_all_data( std::vector host_ptrs; size_t buffer_idx = 0; // buffer to use for async loads ggml_backend_t upload_backend = [&](const char * func) -> ggml_backend_t { - if (use_mmap || check_tensors) { + if (load_from_mmap || check_tensors) { return nullptr; } // When not using mmaped io use async uploads from pinned memory to GPU memory. @@ -1679,7 +1681,7 @@ bool llama_model_loader::load_all_data( // without mmap, tensors in non-host buffers are staged through a temporary buffer sized like the tensor // load them biggest-first so the largest staging buffer is allocated while the fewest weights are resident - if (!use_mmap) { + if (!load_from_mmap) { std::stable_sort(tensors.begin(), tensors.end(), [](const ggml_tensor * a, const ggml_tensor * b) { const bool staged_a = a->buffer && !ggml_backend_buffer_is_host(a->buffer); const bool staged_b = b->buffer && !ggml_backend_buffer_is_host(b->buffer); @@ -1705,7 +1707,7 @@ bool llama_model_loader::load_all_data( size_t n_size = ggml_nbytes(cur); - const bool from_mapping = use_mmap || lazy.has(cur); + const bool from_mapping = load_from_mmap || lazy.has(cur); if (from_mapping) { const auto & mapping = mappings.at(weight->idx); @@ -1813,6 +1815,9 @@ bool llama_model_loader::load_all_data( } } } + if (discard_file_cache) { + file->discard_cache(weight->offs, n_size); + } } size_done += n_size; diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index 5edb02555e7b..5381162065f3 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -295,6 +295,8 @@ struct llama_model_loader { bool load_all_data( struct ggml_context * ctx, llama_buf_map & bufs, + bool load_from_mmap, + bool discard_file_cache, llama_mlocks * lmlocks, llama_progress_callback progress_callback, void * progress_callback_user_data); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index e1392d212b23..1e53015d2996 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1704,14 +1704,20 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { } } - // With the n-gram table left on disk, a populated mapping would pull the table's - // third of the file resident for nothing; readahead alone carries the sequential load. + // Do not prefetch files with disk-owned tensor holes. Unsafe contexts load their + // resident tensors through bounded staging and discard copied source pages. llama_mlocks * mmap_locks = use_mlock && !ml.external.any() ? &pimpl->mlock_mmaps : nullptr; - ml.init_mappings(!params.ple_on_disk, mmap_locks); + ml.init_mappings(!params.ple_on_disk && !ml.external.any(), mmap_locks); pimpl->mappings.reserve(ml.mappings.size()); // create the backend buffers - std::vector> ctx_buf_maps; + struct ctx_buf_map { + ggml_context * ctx; + llama_buf_map bufs; + bool load_from_mmap; + bool discard_file_cache; + }; + std::vector ctx_buf_maps; ctx_buf_maps.reserve(ml.ctx_map.size()); bool keep_mappings = false; @@ -1783,6 +1789,15 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { if (buf == nullptr) { throw std::runtime_error(format("unable to allocate %s buffer", ggml_backend_buft_name(buft))); } + if (use_mlock && ml.external.any() && !is_lazy_mapped) { + size_t lock_first = first; + size_t lock_last = last; + llama_mlock::align_range(&lock_first, &lock_last); + pimpl->mlock_mmaps.emplace_back(new llama_mlock); + auto & mlock_mmap = pimpl->mlock_mmaps.back(); + mlock_mmap->init((char *) addr + lock_first); + mlock_mmap->grow_to(lock_last - lock_first); + } bufs.emplace_back(buf); buf_map.emplace(idx, buf); } @@ -1819,7 +1834,8 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { pimpl->ctxs_bufs.emplace_back(std::move(ctx_ptr), std::move(bufs)); - ctx_buf_maps.emplace_back(ctx, buf_map); + const bool load_from_mmap = context_mmap_safe && (ml.use_mmap || is_lazy_mapped); + ctx_buf_maps.push_back({ ctx, std::move(buf_map), load_from_mmap, ml.use_mmap && !load_from_mmap }); } if (llama_supports_gpu_offload()) { @@ -1851,16 +1867,19 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) { } // without mmap, load non-host buffers first: their tensors go through a staging buffer, which is cheapest while the fewest weights are resident - if (!ml.use_mmap) { + if (!ml.use_mmap || ml.external.any()) { std::stable_partition(ctx_buf_maps.begin(), ctx_buf_maps.end(), [](const auto & ctx_buf_map) { - const auto & buf_map = ctx_buf_map.second; - return !buf_map.empty() && !ggml_backend_buffer_is_host(buf_map.begin()->second); + const auto & buf_map = ctx_buf_map.bufs; + return !ctx_buf_map.load_from_mmap && !buf_map.empty() && + !ggml_backend_buffer_is_host(buf_map.begin()->second); }); } // load tensor data - for (auto & [ctx, buf_map] : ctx_buf_maps) { - if (!ml.load_all_data(ctx, buf_map, mmap_locks, params.progress_callback, params.progress_callback_user_data)) { + for (auto & ctx_buf_map : ctx_buf_maps) { + if (!ml.load_all_data(ctx_buf_map.ctx, ctx_buf_map.bufs, ctx_buf_map.load_from_mmap, + ctx_buf_map.discard_file_cache, mmap_locks, + params.progress_callback, params.progress_callback_user_data)) { return false; } } diff --git a/tests/test-expert-store.cpp b/tests/test-expert-store.cpp index 4b5847d7f8a6..d597e0a750f6 100644 --- a/tests/test-expert-store.cpp +++ b/tests/test-expert-store.cpp @@ -208,6 +208,19 @@ void test_alignment_and_large_offsets() { require_throws([] { llama_expert_store_align_read(UINT64_MAX - 4, 8, 4096, UINT64_MAX); }); + + size_t granularity_first = 1; + size_t granularity_last = 1; + llama_mlock::align_range(&granularity_first, &granularity_last); + const size_t lock_granularity = granularity_last; + REQUIRE(granularity_first == 0); + REQUIRE(lock_granularity > 1); + + size_t lock_first = lock_granularity + 1; + size_t lock_last = 2 * lock_granularity; + llama_mlock::align_range(&lock_first, &lock_last); + REQUIRE(lock_first == lock_granularity); + REQUIRE(lock_last == 2 * lock_granularity); } void test_published_layout_accounting() { From 9b53142e71a9a7df11d2db7f7eb0bae7c0f465ec Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Mon, 14 Sep 2026 01:01:32 -0700 Subject: [PATCH 5/8] mmap : bound advice around external holes Use random file advice when disk-owned tensor ranges are excluded so Linux readahead cannot cross into the expert corpus. Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/llama-mmap.cpp | 13 ++++++++++--- src/llama-mmap.h | 1 + tests/test-expert-store.cpp | 6 ++++++ 3 files changed, 17 insertions(+), 3 deletions(-) diff --git a/src/llama-mmap.cpp b/src/llama-mmap.cpp index fc6651043c69..aefd9340aab7 100644 --- a/src/llama-mmap.cpp +++ b/src/llama-mmap.cpp @@ -493,9 +493,12 @@ struct llama_mmap::impl { int flags = MAP_SHARED; if (numa) { prefetch = 0; } #ifdef __linux__ - if (posix_fadvise(fd, 0, 0, POSIX_FADV_SEQUENTIAL)) { - LLAMA_LOG_WARN("warning: posix_fadvise(.., POSIX_FADV_SEQUENTIAL) failed: %s\n", - strerror(errno)); + const bool sequential = llama_mmap::use_sequential_file_advice(excluded_ranges); + const int file_advice = sequential ? POSIX_FADV_SEQUENTIAL : POSIX_FADV_RANDOM; + const int advice_error = posix_fadvise(fd, 0, 0, file_advice); + if (advice_error) { + LLAMA_LOG_WARN("warning: posix_fadvise(.., %s) failed: %s\n", + sequential ? "POSIX_FADV_SEQUENTIAL" : "POSIX_FADV_RANDOM", strerror(advice_error)); } // MAP_POPULATE would fault in excluded ranges too if (prefetch && excluded_ranges.empty()) { flags |= MAP_POPULATE; } @@ -688,6 +691,10 @@ llama_mmap::llama_mmap(struct llama_file * file, size_t prefetch, bool numa, const ranges & excluded_ranges) : pimpl(std::make_unique(file, prefetch, numa, excluded_ranges)) {} llama_mmap::~llama_mmap() = default; +bool llama_mmap::use_sequential_file_advice(const ranges & excluded_ranges) { + return excluded_ranges.empty(); +} + size_t llama_mmap::size() const { return pimpl->size; } void * llama_mmap::addr() const { return pimpl->addr; } diff --git a/src/llama-mmap.h b/src/llama-mmap.h index a99bc716c5c0..7fbcb0e1c107 100644 --- a/src/llama-mmap.h +++ b/src/llama-mmap.h @@ -58,6 +58,7 @@ struct llama_mmap { void unmap_fragment(size_t first, size_t last); static const bool SUPPORTED; + static bool use_sequential_file_advice(const ranges & excluded_ranges); private: struct impl; diff --git a/tests/test-expert-store.cpp b/tests/test-expert-store.cpp index d597e0a750f6..c91b8493d0a9 100644 --- a/tests/test-expert-store.cpp +++ b/tests/test-expert-store.cpp @@ -223,6 +223,11 @@ void test_alignment_and_large_offsets() { REQUIRE(lock_last == 2 * lock_granularity); } +void test_external_mapping_access_policy() { + REQUIRE(llama_mmap::use_sequential_file_advice({})); + REQUIRE(!llama_mmap::use_sequential_file_advice({ { 4096, 8192 } })); +} + void test_published_layout_accounting() { const int64_t n_embd = 7680; const int64_t n_ff = 1536; @@ -543,6 +548,7 @@ int main() { fixture f; test_layout_and_offsets(f); test_alignment_and_large_offsets(); + test_external_mapping_access_policy(); test_published_layout_accounting(); test_large_offset_read(); test_cache_and_remapping(f); From c1a00d2784cb38ae1e0fc42aecd80ed9cb28c939 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Mon, 14 Sep 2026 01:04:39 -0700 Subject: [PATCH 6/8] mmap : suppress prefetch across expert holes Keep lazy-table prefetch behavior unchanged, but give external routed tensors a strict policy that disables positive prefetch ranges on every platform. Assisted-by: GPT-5.6 Sol Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/llama-mmap.cpp | 40 +++++++++++++++++++++++-------------- src/llama-mmap.h | 6 ++++-- src/llama-model-loader.cpp | 3 ++- tests/test-expert-store.cpp | 20 +++++++++++++++++-- 4 files changed, 49 insertions(+), 20 deletions(-) diff --git a/src/llama-mmap.cpp b/src/llama-mmap.cpp index aefd9340aab7..7c85562dafa0 100644 --- a/src/llama-mmap.cpp +++ b/src/llama-mmap.cpp @@ -460,7 +460,6 @@ void llama_file::write_u32(uint32_t val) const { pimpl->write_u32(val); } // llama_mmap -#if defined(_POSIX_MAPPED_FILES) || defined(_WIN32) // 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; @@ -478,22 +477,21 @@ static llama_mmap::ranges ranges_complement(llama_mmap::ranges ranges, size_t li if (pos < limit) { res.emplace_back(pos, limit); } - return res; } -#endif struct llama_mmap::impl { #ifdef _POSIX_MAPPED_FILES std::vector> mapped_fragments; - impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & excluded_ranges) { + impl(struct llama_file * file, size_t prefetch, bool numa, + const llama_mmap::ranges & excluded_ranges, bool strict_exclusion) { size = file->size(); int fd = file->file_id(); int flags = MAP_SHARED; if (numa) { prefetch = 0; } #ifdef __linux__ - const bool sequential = llama_mmap::use_sequential_file_advice(excluded_ranges); + const bool sequential = llama_mmap::use_sequential_file_advice(strict_exclusion); const int file_advice = sequential ? POSIX_FADV_SEQUENTIAL : POSIX_FADV_RANDOM; const int advice_error = posix_fadvise(fd, 0, 0, file_advice); if (advice_error) { @@ -521,10 +519,9 @@ struct llama_mmap::impl { } }; - if (prefetch > 0) { - for (const auto & range : ranges_complement(excluded_ranges, std::min(file->size(), prefetch))) { - advise(range.first, range.second, POSIX_MADV_WILLNEED, "POSIX_MADV_WILLNEED"); - } + for (const auto & range : + llama_mmap::planned_prefetch_ranges(file->size(), prefetch, excluded_ranges, strict_exclusion)) { + advise(range.first, range.second, POSIX_MADV_WILLNEED, "POSIX_MADV_WILLNEED"); } for (const auto & range : excluded_ranges) { advise(range.first, range.second, POSIX_MADV_RANDOM, "POSIX_MADV_RANDOM"); @@ -597,7 +594,8 @@ struct llama_mmap::impl { #elif defined(_WIN32) HANDLE hMapping = nullptr; - impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & excluded_ranges) { + impl(struct llama_file * file, size_t prefetch, bool numa, + const llama_mmap::ranges & excluded_ranges, bool strict_exclusion) { GGML_UNUSED(numa); size = file->size(); @@ -628,7 +626,8 @@ struct llama_mmap::impl { if (pPrefetchVirtualMemory) { std::vector entries; - for (const auto & range : ranges_complement(excluded_ranges, std::min(size, prefetch))) { + for (const auto & range : + llama_mmap::planned_prefetch_ranges(size, prefetch, excluded_ranges, strict_exclusion)) { WIN32_MEMORY_RANGE_ENTRY entry; entry.VirtualAddress = (char *) addr + range.first; entry.NumberOfBytes = (SIZE_T) (range.second - range.first); @@ -666,11 +665,13 @@ struct llama_mmap::impl { } } #else - impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & excluded_ranges) { + impl(struct llama_file * file, size_t prefetch, bool numa, + const llama_mmap::ranges & excluded_ranges, bool strict_exclusion) { GGML_UNUSED(file); GGML_UNUSED(prefetch); GGML_UNUSED(numa); GGML_UNUSED(excluded_ranges); + GGML_UNUSED(strict_exclusion); throw std::runtime_error("mmap not supported"); } @@ -688,11 +689,20 @@ struct llama_mmap::impl { }; llama_mmap::llama_mmap(struct llama_file * file, size_t prefetch, bool numa, - const ranges & excluded_ranges) : pimpl(std::make_unique(file, prefetch, numa, excluded_ranges)) {} + const ranges & excluded_ranges, bool strict_exclusion) : + pimpl(std::make_unique(file, prefetch, numa, excluded_ranges, strict_exclusion)) {} llama_mmap::~llama_mmap() = default; -bool llama_mmap::use_sequential_file_advice(const ranges & excluded_ranges) { - return excluded_ranges.empty(); +bool llama_mmap::use_sequential_file_advice(bool strict_exclusion) { + return !strict_exclusion; +} + +llama_mmap::ranges llama_mmap::planned_prefetch_ranges( + size_t file_size, size_t prefetch, const ranges & excluded_ranges, bool strict_exclusion) { + if (strict_exclusion || prefetch == 0) { + return {}; + } + return ranges_complement(excluded_ranges, std::min(file_size, prefetch)); } size_t llama_mmap::size() const { return pimpl->size; } diff --git a/src/llama-mmap.h b/src/llama-mmap.h index 7fbcb0e1c107..6d146a61f84e 100644 --- a/src/llama-mmap.h +++ b/src/llama-mmap.h @@ -49,7 +49,7 @@ struct llama_mmap { llama_mmap(const llama_mmap &) = delete; llama_mmap(struct llama_file * file, size_t prefetch = (size_t) -1, bool numa = false, - const ranges & excluded_ranges = {}); + const ranges & excluded_ranges = {}, bool strict_exclusion = false); ~llama_mmap(); size_t size() const; @@ -58,7 +58,9 @@ struct llama_mmap { void unmap_fragment(size_t first, size_t last); static const bool SUPPORTED; - static bool use_sequential_file_advice(const ranges & excluded_ranges); + static bool use_sequential_file_advice(bool strict_exclusion); + static ranges planned_prefetch_ranges( + size_t file_size, size_t prefetch, const ranges & excluded_ranges, bool strict_exclusion); private: struct impl; diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 3558f623132e..c703d1aa797f 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1492,7 +1492,8 @@ void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps const auto & external_ranges = external.for_file(idx); excluded.insert(excluded.end(), external_ranges.begin(), external_ranges.end()); - std::unique_ptr mapping = std::make_unique(file.get(), prefetch_size, is_numa, excluded); + std::unique_ptr mapping = std::make_unique( + file.get(), prefetch_size, is_numa, excluded, !external_ranges.empty()); for (const auto & range : external_ranges) { mapping->unmap_fragment(range.first, range.second); } diff --git a/tests/test-expert-store.cpp b/tests/test-expert-store.cpp index c91b8493d0a9..b84bb7738bfc 100644 --- a/tests/test-expert-store.cpp +++ b/tests/test-expert-store.cpp @@ -224,8 +224,24 @@ void test_alignment_and_large_offsets() { } void test_external_mapping_access_policy() { - REQUIRE(llama_mmap::use_sequential_file_advice({})); - REQUIRE(!llama_mmap::use_sequential_file_advice({ { 4096, 8192 } })); + const size_t page = 4096; + const llama_mmap::ranges external = { + { 0, page }, + { 2 * page + 1, 4 * page - 1 }, + { 5 * page, 6 * page }, + { 7 * page, 8 * page }, + { 9 * page, 10 * page }, + }; + + REQUIRE(llama_mmap::use_sequential_file_advice(false)); + REQUIRE(!llama_mmap::use_sequential_file_advice(true)); + REQUIRE(llama_mmap::planned_prefetch_ranges(10 * page, 10 * page, external, true).empty()); + REQUIRE(llama_mmap::planned_prefetch_ranges(10 * page, 10 * page, external, true).empty()); + + const auto lazy_ranges = llama_mmap::planned_prefetch_ranges(10 * page, 10 * page, external, false); + REQUIRE(lazy_ranges.size() == 4); + REQUIRE(lazy_ranges.front() == std::make_pair(page, 2 * page + 1)); + REQUIRE(lazy_ranges.back() == std::make_pair(8 * page, 9 * page)); } void test_published_layout_accounting() { From 8805ae2d415e2ecc155a105659a9d49072550807 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Mon, 14 Sep 2026 01:18:20 -0700 Subject: [PATCH 7/8] mmap : fail closed on strict advice errors Throw before mmap when POSIX_FADV_RANDOM cannot be applied to an external-hole file. Keep legacy sequential advice warning-only and cover both paths through an injectable constructor seam. This successor binds corrective commits 9b53142e71a9a7df11d2db7f7eb0bae7c0f465ec and c1a00d2784cb38ae1e0fc42aecd80ed9cb28c939, including c1 tree 56460c30cd3c4dc46bab0e3d0939290e1a140b34. Assisted-by: GPT-5.6 Sol Copilot-Session: 6df503f4-c9eb-436d-a324-e8c5f8cabed9 Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- src/llama-mmap.cpp | 25 ++++++++++++++++++------ src/llama-mmap.h | 4 +++- tests/test-expert-store.cpp | 38 ++++++++++++++++++++++++++++++++++++- 3 files changed, 59 insertions(+), 8 deletions(-) diff --git a/src/llama-mmap.cpp b/src/llama-mmap.cpp index 7c85562dafa0..141640f9d192 100644 --- a/src/llama-mmap.cpp +++ b/src/llama-mmap.cpp @@ -485,7 +485,8 @@ struct llama_mmap::impl { std::vector> mapped_fragments; impl(struct llama_file * file, size_t prefetch, bool numa, - const llama_mmap::ranges & excluded_ranges, bool strict_exclusion) { + const llama_mmap::ranges & excluded_ranges, bool strict_exclusion, + llama_mmap::file_advice_override file_advice_override) { size = file->size(); int fd = file->file_id(); int flags = MAP_SHARED; @@ -493,13 +494,21 @@ struct llama_mmap::impl { #ifdef __linux__ const bool sequential = llama_mmap::use_sequential_file_advice(strict_exclusion); const int file_advice = sequential ? POSIX_FADV_SEQUENTIAL : POSIX_FADV_RANDOM; - const int advice_error = posix_fadvise(fd, 0, 0, file_advice); + const int advice_error = file_advice_override ? + file_advice_override(fd, file_advice) : posix_fadvise(fd, 0, 0, file_advice); if (advice_error) { + if (strict_exclusion) { + throw std::runtime_error(format( + "posix_fadvise(.., POSIX_FADV_RANDOM) failed for external tensor mapping: %s", + strerror(advice_error))); + } LLAMA_LOG_WARN("warning: posix_fadvise(.., %s) failed: %s\n", sequential ? "POSIX_FADV_SEQUENTIAL" : "POSIX_FADV_RANDOM", strerror(advice_error)); } // MAP_POPULATE would fault in excluded ranges too if (prefetch && excluded_ranges.empty()) { flags |= MAP_POPULATE; } +#else + GGML_UNUSED(file_advice_override); #endif addr = mmap(NULL, file->size(), PROT_READ, flags, fd, 0); if (addr == MAP_FAILED) { @@ -595,8 +604,10 @@ struct llama_mmap::impl { HANDLE hMapping = nullptr; impl(struct llama_file * file, size_t prefetch, bool numa, - const llama_mmap::ranges & excluded_ranges, bool strict_exclusion) { + const llama_mmap::ranges & excluded_ranges, bool strict_exclusion, + llama_mmap::file_advice_override file_advice_override) { GGML_UNUSED(numa); + GGML_UNUSED(file_advice_override); size = file->size(); @@ -666,12 +677,14 @@ struct llama_mmap::impl { } #else impl(struct llama_file * file, size_t prefetch, bool numa, - const llama_mmap::ranges & excluded_ranges, bool strict_exclusion) { + const llama_mmap::ranges & excluded_ranges, bool strict_exclusion, + llama_mmap::file_advice_override file_advice_override) { GGML_UNUSED(file); GGML_UNUSED(prefetch); GGML_UNUSED(numa); GGML_UNUSED(excluded_ranges); GGML_UNUSED(strict_exclusion); + GGML_UNUSED(file_advice_override); throw std::runtime_error("mmap not supported"); } @@ -689,8 +702,8 @@ struct llama_mmap::impl { }; llama_mmap::llama_mmap(struct llama_file * file, size_t prefetch, bool numa, - const ranges & excluded_ranges, bool strict_exclusion) : - pimpl(std::make_unique(file, prefetch, numa, excluded_ranges, strict_exclusion)) {} + const ranges & excluded_ranges, bool strict_exclusion, file_advice_override file_advice) : + pimpl(std::make_unique(file, prefetch, numa, excluded_ranges, strict_exclusion, file_advice)) {} llama_mmap::~llama_mmap() = default; bool llama_mmap::use_sequential_file_advice(bool strict_exclusion) { diff --git a/src/llama-mmap.h b/src/llama-mmap.h index 6d146a61f84e..c64c5acfb6a1 100644 --- a/src/llama-mmap.h +++ b/src/llama-mmap.h @@ -46,10 +46,12 @@ struct llama_file { struct llama_mmap { // list of [first, last) byte ranges within a file using ranges = std::vector>; + using file_advice_override = int (*)(int fd, int advice); llama_mmap(const llama_mmap &) = delete; llama_mmap(struct llama_file * file, size_t prefetch = (size_t) -1, bool numa = false, - const ranges & excluded_ranges = {}, bool strict_exclusion = false); + const ranges & excluded_ranges = {}, bool strict_exclusion = false, + file_advice_override file_advice = nullptr); ~llama_mmap(); size_t size() const; diff --git a/tests/test-expert-store.cpp b/tests/test-expert-store.cpp index b84bb7738bfc..340390971ce3 100644 --- a/tests/test-expert-store.cpp +++ b/tests/test-expert-store.cpp @@ -5,6 +5,7 @@ #include "gguf.h" #include +#include #include #include #include @@ -236,7 +237,6 @@ void test_external_mapping_access_policy() { REQUIRE(llama_mmap::use_sequential_file_advice(false)); REQUIRE(!llama_mmap::use_sequential_file_advice(true)); REQUIRE(llama_mmap::planned_prefetch_ranges(10 * page, 10 * page, external, true).empty()); - REQUIRE(llama_mmap::planned_prefetch_ranges(10 * page, 10 * page, external, true).empty()); const auto lazy_ranges = llama_mmap::planned_prefetch_ranges(10 * page, 10 * page, external, false); REQUIRE(lazy_ranges.size() == 4); @@ -244,6 +244,39 @@ void test_external_mapping_access_policy() { REQUIRE(lazy_ranges.back() == std::make_pair(8 * page, 9 * page)); } +#if defined(__linux__) +int file_advice_calls = 0; + +int fail_file_advice(int, int) { + ++file_advice_calls; + return EIO; +} + +void test_external_mapping_advice_failure() { + temp_file file { ".bin" }; + { + std::ofstream out(file.path, std::ios::binary); + REQUIRE(out.good()); + out.seekp(8191); + out.put('\0'); + } + + llama_file input(file.path.string(), "rb"); + file_advice_calls = 0; + bool continued_after_advice = false; + require_throws([&] { + llama_mmap mapping(&input, 0, false, { { 4096, 8192 } }, true, fail_file_advice); + continued_after_advice = true; + }); + REQUIRE(file_advice_calls == 1); + REQUIRE(!continued_after_advice); + + llama_mmap legacy_mapping(&input, 0, false, {}, false, fail_file_advice); + REQUIRE(file_advice_calls == 2); + REQUIRE(legacy_mapping.addr() != nullptr); +} +#endif + void test_published_layout_accounting() { const int64_t n_embd = 7680; const int64_t n_ff = 1536; @@ -565,6 +598,9 @@ int main() { test_layout_and_offsets(f); test_alignment_and_large_offsets(); test_external_mapping_access_policy(); +#if defined(__linux__) + test_external_mapping_advice_failure(); +#endif test_published_layout_accounting(); test_large_offset_read(); test_cache_and_remapping(f); From f8d537e3642eb03539c5e840a929e86a76dde134 Mon Sep 17 00:00:00 2001 From: Jerome Coste Date: Mon, 14 Sep 2026 01:18:57 -0700 Subject: [PATCH 8/8] mmap : attest expert store successor Binds fail-closed corrective head 8805ae2d415e2ecc155a105659a9d49072550807 and its tree a85d80bc4356981a9e5da1ab236f766961d082dc without source changes. Assisted-by: GPT-5.6 Sol Copilot-Session: 6df503f4-c9eb-436d-a324-e8c5f8cabed9 Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>