Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Allow exporting a view of the KV cache #4180

Merged
merged 6 commits into from
Nov 23, 2023
Merged

Allow exporting a view of the KV cache #4180

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
71fcb7e
Allow exporting a view of the KV cache
KerfuffleV2 Nov 23, 2023
cb137d8
Allow dumping the sequences per cell in common
KerfuffleV2 Nov 23, 2023
22d0485
Track max contiguous cells value and position as well
KerfuffleV2 Nov 23, 2023
7688d72
Fix max contiguous empty cells index calculation
KerfuffleV2 Nov 23, 2023
bc1c346
Fix off by one error in dump_kv_cache_view
KerfuffleV2 Nov 23, 2023
aa21e6d
Add doc comments for KV cache view functions
KerfuffleV2 Nov 23, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Fix max contiguous empty cells index calculation 
Make dump functions deal with lengths or sequences counts > 10 better
  • Loading branch information
@KerfuffleV2
KerfuffleV2 committed Nov 23, 2023
commit 7688d7204ff13648e2c7555c05cd4bae93995a3c
16 changes: 10 additions & 6 deletions common/common.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1393,7 +1393,9 @@ void dump_non_result_info_yaml(FILE * stream, const gpt_params & params, const l
//

void dump_kv_cache_view(const llama_kv_cache_view & view, int row_size) {
printf("=== Dumping KV cache. total cells %d, max sequences per cell %d, populated cells %d, total tokens in cache %d, max contiguous cells=%d @ %d\n",
static const char slot_chars[] = ".123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz+";

printf("=== Dumping KV cache. total cells %d, max sequences per cell %d, populated cells %d, total tokens in cache %d, largest empty slot=%d @ %d",
view.n_cells, view.n_max_seq, view.used_cells, view.token_count, view.max_contiguous_cells, view.max_contiguous_cells_idx);
llama_kv_cache_view_cell * c_curr = view.cells;
struct llama_kv_cache_view_cell_sequence * cs_curr = view.cells_sequences;
Expand All @@ -1405,13 +1407,14 @@ void dump_kv_cache_view(const llama_kv_cache_view & view, int row_size) {
for (int j = 0; j < view.n_max_seq; j++) {
if (cs_curr[j].seq_id >= 0) { seq_count++; }
}
putchar(seq_count == 0 ? '.' : ('0' + (std::min(9, seq_count))));
putchar(slot_chars[std::min(sizeof(slot_chars) - 1, size_t(seq_count))]);
}
printf("\n=== Done dumping\n");
}

void dump_kv_cache_view_seqs(const llama_kv_cache_view & view, int row_size) {
printf("=== Dumping KV cache. total cells %d, max sequences per cell %d, populated cells %d, total tokens in cache %d, max contiguous cells=%d @ %d\n",
static const char slot_chars[] = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz";
printf("=== Dumping KV cache. total cells %d, max sequences per cell %d, populated cells %d, total tokens in cache %d, largest empty slot=%d @ %d\n",
view.n_cells, view.n_max_seq, view.used_cells, view.token_count, view.max_contiguous_cells, view.max_contiguous_cells_idx);

std::unordered_map<llama_seq_id, size_t> seqs;
Expand All @@ -1421,16 +1424,17 @@ void dump_kv_cache_view_seqs(const llama_kv_cache_view & view, int row_size) {
for (int j = 0; j < view.n_max_seq; j++) {
if (cs_curr[j].seq_id < 0) { continue; }
if (seqs.find(cs_curr[j].seq_id) == seqs.end()) {
if (seqs.size() + 1 >= sizeof(slot_chars)) { break; }
seqs[cs_curr[j].seq_id] = seqs.size();
if (seqs.size() >= 10) { break; }
}
}
if (seqs.size() >= 10) { break; }
if (seqs.size() + 1 >= sizeof(slot_chars)) { break; }
}
printf("=== Sequence legend: ");
for (const auto & it : seqs) {
printf("%zu=%d, ", it.second, it.first);
}
printf("'+'=other sequence ids");

c_curr = view.cells;
cs_curr = view.cells_sequences;
Expand All @@ -1441,7 +1445,7 @@ void dump_kv_cache_view_seqs(const llama_kv_cache_view & view, int row_size) {
for (int j = 0; j < view.n_max_seq; j++) {
if (cs_curr[j].seq_id >= 0) {
const auto & it = seqs.find(cs_curr[j].seq_id);
putchar(it != seqs.end() ? int('0' + it->second) : '+');
putchar(it != seqs.end() ? int(slot_chars[it->second]) : '+');
} else {
putchar('.');
}
Expand Down
2 changes: 1 addition & 1 deletion llama.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8857,8 +8857,8 @@ void llama_kv_cache_view_update(const struct llama_context * ctx, struct llama_k

if (curr_size > 0) {
if (curr_contig_idx >= 0 && uint32_t(i - curr_contig_idx) > max_contig) {
max_contig_idx = i;
max_contig = i - curr_contig_idx;
max_contig_idx = curr_contig_idx;
}
curr_contig_idx = -1;
} else if (curr_contig_idx < 0) {
Expand Down
Loading