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
#pragma once
#include "ggml-backend.h"
#include <map>
#include <memory>
#include <vector>
struct llama_model;
// keeps recently and repeatedly used experts of host-resident MoE layers in a device buffer
// each layer has a slot map in host memory: when the scheduler copies it to the device, the copy callback uploads the missing experts
class llama_moe_cache {
public:
// backends are all the backends of the context, each GPU gets its own cache of the given size for the layers assigned to it
llama_moe_cache(const llama_model & model, const std::vector<ggml_backend_t> & backends, const std::vector<ggml_backend_buffer_type_t> & bufts, size_t size);
~llama_moe_cache();
// the device that caches layer il
ggml_backend_t backend(int32_t il) const;
// the slot map of layer il, if its experts can be read from the cache for n_tokens tokens, nullptr otherwise
ggml_tensor * get_slot_map(int32_t il, int64_t n_tokens, int64_t n_expert_used) const;
// the experts of w in the cache, nullptr if w is not cached
ggml_tensor * get_experts(const ggml_tensor * w) const;
// ggml_backend_sched copy callback, returns false if src is not a slot map
bool copy(ggml_backend_t backend, const ggml_tensor * src, ggml_tensor * dst, ggml_cgraph * graph);
// for large batches: copy the experts of w that are in the cache, starting at expert e and up to expert last, to the copy dst of w
// returns the number of experts copied, 0 if expert e is not in the cache
int64_t copy_experts(ggml_backend_t backend, const ggml_tensor * w, ggml_tensor * dst, int64_t e, int64_t last);
std::map<ggml_backend_buffer_type_t, size_t> memory_breakdown() const;
private:
struct impl;
std::unique_ptr<impl> pimpl;
};
using llama_moe_cache_ptr = std::unique_ptr<llama_moe_cache>;