98 lines
2.7 KiB
C++
98 lines
2.7 KiB
C++
#pragma once
|
|
|
|
#include "llama-mmap.h"
|
|
#include "llama-model-loader.h"
|
|
|
|
#include <condition_variable>
|
|
#include <cstdint>
|
|
#include <deque>
|
|
#include <map>
|
|
#include <memory>
|
|
#include <mutex>
|
|
#include <string>
|
|
#include <thread>
|
|
#include <utility>
|
|
#include <vector>
|
|
|
|
struct 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);
|
|
~llama_longhaul_cache();
|
|
|
|
bool remap(int layer, ggml_tensor * ids);
|
|
void release(int layer);
|
|
|
|
uint32_t capacity() const;
|
|
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 direct;
|
|
std::vector<uint8_t> staging;
|
|
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;
|
|
|
|
std::vector<uint8_t> requested;
|
|
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<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;
|
|
|
|
std::map<ggml_backend_dev_t, ggml_backend_ptr> upload_backends;
|
|
|
|
bool source_is_direct(const llama_model_loader::longhaul_source & source) const;
|
|
ggml_backend_t upload_backend(const llama_model_loader::longhaul_source & source);
|
|
bool load_plan_sources();
|
|
void invalidate_plan(int layer);
|
|
void io_worker();
|
|
};
|