vrr/prima.cpp
2
0
Fork 0
mirror of https://github.com/Lizonghang/prima.cpp.git synced 2025-09-10 11:24:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

add gpu support in llama_model_kvcache_size and llama_model_compute_buf_size

Browse source
This commit is contained in:
Lizonghang 2024-11-29 21:06:32 +04:00
parent f8e9dc2713
commit 6f54a12c7d
3 changed files with 58 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

14
include/llama.h
View file

@ -525,10 +525,20 @@ extern "C" {
LLAMA_API uint64_t llama_model_n_params(const struct llama_model * model);
// Return the size of compute buffer size, including input tensors and activations
LLAMA_API uint64_t llama_model_compute_buf_size(const struct llama_model * model, const struct llama_context_params cparams, bool compress_memory);
LLAMA_API void llama_model_compute_buf_size(
uint64_t * cpu_buf,
uint64_t * gpu_buf,
const struct llama_model * model,
const struct llama_context_params cparams,
bool use_gpu);
// Return the size of KV cache in the model
LLAMA_API uint64_t llama_model_kvcache_size(const struct llama_model * model, const struct llama_context_params cparams);
LLAMA_API void llama_model_kvcache_size(
uint64_t * cpu_cache,
uint64_t * gpu_cache,
const struct llama_model * model,
const struct llama_context_params cparams,
bool use_gpu);
// Return the total number of float operations in the model
LLAMA_API void llama_model_n_flops(