131 lines
3.6 KiB
C++
131 lines
3.6 KiB
C++
#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::capacity() const {
|
|
return n_slots;
|
|
}
|
|
|
|
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();
|
|
}
|