cache.h 459 Bytes
Newer Older
blkmjsian's avatar
blkmjsian committed
1
2
3
4
5
#ifndef CACHE_H
#define CACHE_H

#include <infinirt.h>

6
__INFINI_C __export struct KVCache *createKVCache(
blkmjsian's avatar
blkmjsian committed
7
8
9
10
11
12
13
14
15
16
    size_t nlayers,
    size_t max_len,
    size_t nkvh_,
    size_t dk,
    size_t dv,
    infiniDtype_t dtype,
    infiniDevice_t device,
    int *dev_ids,
    size_t ndev);

17
__INFINI_C __export struct KVCache *duplicateKVCache(const KVCache *kv_cache, size_t seq_len);
blkmjsian's avatar
blkmjsian committed
18

19
__INFINI_C __export void dropKVCache(KVCache *kv_cache);
blkmjsian's avatar
blkmjsian committed
20
21

#endif /* CACHE_H */