MPI support, first cut

This commit is contained in:
Evan Miller 2023-07-03 21:51:05 -04:00
parent d7d2e6a0f0
commit f85785f650
6 changed files with 220 additions and 33 deletions

View file

@ -145,6 +145,8 @@ extern "C" {
// If numa is true, use NUMA optimizations
// Call once at the start of the program
LLAMA_API void llama_init_backend(bool numa);
// Call once at the end of the program - currently only used for MPI
LLAMA_API void llama_finalize_backend();
LLAMA_API int64_t llama_time_us();
@ -257,6 +259,8 @@ extern "C" {
LLAMA_API int llama_n_vocab(const struct llama_context * ctx);
LLAMA_API int llama_n_ctx (const struct llama_context * ctx);
LLAMA_API int llama_n_embd (const struct llama_context * ctx);
LLAMA_API int llama_mpi_rank (const struct llama_context * ctx);
LLAMA_API int llama_mpi_size (const struct llama_context * ctx);
// Get the vocabulary as output parameters.
// Returns number of results.