perf: accelerate longhaul prompt loading

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) {
-    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<uint8_t> 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<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 {
+            const auto & source = *job->source;
+            void * destination =
+                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(
-                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<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));
-    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;
+}