From fdd9fcf85edff4316889e3b1d0a781b5d104f319 Mon Sep 17 00:00:00 2001 From: Owen Qwen Date: Wed, 29 Jul 2026 12:44:50 -0500 Subject: [PATCH] perf: accelerate longhaul prompt loading --- docs/longhaul.md | 22 +++ src/llama-longhaul.cpp | 291 +++++++++++++++++++++++++----- src/llama-longhaul.h | 46 ++++- tests/CMakeLists.txt | 1 + tests/test-longhaul.cpp | 176 ++++++++++++++++++ tools/llama-bench/README.md | 1 + tools/llama-bench/llama-bench.cpp | 43 ++++- 7 files changed, 525 insertions(+), 55 deletions(-) create mode 100644 tests/test-longhaul.cpp diff --git a/docs/longhaul.md b/docs/longhaul.md index 63add34..021b93c 100644 --- a/docs/longhaul.md +++ b/docs/longhaul.md @@ -36,3 +36,25 @@ MTP/speculative decoding, tensor validation during loading, vocabulary-only load The cache budget covers the compact routed-expert tensors. It does not include dense weights, attention weights, the KV cache, graph allocations, or the temporary buffer used for one expert read. Disk reads bypass the macOS unified file cache where supported, avoiding a second long-lived copy of streamed weights in system RAM. This implementation synchronizes at each routed MoE layer to discover the selected experts, populate missing cache slots, and continue execution with cache-local expert IDs. Storage speed and expert reuse therefore have a large effect on generation speed. + +Expert IDs are planned as a batch at each synchronization point. Experts already +needed by that batch are protected from eviction, duplicate IDs are loaded only +once, and independent expert slices are read concurrently on shared-memory Metal +devices. Private Metal buffers use a staged fallback. + +## Benchmarking prompt processing + +`llama-bench` accepts the longhaul load mode and cache budget: + +```sh +llama-bench \ + --model /path/to/model.gguf \ + --load-mode longhaul \ + --longhaul-cache 2 \ + --n-gpu-layers 99 \ + --n-prompt 2048 \ + --n-gen 0 +``` + +Use `--no-warmup --repetitions 1` in separate processes to measure a cold expert +cache. Leave warmup enabled to measure steady-state cache reuse. diff --git a/src/llama-longhaul.cpp b/src/llama-longhaul.cpp index 96d78b8..1b9ea98 100644 --- a/src/llama-longhaul.cpp +++ b/src/llama-longhaul.cpp @@ -14,94 +14,279 @@ llama_longhaul_cache::llama_longhaul_cache( uint32_t n_layers) : files(std::move(files)), sources(std::move(sources)), + sources_by_layer(n_layers), layers(n_layers), n_slots(n_slots), n_experts(n_experts) { for (auto & file : this->files) { file->set_no_cache(); } + for (const auto & source : this->sources) { + sources_by_layer.at(source.layer).push_back(&source); + } for (auto & layer : layers) { layer.expert_ids.assign(n_slots, -1); + layer.expert_slots.assign(n_experts, -1); layer.last_used.assign(n_slots, 0); } + + requested.resize(n_experts); + requested_experts.reserve(n_slots); + missing_experts.reserve(n_slots); + available_slots.reserve(n_slots); + load_plan.reserve(n_slots); + id_buffer.reserve(n_slots); + + const unsigned int n_threads = std::max(1u, std::min(4u, std::thread::hardware_concurrency())); + io_workers.reserve(n_threads); + for (unsigned int i = 0; i < n_threads; ++i) { + io_workers.emplace_back(&llama_longhaul_cache::io_worker, this); + } } llama_longhaul_cache::~llama_longhaul_cache() { - LLAMA_LOG_INFO("%s: hits = %" PRIu64 ", misses = %" PRIu64 ", read = %.2f MiB\n", - __func__, n_hits, n_misses, bytes_read / 1024.0 / 1024.0); + { + std::lock_guard lock(io_mutex); + io_stopping = true; + } + io_ready.notify_all(); + for (auto & worker : io_workers) { + worker.join(); + } + + LLAMA_LOG_INFO( + "%s: batches = %" PRIu64 ", ids = %" PRIu64 ", unique = %" PRIu64 + ", duplicates = %" PRIu64 ", hits = %" PRIu64 ", misses = %" PRIu64 + ", read = %.2f MiB, io = %.2f ms, remap = %.2f ms\n", + __func__, n_batches, n_ids, n_unique, n_duplicates, n_hits, n_misses, + bytes_read / 1024.0 / 1024.0, io_wall_us / 1000.0, remap_us / 1000.0); } -int llama_longhaul_cache::find_slot(int layer, int32_t expert_id) { - auto & state = layers.at(layer); - for (size_t i = 0; i < n_slots; ++i) { - if (state.expert_ids[i] == expert_id) { - state.last_used[i] = ++tick; - ++n_hits; - return i; - } +bool llama_longhaul_cache::source_is_direct( + const llama_model_loader::longhaul_source & source) const { + if (ggml_backend_buffer_is_host(source.tensor->buffer)) { + return true; } - size_t slot = 0; - for (size_t i = 0; i < n_slots; ++i) { - if (state.expert_ids[i] == -1) { - slot = i; - break; - } - if (state.last_used[i] < state.last_used[slot]) { - slot = i; - } - } - if (!load_slot(layer, slot, expert_id)) { - return -1; - } - state.expert_ids[slot] = expert_id; - state.last_used[slot] = ++tick; - ++n_misses; - return slot; + const char * name = ggml_backend_buft_name(ggml_backend_buffer_get_type(source.tensor->buffer)); + return name != nullptr && + std::strncmp(name, "MTL", 3) == 0 && + std::strstr(name, "_Private") == nullptr; } -bool llama_longhaul_cache::load_slot(int layer, int slot, int32_t expert_id) { - try { - std::vector buffer; - for (const auto & source : sources) { - if (source.layer != layer) { +void llama_longhaul_cache::io_worker() { + while (true) { + io_job * job = nullptr; + { + std::unique_lock lock(io_mutex); + io_ready.wait(lock, [&] { return io_stopping || !io_queue.empty(); }); + if (io_stopping && io_queue.empty()) { + return; + } + job = io_queue.front(); + io_queue.pop_front(); + } + + try { + const auto & source = *job->source; + void * destination = + static_cast(source.tensor->data) + size_t(job->slot) * source.tensor->nb[2]; + files.at(source.file_idx)->read_at( + destination, source.expert_size, + source.offset + size_t(job->expert_id) * source.expert_size); + job->ok = true; + } catch (const std::exception & e) { + job->error = e.what(); + } + + { + std::lock_guard lock(io_mutex); + --io_pending; + } + io_done.notify_one(); + } +} + +bool llama_longhaul_cache::load_plan_sources() { + const size_t n_direct_sources = std::count_if( + sources_by_layer.at(locked_layer).begin(), + sources_by_layer.at(locked_layer).end(), + [&](const auto * source) { return source_is_direct(*source); }); + const size_t n_jobs = load_plan.size() * n_direct_sources; + io_jobs.clear(); + io_jobs.reserve(n_jobs); + + for (const auto & item : load_plan) { + for (const auto * source : sources_by_layer.at(locked_layer)) { + if (source_is_direct(*source)) { + io_jobs.push_back({ source, item.first, item.second, false, {} }); + } + } + } + + if (!io_jobs.empty()) { + std::lock_guard lock(io_mutex); + io_pending = io_jobs.size(); + for (auto & job : io_jobs) { + io_queue.push_back(&job); + } + io_ready.notify_all(); + } + + bool staged_ok = true; + for (const auto & item : load_plan) { + const int32_t expert_id = item.first; + const int slot = item.second; + for (const auto * source : sources_by_layer.at(locked_layer)) { + if (source_is_direct(*source)) { continue; } - buffer.resize(source.expert_size); - files.at(source.file_idx)->read_at( - buffer.data(), buffer.size(), source.offset + size_t(expert_id) * source.expert_size); - ggml_backend_tensor_set(source.tensor, buffer.data(), size_t(slot) * source.tensor->nb[2], buffer.size()); - bytes_read += buffer.size(); + + read_buffer.resize(source->expert_size); + try { + files.at(source->file_idx)->read_at( + read_buffer.data(), read_buffer.size(), + source->offset + size_t(expert_id) * source->expert_size); + ggml_backend_tensor_set( + source->tensor, read_buffer.data(), + size_t(slot) * source->tensor->nb[2], read_buffer.size()); + bytes_read += read_buffer.size(); + } catch (const std::exception & e) { + last_error = e.what(); + staged_ok = false; + break; + } } - return true; - } catch (const std::exception & e) { - last_error = e.what(); - return false; + if (!staged_ok) { + break; + } + } + + if (!io_jobs.empty()) { + std::unique_lock lock(io_mutex); + io_done.wait(lock, [&] { return io_pending == 0; }); + } + + for (const auto & job : io_jobs) { + if (!job.ok) { + last_error = job.error; + return false; + } + bytes_read += job.source->expert_size; + } + return staged_ok; +} + +void llama_longhaul_cache::invalidate_plan(int layer) { + auto & state = layers.at(layer); + for (const auto & item : load_plan) { + const int slot = item.second; + const int32_t old_expert = state.expert_ids.at(slot); + if (old_expert >= 0) { + state.expert_slots.at(old_expert) = -1; + } + state.expert_ids.at(slot) = -1; + state.last_used.at(slot) = 0; } } bool llama_longhaul_cache::remap(int layer, ggml_tensor * ids) { + const int64_t t_start_us = ggml_time_us(); last_error.clear(); mutex.lock(); locked = true; locked_layer = layer; - std::vector values(ggml_nbytes(ids) / sizeof(int32_t)); - ggml_backend_tensor_get(ids, values.data(), 0, ggml_nbytes(ids)); - for (int32_t & id : values) { + const size_t n_values = ggml_nbytes(ids) / sizeof(int32_t); + id_buffer.resize(n_values); + ggml_backend_tensor_get(ids, id_buffer.data(), 0, ggml_nbytes(ids)); + + std::fill(requested.begin(), requested.end(), 0); + requested_experts.clear(); + missing_experts.clear(); + available_slots.clear(); + load_plan.clear(); + + auto & state = layers.at(layer); + for (const int32_t id : id_buffer) { if (id < 0 || id >= (int32_t) n_experts) { last_error = "router produced an out-of-range expert ID"; release(layer); return false; } - const int slot = find_slot(layer, id); - if (slot < 0) { - release(layer); - return false; + if (!requested.at(id)) { + requested.at(id) = 1; + requested_experts.push_back(id); + if (state.expert_slots.at(id) < 0) { + missing_experts.push_back(id); + } } + } + + if (requested_experts.size() > n_slots) { + last_error = "router requested more unique experts than the longhaul cache can hold"; + release(layer); + return false; + } + + for (size_t slot = 0; slot < n_slots; ++slot) { + const int32_t expert_id = state.expert_ids.at(slot); + if (expert_id < 0 || !requested.at(expert_id)) { + available_slots.push_back(slot); + } + } + std::stable_sort(available_slots.begin(), available_slots.end(), [&](int32_t a, int32_t b) { + const bool a_empty = state.expert_ids.at(a) < 0; + const bool b_empty = state.expert_ids.at(b) < 0; + if (a_empty != b_empty) { + return a_empty; + } + return state.last_used.at(a) < state.last_used.at(b); + }); + + if (missing_experts.size() > available_slots.size()) { + last_error = "longhaul could not reserve enough cache slots for the routed expert batch"; + release(layer); + return false; + } + for (size_t i = 0; i < missing_experts.size(); ++i) { + load_plan.emplace_back(missing_experts.at(i), available_slots.at(i)); + } + + const int64_t io_start_us = ggml_time_us(); + if (!load_plan_sources()) { + invalidate_plan(layer); + release(layer); + return false; + } + io_wall_us += ggml_time_us() - io_start_us; + + for (const auto & item : load_plan) { + const int32_t expert_id = item.first; + const int slot = item.second; + const int32_t old_expert = state.expert_ids.at(slot); + if (old_expert >= 0) { + state.expert_slots.at(old_expert) = -1; + } + state.expert_ids.at(slot) = expert_id; + state.expert_slots.at(expert_id) = slot; + } + + for (int32_t & id : id_buffer) { + const int slot = state.expert_slots.at(id); + GGML_ASSERT(slot >= 0); + state.last_used.at(slot) = ++tick; id = slot; } - ggml_backend_tensor_set(ids, values.data(), 0, ggml_nbytes(ids)); + ggml_backend_tensor_set(ids, id_buffer.data(), 0, ggml_nbytes(ids)); + + ++n_batches; + n_ids += n_values; + n_unique += requested_experts.size(); + n_duplicates += n_values - requested_experts.size(); + n_misses += missing_experts.size(); + n_hits += n_values - missing_experts.size(); + remap_us += ggml_time_us() - t_start_us; return true; } @@ -128,3 +313,11 @@ const std::string & llama_longhaul_cache::error() const { bool llama_longhaul_cache::failed() const { return !last_error.empty(); } + +uint64_t llama_longhaul_cache::misses() const { + return n_misses; +} + +uint64_t llama_longhaul_cache::bytes_read_count() const { + return bytes_read; +} diff --git a/src/llama-longhaul.h b/src/llama-longhaul.h index 5206eaa..cf33ec4 100644 --- a/src/llama-longhaul.h +++ b/src/llama-longhaul.h @@ -3,9 +3,14 @@ #include "llama-mmap.h" #include "llama-model-loader.h" +#include #include +#include #include #include +#include +#include +#include #include struct llama_longhaul_cache { @@ -24,27 +29,64 @@ struct llama_longhaul_cache { uint32_t max_ubatch(uint32_t n_expert_used) const; const std::string & error() const; bool failed() const; + uint64_t misses() const; + uint64_t bytes_read_count() const; private: struct layer_state { std::vector expert_ids; + std::vector expert_slots; std::vector last_used; }; + struct io_job { + const llama_model_loader::longhaul_source * source; + int32_t expert_id; + int slot; + bool ok = false; + std::string error; + }; + llama_files files; std::vector sources; + std::vector> sources_by_layer; std::vector layers; size_t n_slots; uint32_t n_experts; uint64_t tick = 0; + uint64_t n_batches = 0; + uint64_t n_ids = 0; + uint64_t n_unique = 0; + uint64_t n_duplicates = 0; uint64_t n_hits = 0; uint64_t n_misses = 0; uint64_t bytes_read = 0; + uint64_t io_wall_us = 0; + uint64_t remap_us = 0; std::mutex mutex; bool locked = false; int locked_layer = -1; std::string last_error; - int find_slot(int layer, int32_t expert_id); - bool load_slot(int layer, int slot, int32_t expert_id); + std::vector requested; + std::vector requested_experts; + std::vector missing_experts; + std::vector available_slots; + std::vector> load_plan; + std::vector id_buffer; + std::vector read_buffer; + + std::vector io_workers; + std::vector io_jobs; + std::deque io_queue; + std::mutex io_mutex; + std::condition_variable io_ready; + std::condition_variable io_done; + size_t io_pending = 0; + bool io_stopping = false; + + bool source_is_direct(const llama_model_loader::longhaul_source & source) const; + bool load_plan_sources(); + void invalidate_plan(int layer); + void io_worker(); }; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 7a93b19..66fd403 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -258,6 +258,7 @@ llama_build_and_test(test-thread-safety.cpp ARGS -m "${MODEL_DEST}" -ngl 99 -p " set_tests_properties(test-thread-safety PROPERTIES FIXTURES_REQUIRED test-download-model) llama_build_and_test(test-arg-parser.cpp) +llama_build_and_test(test-longhaul.cpp) if (NOT LLAMA_SANITIZE_ADDRESS AND NOT GGML_SCHED_NO_REALLOC) # TODO: repair known memory leaks diff --git a/tests/test-longhaul.cpp b/tests/test-longhaul.cpp new file mode 100644 index 0000000..39a2ebe --- /dev/null +++ b/tests/test-longhaul.cpp @@ -0,0 +1,176 @@ +#include "testing.h" + +#include "../src/llama-longhaul.h" + +#include "ggml-backend.h" + +#include +#include +#include + +struct longhaul_fixture { + static constexpr size_t expert_size = 32; + + FILE * file = nullptr; + ggml_backend_ptr backend; + ggml_context_ptr ctx; + ggml_backend_buffer_ptr buffer; + ggml_tensor * weights_a = nullptr; + ggml_tensor * weights_b = nullptr; + ggml_tensor * ids = nullptr; + std::unique_ptr cache; + + longhaul_fixture( + size_t n_slots, + uint32_t n_experts, + bool two_sources = false, + uint32_t written_experts = 0) { + file = tmpfile(); + GGML_ASSERT(file != nullptr); + write_experts(written_experts == 0 ? n_experts : written_experts, two_sources); + + backend.reset(ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr)); + GGML_ASSERT(backend); + + ggml_init_params params = { + /*.mem_size =*/ 16 * 1024, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ctx.reset(ggml_init(params)); + GGML_ASSERT(ctx); + + weights_a = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_I8, expert_size, 1, n_slots); + if (two_sources) { + weights_b = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_I8, expert_size, 1, n_slots); + } + ids = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, n_slots); + buffer.reset(ggml_backend_alloc_ctx_tensors(ctx.get(), backend.get())); + GGML_ASSERT(buffer); + + llama_files files; + files.emplace_back(std::make_unique(file)); + + std::vector sources = { + { weights_a, 0, 0, expert_size, 0 }, + }; + if (two_sources) { + sources.push_back({ weights_b, 0, n_experts * expert_size, expert_size, 0 }); + } + cache = std::make_unique( + std::move(files), std::move(sources), n_slots, n_experts, 1); + } + + ~longhaul_fixture() { + cache.reset(); + if (file) { + fclose(file); + } + } + + void write_experts(uint32_t n_experts, bool two_sources) { + GGML_ASSERT(fseek(file, 0, SEEK_END) == 0); + for (int source = 0; source < (two_sources ? 2 : 1); ++source) { + for (uint32_t expert = 0; expert < n_experts; ++expert) { + std::vector data(expert_size, uint8_t(1 + expert + source * 32)); + GGML_ASSERT(fwrite(data.data(), 1, data.size(), file) == data.size()); + } + } + GGML_ASSERT(fflush(file) == 0); + } + + std::vector remap(std::vector values) { + GGML_ASSERT(values.size() <= (size_t) ids->ne[0]); + ggml_backend_tensor_set(ids, values.data(), 0, values.size() * sizeof(int32_t)); + ids->ne[0] = values.size(); + const bool ok = cache->remap(0, ids); + if (ok) { + ggml_backend_tensor_get(ids, values.data(), 0, values.size() * sizeof(int32_t)); + } + cache->release(0); + return ok ? values : std::vector{}; + } + + uint8_t slot_value(ggml_tensor * tensor, int32_t slot) { + uint8_t value = 0; + ggml_backend_tensor_get(tensor, &value, size_t(slot) * tensor->nb[2], 1); + return value; + } +}; + +static void test_batch_planning(testing & t) { + longhaul_fixture fixture(3, 4); + + const auto initial = fixture.remap({0, 1, 2}); + t.assert_equal(3u, initial.size()); + t.assert_equal(3u, fixture.cache->misses()); + + // Expert 0 is the oldest cache entry. A sequential miss for expert 3 used + // to evict it before the later 0 in this same routed batch was observed. + const auto remapped = fixture.remap({3, 0, 3}); + t.assert_equal(3u, remapped.size()); + t.assert_equal("only expert 3 is loaded", 4u, fixture.cache->misses()); + t.assert_true("duplicates map to the same slot", remapped[0] == remapped[2]); + t.assert_true("different experts map to different slots", remapped[0] != remapped[1]); + t.assert_equal(uint8_t(4), fixture.slot_value(fixture.weights_a, remapped[0])); + t.assert_equal(uint8_t(1), fixture.slot_value(fixture.weights_a, remapped[1])); +} + +static void test_multiple_sources(testing & t) { + longhaul_fixture fixture(2, 4, true); + + const auto remapped = fixture.remap({2, 1}); + t.assert_equal(2u, remapped.size()); + t.assert_equal(uint64_t(4 * longhaul_fixture::expert_size), fixture.cache->bytes_read_count()); + t.assert_equal(uint8_t(3), fixture.slot_value(fixture.weights_a, remapped[0])); + t.assert_equal(uint8_t(35), fixture.slot_value(fixture.weights_b, remapped[0])); + t.assert_equal(uint8_t(2), fixture.slot_value(fixture.weights_a, remapped[1])); + t.assert_equal(uint8_t(34), fixture.slot_value(fixture.weights_b, remapped[1])); +} + +static void test_invalid_ids(testing & t) { + longhaul_fixture fixture(2, 4); + + const auto remapped = fixture.remap({0, 4}); + t.assert_true(remapped.empty()); + t.assert_true(fixture.cache->failed()); + t.assert_equal(0u, fixture.cache->misses()); +} + +static void test_read_failure_recovery(testing & t) { + longhaul_fixture fixture(1, 4, false, 3); + + const auto failed = fixture.remap({3}); + t.assert_true(failed.empty()); + t.assert_true(fixture.cache->failed()); + t.assert_equal(0u, fixture.cache->misses()); + + fixture.write_experts(1, false); + const auto remapped = fixture.remap({3}); + t.assert_equal(1u, remapped.size()); + t.assert_equal(1u, fixture.cache->misses()); + t.assert_equal(uint8_t(1), fixture.slot_value(fixture.weights_a, remapped[0])); +} + +int main(int argc, char ** argv) { + testing t; + + const char * verbose = getenv("LLAMA_TEST_VERBOSE"); + if (verbose) { + t.verbose = std::string(verbose) == "1"; + } + if (!t.verbose) { + llama_log_set([](ggml_log_level, const char *, void *) {}, nullptr); + } + + if (argc > 1) { + t.set_filter(argv[1]); + } + + t.test("batch_planning", test_batch_planning); + t.test("multiple_sources", test_multiple_sources); + t.test("invalid_ids", test_invalid_ids); + t.test("read_failure", test_read_failure_recovery); + + return t.summary(); +} diff --git a/tools/llama-bench/README.md b/tools/llama-bench/README.md index d539785..91c3711 100644 --- a/tools/llama-bench/README.md +++ b/tools/llama-bench/README.md @@ -56,6 +56,7 @@ test parameters: -ub, --ubatch-size (default: 512) -ctk, --cache-type-k (default: f16) -ctv, --cache-type-v (default: f16) + --longhaul-cache expert cache size for longhaul mode -t, --threads (default: system dependent) -C, --cpu-mask (default: 0x0) --cpu-strict <0|1> (default: 0) diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index c17a27b..f148704 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -341,6 +341,7 @@ struct cmd_params { std::vector n_cpu_moe; std::vector split_mode; std::vector load_mode; + uint64_t longhaul_cache_bytes; std::vector main_gpu; std::vector no_kv_offload; std::vector flash_attn; @@ -385,6 +386,7 @@ static const cmd_params cmd_params_defaults = { /* n_cpu_moe */ { 0 }, /* split_mode */ { LLAMA_SPLIT_MODE_LAYER }, /* load_mode */ { LLAMA_LOAD_MODE_MMAP }, + /* longhaul_cache_bytes */ 0, /* main_gpu */ { 0 }, /* no_kv_offload */ { false }, /* flash_attn */ { LLAMA_FLASH_ATTN_TYPE_AUTO }, @@ -459,7 +461,8 @@ static void print_usage(int /* argc */, char ** argv) { printf(" -nkvo, --no-kv-offload <0|1> (default: %s)\n", join(cmd_params_defaults.no_kv_offload, ",").c_str()); printf(" -fa, --flash-attn (default: %s)\n", join(transform_to_str(cmd_params_defaults.flash_attn, llama_flash_attn_type_name), ",").c_str()); printf(" -dev, --device (default: auto)\n"); - printf(" -lm, --load-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str()); + printf(" -lm, --load-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str()); + printf(" --longhaul-cache expert cache size for longhaul mode\n"); printf(" -mmp, --mmap <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n"); printf(" -dio, --direct-io <0|1> (DEPRECATED IN FAVOUR OF --load-mode)\n"); printf(" -embd, --embeddings <0|1> (default: %s)\n", join(cmd_params_defaults.embeddings, ",").c_str()); @@ -521,6 +524,7 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { params.progress = cmd_params_defaults.progress; params.no_warmup = cmd_params_defaults.no_warmup; params.offline = cmd_params_defaults.offline; + params.longhaul_cache_bytes = cmd_params_defaults.longhaul_cache_bytes; if (const char * env = getenv("HF_TOKEN")) { params.hf_token = env; @@ -774,6 +778,8 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { mode = LLAMA_LOAD_MODE_MMAP_MLOCK; } else if (m == "dio") { mode = LLAMA_LOAD_MODE_DIRECT_IO; + } else if (m == "longhaul") { + mode = LLAMA_LOAD_MODE_LONGHAUL; } else { invalid_param = true; break; @@ -784,6 +790,17 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { break; } params.load_mode.insert(params.load_mode.end(), modes.begin(), modes.end()); + } else if (arg == "--longhaul-cache") { + if (++i >= argc) { + invalid_param = true; + break; + } + const uint64_t gib = std::stoull(argv[i]); + if (gib == 0 || gib > UINT64_MAX / 1024 / 1024 / 1024) { + invalid_param = true; + break; + } + params.longhaul_cache_bytes = gib * 1024 * 1024 * 1024; } else if (arg == "-mg" || arg == "--main-gpu") { if (++i >= argc) { invalid_param = true; @@ -1135,6 +1152,11 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { if (params.load_mode.empty()) { params.load_mode = cmd_params_defaults.load_mode; } + if (std::find(params.load_mode.begin(), params.load_mode.end(), LLAMA_LOAD_MODE_LONGHAUL) != params.load_mode.end() && + params.longhaul_cache_bytes == 0) { + fprintf(stderr, "error: longhaul load mode requires --longhaul-cache N\n"); + exit(1); + } if (params.main_gpu.empty()) { params.main_gpu = cmd_params_defaults.main_gpu; } @@ -1201,6 +1223,7 @@ struct cmd_params_instance { int n_cpu_moe; llama_split_mode split_mode; llama_load_mode load_mode; + uint64_t longhaul_cache_bytes; int main_gpu; bool no_kv_offload; llama_flash_attn_type flash_attn; @@ -1222,6 +1245,7 @@ struct cmd_params_instance { } mparams.split_mode = split_mode; mparams.load_mode = load_mode; + mparams.longhaul_cache_bytes = longhaul_cache_bytes; mparams.main_gpu = main_gpu; mparams.tensor_split = tensor_split.data(); mparams.no_host = no_host; @@ -1269,7 +1293,8 @@ struct cmd_params_instance { return model == other.model && n_gpu_layers == other.n_gpu_layers && n_cpu_moe == other.n_cpu_moe && split_mode == other.split_mode && main_gpu == other.main_gpu && tensor_split == other.tensor_split && - load_mode == other.load_mode && devices == other.devices && no_host == other.no_host && + load_mode == other.load_mode && longhaul_cache_bytes == other.longhaul_cache_bytes && + devices == other.devices && no_host == other.no_host && vec_tensor_buft_override_equal(tensor_buft_overrides, other.tensor_buft_overrides); } @@ -1342,6 +1367,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .n_cpu_moe = */ ncmoe, /* .split_mode = */ sm, /* .load_mode = */ lm, + /* .longhaul_cache_bytes = */ params.longhaul_cache_bytes, /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, /* .flash_attn = */ fa, @@ -1378,6 +1404,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .n_cpu_moe = */ ncmoe, /* .split_mode = */ sm, /* .load_mode = */ lm, + /* .longhaul_cache_bytes = */ params.longhaul_cache_bytes, /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, /* .flash_attn = */ fa, @@ -1414,6 +1441,7 @@ static std::vector get_cmd_params_instances(const cmd_param /* .n_cpu_moe = */ ncmoe, /* .split_mode = */ sm, /* .load_mode = */ lm, + /* .longhaul_cache_bytes = */ params.longhaul_cache_bytes, /* .main_gpu = */ mg, /* .no_kv_offload = */ nkvo, /* .flash_attn = */ fa, @@ -1455,6 +1483,7 @@ struct test { int n_cpu_moe; llama_split_mode split_mode; llama_load_mode load_mode; + uint64_t longhaul_cache_bytes; int main_gpu; bool no_kv_offload; llama_flash_attn_type flash_attn; @@ -1494,6 +1523,7 @@ struct test { n_cpu_moe = inst.n_cpu_moe; split_mode = inst.split_mode; load_mode = inst.load_mode; + longhaul_cache_bytes = inst.longhaul_cache_bytes; main_gpu = inst.main_gpu; no_kv_offload = inst.no_kv_offload; flash_attn = inst.flash_attn; @@ -1561,7 +1591,7 @@ struct test { "n_ubatch", "n_threads", "cpu_mask", "cpu_strict", "poll", "type_k", "type_v", "n_gpu_layers", "n_cpu_moe", "split_mode", "main_gpu", "no_kv_offload", "flash_attn", "devices", "tensor_split", - "tensor_buft_overrides", "load_mode", "embeddings", + "tensor_buft_overrides", "load_mode", "longhaul_cache_bytes", "embeddings", "no_op_offload", "no_host", "fit_target", "fit_min_ctx", "n_prompt", "n_gen", "n_depth", "test_time", "avg_ns", "stddev_ns", "avg_ts", "stddev_ts" @@ -1576,7 +1606,8 @@ struct test { field == "poll" || field == "model_size" || field == "model_n_params" || field == "n_gpu_layers" || field == "main_gpu" || field == "n_prompt" || field == "n_gen" || field == "n_depth" || field == "avg_ns" || field == "stddev_ns" || field == "no_op_offload" || field == "n_cpu_moe" || - field == "fit_target" || field == "fit_min_ctx" || field == "flash_attn") { + field == "fit_target" || field == "fit_min_ctx" || field == "flash_attn" || + field == "longhaul_cache_bytes") { return INT; } if (field == "f16_kv" || field == "no_kv_offload" || field == "cpu_strict" || @@ -1656,6 +1687,7 @@ struct test { tensor_split_str, tensor_buft_overrides_str, llama_load_mode_name(load_mode), + std::to_string(longhaul_cache_bytes), std::to_string(embeddings), std::to_string(no_op_offload), std::to_string(no_host), @@ -1970,6 +2002,9 @@ struct markdown_printer : public printer { if (params.load_mode.size() > 1 || params.load_mode != cmd_params_defaults.load_mode) { fields.emplace_back("load_mode"); } + if (params.longhaul_cache_bytes != cmd_params_defaults.longhaul_cache_bytes) { + fields.emplace_back("longhaul_cache_bytes"); + } if (params.embeddings.size() > 1 || params.embeddings != cmd_params_defaults.embeddings) { fields.emplace_back("embeddings"); }