File size: 3,576 Bytes
c8ad244 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | #pragma once
#include <cstdint>
#include <cstddef>
#include <mutex>
#include <unordered_map>
#include <list>
#include <vector>
#include <array>
#include <atomic>
#include <utility>
struct ggml_tensor;
// Identifies one expert's one weight matrix
struct llama_expert_key {
int32_t layer; // 0..n_layer-1
int32_t expert_idx; // 0..n_expert-1
int32_t weight_type; // 0=up, 1=gate, 2=down
bool operator==(const llama_expert_key & other) const {
return layer == other.layer
&& expert_idx == other.expert_idx
&& weight_type == other.weight_type;
}
};
struct llama_expert_key_hash {
size_t operator()(const llama_expert_key & k) const {
// Combine the three fields into a single hash
size_t h = 0;
h ^= std::hash<int32_t>{}(k.layer) + 0x9e3779b9 + (h << 6) + (h >> 2);
h ^= std::hash<int32_t>{}(k.expert_idx) + 0x9e3779b9 + (h << 6) + (h >> 2);
h ^= std::hash<int32_t>{}(k.weight_type) + 0x9e3779b9 + (h << 6) + (h >> 2);
return h;
}
};
// Where to find an expert's weights on disk
struct llama_expert_disk_info {
int fd; // open file descriptor for GGUF
size_t file_offset; // absolute byte offset in GGUF file
size_t size_bytes; // bytes for this one expert slice
int ggml_type; // quantization type (enum ggml_type)
};
// One cached expert weight slab
struct llama_expert_entry {
llama_expert_key key;
void * data; // page-aligned buffer
size_t size_bytes;
std::list<llama_expert_key>::iterator lru_it;
};
struct llama_expert_cache_stats {
uint64_t hits;
uint64_t misses;
uint64_t evictions;
size_t bytes_used;
size_t bytes_capacity;
double hit_rate() const {
uint64_t total = hits + misses;
return total > 0 ? (double)hits / total : 0.0;
}
};
class llama_expert_cache {
public:
explicit llama_expert_cache(size_t max_bytes);
~llama_expert_cache();
// Non-copyable, non-movable
llama_expert_cache(const llama_expert_cache &) = delete;
llama_expert_cache & operator=(const llama_expert_cache &) = delete;
// Returns pointer to cached expert data.
// Loads from disk on miss. Thread-safe.
void * ensure(const llama_expert_key & key,
const llama_expert_disk_info & disk_info);
// Get cached data or allocate empty slot for caller to fill.
// Returns {pointer, true} on cache hit, {pointer, false} on miss (caller must fill).
// Returns {nullptr, false} on allocation failure.
std::pair<void *, bool> get_or_alloc(const llama_expert_key & key, size_t size_bytes);
// Update LRU ordering (call when expert is accessed but already loaded)
void touch(const llama_expert_key & key);
// Check if an expert is cached without loading
bool contains(const llama_expert_key & key) const;
// Get statistics
llama_expert_cache_stats get_stats() const;
// Reset stats counters
void reset_stats();
private:
void evict_until_free(size_t needed);
void * alloc_aligned(size_t size);
void free_aligned(void * ptr, size_t size);
void * load_from_disk(const llama_expert_disk_info & info);
size_t max_bytes_;
size_t used_bytes_;
std::list<llama_expert_key> lru_order_; // front = most recent
std::unordered_map<llama_expert_key, llama_expert_entry,
llama_expert_key_hash> cache_;
mutable std::mutex mutex_;
llama_expert_cache_stats stats_;
};
|