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

Merge branch 'dev' into feat/auto-exit

Browse source
This commit is contained in:
Li, Zonghang 2025-05-20 02:04:14 +08:00 • committed by GitHub
commit 7b0ededd24
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 488 additions and 87 deletions
Download patch file Download diff file Expand all files Collapse all files

41
include/llama.h
View file

@ -283,7 +283,7 @@ extern "C" {
uint32_t n_world; // number of nodes
uint32_t rank; // my node rank
uint32_t n_layer_window[32]; // number of layers to kept each time
int32_t n_gpu_layers; // number of layers to store in VRAM
int32_t n_gpu_layers; // number of layers to store in VRAM
enum llama_split_mode split_mode; // how to split the model across multiple GPUs
// main_gpu interpretation depends on split_mode:
@ -453,7 +453,7 @@ extern "C" {
LLAMA_API void llama_free_sockets (struct llama_context * ctx, char ** msg);
LLAMA_API int llama_gather_device_info(struct llama_context * ctx, struct device_info * dev_info_set);
LLAMA_API int llama_send_device_info (struct llama_context * ctx, struct device_info * dev_info);
LLAMA_API int llama_bcast_startup_args(struct llama_context * ctx, uint32_t rank, struct startup_args * args);
LLAMA_API int llama_bcast_startup_args(struct llama_context * ctx, uint32_t rank, struct startup_args * args);
LLAMA_API int llama_bcast_layer_setup (struct llama_context * ctx, uint32_t * n_layer_window, uint32_t * n_gpu_layers);
LLAMA_API int llama_rebuild_topo (struct llama_context * ctx, uint32_t * n_layer_window, struct device_info * dev_info_set);
LLAMA_API int llama_recv_layer_setup (struct llama_context * ctx, uint32_t * n_layer_window, uint32_t * n_gpu_layers);
@ -712,8 +712,10 @@ extern "C" {
LLAMA_API int32_t llama_get_kv_cache_used_cells(const struct llama_context * ctx);
// Clear the KV cache - both cell info is erased and KV data is zeroed
LLAMA_API void llama_kv_cache_clear(
struct llama_context * ctx);
LLAMA_API void llama_kv_cache_clear(struct llama_context * ctx);
// Notify other devices to clear their KV cache
LLAMA_API void llama_send_kv_cache_clear(struct llama_context * ctx);
// Removes all tokens that belong to the specified sequence and have positions in [p0, p1)
// Returns false if a partial sequence cannot be removed. Removing a whole sequence never fails
@ -725,6 +727,13 @@ extern "C" {
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1);
// Notify other nodes to remove a range from their KV cache
LLAMA_API void llama_send_kv_cache_seq_rm(
struct llama_context * ctx,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1);
// Copy all tokens that belong to the specified sequence to another sequence
// Note that this does not allocate extra KV cache memory - it simply assigns the tokens to the new sequence
@ -736,6 +745,14 @@ extern "C" {
llama_seq_id seq_id_dst,
llama_pos p0,
llama_pos p1);
// Notify other nodes to copy a range of KV entries
LLAMA_API void llama_send_kv_cache_seq_cp(
struct llama_context * ctx,
llama_seq_id seq_id_src,
llama_seq_id seq_id_dst,
llama_pos p0,
llama_pos p1);
// Removes all tokens that do not belong to the specified sequence
LLAMA_API void llama_kv_cache_seq_keep(
@ -755,6 +772,14 @@ extern "C" {
llama_pos p1,
llama_pos delta);
// Notify other nodes to shift (add) their KV cache entries
LLAMA_API void llama_send_kv_cache_seq_add(
struct llama_context * ctx,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1,
llama_pos delta);
// Integer division of the positions by factor of `d > 1`
// If the KV cache is RoPEd, the KV data is updated accordingly:
// - lazily on next llama_decode()
@ -767,6 +792,14 @@ extern "C" {
llama_pos p0,
llama_pos p1,
int d);
// Notify other nodes to perform a division operation on a KV cache range
LLAMA_API void llama_send_kv_cache_seq_div(
struct llama_context * ctx,
llama_seq_id seq_id,
llama_pos p0,
llama_pos p1,
int d);
// Returns the largest position present in the KV cache for the specified sequence
LLAMA_API llama_pos llama_kv_cache_seq_pos_max(