diff --git a/common/arg.cpp b/common/arg.cpp index be5f8ebe34bd..c55b9a58aef5 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2806,6 +2806,26 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.ple_direct_io = value; } ).set_env("LLAMA_ARG_NGRAM_DIRECT_IO")); + add_opt(common_arg( + {"--expert-cache-slots"}, "N", + "DeepSeek V4.1 routed experts resident per layer; requires --expert-cache-mib", + [](common_params & params, int value) { + if (value <= 0) { + throw std::invalid_argument("invalid value"); + } + params.expert_cache_slots = value; + } + ).set_env("LLAMA_ARG_EXPERT_CACHE_SLOTS")); + add_opt(common_arg( + {"--expert-cache-mib"}, "MiB", + "aggregate DeepSeek V4.1 fixed expert slot-tensor capacity; requires --expert-cache-slots", + [](common_params & params, int value) { + if (value <= 0) { + throw std::invalid_argument("invalid value"); + } + params.expert_cache_mib = value; + } + ).set_env("LLAMA_ARG_EXPERT_CACHE_MIB")); add_opt(common_arg( {"-cmoe", "--cpu-moe"}, "keep all Mixture of Experts (MoE) weights in the CPU", diff --git a/common/common.cpp b/common/common.cpp index a6a5364163bc..b80b7fa5f809 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1697,6 +1697,8 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.ple_direct_io = params.ple_direct_io; mparams.ple_io_threads = params.ple_io_threads; mparams.ple_cache_mb = params.ple_cache_mb; + mparams.expert_cache_slots = params.expert_cache_slots; + mparams.expert_cache_bytes = params.expert_cache_mib > 0 ? (size_t) params.expert_cache_mib << 20 : 0; if (params.kv_overrides.empty()) { mparams.kv_overrides = NULL; diff --git a/common/common.h b/common/common.h index 8c4427e0c719..a838f90b52a8 100644 --- a/common/common.h +++ b/common/common.h @@ -626,6 +626,8 @@ struct common_params { bool ple_direct_io = true; // ... read with O_DIRECT int32_t ple_io_threads = 64; // ... parallel readers (random 4 KiB reads: this NVMe gives 62k IOPS at 16, 130k at 64, ~160k at 128+) int32_t ple_cache_mb = 256; // ... row cache, 0 disables + int32_t expert_cache_slots = 0; // DeepSeek V4.1 routed experts resident per layer + int32_t expert_cache_mib = 0; // aggregate fixed slot-tensor capacity bool single_turn = false; // single turn chat conversation diff --git a/include/llama.h b/include/llama.h index 429fc1cb8480..4720bf625606 100644 --- a/include/llama.h +++ b/include/llama.h @@ -348,6 +348,10 @@ extern "C" { int32_t ple_io_threads; // parallel pread workers int32_t ple_cache_mb; // in-memory cache of recently read rows, 0 disables + // DeepSeek V4.1 routed-expert cache. Both values must be non-zero. + size_t expert_cache_bytes; + int32_t expert_cache_slots; + // proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices() const float * tensor_split; diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 49a78bd220ed..7c403a37e0d9 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -18,6 +18,8 @@ set(LLAMA_CORE_SOURCES llama-cparams.cpp llama-dsv41.cpp llama-dsv41-engram.cpp + llama-dsv41-expert.cpp + llama-expert-store.cpp llama-grammar.cpp llama-graph.cpp llama-hparams.cpp diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 5fee486c172d..bf77844264c0 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -481,11 +481,15 @@ llama_context::llama_context( sampling.token_ids_full_vocab[i] = i; } } + + model.acquire_runtime_context(); } llama_context::~llama_context() { // wait for any pending asynchronous copies into the output buffers before they are freed synchronize(); + model.release_runtime_work(); + model.release_runtime_context(); // when training, ggml_opt allocates extra buffers through the scheduler, so the sizes no longer match the expectation if (!model.hparams.no_alloc && !opt_ctx) { @@ -1422,10 +1426,21 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll const auto status = graph_compute(res->get_gf(), ubatch.n_tokens > 1); if (status != GGML_STATUS_SUCCESS) { + model.release_runtime_work_after_sync(sched.get()); LLAMA_LOG_ERROR("%s: failed to compute graph, compute status: %d\n", __func__, status); ret = status; return nullptr; } + if (model.requires_synchronous_graph()) { + synchronize(); + const std::string error = model.consume_runtime_error(); + if (!error.empty()) { + model.release_runtime_work(); + LLAMA_LOG_ERROR("%s: model runtime failed: %s\n", __func__, error.c_str()); + ret = GGML_STATUS_FAILED; + return nullptr; + } + } ret = GGML_STATUS_SUCCESS; diff --git a/src/llama-dsv41-expert.cpp b/src/llama-dsv41-expert.cpp new file mode 100644 index 000000000000..3d82f29b4739 --- /dev/null +++ b/src/llama-dsv41-expert.cpp @@ -0,0 +1,552 @@ +#include "llama-dsv41-expert.h" + +#include "llama-dsv41.h" +#include "llama-impl.h" + +#include "ggml-cpp.h" + +#include +#include +#include +#include +#include +#include +#include + +// The remap node follows ggml-org/llama.cpp#25294 commit 4260e4608. +// Lease publication and release follow ggml-org/llama.cpp#27861 commit bccbacdb8. + +namespace { + +size_t checked_add(size_t a, size_t b, const char * message) { + if (b > std::numeric_limits::max() - a) { + throw std::runtime_error(message); + } + return a + b; +} + +size_t checked_mul(size_t a, size_t b, const char * message) { + if (a != 0 && b > std::numeric_limits::max()/a) { + throw std::runtime_error(message); + } + return a*b; +} + +size_t projection_index(llama_expert_projection projection) { + if (projection < LLAMA_EXPERT_PROJECTION_GATE || projection > LLAMA_EXPERT_PROJECTION_DOWN) { + throw std::invalid_argument("DeepSeek V4.1 expert projection is invalid"); + } + return static_cast(projection); +} + +void require_local_cpu(ggml_backend_sched_t sched, ggml_backend_t backend_cpu) { + if (sched == nullptr || backend_cpu == nullptr || !ggml_backend_is_cpu(backend_cpu)) { + throw std::invalid_argument("DeepSeek V4.1 expert callback requires a local CPU backend"); + } + for (int i = 0; i < ggml_backend_sched_get_n_backends(sched); ++i) { + if (ggml_backend_sched_get_backend(sched, i) == backend_cpu) { + return; + } + } + throw std::invalid_argument("DeepSeek V4.1 expert CPU backend is not in the scheduler"); +} + +struct dsv41_expert_callback_state { + llama_dsv41_expert_runtime * runtime = nullptr; + int32_t layer = -1; +}; + +} + +struct llama_dsv41_expert_runtime::impl { + struct logical_slot { + int32_t expert_id = -1; + uint32_t pins = 0; + uint64_t last_use = 0; + }; + + struct layer_state { + std::array tensors = {}; + std::vector slots; + std::unique_ptr lease; + std::vector pinned_slots; + std::vector active_ids; + std::vector active_remap; + }; + + llama_dsv41_expert_runtime_params params; + std::unique_ptr store; + std::vector layers; + std::vector callbacks; + std::vector> contexts; + std::vector buffers; + upload_fn upload; + publish_fn before_publish; + size_t bytes_cache = 0; + size_t bytes_staging = 0; + uint64_t use_clock = 0; + std::string error; + bool context_active = false; + mutable std::mutex mutex; + + impl( + llama_dsv41_expert_runtime * owner, + std::vector tensors, + const llama_dsv41_expert_runtime_params & params, + buft_selector select_buft, + upload_fn upload, + publish_fn before_publish) + : params(params), upload(std::move(upload)), before_publish(std::move(before_publish)) { + if (params.cache_bytes == 0 || params.cache_slots == 0) { + throw std::runtime_error("DeepSeek V4.1 expert cache byte and slot capacity must be non-zero"); + } + if (!select_buft) { + throw std::invalid_argument("DeepSeek V4.1 expert cache buffer selector is empty"); + } + if (tensors.size() != LLAMA_DSV41_N_LAYER*3) { + throw std::runtime_error("DeepSeek V4.1 must register 40 gate/up/down routed tensor sets"); + } + + layers.resize(LLAMA_DSV41_N_LAYER); + callbacks.resize(LLAMA_DSV41_N_LAYER); + std::vector layer_plane_bytes(LLAMA_DSV41_N_LAYER, 0); + for (const auto & tensor : tensors) { + llama_expert_store_validate_tensor(tensor); + if (tensor.layer < 0 || tensor.layer >= (int32_t) LLAMA_DSV41_N_LAYER) { + throw std::runtime_error("DeepSeek V4.1 routed tensor layer is invalid"); + } + if (tensor.ne[2] != LLAMA_DSV41_N_EXPERT) { + throw std::runtime_error("DeepSeek V4.1 routed tensor expert count mismatch"); + } + layer_plane_bytes[tensor.layer] = checked_add( + layer_plane_bytes[tensor.layer], + tensor.nb[2], + "DeepSeek V4.1 expert plane byte count overflow"); + bytes_cache = checked_add( + bytes_cache, + checked_mul(tensor.nb[2], params.cache_slots, "DeepSeek V4.1 expert cache byte count overflow"), + "DeepSeek V4.1 expert cache byte count overflow"); + + ggml_backend_buffer_type_t buft = select_buft(tensor); + if (buft == nullptr) { + throw std::runtime_error("DeepSeek V4.1 expert cache has no backend buffer type"); + } + ggml_context * ctx = nullptr; + for (auto & item : contexts) { + if (item.first == buft) { + ctx = item.second.get(); + break; + } + } + if (ctx == nullptr) { + ggml_init_params ctx_params = { + /*.mem_size =*/ ggml_tensor_overhead()*(LLAMA_DSV41_N_LAYER*3 + 1), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ctx = ggml_init(ctx_params); + if (ctx == nullptr) { + throw std::runtime_error("DeepSeek V4.1 failed to create expert cache tensor context"); + } + contexts.emplace_back(buft, ctx); + } + + ggml_tensor * cache = ggml_new_tensor_3d( + ctx, tensor.type, tensor.ne[0], tensor.ne[1], params.cache_slots); + ggml_format_name(cache, "%s.cache", tensor.name.c_str()); + layers[tensor.layer].tensors[projection_index(tensor.projection)] = cache; + } + if (bytes_cache > params.cache_bytes) { + throw std::runtime_error(format( + "DeepSeek V4.1 expert cache requires %zu bytes for %zu slots per layer, configured %zu", + bytes_cache, params.cache_slots, params.cache_bytes)); + } + + for (int32_t il = 0; il < (int32_t) LLAMA_DSV41_N_LAYER; ++il) { + for (ggml_tensor * tensor : layers[il].tensors) { + if (tensor == nullptr) { + throw std::runtime_error(format("DeepSeek V4.1 layer %d is missing an expert cache tensor", il)); + } + } + layers[il].slots.resize(params.cache_slots); + callbacks[il] = { owner, il }; + bytes_staging = std::max( + bytes_staging, + checked_mul(layer_plane_bytes[il], params.cache_slots, "DeepSeek V4.1 expert staging byte count overflow")); + } + + for (auto & item : contexts) { + ggml_backend_buffer_t buffer = nullptr; + if (params.no_alloc) { + buffer = ggml_backend_buft_alloc_buffer(item.first, 0); + for (ggml_tensor * tensor = ggml_get_first_tensor(item.second.get()); + tensor != nullptr; + tensor = ggml_get_next_tensor(item.second.get(), tensor)) { + tensor->buffer = buffer; + } + } else { + buffer = ggml_backend_alloc_ctx_tensors_from_buft(item.second.get(), item.first); + } + if (buffer == nullptr) { + throw std::runtime_error(format( + "DeepSeek V4.1 failed to allocate %s expert cache buffer", + ggml_backend_buft_name(item.first))); + } + ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + buffers.emplace_back(buffer); + } + + if (!params.no_alloc) { + llama_expert_store_params store_params; + store_params.cache_bytes = bytes_staging; + store_params.cache_slots = checked_mul(params.cache_slots, 3, "DeepSeek V4.1 expert staging slot count overflow"); + store_params.direct_io = params.direct_io; + store_params.allow_buffered_io = params.allow_buffered_io; + store = std::make_unique(std::move(tensors), store_params); + } + + if (!this->upload) { + this->upload = [](ggml_tensor * tensor, size_t offset, const void * data, size_t size) { + ggml_backend_tensor_set(tensor, data, offset, size); + }; + } + } + + const llama_expert_store::payload & payload( + const std::vector & payloads, + int32_t layer, + llama_expert_projection projection, + int32_t expert_id) const { + for (const auto & payload : payloads) { + if (payload.layer == layer && payload.projection == projection && payload.expert_id == expert_id) { + return payload; + } + } + throw std::runtime_error("DeepSeek V4.1 expert store did not return a requested payload"); + } + + std::vector remap(int32_t layer, const std::vector & expert_ids) { + std::lock_guard lock(mutex); + if (store == nullptr) { + throw std::runtime_error("DeepSeek V4.1 expert cache is metadata-only"); + } + if (layer < 0 || layer >= (int32_t) layers.size()) { + throw std::invalid_argument("DeepSeek V4.1 expert layer is invalid"); + } + layer_state & state = layers[layer]; + if (state.lease) { + if (state.active_ids == expert_ids) { + return state.active_remap; + } + throw std::runtime_error(format("DeepSeek V4.1 expert layer %d still has an in-flight lease", layer)); + } + + std::vector unique = expert_ids; + for (int32_t expert_id : unique) { + if (expert_id < 0 || expert_id >= (int32_t) LLAMA_DSV41_N_EXPERT) { + throw std::runtime_error("DeepSeek V4.1 selected expert ID is out of range"); + } + } + std::sort(unique.begin(), unique.end()); + unique.erase(std::unique(unique.begin(), unique.end()), unique.end()); + if (unique.size() > params.cache_slots) { + throw std::runtime_error(format( + "DeepSeek V4.1 selected expert union has %zu entries, cache has %zu slots", + unique.size(), params.cache_slots)); + } + + std::map resident; + std::vector empty; + std::vector victims; + for (uint32_t slot = 0; slot < state.slots.size(); ++slot) { + const logical_slot & entry = state.slots[slot]; + if (entry.expert_id >= 0) { + resident.emplace(entry.expert_id, slot); + if (entry.pins == 0 && !std::binary_search(unique.begin(), unique.end(), entry.expert_id)) { + victims.push_back(slot); + } + } else { + empty.push_back(slot); + } + } + std::sort(victims.begin(), victims.end(), [&](uint32_t a, uint32_t b) { + if (state.slots[a].last_use != state.slots[b].last_use) { + return state.slots[a].last_use < state.slots[b].last_use; + } + return a < b; + }); + + std::vector misses; + for (int32_t expert_id : unique) { + if (resident.count(expert_id) == 0) { + misses.push_back(expert_id); + } + } + if (misses.size() > empty.size() + victims.size()) { + throw std::runtime_error("DeepSeek V4.1 expert cache capacity is pinned"); + } + + std::vector targets = empty; + targets.insert(targets.end(), victims.begin(), victims.end()); + targets.resize(misses.size()); + std::sort(targets.begin(), targets.end()); + + auto lease = std::make_unique(); + if (!misses.empty()) { + *lease = store->acquire({ + { layer, LLAMA_EXPERT_PROJECTION_GATE, misses }, + { layer, LLAMA_EXPERT_PROJECTION_UP, misses }, + { layer, LLAMA_EXPERT_PROJECTION_DOWN, misses }, + }); + const auto payloads = lease->payloads(); + try { + for (size_t i = 0; i < misses.size(); ++i) { + const uint32_t slot = targets[i]; + for (llama_expert_projection projection : { + LLAMA_EXPERT_PROJECTION_GATE, + LLAMA_EXPERT_PROJECTION_UP, + LLAMA_EXPERT_PROJECTION_DOWN }) { + const auto & item = payload(payloads, layer, projection, misses[i]); + ggml_tensor * tensor = state.tensors[projection_index(projection)]; + upload(tensor, slot*tensor->nb[2], item.data, item.size); + } + } + } catch (...) { + for (uint32_t slot : targets) { + state.slots[slot] = {}; + } + throw; + } + + for (size_t i = 0; i < misses.size(); ++i) { + const uint32_t slot = targets[i]; + state.slots[slot].expert_id = misses[i]; + resident[misses[i]] = slot; + } + } + + std::vector pinned_slots; + pinned_slots.reserve(unique.size()); + for (int32_t expert_id : unique) { + const uint32_t slot = resident.at(expert_id); + pinned_slots.push_back(slot); + } + + std::vector result; + result.reserve(expert_ids.size()); + for (int32_t expert_id : expert_ids) { + result.push_back((int32_t) resident.at(expert_id)); + } + std::vector active_ids = expert_ids; + std::vector active_remap = result; + + if (before_publish) { + before_publish(); + } + for (uint32_t slot : pinned_slots) { + logical_slot & entry = state.slots[slot]; + entry.last_use = ++use_clock; + entry.pins++; + } + state.pinned_slots = std::move(pinned_slots); + state.lease = std::move(lease); + state.active_ids = std::move(active_ids); + state.active_remap = std::move(active_remap); + return result; + } + + void release(int32_t layer) { + std::lock_guard lock(mutex); + if (layer < 0 || layer >= (int32_t) layers.size()) { + return; + } + layer_state & state = layers[layer]; + for (uint32_t slot : state.pinned_slots) { + if (state.slots[slot].pins == 0) { + GGML_ABORT("DeepSeek V4.1 expert slot pin underflow"); + } + state.slots[slot].pins--; + } + state.pinned_slots.clear(); + state.lease.reset(); + state.active_ids.clear(); + state.active_remap.clear(); + } +}; + +llama_dsv41_expert_runtime::llama_dsv41_expert_runtime( + std::vector tensors, + const llama_dsv41_expert_runtime_params & params, + buft_selector select_buft, + upload_fn upload, + publish_fn before_publish) + : pimpl(std::make_unique( + this, std::move(tensors), params, std::move(select_buft), std::move(upload), std::move(before_publish))) { +} + +llama_dsv41_expert_runtime::~llama_dsv41_expert_runtime() = default; + +std::vector llama_dsv41_expert_runtime::remap( + int32_t layer, const std::vector & expert_ids) { + return pimpl->remap(layer, expert_ids); +} + +void llama_dsv41_expert_runtime::release(int32_t layer) { + pimpl->release(layer); +} + +void llama_dsv41_expert_runtime::release_all() { + for (int32_t il = 0; il < (int32_t) LLAMA_DSV41_N_LAYER; ++il) { + pimpl->release(il); + } +} + +void llama_dsv41_expert_runtime::release_all_after_sync(ggml_backend_sched_t sched) { + if (sched != nullptr) { + ggml_backend_sched_synchronize(sched); + } + release_all(); +} + +ggml_tensor * llama_dsv41_expert_runtime::cache_tensor( + int32_t layer, llama_expert_projection projection) const { + if (layer < 0 || layer >= (int32_t) pimpl->layers.size()) { + throw std::invalid_argument("DeepSeek V4.1 expert layer is invalid"); + } + return pimpl->layers[layer].tensors[projection_index(projection)]; +} + +size_t llama_dsv41_expert_runtime::cache_slots() const { + return pimpl->params.cache_slots; +} + +size_t llama_dsv41_expert_runtime::cache_bytes() const { + return pimpl->bytes_cache; +} + +size_t llama_dsv41_expert_runtime::staging_bytes() const { + return pimpl->bytes_staging; +} + +void llama_dsv41_expert_runtime::set_error(const std::string & error) { + std::lock_guard lock(pimpl->mutex); + if (pimpl->error.empty()) { + pimpl->error = error; + } +} + +std::string llama_dsv41_expert_runtime::consume_error() { + std::lock_guard lock(pimpl->mutex); + std::string result; + result.swap(pimpl->error); + return result; +} + +void llama_dsv41_expert_runtime::acquire_context() { + std::lock_guard lock(pimpl->mutex); + if (pimpl->context_active) { + throw std::runtime_error("DeepSeek V4.1 bounded expert runtime supports one context per model"); + } + pimpl->context_active = true; +} + +void llama_dsv41_expert_runtime::release_context() { + std::lock_guard lock(pimpl->mutex); + pimpl->context_active = false; +} + +std::vector llama_dsv41_register_expert_tensors( + const std::function & ne)> & register_tensor) { + if (!register_tensor) { + throw std::invalid_argument("DeepSeek V4.1 expert tensor registrar is empty"); + } + std::vector result; + result.reserve(LLAMA_DSV41_N_LAYER*3); + for (int32_t il = 0; il < (int32_t) LLAMA_DSV41_N_LAYER; ++il) { + result.push_back(register_tensor( + "blk." + std::to_string(il) + ".ffn_gate_exps.weight", + il, LLAMA_EXPERT_PROJECTION_GATE, + { LLAMA_DSV41_N_EMBD, LLAMA_DSV41_N_FF_EXP, LLAMA_DSV41_N_EXPERT })); + result.push_back(register_tensor( + "blk." + std::to_string(il) + ".ffn_up_exps.weight", + il, LLAMA_EXPERT_PROJECTION_UP, + { LLAMA_DSV41_N_EMBD, LLAMA_DSV41_N_FF_EXP, LLAMA_DSV41_N_EXPERT })); + result.push_back(register_tensor( + "blk." + std::to_string(il) + ".ffn_down_exps.weight", + il, LLAMA_EXPERT_PROJECTION_DOWN, + { LLAMA_DSV41_N_FF_EXP, LLAMA_DSV41_N_EMBD, LLAMA_DSV41_N_EXPERT })); + } + return result; +} + +static void dsv41_expert_remap_callback( + ggml_tensor * dst, + const ggml_tensor * src, + int, + int, + void * userdata) { + auto * state = static_cast(userdata); + GGML_ASSERT(dst->type == GGML_TYPE_I32 && src->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_is_contiguous(dst) && ggml_is_contiguous(src)); + const int32_t * ids = static_cast(src->data); + const size_t count = ggml_nelements(src); + try { + const std::vector remapped = state->runtime->remap( + state->layer, std::vector(ids, ids + count)); + memcpy(dst->data, remapped.data(), remapped.size()*sizeof(int32_t)); + } catch (const std::exception & error) { + std::fill_n(static_cast(dst->data), count, 0); + state->runtime->set_error(error.what()); + } +} + +static void dsv41_expert_release_callback( + ggml_tensor * dst, + const ggml_tensor * src, + int, + int, + void * userdata) { + auto * state = static_cast(userdata); + memcpy(dst->data, src->data, ggml_nbytes(src)); + state->runtime->release(state->layer); +} + +ggml_tensor * llama_dsv41_build_expert_remap( + ggml_context * ctx, + ggml_tensor * selected_experts, + llama_dsv41_expert_runtime & runtime, + int32_t layer, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu) { + if (ctx == nullptr || selected_experts == nullptr || selected_experts->type != GGML_TYPE_I32) { + throw std::invalid_argument("DeepSeek V4.1 expert remap input is invalid"); + } + require_local_cpu(sched, backend_cpu); + ggml_tensor * original = ggml_cont(ctx, selected_experts); + ggml_tensor * remapped = ggml_map_custom1( + ctx, original, dsv41_expert_remap_callback, 1, &runtime.pimpl->callbacks.at(layer)); + ggml_backend_sched_set_tensor_backend(sched, remapped, backend_cpu); + return remapped; +} + +ggml_tensor * llama_dsv41_build_expert_release( + ggml_context * ctx, + ggml_tensor * experts, + llama_dsv41_expert_runtime & runtime, + int32_t layer, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu) { + if (ctx == nullptr || experts == nullptr) { + throw std::invalid_argument("DeepSeek V4.1 expert release input is invalid"); + } + require_local_cpu(sched, backend_cpu); + ggml_tensor * completion = ggml_sum(ctx, experts); + completion = ggml_map_custom1( + ctx, completion, dsv41_expert_release_callback, 1, &runtime.pimpl->callbacks.at(layer)); + ggml_backend_sched_set_tensor_backend(sched, completion, backend_cpu); + return completion; +} diff --git a/src/llama-dsv41-expert.h b/src/llama-dsv41-expert.h new file mode 100644 index 000000000000..82ce5037e8fc --- /dev/null +++ b/src/llama-dsv41-expert.h @@ -0,0 +1,89 @@ +#pragma once + +#include "llama-expert-store.h" + +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include +#include + +struct ggml_cgraph; +struct ggml_context; +struct ggml_tensor; + +struct llama_dsv41_expert_runtime_params { + size_t cache_bytes = 0; + size_t cache_slots = 0; + bool direct_io = true; + bool allow_buffered_io = false; + bool no_alloc = false; +}; + +struct llama_dsv41_expert_runtime { + using buft_selector = std::function; + using upload_fn = std::function; + using publish_fn = std::function; + + llama_dsv41_expert_runtime( + std::vector tensors, + const llama_dsv41_expert_runtime_params & params, + buft_selector select_buft, + upload_fn upload = {}, + publish_fn before_publish = {}); + ~llama_dsv41_expert_runtime(); + + llama_dsv41_expert_runtime(const llama_dsv41_expert_runtime &) = delete; + llama_dsv41_expert_runtime & operator=(const llama_dsv41_expert_runtime &) = delete; + + std::vector remap(int32_t layer, const std::vector & expert_ids); + void acquire_context(); + void release_context(); + void release(int32_t layer); + void release_all(); + void release_all_after_sync(ggml_backend_sched_t sched); + + ggml_tensor * cache_tensor(int32_t layer, llama_expert_projection projection) const; + size_t cache_slots() const; + size_t cache_bytes() const; + size_t staging_bytes() const; + + void set_error(const std::string & error); + std::string consume_error(); + +private: + friend ggml_tensor * llama_dsv41_build_expert_remap( + ggml_context *, ggml_tensor *, llama_dsv41_expert_runtime &, int32_t, ggml_backend_sched_t, ggml_backend_t); + friend ggml_tensor * llama_dsv41_build_expert_release( + ggml_context *, ggml_tensor *, llama_dsv41_expert_runtime &, int32_t, ggml_backend_sched_t, ggml_backend_t); + + struct impl; + std::unique_ptr pimpl; +}; + +std::vector llama_dsv41_register_expert_tensors( + const std::function & ne)> & register_tensor); + +ggml_tensor * llama_dsv41_build_expert_remap( + ggml_context * ctx, + ggml_tensor * selected_experts, + llama_dsv41_expert_runtime & runtime, + int32_t layer, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu); + +ggml_tensor * llama_dsv41_build_expert_release( + ggml_context * ctx, + ggml_tensor * experts, + llama_dsv41_expert_runtime & runtime, + int32_t layer, + ggml_backend_sched_t sched, + ggml_backend_t backend_cpu); diff --git a/src/llama-expert-store.cpp b/src/llama-expert-store.cpp new file mode 100644 index 000000000000..c8597e52845f --- /dev/null +++ b/src/llama-expert-store.cpp @@ -0,0 +1,646 @@ +#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, 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; + 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(); + } + + 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 < need) { +#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, 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())); + } + 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) { + 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); + *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()); + + uint64_t bytes_read = 0; + std::vector staged; + staged.reserve(misses.size()); + for (const expert_key & key : misses) { + const auto & tensor = pimpl->get_tensor(key); + staged.push_back(pimpl->read_expert(tensor, key.expert_id, &bytes_read)); + } + + for (uint32_t victim : victims) { + auto & entry = pimpl->slots[victim]; + resident.erase(entry.key); + pimpl->bytes_resident -= entry.bytes.size; + entry.bytes.clear(); + entry.occupied = false; + entry.tensor = nullptr; + entry.last_use = 0; + pimpl->counters.evictions++; + } + + for (size_t i = 0; i < misses.size(); ++i) { + const expert_key & key = misses[i]; + const auto & tensor = pimpl->get_tensor(key); + auto & entry = pimpl->slots[target_slots[i]]; + entry.bytes = std::move(staged[i]); + entry.occupied = true; + entry.key = key; + entry.tensor = &tensor; + entry.pins = 0; + pimpl->bytes_resident += entry.bytes.size; + resident[entry.key] = target_slots[i]; + } + + auto lease_impl = std::make_unique(); + 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)); + } + lease_impl->pinned_slots.reserve(unique_keys.size()); + for (const expert_key & key : unique_keys) { + 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->store = pimpl; + + 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; +} + +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 new file mode 100644 index 000000000000..10ce3a1a19d7 --- /dev/null +++ b/src/llama-expert-store.h @@ -0,0 +1,106 @@ +#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; + size_t file_index = 0; + 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; + bool allow_buffered_io = false; // opt-in only; page-cache bytes are outside cache_bytes +}; + +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; + bool direct_io_active() const; + +private: + struct impl; + std::shared_ptr pimpl; +}; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 010c40bcc500..9a7219184874 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1547,15 +1547,17 @@ ggml_tensor * llm_graph_context::build_lora_mm_id( ggml_tensor * w, // ggml_tensor * as ggml_tensor * cur, // ggml_tensor * b ggml_tensor * ids, - ggml_tensor * w_s) const { + ggml_tensor * w_s, + ggml_tensor * ids_scale) const { ggml_tensor * res = ggml_mul_mat_id(ctx0, w, cur, ids); if (w_s) { + ids_scale = ids_scale ? ids_scale : ids; const int64_t n_expert = w_s->ne[0]; const int64_t n_tokens = cur->ne[2]; ggml_tensor * s = ggml_reshape_3d(ctx0, w_s, 1, n_expert, 1); s = ggml_repeat_4d(ctx0, s, 1, n_expert, n_tokens, 1); - s = ggml_get_rows(ctx0, s, ids); + s = ggml_get_rows(ctx0, s, ids_scale); res = ggml_mul(ctx0, res, s); } for (const auto & lora : *loras) { @@ -1947,6 +1949,26 @@ ggml_tensor * llm_graph_context::build_ffn( return cur; } +llm_moe_expert_ids llm_build_moe_expert_ids( + ggml_context * ctx, + llm_arch arch, + ggml_tensor * selected, + ggml_tensor * lookup, + int64_t n_expert, + int64_t n_expert_total, + uint32_t n_group_experts) { + if (lookup != nullptr) { + return { selected, lookup }; + } + if (arch == LLM_ARCH_GROVEMOE && n_expert != n_expert_total) { + GGML_ASSERT(n_group_experts > 0); + // TODO: Use scalar div instead when/if implemented + ggml_tensor * f_sel = ggml_cast(ctx, selected, GGML_TYPE_F32); + lookup = ggml_cast(ctx, ggml_scale(ctx, f_sel, 1.0f / float(n_group_experts)), GGML_TYPE_I32); + } + return { selected, lookup ? lookup : selected }; +} + ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * cur, ggml_tensor * gate_inp, @@ -1966,7 +1988,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * up_exps_s, ggml_tensor * gate_exps_s, ggml_tensor * down_exps_s, - ggml_tensor * selected_experts_in) const { + ggml_tensor * selected_experts_in, + ggml_tensor * selected_experts_lookup_in) const { return build_moe_ffn( cur, gate_inp, /* gate_inp_b */ nullptr, @@ -1987,7 +2010,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( up_exps_s, gate_exps_s, down_exps_s, - selected_experts_in + selected_experts_in, + selected_experts_lookup_in ); } @@ -2015,7 +2039,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( ggml_tensor * up_exps_s, ggml_tensor * gate_exps_s, ggml_tensor * down_exps_s, - ggml_tensor * selected_experts_in) const { + ggml_tensor * selected_experts_in, + ggml_tensor * selected_experts_lookup_in) const { const int64_t n_embd = cur->ne[0]; const int64_t n_tokens = cur->ne[1]; const bool weight_before_ffn = arch == LLM_ARCH_LLAMA4; // for llama4, we apply the sigmoid-ed weights before the FFN @@ -2112,10 +2137,13 @@ ggml_tensor * llm_graph_context::build_moe_ffn( } cb(selected_experts, "ffn_moe_topk", il); + const llm_moe_expert_ids expert_ids = llm_build_moe_expert_ids( + ctx0, arch, selected_experts, selected_experts_lookup_in, + n_expert, hparams.n_expert, hparams.n_group_experts); + selected_experts = expert_ids.routing; + ggml_tensor * selected_experts_lookup = expert_ids.lookup; + if (arch == LLM_ARCH_GROVEMOE && n_expert != hparams.n_expert) { - // TODO: Use scalar div instead when/if implemented - ggml_tensor * f_sel = ggml_cast(ctx0, selected_experts, GGML_TYPE_F32); - selected_experts = ggml_cast(ctx0, ggml_scale(ctx0, f_sel, 1.0f / float(hparams.n_group_experts)), GGML_TYPE_I32); probs = ggml_reshape_3d(ctx0, probs, 1, hparams.n_expert, n_tokens); } else { probs = ggml_reshape_3d(ctx0, probs, 1, n_expert, n_tokens); @@ -2169,7 +2197,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( if (gate_up_exps) { // merged gate_up path: one mul_mat_id, then split into gate and up views - ggml_tensor * gate_up = build_lora_mm_id(gate_up_exps, cur, selected_experts, up_exps_s); // [n_ff*2, n_expert_used, n_tokens] + ggml_tensor * gate_up = build_lora_mm_id( + gate_up_exps, cur, selected_experts_lookup, up_exps_s, selected_experts); // [n_ff*2, n_expert_used, n_tokens] cb(gate_up, "ffn_moe_gate_up", il); if (up_exps_s) { @@ -2188,7 +2217,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( cb(up, "ffn_moe_up", il); } else { // separate gate and up path - up = build_lora_mm_id(up_exps, cur, selected_experts, up_exps_s); // [n_ff, n_expert_used, n_tokens] + up = build_lora_mm_id( + up_exps, cur, selected_experts_lookup, up_exps_s, selected_experts); // [n_ff, n_expert_used, n_tokens] cb(up, "ffn_moe_up", il); if (up_exps_s) { @@ -2201,7 +2231,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( } if (gate_exps) { - cur = build_lora_mm_id(gate_exps, cur, selected_experts, gate_exps_s); // [n_ff, n_expert_used, n_tokens] + cur = build_lora_mm_id( + gate_exps, cur, selected_experts_lookup, gate_exps_s, selected_experts); // [n_ff, n_expert_used, n_tokens] cb(cur, "ffn_moe_gate", il); } else { cur = up; @@ -2302,7 +2333,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( GGML_ABORT("fatal error"); } - experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s); // [n_embd, n_expert_used, n_tokens] + experts = build_lora_mm_id( + down_exps, cur, selected_experts_lookup, down_exps_s, selected_experts); // [n_embd, n_expert_used, n_tokens] cb(experts, "ffn_moe_down", il); if (down_exps_s) { diff --git a/src/llama-graph.h b/src/llama-graph.h index ce68832ce47a..e885bf4a8513 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -75,6 +75,20 @@ enum llm_norm_type { LLM_NORM_GROUP, }; +struct llm_moe_expert_ids { + ggml_tensor * routing; + ggml_tensor * lookup; +}; + +llm_moe_expert_ids llm_build_moe_expert_ids( + ggml_context * ctx, + llm_arch arch, + ggml_tensor * selected, + ggml_tensor * lookup, + int64_t n_expert, + int64_t n_expert_total, + uint32_t n_group_experts); + // TODO: tmp - need something better to pass the data from the encoder to the decoder struct llama_cross { // the output embeddings from the encoder as a ggml tensor @@ -1060,7 +1074,8 @@ struct llm_graph_context { ggml_tensor * w, // ggml_tensor * as ggml_tensor * cur, // ggml_tensor * b ggml_tensor * ids, - ggml_tensor * w_s = nullptr) const; + ggml_tensor * w_s = nullptr, + ggml_tensor * ids_scale = nullptr) const; ggml_tensor * build_norm( ggml_tensor * cur, @@ -1129,7 +1144,8 @@ struct llm_graph_context { ggml_tensor * up_exps_s = nullptr, ggml_tensor * gate_exps_s = nullptr, ggml_tensor * down_exps_s = nullptr, - ggml_tensor * selected_experts_in = nullptr) const; + ggml_tensor * selected_experts_in = nullptr, + ggml_tensor * selected_experts_lookup_in = nullptr) const; ggml_tensor * build_moe_ffn( ggml_tensor * cur, @@ -1155,7 +1171,8 @@ struct llm_graph_context { ggml_tensor * up_exps_s = nullptr, ggml_tensor * gate_exps_s = nullptr, ggml_tensor * down_exps_s = nullptr, - ggml_tensor * selected_experts_in = nullptr) const; + ggml_tensor * selected_experts_in = nullptr, + ggml_tensor * selected_experts_lookup_in = nullptr) const; // // inputs diff --git a/src/llama-mmap.cpp b/src/llama-mmap.cpp index 715a6e3548e6..141640f9d192 100644 --- a/src/llama-mmap.cpp +++ b/src/llama-mmap.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -172,7 +173,13 @@ 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; + } + + void discard_cache(size_t offset, size_t length) const { + GGML_UNUSED(offset); + GGML_UNUSED(length); } ~impl() { @@ -374,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); @@ -408,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 @@ -439,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; @@ -457,27 +477,38 @@ 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 & lazy_ranges) { + impl(struct llama_file * file, size_t prefetch, bool numa, + 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; 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(strict_exclusion); + const int file_advice = sequential ? POSIX_FADV_SEQUENTIAL : POSIX_FADV_RANDOM; + 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 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; } +#else + GGML_UNUSED(file_advice_override); #endif addr = mmap(NULL, file->size(), PROT_READ, flags, fd, 0); if (addr == MAP_FAILED) { @@ -497,12 +528,11 @@ struct llama_mmap::impl { } }; - if (prefetch > 0) { - for (const auto & range : ranges_complement(lazy_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 : lazy_ranges) { + for (const auto & range : excluded_ranges) { advise(range.first, range.second, POSIX_MADV_RANDOM, "POSIX_MADV_RANDOM"); } if (numa) { @@ -573,8 +603,11 @@ 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, bool strict_exclusion, + llama_mmap::file_advice_override file_advice_override) { GGML_UNUSED(numa); + GGML_UNUSED(file_advice_override); size = file->size(); @@ -604,7 +637,8 @@ 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 : + 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); @@ -642,11 +676,15 @@ 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, bool strict_exclusion, + llama_mmap::file_advice_override file_advice_override) { GGML_UNUSED(file); GGML_UNUSED(prefetch); GGML_UNUSED(numa); - GGML_UNUSED(lazy_ranges); + GGML_UNUSED(excluded_ranges); + GGML_UNUSED(strict_exclusion); + GGML_UNUSED(file_advice_override); throw std::runtime_error("mmap not supported"); } @@ -664,9 +702,22 @@ 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, 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) { + 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; } void * llama_mmap::addr() const { return pimpl->addr; } @@ -782,6 +833,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; @@ -814,6 +877,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 cc28c8a73fa5..c64c5acfb6a1 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; @@ -44,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 & lazy_ranges = {}); + const ranges & excluded_ranges = {}, bool strict_exclusion = false, + file_advice_override file_advice = nullptr); ~llama_mmap(); size_t size() const; @@ -56,6 +60,9 @@ struct llama_mmap { void unmap_fragment(size_t first, size_t last); static const bool SUPPORTED; + 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; @@ -69,6 +76,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 0c609c2486b0..d8cb220b9dfd 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1132,6 +1132,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) { @@ -1408,6 +1416,41 @@ 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.file_index = 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)); @@ -1447,10 +1490,17 @@ 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, !external_ranges.empty()); + 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)); @@ -1461,7 +1511,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); + } } } @@ -1512,6 +1564,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) { @@ -1538,12 +1592,44 @@ bool llama_model_loader::load_all_data( // 64MB works well for NVMe drives const size_t buffer_size = alignment != 1 ? 64 * 1024 * 1024 + 2 * alignment : 1 * 1024 * 1024; - std::vector host_buffers; - std::vector events; - std::vector host_ptrs; + struct async_upload_resources { + std::vector host_buffers; + std::vector events; + std::vector host_ptrs; + ggml_backend_t backend = nullptr; + + void reset() { + for (auto * event : events) { + if (backend != nullptr) { + ggml_backend_event_synchronize(event); + } + ggml_backend_event_free(event); + } + events.clear(); + for (auto * buffer : host_buffers) { + ggml_backend_buffer_free(buffer); + } + host_buffers.clear(); + host_ptrs.clear(); + ggml_backend_free(backend); + backend = nullptr; + } + + ~async_upload_resources() { + reset(); + } + } async_upload; + async_upload.host_buffers.reserve(n_buffers); + async_upload.events.reserve(n_buffers); + async_upload.host_ptrs.reserve(n_buffers); + + auto & host_buffers = async_upload.host_buffers; + auto & events = async_upload.events; + auto & host_ptrs = async_upload.host_ptrs; + auto & upload_backend = async_upload.backend; 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) { + upload_backend = [&](const char * func) -> ggml_backend_t { + if (load_from_mmap || check_tensors) { return nullptr; } // When not using mmaped io use async uploads from pinned memory to GPU memory. @@ -1590,6 +1676,7 @@ bool llama_model_loader::load_all_data( if (!buf) { LLAMA_LOG_DEBUG("%s: failed to allocate host buffer for async uploads for device %s\n", func, ggml_backend_dev_name(dev)); + async_upload.reset(); return nullptr; } @@ -1600,6 +1687,7 @@ bool llama_model_loader::load_all_data( if (!event) { LLAMA_LOG_DEBUG("%s: failed to create event for async uploads for device %s\n", func, ggml_backend_dev_name(dev)); + async_upload.reset(); return nullptr; } @@ -1610,6 +1698,7 @@ bool llama_model_loader::load_all_data( if (!backend) { LLAMA_LOG_DEBUG("%s: failed to initialize backend for device %s for async uploads\n", func, ggml_backend_dev_name(dev)); + async_upload.reset(); return nullptr; } @@ -1630,7 +1719,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); @@ -1656,7 +1745,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); @@ -1764,20 +1853,15 @@ bool llama_model_loader::load_all_data( } } } + if (discard_file_cache) { + file->discard_cache(weight->offs, n_size); + } } size_done += n_size; } - // free temporary resources used for async uploads - for (auto * event : events) { - ggml_backend_event_synchronize(event); - ggml_backend_event_free(event); - } - for (auto * buf : host_buffers) { - ggml_backend_buffer_free(buf); - } - ggml_backend_free(upload_backend); + async_upload.reset(); // check validation results bool validation_failed = false; diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index 72bbd53e7d03..5381162065f3 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); @@ -256,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 9da93a11151e..e02446d1493c 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -1735,14 +1735,22 @@ 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); + // 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 && !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; // 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(); @@ -1778,9 +1786,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, @@ -1797,6 +1820,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); } @@ -1833,7 +1865,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()) { @@ -1865,21 +1898,24 @@ 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, use_mlock ? &pimpl->mlock_mmaps : NULL, 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; } } - if (use_mmap_buffer) { + if (keep_mappings) { for (auto & mapping : ml.mappings) { pimpl->mappings.emplace_back(std::move(mapping)); } @@ -2238,6 +2274,20 @@ ggml_backend_buffer_type_t llama_model::select_buft(int il) const { }); } +ggml_backend_buffer_type_t llama_model::select_moe_buft( + int il, enum ggml_type type, int64_t ne0, int64_t ne1, int64_t ne2) const { + return ::select_buft( + *pimpl->dev_layer.at(il).buft_list, + [&](ggml_context * ctx) { + const int64_t n_expert_used = hparams.n_expert_used_max(); + GGML_ASSERT(n_expert_used > 0); + ggml_tensor * weight = ggml_new_tensor_3d(ctx, type, ne0, ne1, ne2); + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, ne0, n_expert_used, 512); + ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_expert_used, 512); + return ggml_mul_mat_id(ctx, weight, input, ids); + }); +} + bool llama_model::has_tensor_overrides() const { return pimpl->has_tensor_overrides; } @@ -2803,6 +2853,8 @@ llama_model_params llama_model_default_params() { /*.main_gpu =*/ 0, /*.ple_io_threads =*/ 64, /*.ple_cache_mb =*/ 256, + /*.expert_cache_bytes =*/ 0, + /*.expert_cache_slots =*/ 0, /*.tensor_split =*/ nullptr, /*.progress_callback =*/ nullptr, /*.progress_callback_user_data =*/ nullptr, diff --git a/src/llama-model.h b/src/llama-model.h index c5895ab9b41f..2bc5b96e4d92 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -748,6 +748,8 @@ struct llama_model { ggml_backend_dev_t dev_output() const; ggml_backend_buffer_type_t select_buft(int il) const; + ggml_backend_buffer_type_t select_moe_buft( + int il, enum ggml_type type, int64_t ne0, int64_t ne1, int64_t ne2) const; bool has_tensor_overrides() const; @@ -767,6 +769,13 @@ struct llama_model { virtual void load_vocab (llama_model_loader & ml) = 0; virtual bool load_tensors(llama_model_loader & ml) = 0; // returns false if cancelled by progress_callback + virtual bool requires_synchronous_graph() const { return false; } + virtual std::string consume_runtime_error() const { return {}; } + virtual void release_runtime_work() const {} + virtual void release_runtime_work_after_sync(ggml_backend_sched_t) const { release_runtime_work(); } + virtual void acquire_runtime_context() const {} + virtual void release_runtime_context() const {} + // model must define these virtual void load_arch_hparams(llama_model_loader & ml) = 0; virtual void load_arch_tensors(llama_model_loader & ml) = 0; diff --git a/src/models/deepseek41.cpp b/src/models/deepseek41.cpp index 4312290507c0..ba7923663f53 100644 --- a/src/models/deepseek41.cpp +++ b/src/models/deepseek41.cpp @@ -1,5 +1,6 @@ #include "llama-dsv41.h" #include "llama-dsv41-engram.h" +#include "llama-dsv41-expert.h" #include "llama-hparams.h" #include "models.h" @@ -148,7 +149,7 @@ void llama_model_deepseek41::load_arch_hparams(llama_model_loader & ml) { type = LLM_TYPE_UNKNOWN; } -[[noreturn]] void llama_model_deepseek41::load_arch_tensors(llama_model_loader & ml) { +void llama_model_deepseek41::load_arch_tensors(llama_model_loader & ml) { LLAMA_LOAD_LOCALS; const int64_t q_lora_rank = hparams.n_lora_q; @@ -161,6 +162,31 @@ void llama_model_deepseek41::load_arch_hparams(llama_model_loader & ml) { const int64_t hc_dim = hc_mult*n_embd; const int64_t hc_mix_dim = (2 + hc_mult)*hc_mult; + const std::vector expert_tensors = + llama_dsv41_register_expert_tensors([&](const std::string & name, + int32_t layer, + llama_expert_projection projection, + const std::initializer_list & ne) { + return ml.register_external_tensor(name, layer, projection, ne); + }); + if (params.expert_cache_bytes == 0 || params.expert_cache_slots <= 0) { + throw std::runtime_error( + "DeepSeek V4.1 requires non-zero expert_cache_bytes and expert_cache_slots before tensor allocation"); + } + llama_dsv41_expert_runtime_params expert_params; + expert_params.cache_bytes = params.expert_cache_bytes; + expert_params.cache_slots = params.expert_cache_slots; + expert_params.direct_io = true; + expert_params.allow_buffered_io = false; + expert_params.no_alloc = ml.no_alloc; + experts = std::make_shared( + expert_tensors, + expert_params, + [this](const llama_expert_store_tensor & tensor) { + return select_moe_buft( + tensor.layer, tensor.type, tensor.ne[0], tensor.ne[1], params.expert_cache_slots); + }); + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0); output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, 0); @@ -203,9 +229,9 @@ void llama_model_deepseek41::load_arch_hparams(llama_model_loader & ml) { layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", il), { n_expert }, 0); layer.ffn_exp_probs_b_vl = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B_VL, "bias", il), { n_expert }, TENSOR_NOT_REQUIRED); layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", il), { n_embd }, 0); - layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", il), { n_embd, n_ff_exp, n_expert }, 0); - layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, 0); - layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", il), { n_embd, n_ff_exp, n_expert }, 0); + layer.ffn_gate_exps = experts->cache_tensor(il, LLAMA_EXPERT_PROJECTION_GATE); + layer.ffn_down_exps = experts->cache_tensor(il, LLAMA_EXPERT_PROJECTION_DOWN); + layer.ffn_up_exps = experts->cache_tensor(il, LLAMA_EXPERT_PROJECTION_UP); layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_exp*n_expert_shared }, 0); layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_exp*n_expert_shared, n_embd }, 0); layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_exp*n_expert_shared }, 0); @@ -245,9 +271,38 @@ void llama_model_deepseek41::load_arch_hparams(llama_model_loader & ml) { } } - throw std::runtime_error( - std::string("DeepSeek V4.1 Engram metadata and disk extents are valid, but execution is blocked: ") + - llama_dsv41_runtime_dependency_error()); +} + +bool llama_model_deepseek41::requires_synchronous_graph() const { + return experts != nullptr; +} + +std::string llama_model_deepseek41::consume_runtime_error() const { + return experts ? experts->consume_error() : std::string(); +} + +void llama_model_deepseek41::release_runtime_work() const { + if (experts) { + experts->release_all(); + } +} + +void llama_model_deepseek41::release_runtime_work_after_sync(ggml_backend_sched_t sched) const { + if (experts) { + experts->release_all_after_sync(sched); + } +} + +void llama_model_deepseek41::acquire_runtime_context() const { + if (experts) { + experts->acquire_context(); + } +} + +void llama_model_deepseek41::release_runtime_context() const { + if (experts) { + experts->release_context(); + } } [[noreturn]] std::unique_ptr llama_model_deepseek41::build_arch_graph(const llm_graph_params &) const { diff --git a/src/models/models.h b/src/models/models.h index 8f106aa98753..ef05debe660c 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -9,6 +9,7 @@ #include class llama_memory_hybrid_idx_context; +struct llama_dsv41_expert_runtime; // ref: https://github.com/ggml-org/llama.cpp/pull/28068 static inline ggml_tensor * build_gdn_l2_norm(ggml_context * ctx, ggml_tensor * x, float eps) { @@ -1318,9 +1319,16 @@ struct llama_model_deepseek41 : public llama_model_deepseek4 { struct engram_model; std::shared_ptr engram; + std::shared_ptr experts; void load_arch_hparams(llama_model_loader & ml) override; - [[noreturn]] void load_arch_tensors(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + bool requires_synchronous_graph() const override; + std::string consume_runtime_error() const override; + void release_runtime_work() const override; + void release_runtime_work_after_sync(ggml_backend_sched_t sched) const override; + void acquire_runtime_context() const override; + void release_runtime_context() const override; [[noreturn]] std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 40423d08ca1b..2c4242b749c9 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -198,6 +198,8 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) llama_build_and_test(test-deepseek41-schema.cpp) llama_build_and_test(test-deepseek41-engram.cpp) + llama_build_and_test(test-deepseek41-expert.cpp) + llama_build_and_test(test-expert-store.cpp) llama_build_and_test(test-deepseek41-runtime.cpp) llama_build_and_test(test-engram.cpp) llama_build(test-llama-archs.cpp) diff --git a/tests/test-deepseek41-expert.cpp b/tests/test-deepseek41-expert.cpp new file mode 100644 index 000000000000..269922196301 --- /dev/null +++ b/tests/test-deepseek41-expert.cpp @@ -0,0 +1,583 @@ +#include "../src/llama-dsv41-expert.h" +#include "../src/llama-dsv41.h" +#include "../src/llama-graph.h" +#include "../src/llama-model-loader.h" + +#include "ggml-backend.h" +#include "ggml-cpu.h" +#include "ggml.h" +#include "gguf.h" +#include "../ggml/src/ggml-backend-impl.h" + +#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; + + temp_file() { + static uint64_t sequence = 0; + const auto stamp = std::chrono::steady_clock::now().time_since_epoch().count(); + path = std::filesystem::temp_directory_path() / + ("llama-dsv41-expert-" + std::to_string(stamp) + "-" + std::to_string(++sequence) + ".gguf"); + } + + ~temp_file() { + std::error_code ec; + std::filesystem::remove(path, ec); + } +}; + +template +void require_throws(F && fn) { + bool threw = false; + try { + fn(); + } catch (const std::exception &) { + threw = true; + } + REQUIRE(threw); +} + +struct fixture { + static constexpr int64_t n_embd = 256; + static constexpr int64_t n_ff = 256; + static constexpr int64_t n_expert = LLAMA_DSV41_N_EXPERT; + + temp_file file; + 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 + down_plane)*n_expert; + ggml_init_params params = { + /*.mem_size =*/ data_size + 8*ggml_tensor_overhead() + 4096, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ false, + }; + ggml_context * ctx = ggml_init(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(gate, 0x10); + fill(up, 0x20); + fill(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); + const auto gate_extent = loader.register_external_tensor( + "blk.0.ffn_gate_exps.weight", 0, LLAMA_EXPERT_PROJECTION_GATE, { n_embd, n_ff, n_expert }); + const auto up_extent = loader.register_external_tensor( + "blk.0.ffn_up_exps.weight", 0, LLAMA_EXPERT_PROJECTION_UP, { n_embd, n_ff, n_expert }); + const auto down_extent = loader.register_external_tensor( + "blk.0.ffn_down_exps.weight", 0, LLAMA_EXPERT_PROJECTION_DOWN, { n_ff, n_embd, n_expert }); + for (int32_t il = 0; il < (int32_t) LLAMA_DSV41_N_LAYER; ++il) { + for (auto extent : { gate_extent, up_extent, down_extent }) { + extent.layer = il; + extent.name = "blk." + std::to_string(il) + extent.name.substr(5); + tensors.push_back(std::move(extent)); + } + } + } + + static void fill(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%32; + } + } + + size_t cache_bytes(size_t slots) const { + size_t result = 0; + for (const auto & tensor : tensors) { + result += tensor.nb[2]*slots; + } + return result; + } + + llama_dsv41_expert_runtime make_runtime( + size_t slots, + llama_dsv41_expert_runtime::upload_fn upload = {}, + llama_dsv41_expert_runtime::publish_fn before_publish = {}) const { + llama_dsv41_expert_runtime_params params; + params.cache_slots = slots; + params.cache_bytes = cache_bytes(slots); + params.direct_io = false; + return llama_dsv41_expert_runtime( + tensors, + params, + [](const llama_expert_store_tensor &) { return ggml_backend_cpu_buffer_type(); }, + std::move(upload), + std::move(before_publish)); + } +}; + +void test_registration() { + std::vector names; + const auto tensors = llama_dsv41_register_expert_tensors( + [&](const std::string & name, + int32_t layer, + llama_expert_projection projection, + const std::initializer_list & ne) { + names.push_back(name); + llama_expert_store_tensor tensor; + tensor.name = name; + tensor.fname = "unused"; + tensor.file_index = layer % 3; + tensor.layer = layer; + tensor.projection = projection; + tensor.type = projection == LLAMA_EXPERT_PROJECTION_DOWN ? GGML_TYPE_Q2_K : GGML_TYPE_IQ2_XXS; + std::copy(ne.begin(), ne.end(), tensor.ne); + tensor.nb[0] = ggml_type_size(tensor.type); + tensor.nb[1] = ggml_row_size(tensor.type, tensor.ne[0]); + tensor.nb[2] = tensor.nb[1]*tensor.ne[1]; + tensor.file_offset = 4096 + (size_t) layer*3*1024 + (size_t) projection*1024; + tensor.file_size = tensor.nb[2]*tensor.ne[2]; + return tensor; + }); + REQUIRE(tensors.size() == LLAMA_DSV41_N_LAYER*3); + REQUIRE(names.front() == "blk.0.ffn_gate_exps.weight"); + REQUIRE(names.back() == "blk.39.ffn_down_exps.weight"); + for (int32_t il = 0; il < (int32_t) LLAMA_DSV41_N_LAYER; ++il) { + REQUIRE(tensors[3*il + 0].layer == il); + REQUIRE(tensors[3*il + 1].layer == il); + REQUIRE(tensors[3*il + 2].layer == il); + REQUIRE(tensors[3*il + 0].file_index == (size_t) il % 3); + REQUIRE(tensors[3*il + 0].file_offset == 4096 + (size_t) il*3*1024); + REQUIRE(tensors[3*il + 0].ne[2] == LLAMA_DSV41_N_EXPERT); + REQUIRE(tensors[3*il + 0].nb[2] == 3041280); + REQUIRE(tensors[3*il + 1].nb[2] == 3041280); + REQUIRE(tensors[3*il + 2].nb[2] == 3870720); + } + REQUIRE(std::none_of(names.begin(), names.end(), [](const std::string & name) { + return name.find("shexp") != std::string::npos; + })); + size_t one_slot_bytes = 0; + for (const auto & tensor : tensors) { + one_slot_bytes += tensor.nb[2]; + } + REQUIRE(one_slot_bytes == 398131200); +} + +void test_configuration(const fixture & f) { + llama_dsv41_expert_runtime_params params; + params.direct_io = false; + require_throws([&] { + llama_dsv41_expert_runtime runtime( + f.tensors, params, [](const llama_expert_store_tensor &) { return ggml_backend_cpu_buffer_type(); }); + }); + + params.cache_slots = 1; + params.cache_bytes = f.cache_bytes(1) - 1; + require_throws([&] { + llama_dsv41_expert_runtime runtime( + f.tensors, params, [](const llama_expert_store_tensor &) { return ggml_backend_cpu_buffer_type(); }); + }); + + auto runtime = f.make_runtime(1); + runtime.acquire_context(); + require_throws([&] { runtime.acquire_context(); }); + runtime.release_context(); + runtime.acquire_context(); + runtime.release_context(); +} + +void test_remap_upload_and_eviction(const fixture & f) { + struct upload_record { + std::string name; + size_t offset; + std::vector bytes; + }; + std::vector uploads; + auto runtime = f.make_runtime(3, [&](ggml_tensor * tensor, size_t offset, const void * data, size_t size) { + uploads.push_back({ tensor->name, offset, std::vector( + static_cast(data), static_cast(data) + size) }); + ggml_backend_tensor_set(tensor, data, offset, size); + }); + + const auto remapped = runtime.remap(0, { 3, 3, 1, 2 }); + REQUIRE(remapped == std::vector({ 2, 2, 0, 1 })); + REQUIRE(uploads.size() == 9); + REQUIRE(uploads[0].bytes.back() == 0x11); + REQUIRE(uploads[1].bytes.back() == 0x21); + REQUIRE(uploads[2].bytes.back() == 0x31); + require_throws([&] { runtime.remap(0, { 0 }); }); + runtime.release(0); + + uploads.clear(); + REQUIRE(runtime.remap(0, { 3 }).front() == 2); + REQUIRE(uploads.empty()); + runtime.release(0); + + REQUIRE(runtime.remap(0, { 4 }).front() == 0); + runtime.release(0); +} + +void test_capacity_and_upload_failure(const fixture & f) { + auto runtime = f.make_runtime(2); + require_throws([&] { runtime.remap(0, { 0, 1, 2 }); }); + REQUIRE(runtime.remap(0, { 0, 0, 1 }).size() == 3); + runtime.release(0); + + size_t calls = 0; + auto failing = f.make_runtime(1, [&](ggml_tensor * tensor, size_t offset, const void * data, size_t size) { + if (++calls == 5) { + throw std::runtime_error("synthetic upload failure"); + } + ggml_backend_tensor_set(tensor, data, offset, size); + }); + REQUIRE(failing.remap(0, { 0 }).front() == 0); + failing.release(0); + require_throws([&] { failing.remap(0, { 1 }); }); + const size_t after_failure = calls; + REQUIRE(failing.remap(0, { 0 }).front() == 0); + REQUIRE(calls == after_failure + 3); + failing.release(0); +} + +void test_publication_allocation_failure(const fixture & f) { + size_t attempts = 0; + auto runtime = f.make_runtime(1, {}, [&] { + if (attempts++ == 0) { + throw std::bad_alloc(); + } + }); + + require_throws([&] { runtime.remap(0, { 0 }); }); + REQUIRE(runtime.remap(0, { 1 }).front() == 0); + runtime.release(0); +} + +void test_graph_callbacks(const fixture & f) { + auto runtime = f.make_runtime(2); + ggml_init_params params = { + /*.mem_size =*/ 2*1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml_backend_t backend = ggml_backend_cpu_init(); + REQUIRE(backend != nullptr); + ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + ggml_backend_sched_t sched = ggml_backend_sched_new(&backend, &buft, 1, 64, false, true); + REQUIRE(sched != nullptr); + + ggml_tensor * selected = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 3, 1); + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, fixture::n_embd, 3, 1); + ggml_set_input(selected); + ggml_set_input(input); + ggml_tensor * remapped = llama_dsv41_build_expert_remap(ctx, selected, runtime, 0, sched, backend); + ggml_tensor * expert_values = ggml_mul_mat_id( + ctx, runtime.cache_tensor(0, LLAMA_EXPERT_PROJECTION_GATE), input, remapped); + ggml_tensor * release = llama_dsv41_build_expert_release(ctx, expert_values, runtime, 0, sched, backend); + ggml_set_output(remapped); + ggml_set_output(release); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, release); + REQUIRE(ggml_backend_sched_alloc_graph(sched, graph)); + const int32_t original[] = { 3, 1, 1 }; + std::vector input_data(fixture::n_embd*3, 1.0f); + ggml_backend_tensor_set(selected, original, 0, sizeof(original)); + ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size()*sizeof(float)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + const std::string error = runtime.consume_error(); + if (!error.empty()) { + throw std::runtime_error(error); + } + REQUIRE(ggml_backend_sched_get_tensor_backend(sched, remapped) == backend); + REQUIRE(ggml_backend_sched_get_tensor_backend(sched, release) == backend); + REQUIRE(runtime.remap(0, { 0 }).front() >= 0); + runtime.release(0); + runtime.release_all(); + + const int32_t over_capacity[] = { 3, 1, 2 }; + ggml_backend_tensor_set(selected, over_capacity, 0, sizeof(over_capacity)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + REQUIRE(!runtime.consume_error().empty()); + int32_t safe_ids[3] = { -1, -1, -1 }; + ggml_backend_tensor_get(remapped, safe_ids, 0, sizeof(safe_ids)); + REQUIRE(safe_ids[0] == 0 && safe_ids[1] == 0 && safe_ids[2] == 0); + REQUIRE(runtime.remap(0, { 0 }).front() >= 0); + runtime.release(0); + + ggml_backend_sched_free(sched); + ggml_backend_free(backend); + ggml_free(ctx); +} + +void test_graph_upload_failure_sentinel(const fixture & f) { + size_t calls = 0; + auto runtime = f.make_runtime(1, [&](ggml_tensor * tensor, size_t offset, const void * data, size_t size) { + if (++calls == 5) { + throw std::runtime_error("synthetic graph upload failure"); + } + ggml_backend_tensor_set(tensor, data, offset, size); + }); + ggml_init_params params = { + /*.mem_size =*/ 2*1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml_backend_t backend = ggml_backend_cpu_init(); + REQUIRE(backend != nullptr); + ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + ggml_backend_sched_t sched = ggml_backend_sched_new(&backend, &buft, 1, 64, false, true); + REQUIRE(sched != nullptr); + + ggml_tensor * selected = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 1, 1); + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, fixture::n_embd, 1, 1); + ggml_set_input(selected); + ggml_set_input(input); + ggml_tensor * remapped = llama_dsv41_build_expert_remap(ctx, selected, runtime, 0, sched, backend); + ggml_tensor * expert_values = ggml_mul_mat_id( + ctx, runtime.cache_tensor(0, LLAMA_EXPERT_PROJECTION_GATE), input, remapped); + ggml_tensor * release = llama_dsv41_build_expert_release(ctx, expert_values, runtime, 0, sched, backend); + ggml_set_output(remapped); + ggml_set_output(release); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, release); + REQUIRE(ggml_backend_sched_alloc_graph(sched, graph)); + + std::vector input_data(fixture::n_embd, 1.0f); + ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size()*sizeof(float)); + const int32_t first[] = { 0 }; + ggml_backend_tensor_set(selected, first, 0, sizeof(first)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + REQUIRE(runtime.consume_error().empty()); + + const int32_t failed[] = { 1 }; + ggml_backend_tensor_set(selected, failed, 0, sizeof(failed)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + REQUIRE(runtime.consume_error().find("synthetic graph upload failure") != std::string::npos); + int32_t safe_id = -1; + ggml_backend_tensor_get(remapped, &safe_id, 0, sizeof(safe_id)); + REQUIRE(safe_id == 0); + + REQUIRE(runtime.remap(0, { 2 }).front() == 0); + runtime.release(0); + ggml_backend_sched_free(sched); + ggml_backend_free(backend); + ggml_free(ctx); +} + +void test_grovemoe_lookup_ids() { + ggml_init_params params = { + /*.mem_size =*/ 1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml_backend_t backend = ggml_backend_cpu_init(); + REQUIRE(backend != nullptr); + ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + ggml_backend_sched_t sched = ggml_backend_sched_new(&backend, &buft, 1, 64, false, true); + REQUIRE(sched != nullptr); + + ggml_tensor * selected = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 2, 1); + ggml_tensor * explicit_slots = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 2, 1); + ggml_set_input(selected); + ggml_set_input(explicit_slots); + const llm_moe_expert_ids grovemoe = llm_build_moe_expert_ids( + ctx, LLM_ARCH_GROVEMOE, selected, nullptr, 2, 8, 4); + const llm_moe_expert_ids deepseek = llm_build_moe_expert_ids( + ctx, LLM_ARCH_DEEPSEEK41, selected, explicit_slots, 8, 8, 0); + REQUIRE(grovemoe.routing == selected); + REQUIRE(grovemoe.lookup != selected); + REQUIRE(deepseek.routing == selected); + REQUIRE(deepseek.lookup == explicit_slots); + ggml_set_output(grovemoe.routing); + ggml_set_output(grovemoe.lookup); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, grovemoe.routing); + ggml_build_forward_expand(graph, grovemoe.lookup); + REQUIRE(ggml_backend_sched_alloc_graph(sched, graph)); + + const int32_t original[] = { 7, 1 }; + ggml_backend_tensor_set(selected, original, 0, sizeof(original)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + int32_t routing_ids[2] = {}; + int32_t chunk_ids[2] = {}; + ggml_backend_tensor_get(grovemoe.routing, routing_ids, 0, sizeof(routing_ids)); + ggml_backend_tensor_get(grovemoe.lookup, chunk_ids, 0, sizeof(chunk_ids)); + REQUIRE(routing_ids[0] == 7 && routing_ids[1] == 1); + REQUIRE(chunk_ids[0] == 1 && chunk_ids[1] == 0); + + ggml_backend_sched_free(sched); + ggml_backend_free(backend); + ggml_free(ctx); +} + +struct sync_test_context { + llama_dsv41_expert_runtime * runtime = nullptr; + int synchronize_count = 0; + bool saw_pinned = false; +}; + +const char * sync_test_backend_name(ggml_backend_t) { + return "dsv41-sync-test"; +} + +void sync_test_backend_synchronize(ggml_backend_t backend) { + auto * state = static_cast(backend->context); + state->synchronize_count++; + try { + state->runtime->remap(0, { 2 }); + } catch (const std::exception &) { + state->saw_pinned = true; + } +} + +const char * sync_test_device_name(ggml_backend_dev_t) { + return "dsv41-sync-test"; +} + +enum ggml_backend_dev_type sync_test_device_type(ggml_backend_dev_t) { + return GGML_BACKEND_DEVICE_TYPE_CPU; +} + +bool sync_test_device_supports_op(ggml_backend_dev_t, const ggml_tensor *) { + return true; +} + +bool sync_test_device_supports_buft(ggml_backend_dev_t, ggml_backend_buffer_type_t buft) { + return buft == ggml_backend_cpu_buffer_type(); +} + +void test_release_after_sync(const fixture & f) { + auto runtime = f.make_runtime(1); + REQUIRE(runtime.remap(0, { 1 }).front() == 0); + + sync_test_context state = { &runtime }; + ggml_backend_device device = {}; + device.iface.get_name = sync_test_device_name; + device.iface.get_type = sync_test_device_type; + device.iface.supports_op = sync_test_device_supports_op; + device.iface.supports_buft = sync_test_device_supports_buft; + ggml_backend backend = {}; + backend.iface.get_name = sync_test_backend_name; + backend.iface.synchronize = sync_test_backend_synchronize; + backend.device = &device; + backend.context = &state; + ggml_backend_t backends[] = { &backend }; + ggml_backend_buffer_type_t bufts[] = { ggml_backend_cpu_buffer_type() }; + ggml_backend_sched_t sched = ggml_backend_sched_new(backends, bufts, 1, 16, false, false); + REQUIRE(sched != nullptr); + + runtime.release_all_after_sync(sched); + REQUIRE(state.synchronize_count == 1); + REQUIRE(state.saw_pinned); + REQUIRE(runtime.remap(0, { 2 }).front() == 0); + runtime.release(0); + ggml_backend_sched_free(sched); +} + +void test_original_and_slot_ids() { + ggml_init_params params = { + /*.mem_size =*/ 1024*1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml_backend_t backend = ggml_backend_cpu_init(); + REQUIRE(backend != nullptr); + ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); + ggml_backend_sched_t sched = ggml_backend_sched_new(&backend, &buft, 1, 64, false, true); + REQUIRE(sched != nullptr); + + ggml_tensor * cache = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 1, 2); + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 2, 1); + ggml_tensor * slots = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 2, 1); + ggml_tensor * probs = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 4, 1); + ggml_tensor * original = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 2, 1); + for (ggml_tensor * tensor : { cache, input, slots, probs, original }) { + ggml_set_input(tensor); + } + + ggml_tensor * expert_values = ggml_mul_mat_id(ctx, cache, input, slots); + ggml_tensor * routing_weights = ggml_get_rows(ctx, probs, original); + ggml_set_output(expert_values); + ggml_set_output(routing_weights); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, expert_values); + ggml_build_forward_expand(graph, routing_weights); + REQUIRE(ggml_backend_sched_alloc_graph(sched, graph)); + + const float cache_data[] = { 10.0f, 20.0f }; + const float input_data[] = { 1.0f, 1.0f }; + const int32_t slot_data[] = { 1, 0 }; + const float prob_data[] = { 1.0f, 2.0f, 3.0f, 4.0f }; + const int32_t original_data[] = { 3, 1 }; + ggml_backend_tensor_set(cache, cache_data, 0, sizeof(cache_data)); + ggml_backend_tensor_set(input, input_data, 0, sizeof(input_data)); + ggml_backend_tensor_set(slots, slot_data, 0, sizeof(slot_data)); + ggml_backend_tensor_set(probs, prob_data, 0, sizeof(prob_data)); + ggml_backend_tensor_set(original, original_data, 0, sizeof(original_data)); + REQUIRE(ggml_backend_sched_graph_compute(sched, graph) == GGML_STATUS_SUCCESS); + + float expert_result[2] = {}; + float weight_result[2] = {}; + ggml_backend_tensor_get(expert_values, expert_result, 0, sizeof(expert_result)); + ggml_backend_tensor_get(routing_weights, weight_result, 0, sizeof(weight_result)); + REQUIRE(expert_result[0] == 20.0f && expert_result[1] == 10.0f); + REQUIRE(weight_result[0] == 4.0f && weight_result[1] == 2.0f); + + ggml_backend_sched_free(sched); + ggml_backend_free(backend); + ggml_free(ctx); +} + +} + +int main() { + try { + test_registration(); + fixture f; + test_configuration(f); + test_remap_upload_and_eviction(f); + test_capacity_and_upload_failure(f); + test_publication_allocation_failure(f); + test_graph_callbacks(f); + test_graph_upload_failure_sentinel(f); + test_grovemoe_lookup_ids(); + test_release_after_sync(f); + test_original_and_slot_ids(); + } catch (const std::exception & error) { + std::fprintf(stderr, "test-deepseek41-expert: %s\n", error.what()); + return 1; + } + return 0; +} diff --git a/tests/test-expert-store.cpp b/tests/test-expert-store.cpp new file mode 100644 index 000000000000..fbf5f0b10954 --- /dev/null +++ b/tests/test-expert-store.cpp @@ -0,0 +1,651 @@ +#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 +#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(tensor.file_index == 0); + 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); + }); + + 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_external_mapping_access_policy() { + 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()); + + 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)); +} + +#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; + 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; + params.allow_buffered_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); +} + +#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 + +#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 } } }); + 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); +} + +void test_failed_replacement_keeps_resident_entry(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.max_plane_size(), 1, 4096, false }; + llama_expert_store store(std::move(tensors), params); + { + auto resident = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + REQUIRE(resident.payloads()[0].expert_id == 0); + } + const llama_expert_store_stats before = store.stats(); + std::filesystem::resize_file(copy.path, f.tensors[0].file_offset + 2*f.tensors[0].nb[2] - 1); + require_throws([&] { + store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 1 } } }); + }); + REQUIRE(store.resident_entries() == 1); + REQUIRE(store.stats().evictions == before.evictions); + auto hit = store.acquire({ { 0, LLAMA_EXPERT_PROJECTION_GATE, { 0 } } }); + REQUIRE(hit.payloads()[0].expert_id == 0); +} + +} + +int main() { + try { + fixture f; + 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); + 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); + test_payload_validation(f); + test_truncated_file(f); + test_failed_replacement_keeps_resident_entry(f); + } catch (const std::exception & e) { + fprintf(stderr, "test-expert-store: %s\n", e.what()); + return 1; + } + return 0; +}