Inital commit

src/llama-longhaul.cpp

@@ -0,0 +1,126 @@
+#include "llama-longhaul.h"
+
+#include "llama-impl.h"
+
+#include <algorithm>
+#include <cstring>
+#include <inttypes.h>
+
+llama_longhaul_cache::llama_longhaul_cache(
+        llama_files files,
+        std::vector<llama_model_loader::longhaul_source> sources,
+        size_t n_slots,
+        uint32_t n_experts,
+        uint32_t n_layers) :
+    files(std::move(files)),
+    sources(std::move(sources)),
+    layers(n_layers),
+    n_slots(n_slots),
+    n_experts(n_experts) {
+    for (auto & file : this->files) {
+        file->set_no_cache();
+    }
+    for (auto & layer : layers) {
+        layer.expert_ids.assign(n_slots, -1);
+        layer.last_used.assign(n_slots, 0);
+    }
+}
+
+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);
+}
+
+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;
+        }
+    }
+
+    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;
+}
+
+bool llama_longhaul_cache::load_slot(int layer, int slot, int32_t expert_id) {
+    try {
+        std::vector<uint8_t> buffer;
+        for (const auto & source : sources) {
+            if (source.layer != layer) {
+                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();
+        }
+        return true;
+    } catch (const std::exception & e) {
+        last_error = e.what();
+        return false;
+    }
+}
+
+bool llama_longhaul_cache::remap(int layer, ggml_tensor * ids) {
+    last_error.clear();
+    mutex.lock();
+    locked = true;
+    locked_layer = layer;
+
+    std::vector<int32_t> values(ggml_nbytes(ids) / sizeof(int32_t));
+    ggml_backend_tensor_get(ids, values.data(), 0, ggml_nbytes(ids));
+    for (int32_t & id : values) {
+        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;
+        }
+        id = slot;
+    }
+    ggml_backend_tensor_set(ids, values.data(), 0, ggml_nbytes(ids));
+    return true;
+}
+
+void llama_longhaul_cache::release(int layer) {
+    if (locked && (locked_layer == layer || layer < 0)) {
+        locked = false;
+        locked_layer = -1;
+        mutex.unlock();
+    }
+}
+
+uint32_t llama_longhaul_cache::max_ubatch(uint32_t n_expert_used) const {
+    return std::max<uint32_t>(1, n_slots / n_expert_used);
+}
+
+const std::string & llama_longhaul_cache::error() const {
+    return last_error;
+}
+
+bool llama_longhaul_cache::failed() const {
+    return !last_error.empty();
+}