perf: accelerate longhaul prompt loading

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.

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<std::mutex> 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) {
+bool llama_longhaul_cache::source_is_direct(
-    auto & state = layers.at(layer);
+        const llama_model_loader::longhaul_source & source) const {
-    for (size_t i = 0; i < n_slots; ++i) {
+    if (ggml_backend_buffer_is_host(source.tensor->buffer)) {
-        if (state.expert_ids[i] == expert_id) {
+        return true;
-            state.last_used[i] = ++tick;
-            ++n_hits;
-            return i;
-        }
     }
 
-    size_t slot = 0;
+    const char * name = ggml_backend_buft_name(ggml_backend_buffer_get_type(source.tensor->buffer));
-    for (size_t i = 0; i < n_slots; ++i) {
+    return name != nullptr &&
-        if (state.expert_ids[i] == -1) {
+        std::strncmp(name, "MTL", 3) == 0 &&
-            slot = i;
+        std::strstr(name, "_Private") == nullptr;
-            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;
 }
 
-bool llama_longhaul_cache::load_slot(int layer, int slot, int32_t expert_id) {
+void llama_longhaul_cache::io_worker() {
+    while (true) {
+        io_job * job = nullptr;
+        {
+            std::unique_lock<std::mutex> 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 {
-        std::vector<uint8_t> buffer;
+            const auto & source = *job->source;
-        for (const auto & source : sources) {
+            void * destination =
-            if (source.layer != layer) {
+                static_cast<uint8_t *>(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<std::mutex> 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<std::mutex> 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(
+            read_buffer.resize(source->expert_size);
-                buffer.data(), buffer.size(), source.offset + size_t(expert_id) * source.expert_size);
+            try {
-            ggml_backend_tensor_set(source.tensor, buffer.data(), size_t(slot) * source.tensor->nb[2], buffer.size());
+                files.at(source->file_idx)->read_at(
-            bytes_read += buffer.size();
+                    read_buffer.data(), read_buffer.size(),
-        }
+                    source->offset + size_t(expert_id) * source->expert_size);
-        return true;
+                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;
+            }
+        }
+        if (!staged_ok) {
+            break;
+        }
+    }
+
+    if (!io_jobs.empty()) {
+        std::unique_lock<std::mutex> 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<int32_t> values(ggml_nbytes(ids) / sizeof(int32_t));
+    const size_t n_values = ggml_nbytes(ids) / sizeof(int32_t);
-    ggml_backend_tensor_get(ids, values.data(), 0, ggml_nbytes(ids));
+    id_buffer.resize(n_values);
-    for (int32_t & id : 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 (!requested.at(id)) {
-        if (slot < 0) {
+            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;
+}

src/llama-longhaul.h

@@ -3,9 +3,14 @@
 #include "llama-mmap.h"
 #include "llama-model-loader.h"
 
+#include <condition_variable>
 #include <cstdint>
+#include <deque>
 #include <memory>
 #include <mutex>
+#include <string>
+#include <thread>
+#include <utility>
 #include <vector>
 
 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<int32_t> expert_ids;
+        std::vector<int32_t> expert_slots;
         std::vector<uint64_t> 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<llama_model_loader::longhaul_source> sources;
+    std::vector<std::vector<const llama_model_loader::longhaul_source *>> sources_by_layer;
     std::vector<layer_state> 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);
+    std::vector<uint8_t> requested;
-    bool load_slot(int layer, int slot, int32_t expert_id);
+    std::vector<int32_t> requested_experts;
+    std::vector<int32_t> missing_experts;
+    std::vector<int32_t> available_slots;
+    std::vector<std::pair<int32_t, int32_t>> load_plan;
+    std::vector<int32_t> id_buffer;
+    std::vector<uint8_t> read_buffer;
+
+    std::vector<std::thread> io_workers;
+    std::vector<io_job> io_jobs;
+    std::deque<io_job *> 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();
 };

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

tests/test-longhaul.cpp

@@ -0,0 +1,176 @@
+#include "testing.h"
+
+#include "../src/llama-longhaul.h"
+
+#include "ggml-backend.h"
+
+#include <cstdio>
+#include <memory>
+#include <vector>
+
+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<llama_longhaul_cache> 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<llama_file>(file));
+
+        std::vector<llama_model_loader::longhaul_source> 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<llama_longhaul_cache>(
+            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<uint8_t> 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<int32_t> remap(std::vector<int32_t> 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<int32_t>{};
+    }
+
+    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();
+}

tools/llama-bench/README.md

@@ -56,6 +56,7 @@ test parameters:
   -ub, --ubatch-size <n>                    (default: 512)
   -ctk, --cache-type-k <t>                  (default: f16)
   -ctv, --cache-type-v <t>                  (default: f16)
+  --longhaul-cache <GiB>                    expert cache size for longhaul mode
   -t, --threads <n>                         (default: system dependent)
   -C, --cpu-mask <hex,hex>                  (default: 0x0)
   --cpu-strict <0|1>                        (default: 0)

tools/llama-bench/llama-bench.cpp

@@ -341,6 +341,7 @@ struct cmd_params {
     std::vector<int>                 n_cpu_moe;
     std::vector<llama_split_mode>    split_mode;
     std::vector<llama_load_mode>     load_mode;
+    uint64_t                         longhaul_cache_bytes;
     std::vector<int>                 main_gpu;
     std::vector<bool>                no_kv_offload;
     std::vector<llama_flash_attn_type> 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 <on|off|auto>                   (default: %s)\n", join(transform_to_str(cmd_params_defaults.flash_attn, llama_flash_attn_type_name), ",").c_str());
     printf("  -dev, --device <dev0/dev1/...>                    (default: auto)\n");
-    printf("  -lm, --load-mode <none|mmap|mlock|mmap+mlock|dio> (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str());
+    printf("  -lm, --load-mode <none|mmap|mlock|mmap+mlock|dio|longhaul> (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str());
+    printf("  --longhaul-cache <GiB>                           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<cmd_params_instance> 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<cmd_params_instance> 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<cmd_params_instance> 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");
         }