Skip to content

Commit d40bdb0

Browse files
ggerganovggerganov
committed
llama : add struct llama_kv_cache (wip) [no ci]
1 parent 4dbc8b9 commit d40bdb0

File tree

8 files changed

+428
-415
lines changed

8 files changed

+428
-415
lines changed

‎common/common.cpp‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -909,7 +909,9 @@ struct common_init_result common_init_from_params(common_params & params) {
909909
return iparams;
910910
}
911911

912-
if (params.ctx_shift && !llama_kv_cache_can_shift(lctx)) {
912+
llama_kv_cache * kv = llama_get_kv_cache(lctx);
913+
914+
if (params.ctx_shift && !llama_kv_cache_can_shift(kv)) {
913915
LOG_WRN("%s: KV cache shifting is not supported for this model, disabling KV cache shifting\n", __func__);
914916
params.ctx_shift = false;
915917
}
@@ -1014,7 +1016,7 @@ struct common_init_result common_init_from_params(common_params & params) {
10141016
if (llama_model_has_decoder(model)) {
10151017
llama_decode(lctx, llama_batch_get_one(tmp.data(), std::min(tmp.size(), (size_t) params.n_batch)));
10161018
}
1017-
llama_kv_cache_clear(lctx);
1019+
llama_kv_cache_clear(kv);
10181020
llama_synchronize(lctx);
10191021
llama_perf_context_reset(lctx);
10201022
}

‎common/speculative.cpp‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -171,8 +171,10 @@ llama_tokens common_speculative_gen_draft(
171171
llama_tokens result;
172172
result.reserve(params.n_draft);
173173

174+
llama_kv_cache * kv = llama_get_kv_cache(ctx);
175+
174176
if (reuse_n == 0) {
175-
llama_kv_cache_clear(ctx);
177+
llama_kv_cache_clear(kv);
176178

177179
prompt.clear();
178180
} else {
@@ -191,14 +193,14 @@ llama_tokens common_speculative_gen_draft(
191193
}
192194

193195
if (reuse_i > 0) {
194-
llama_kv_cache_seq_rm (ctx, 0, 0, reuse_i);
195-
llama_kv_cache_seq_add(ctx, 0, reuse_i, -1, -reuse_i);
196+
llama_kv_cache_seq_rm (kv, 0, 0, reuse_i);
197+
llama_kv_cache_seq_add(kv, 0, reuse_i, -1, -reuse_i);
196198

197199
prompt.erase(prompt.begin(), prompt.begin() + reuse_i);
198200
}
199201

200202
if (reuse_n < (int) prompt.size()) {
201-
llama_kv_cache_seq_rm (ctx, 0, reuse_n, -1);
203+
llama_kv_cache_seq_rm (kv, 0, reuse_n, -1);
202204

203205
prompt.erase(prompt.begin() + reuse_n, prompt.end());
204206
}

‎examples/embedding/embedding.cpp‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,10 +34,11 @@ static void batch_add_seq(llama_batch & batch, const std::vector<int32_t> & toke
3434

3535
static void batch_decode(llama_context * ctx, llama_batch & batch, float * output, int n_seq, int n_embd, int embd_norm) {
3636
const enum llama_pooling_type pooling_type = llama_pooling_type(ctx);
37-
const struct llama_model * model = llama_get_model(ctx);
37+
const llama_model * model = llama_get_model(ctx);
38+
llama_kv_cache * kv = llama_get_kv_cache(ctx);
3839

3940
// clear previous kv_cache values (irrelevant for embeddings)
40-
llama_kv_cache_clear(ctx);
41+
llama_kv_cache_clear(kv);
4142

4243
// run model
4344
LOG_INF("%s: n_tokens = %d, n_seq = %d\n", __func__, batch.n_tokens, n_seq);

‎include/llama.h‎

Lines changed: 42 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ extern "C" {
6060
struct llama_model;
6161
struct llama_context;
6262
struct llama_sampler;
63+
struct llama_kv_cache;
6364

6465
typedef int32_t llama_pos;
6566
typedef int32_t llama_token;
@@ -470,8 +471,9 @@ extern "C" {
470471

471472
DEPRECATED(LLAMA_API int32_t llama_n_vocab (const struct llama_vocab * vocab), "use llama_vocab_n_tokens instead");
472473

473-
LLAMA_API const struct llama_model * llama_get_model (const struct llama_context * ctx);
474-
LLAMA_API enum llama_pooling_type llama_pooling_type(const struct llama_context * ctx);
474+
LLAMA_API const struct llama_model * llama_get_model (const struct llama_context * ctx); // TODO: remove const?
475+
LLAMA_API struct llama_kv_cache * llama_get_kv_cache( struct llama_context * ctx);
476+
LLAMA_API enum llama_pooling_type llama_pooling_type(const struct llama_context * ctx);
475477

476478
LLAMA_API const struct llama_vocab * llama_model_get_vocab(const struct llama_model * model);
477479
LLAMA_API enum llama_rope_type llama_model_rope_type(const struct llama_model * model);
@@ -586,7 +588,7 @@ extern "C" {
586588
// KV cache
587589
//
588590

589-
// TODO: remove llama_kv_cache_view_* API
591+
// TODO: start using struct llama_kv_cache
590592

591593
// Information associated with an individual cell in the KV cache view.
592594
struct llama_kv_cache_view_cell {
@@ -641,41 +643,47 @@ extern "C" {
641643

642644
// Returns the number of tokens in the KV cache (slow, use only for debug)
643645
// If a KV cell has multiple sequences assigned to it, it will be counted multiple times
644-
LLAMA_API int32_t llama_get_kv_cache_token_count(const struct llama_context * ctx);
646+
LLAMA_API int32_t llama_kv_cache_n_tokens(const struct llama_kv_cache * kv);
647+
648+
DEPRECATED(LLAMA_API int32_t llama_get_kv_cache_token_count(const struct llama_context * ctx),
649+
"use llama_kv_cache_n_tokens instead");
645650

646651
// Returns the number of used KV cells (i.e. have at least one sequence assigned to them)
647-
LLAMA_API int32_t llama_get_kv_cache_used_cells(const struct llama_context * ctx);
652+
LLAMA_API int32_t llama_kv_cache_used_cells(const struct llama_kv_cache * kv);
653+
654+
DEPRECATED(LLAMA_API int32_t llama_get_kv_cache_used_cells(const struct llama_context * ctx),
655+
"use llama_kv_cache_used_cells instead");
648656

649657
// Clear the KV cache - both cell info is erased and KV data is zeroed
650658
LLAMA_API void llama_kv_cache_clear(
651-
struct llama_context * ctx);
659+
struct llama_kv_cache * kv);
652660

653661
// Removes all tokens that belong to the specified sequence and have positions in [p0, p1)
654662
// Returns false if a partial sequence cannot be removed. Removing a whole sequence never fails
655663
// seq_id < 0 : match any sequence
656664
// p0 < 0 : [0, p1]
657665
// p1 < 0 : [p0, inf)
658666
LLAMA_API bool llama_kv_cache_seq_rm(
659-
struct llama_context * ctx,
660-
llama_seq_id seq_id,
661-
llama_pos p0,
662-
llama_pos p1);
667+
struct llama_kv_cache * kv,
668+
llama_seq_id seq_id,
669+
llama_pos p0,
670+
llama_pos p1);
663671

664672
// Copy all tokens that belong to the specified sequence to another sequence
665673
// Note that this does not allocate extra KV cache memory - it simply assigns the tokens to the new sequence
666674
// p0 < 0 : [0, p1]
667675
// p1 < 0 : [p0, inf)
668676
LLAMA_API void llama_kv_cache_seq_cp(
669-
struct llama_context * ctx,
670-
llama_seq_id seq_id_src,
671-
llama_seq_id seq_id_dst,
672-
llama_pos p0,
673-
llama_pos p1);
677+
struct llama_kv_cache * kv,
678+
llama_seq_id seq_id_src,
679+
llama_seq_id seq_id_dst,
680+
llama_pos p0,
681+
llama_pos p1);
674682

675683
// Removes all tokens that do not belong to the specified sequence
676684
LLAMA_API void llama_kv_cache_seq_keep(
677-
struct llama_context * ctx,
678-
llama_seq_id seq_id);
685+
struct llama_kv_cache * kv,
686+
llama_seq_id seq_id);
679687

680688
// Adds relative position "delta" to all tokens that belong to the specified sequence and have positions in [p0, p1)
681689
// If the KV cache is RoPEd, the KV data is updated accordingly:
@@ -684,11 +692,11 @@ extern "C" {
684692
// p0 < 0 : [0, p1]
685693
// p1 < 0 : [p0, inf)
686694
LLAMA_API void llama_kv_cache_seq_add(
687-
struct llama_context * ctx,
688-
llama_seq_id seq_id,
689-
llama_pos p0,
690-
llama_pos p1,
691-
llama_pos delta);
695+
struct llama_kv_cache * kv,
696+
llama_seq_id seq_id,
697+
llama_pos p0,
698+
llama_pos p1,
699+
llama_pos delta);
692700

693701
// Integer division of the positions by factor of `d > 1`
694702
// If the KV cache is RoPEd, the KV data is updated accordingly:
@@ -697,31 +705,28 @@ extern "C" {
697705
// p0 < 0 : [0, p1]
698706
// p1 < 0 : [p0, inf)
699707
LLAMA_API void llama_kv_cache_seq_div(
700-
struct llama_context * ctx,
701-
llama_seq_id seq_id,
702-
llama_pos p0,
703-
llama_pos p1,
704-
int d);
708+
struct llama_kv_cache * kv,
709+
llama_seq_id seq_id,
710+
llama_pos p0,
711+
llama_pos p1,
712+
int d);
705713

706714
// Returns the largest position present in the KV cache for the specified sequence
707715
LLAMA_API llama_pos llama_kv_cache_seq_pos_max(
708-
struct llama_context * ctx,
709-
llama_seq_id seq_id);
710-
711-
// TODO: the llama_kv_cache_defrag and llama_kv_cache_update API tightly couples llama_context with llama_kv_cache
712-
// how to avoid this?
716+
struct llama_kv_cache * kv,
717+
llama_seq_id seq_id);
713718

714719
// Defragment the KV cache
715720
// This will be applied:
716721
// - lazily on next llama_decode()
717722
// - explicitly with llama_kv_cache_update()
718-
LLAMA_API void llama_kv_cache_defrag(struct llama_context * ctx);
719-
720-
// Apply the KV cache updates (such as K-shifts, defragmentation, etc.)
721-
LLAMA_API void llama_kv_cache_update(struct llama_context * ctx);
723+
LLAMA_API void llama_kv_cache_defrag(struct llama_kv_cache * kv);
722724

723725
// Check if the context supports KV cache shifting
724-
LLAMA_API bool llama_kv_cache_can_shift(struct llama_context * ctx);
726+
LLAMA_API bool llama_kv_cache_can_shift(const struct llama_kv_cache * kv);
727+
728+
// Apply the KV cache updates (such as K-shifts, defragmentation, etc.)
729+
LLAMA_API void llama_update_kv_cache(struct llama_context * ctx, struct llama_kv_cache * kv);
725730

726731
//
727732
// State / sessions

‎src/llama-context.cpp‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -602,11 +602,15 @@ uint32_t llama_n_seq_max(const struct llama_context * ctx) {
602602
return ctx->kv_self.size;
603603
}
604604

605-
const struct llama_model * llama_get_model(const struct llama_context * ctx) {
605+
const llama_model * llama_get_model(const llama_context * ctx) {
606606
return &ctx->model;
607607
}
608608

609-
enum llama_pooling_type llama_pooling_type(const struct llama_context * ctx) {
609+
llama_kv_cache * llama_get_kv_cache(llama_context * ctx) {
610+
return &ctx->kv_self;
611+
}
612+
613+
enum llama_pooling_type llama_pooling_type(const llama_context * ctx) {
610614
return ctx->cparams.pooling_type;
611615
}
612616

@@ -1142,7 +1146,7 @@ struct llama_data_read {
11421146
if (dest_seq_id != -1) {
11431147
// single sequence
11441148

1145-
llama_kv_cache_seq_rm(kv_self, dest_seq_id, -1, -1);
1149+
kv_self.seq_rm(dest_seq_id, -1, -1);
11461150

11471151
llama_ubatch batch = ctx->sbatch.reserve_ubatch(cell_count, /* has_embd */ false);
11481152
batch.n_tokens = cell_count;
@@ -1185,7 +1189,7 @@ struct llama_data_read {
11851189
return false;
11861190
}
11871191

1188-
llama_kv_cache_clear(kv_self);
1192+
kv_self.clear();
11891193

11901194
for (uint32_t i = 0; i < cell_count; ++i) {
11911195
llama_kv_cell & cell = kv_self.cells[i];
@@ -1362,9 +1366,9 @@ struct llama_data_read {
13621366

13631367
if (!res) {
13641368
if (seq_id == -1) {
1365-
llama_kv_cache_clear(ctx);
1369+
ctx->kv_self.clear();
13661370
} else {
1367-
llama_kv_cache_seq_rm(ctx, seq_id, -1, -1);
1371+
ctx->kv_self.seq_rm(seq_id, -1, -1);
13681372
}
13691373
throw std::runtime_error("failed to restore kv cache");
13701374
}

0 commit comments

Comments
 (0)