#include "llama-kv-cache.h"

#include "llama-impl.h"

#include "llama-io.h"

#include "llama-model.h"

#include "llama-context.h"

#include <algorithm>

#include <cassert>

#include <cmath>

#include <cstring>

#include <limits>

#include <map>

#include <stdexcept>

static bool ggml_is_power_of_2(int n) {

return (n & (n - 1)) == 0;

}

// orthonormal Walsh-Hadamard rotation matrix

// note: res^2 == I

static void ggml_gen_hadamard(ggml_tensor * tensor) {

assert(tensor->type == GGML_TYPE_F32);

const int n = tensor->ne[0];

assert(ggml_is_power_of_2(n));

assert(tensor->ne[1] == n);

assert(tensor->ne[2] == 1);

assert(tensor->ne[3] == 1);

std::vector<float> data_f32;

float * data = (float *) tensor->data;

if (tensor->type != GGML_TYPE_F32) {

data_f32.resize(n*n);

data = data_f32.data();

}

data[0*n + 0] = 1.0 / sqrtf(n);

for (int s = 1; s < n; s *= 2) {

for (int i = 0; i < s; i++) {

for (int j = 0; j < s; j++) {

const float val = data[i*n + j];

data[(i + s)*n + (j )] = val;

data[(i )*n + (j + s)] = val;

data[(i + s)*n + (j + s)] = -val;

}

}

}

if (tensor->type != GGML_TYPE_F32) {

ggml_quantize_chunk(tensor->type, data, tensor->data, 0, 1, n*n, nullptr);

}

}

//

// llama_kv_cache

//

llama_kv_cache::llama_kv_cache(

const llama_model & model,

const llama_hparams & hparams,

ggml_type type_k,

ggml_type type_v,

bool v_trans,

bool offload,

bool unified,

uint32_t kv_size,

uint32_t n_seq_max,

uint32_t n_pad,

uint32_t n_swa,

llama_swa_type swa_type,

llama_memory_t mem_other,

const layer_filter_cb & filter,

const layer_reuse_cb & reuse,

const layer_share_cb & share) :

model(model), hparams(hparams), v_trans(v_trans),

n_seq_max(n_seq_max), n_stream(unified ? 1 : n_seq_max), n_pad(n_pad), n_swa(n_swa), swa_type(swa_type),

other(static_cast<llama_kv_cache *>(mem_other)),

v_cells_impl(other ? other->v_cells_impl : std::make_shared<llama_kv_cells_vec>()),

v_cells(*v_cells_impl) {

// shared cells view the source cache's K/V tensors, so the cell count

// follows the source allocation: a fitted target can be smaller than the

// draft default and oversized views would overflow the source tensors

if (other) {

const uint32_t size_other = other->get_size();

if (kv_size != size_other) {

LLAMA_LOG_WARN("%s: kv_size = %u overridden to %u to match the shared source cache\n", __func__, kv_size, size_other);

kv_size = size_other;

}

}

GGML_ASSERT(kv_size % n_pad == 0);

const uint32_t n_layer = hparams.n_layer_all;

// define a comparator for the buft -> ctx map to ensure that the order is well-defined:

struct ggml_backend_buft_comparator {

bool operator()(const ggml_backend_buffer_type_t & lhs, const ggml_backend_buffer_type_t & rhs) const {

return strcmp(ggml_backend_buft_name(lhs), ggml_backend_buft_name(rhs)) < 0;

}

};

std::map<ggml_backend_buffer_type_t, ggml_context_ptr, ggml_backend_buft_comparator> ctx_map;

// create a context for each buffer type

auto ctx_for_buft = [&](ggml_backend_buffer_type_t buft) -> ggml_context * {

auto it = ctx_map.find(buft);

if (it == ctx_map.end()) {

ggml_init_params params = {

/*.mem_size =*/ size_t(2u*(1 + n_stream)*n_layer*ggml_tensor_overhead()),

/*.mem_buffer =*/ NULL,

/*.no_alloc =*/ true,

};

ggml_context * ctx = ggml_init(params);

if (!ctx) {

return nullptr;

}

ctx_map.emplace(buft, ctx);

return ctx;

}

return it->second.get();

};

GGML_ASSERT(n_stream == 1 || n_stream == n_seq_max);

v_heads.resize(n_stream);

for (uint32_t s = 0; s < n_stream; ++s) {

v_heads[s] = 0;

}

v_cells.resize(n_stream);

for (uint32_t s = 0; s < n_stream; ++s) {

v_cells[s].resize(kv_size);

}

// by default, all sequence ids are mapped to the 0th stream

seq_to_stream.resize(LLAMA_MAX_SEQ, 0);

if (n_stream > 1) {

seq_to_stream.resize(n_stream, 0);

for (uint32_t s = 0; s < n_stream; ++s) {

seq_to_stream[s] = s;

}

}

// [TAG_V_CACHE_VARIABLE]

if (v_trans && hparams.is_n_embd_v_gqa_variable()) {

LLAMA_LOG_WARN("%s: the V embeddings have different sizes across layers and FA is not enabled - padding V cache to %d\n",

__func__, hparams.n_embd_v_gqa_max());

}

const bool is_mla = hparams.is_mla();

for (uint32_t il = 0; il < n_layer; il++) {

if (!hparams.has_kv(il)) {

LLAMA_LOG_DEBUG("%s: layer %3d: does not have KV cache\n", __func__, il);

continue;

}

if (filter && !filter(il)) {

LLAMA_LOG_DEBUG("%s: layer %3d: filtered\n", __func__, il);

continue;

}

if (share && other) {

const int32_t il_share = share(il);

if (il_share >= 0) {

const auto & layer_share = other->layers[other->map_layer_ids[il_share]];

LLAMA_LOG_WARN("%s: layer %3d: sharing with layer %d. k = %p, v = %p\n", __func__, il, il_share,

layer_share.k->data, layer_share.v->data);

map_layer_ids[il] = layers.size();

layers.push_back(layer_share);

layers.back().il = il;

continue;

}

}

if (n_embd_head_k_all == 0) {

n_embd_head_k_all = (int32_t) hparams.n_embd_head_k(il);

} else if (n_embd_head_k_all > 0 && n_embd_head_k_all != (int32_t) hparams.n_embd_head_k(il)) {

n_embd_head_k_all = -1;

}

if (!is_mla) {

if (n_embd_head_v_all == 0) {

n_embd_head_v_all = (int32_t) hparams.n_embd_head_v(il);

} else if (n_embd_head_v_all > 0 && n_embd_head_v_all != (int32_t) hparams.n_embd_head_v(il)) {

n_embd_head_v_all = -1;

}

}

// [TAG_V_CACHE_VARIABLE]

const uint32_t n_embd_k_gqa = hparams.n_embd_k_gqa(il);

const uint32_t n_embd_v_gqa = !v_trans ? hparams.n_embd_v_gqa(il) : hparams.n_embd_v_gqa_max();

const char * dev_name = "CPU";

ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type();

if (offload) {

auto * dev = model.dev_layer(il);

buft = ggml_backend_dev_buffer_type(dev);

dev_name = ggml_backend_dev_name(dev);

}

LLAMA_LOG_DEBUG("%s: layer %3d: dev = %s\n", __func__, il, dev_name);

ggml_context * ctx = ctx_for_buft(buft);

if (!ctx) {

throw std::runtime_error("failed to create ggml context for kv cache");

}

const bool has_k = true;

const bool has_v = !is_mla;

ggml_tensor * k = has_k ? ggml_new_tensor_3d(ctx, type_k, n_embd_k_gqa, kv_size, n_stream) : nullptr;

ggml_tensor * v = has_v ? ggml_new_tensor_3d(ctx, type_v, n_embd_v_gqa, kv_size, n_stream) : nullptr;

has_k && ggml_format_name(k, "cache_k_l%d", il);

has_v && ggml_format_name(v, "cache_v_l%d", il);

std::vector<ggml_tensor *> k_stream;

std::vector<ggml_tensor *> v_stream;

for (uint32_t s = 0; s < n_stream; ++s) {

k_stream.push_back(has_k ? ggml_view_2d(ctx, k, n_embd_k_gqa, kv_size, k->nb[1], s*k->nb[2]) : nullptr);

v_stream.push_back(has_v ? ggml_view_2d(ctx, v, n_embd_v_gqa, kv_size, v->nb[1], s*v->nb[2]) : nullptr);

}

map_layer_ids[il] = layers.size();

layers.push_back({ il, k, v, k_stream, v_stream, });

}

if (reuse) {

LLAMA_LOG_DEBUG("%s: reusing layers:\n", __func__);

for (uint32_t il = 0; il < n_layer; il++) {

const int32_t il_reuse = reuse(il);

if (il_reuse < 0) {

LLAMA_LOG_DEBUG("%s: - layer %3d: no reuse\n", __func__, il);

continue;

}

if (filter && !filter(il)) {

LLAMA_LOG_DEBUG("%s: - layer %3d: filtered\n", __func__, il);

continue;

}

GGML_ASSERT(map_layer_ids.find(il_reuse) != map_layer_ids.end());

map_layer_ids[il] = map_layer_ids[il_reuse];

LLAMA_LOG_DEBUG("%s: - layer %3d: reuse layer %d, is_swa = %d\n", __func__, il, il_reuse, hparams.is_swa(il));

}

}

// allocate tensors and initialize the buffers to avoid NaNs in the padding

for (auto & [buft, ctx] : ctx_map) {

ggml_backend_buffer_t buf;

if (hparams.no_alloc) {

buf = ggml_backend_buft_alloc_buffer(buft, /*size =*/ 0); // dummy buffer

for (ggml_tensor * t = ggml_get_first_tensor(ctx.get()); t != nullptr; t = ggml_get_next_tensor(ctx.get(), t)) {

t->buffer = buf; // set dummy buffer for KV cache so that the backend scheduler won't try to allocate it

}

} else {

buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); // real buffer

}

if (!buf) {

throw std::runtime_error("failed to allocate buffer for kv cache");

}

LLAMA_LOG_INFO("%s: %10s KV buffer size = %8.2f MiB\n", __func__, ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf)/1024.0/1024.0);

ggml_backend_buffer_clear(buf, 0);

ctxs_bufs.emplace_back(std::move(ctx), buf);

}

{

const size_t memory_size_k = size_k_bytes();

const size_t memory_size_v = size_v_bytes();

LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs), K (%s): %7.2f MiB, V (%s): %7.2f MiB\n", __func__,

(float)(memory_size_k + memory_size_v) / (1024.0f * 1024.0f), kv_size, (int) layers.size(), n_seq_max, n_stream,

ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f),

ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f));

}

// TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]

if (other) {

n_embd_head_k_all = other->n_embd_head_k_all;

n_embd_head_v_all = other->n_embd_head_v_all;

attn_rot_k = other->attn_rot_k;

attn_rot_v = other->attn_rot_v;

} else {

const char * LLAMA_ATTN_ROT_DISABLE = getenv("LLAMA_ATTN_ROT_DISABLE");

const bool attn_rot_disable = LLAMA_ATTN_ROT_DISABLE ? atoi(LLAMA_ATTN_ROT_DISABLE) : false;

if (attn_rot_disable) {

LLAMA_LOG_WARN("%s: attention rotation force disabled (LLAMA_ATTN_ROT_DISABLE)\n", __func__);

}

attn_rot_k =

!attn_rot_disable &&

n_embd_head_k_all > 0 &&

ggml_is_quantized(type_k) &&

hparams.n_embd_head_k() % 64 == 0;

// always create Hadamard rotation tensors for DeepSeek lightning indexers

if ((model.arch == LLM_ARCH_DEEPSEEK32 || model.arch == LLM_ARCH_DEEPSEEK4) &&

hparams.n_embd_head_k_full == hparams.indexer_head_size) {

attn_rot_k = true;

}

attn_rot_v =

!attn_rot_disable &&

n_embd_head_v_all > 0 &&

ggml_is_quantized(type_v) &&

hparams.n_embd_head_v() % 64 == 0;

}

LLAMA_LOG_INFO("%s: attn_rot_k = %d, n_embd_head_k_all = %d\n", __func__, attn_rot_k, n_embd_head_k_all);

LLAMA_LOG_INFO("%s: attn_rot_v = %d, n_embd_head_k_all = %d\n", __func__, attn_rot_v, n_embd_head_v_all);

// pre-compute the haramard matrices and keep them in host memory

// TODO: in the future, we can make copies in the backend buffers to avoid host -> device transfers

if (attn_rot_k || attn_rot_v) {

for (int64_t n = 64; n <= std::max(n_embd_head_k_all, n_embd_head_v_all); n *= 2) {

attn_rot_hadamard[n] = std::vector<float>(n*n);

ggml_init_params params = {

/* .mem_size = */ 1*ggml_tensor_overhead(),

/* .mem_buffer = */ nullptr,

/* .no_alloc = */ true,

};

ggml_context_ptr ctx { ggml_init(params) };

ggml_tensor * tmp = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F32, n, n);

tmp->data = attn_rot_hadamard[n].data();

ggml_gen_hadamard(tmp);

}

}

const char * LLAMA_KV_CACHE_DEBUG = getenv("LLAMA_KV_CACHE_DEBUG");

debug = LLAMA_KV_CACHE_DEBUG ? atoi(LLAMA_KV_CACHE_DEBUG) : 0;

}

void llama_kv_cache::clear(bool data) {

for (uint32_t s = 0; s < n_stream; ++s) {

v_cells[s].reset();

v_heads[s] = 0;

}

if (data) {

for (auto & [_, buf] : ctxs_bufs) {

ggml_backend_buffer_clear(buf.get(), 0);

}

}

}

bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {

// TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]

if (other) {

return true;

}

GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()));

if (p0 < 0) {

p0 = 0;

}

if (p1 < 0) {

p1 = std::numeric_limits<llama_pos>::max();

}

if (seq_id >= 0) {

auto & cells = v_cells[seq_to_stream[seq_id]];

auto & head = v_heads[seq_to_stream[seq_id]];

uint32_t new_head = cells.size();

for (uint32_t i = 0; i < cells.size(); ++i) {

if (!cells.pos_in(i, p0, p1)) {

continue;

}

if (cells.seq_has(i, seq_id) && cells.seq_rm(i, seq_id)) {

if (new_head == cells.size()) {

new_head = i;

}

}

}

// If we freed up a slot, set head to it so searching can start there.

if (new_head != cells.size() && new_head < head) {

head = new_head;

}

} else {

// match any sequence

for (uint32_t s = 0; s < n_stream; ++s) {

auto & cells = v_cells[s];

auto & head = v_heads[s];

uint32_t new_head = cells.size();

for (uint32_t i = 0; i < cells.size(); ++i) {

if (!cells.pos_in(i, p0, p1)) {

continue;

}

cells.rm(i);

if (new_head == cells.size()) {

new_head = i;

}

}

// If we freed up a slot, set head to it so searching can start there.

if (new_head != cells.size() && new_head < head) {

head = new_head;

}

}

}

return true;

}

void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {

// TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]

if (other) {

return;

}

GGML_ASSERT(seq_id_src >= 0 && (size_t) seq_id_src < seq_to_stream.size());

GGML_ASSERT(seq_id_dst >= 0 && (size_t) seq_id_dst < seq_to_stream.size());

const auto s0 = seq_to_stream[seq_id_src];

const auto s1 = seq_to_stream[seq_id_dst];

if (s0 == s1) {

// since both sequences are in the same stream, no data copy is necessary

// we just have to update the cells meta data

auto & cells = v_cells[s0];

if (seq_id_src == seq_id_dst) {

return;

}

if (p0 < 0) {

p0 = 0;

}

if (p1 < 0) {

p1 = std::numeric_limits<llama_pos>::max();

}

for (uint32_t i = 0; i < cells.size(); ++i) {

if (!cells.pos_in(i, p0, p1)) {

continue;

}

if (cells.seq_has(i, seq_id_src)) {

cells.seq_add(i, seq_id_dst);

}

}

return;

}

// cross-stream sequence copies require to copy the actual buffer data

bool is_full = true;

if (p0 > 0 && p0 + 1 < (int) get_size()) {

is_full = false;

}

if (p1 > 0 && p1 + 1 < (int) get_size()) {

is_full = false;

}

GGML_ASSERT(is_full && "seq_cp() is only supported for full KV buffers");

// enqueue the copy operation - the buffer copy will be performed during the next update

sc_info.ssrc.push_back(s0);

sc_info.sdst.push_back(s1);

v_cells[s1].reset();

for (uint32_t i = 0; i < v_cells[s0].size(); ++i) {

if (v_cells[s0].seq_has(i, seq_id_src)) {

llama_pos pos = v_cells[s0].pos_get(i);

llama_pos shift = v_cells[s0].get_shift(i);

llama_kv_cell_ext ext = v_cells[s0].ext_get(i);

if (shift != 0) {

pos -= shift;

assert(pos >= 0);

}

v_cells[s1].pos_set(i, pos);

v_cells[s1].seq_add(i, seq_id_dst);

if (shift != 0) {

v_cells[s1].pos_add(i, shift);

}

v_cells[s1].ext_set(i, ext);

}

}

v_heads[s1] = v_heads[s0];

//for (uint32_t s = 0; s < n_stream; ++s) {

// LLAMA_LOG_WARN("%s: seq %d: min = %d, max = %d\n", __func__, s, v_cells[s].seq_pos_min(s), v_cells[s].seq_pos_max(s));

//}

}

void llama_kv_cache::seq_keep(llama_seq_id seq_id) {

// TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]

if (other) {

return;

}

GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());

auto & cells = v_cells[seq_to_stream[seq_id]];

auto & head = v_heads[seq_to_stream[seq_id]];

uint32_t new_head = cells.size();

for (uint32_t i = 0; i < cells.size(); ++i) {

if (cells.seq_keep(i, seq_id)) {

if (new_head == cells.size()) {

new_head = i;

}

}

}

// If we freed up a slot, set head to it so searching can start there.

if (new_head != cells.size() && new_head < head) {

head = new_head;

}

}

void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {

// TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]

if (other) {

return;

}

GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());

GGML_ASSERT(hparams.n_pos_per_embd() == 1 && "seq_add() is only supported for n_pos_per_embd() == 1");

auto & cells = v_cells[seq_to_stream[seq_id]];

auto & head = v_heads[seq_to_stream[seq_id]];

if (shift == 0) {

return;

}

uint32_t new_head = cells.size();

if (p0 < 0) {

p0 = 0;

}

if (p1 < 0) {

p1 = std::numeric_limits<llama_pos>::max();

}

// If there is no range then return early to avoid looping over all cells.

if (p0 == p1) {

return;

}

for (uint32_t i = 0; i < cells.size(); ++i) {

if (!cells.pos_in(i, p0, p1)) {

continue;

}

if (cells.seq_has(i, seq_id)) {

if (cells.pos_add(i, shift)) {

if (new_head == cells.size()) {

new_head = i;

}

}

}

}

// If we freed up a slot, set head to it so searching can start there.

// Otherwise we just start the next search from the beginning.

head = new_head != cells.size() ? new_head : 0;

}

void llama_kv_cache::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {

// TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]

if (other) {

return;

}

GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());

GGML_ASSERT(hparams.n_pos_per_embd() == 1 && "seq_div() is only supported for n_pos_per_embd() == 1");

auto & cells = v_cells[seq_to_stream[seq_id]];

if (d == 1) {

return;

}

if (p0 < 0) {

p0 = 0;

}

if (p1 < 0) {

p1 = std::numeric_limits<llama_pos>::max();

}

// If there is no range then return early to avoid looping over the cache.

if (p0 == p1) {

return;

}

for (uint32_t i = 0; i < cells.size(); ++i) {

if (!cells.pos_in(i, p0, p1)) {

continue;

}

if (cells.seq_has(i, seq_id)) {

cells.pos_div(i, d);

}

}

}

llama_pos llama_kv_cache::seq_pos_min(llama_seq_id seq_id) const {

// TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]

if (other) {

return other->seq_pos_min(seq_id);

}

GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());

const auto & cells = v_cells[seq_to_stream[seq_id]];

return cells.seq_pos_min(seq_id);

}

llama_pos llama_kv_cache::seq_pos_max(llama_seq_id seq_id) const {

// TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]

if (other) {

return other->seq_pos_max(seq_id);

}

GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());

const auto & cells = v_cells[seq_to_stream[seq_id]];

return cells.seq_pos_max(seq_id);

}

std::map<ggml_backend_buffer_type_t, size_t> llama_kv_cache::memory_breakdown() const {

std::map<ggml_backend_buffer_type_t, size_t> ret;

for (const auto & [ctx, buf] : ctxs_bufs) {

ggml_backend_buffer_type_t buft = ggml_backend_buffer_get_type(buf.get());

if (hparams.no_alloc) {

GGML_ASSERT(ggml_backend_buffer_get_base(buf.get()) == nullptr);

ret[buft] += ggml_backend_alloc_ctx_tensors_from_buft_size(ctx.get(), buft);

} else {

// GGML_ASSERT(ggml_backend_buffer_get_base(buf.get()) != nullptr); // multi_buffer does not have a defined base

ret[buft] += ggml_backend_buffer_get_size(buf.get());

}

}

return ret;

}

llama_memory_context_ptr llama_kv_cache::init_batch(

llama_batch_allocr & balloc,

uint32_t n_ubatch,

bool embd_all) {

GGML_UNUSED(embd_all);

do {

balloc.split_reset();

std::vector<llama_ubatch> ubatches;

while (true) {

auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true, 0);

if (ubatch.n_tokens == 0) {

break;

}

ubatches.push_back(std::move(ubatch)); // NOLINT

}

if (balloc.get_n_used() < balloc.get_n_tokens()) {

// failed to find a suitable split

break;

}

auto sinfos = prepare(ubatches);

if (sinfos.empty()) {

break;

}

return std::make_unique<llama_kv_cache_context>(

this, std::move(sinfos), std::move(ubatches));

} while (false);

return std::make_unique<llama_kv_cache_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);

}

llama_memory_context_ptr llama_kv_cache::init_full() {

return std::make_unique<llama_kv_cache_context>(this);

}

llama_memory_context_ptr llama_kv_cache::init_update(llama_context * lctx, bool optimize) {

GGML_UNUSED(optimize);

bool do_shift = get_has_shift();

return std::make_unique<llama_kv_cache_context>(this, lctx, do_shift, std::move(sc_info));

}

llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vector<llama_ubatch> & ubatches) {

llama_kv_cache::slot_info_vec_t res;

struct state_t {

slot_info sinfo; // slot info for the ubatch

std::vector<uint32_t> v_heads_old; // old positions of the heads, before placing the ubatch

std::vector<llama_kv_cells> v_cells; // copy of the old cells, before placing the ubatch

};

// remember the old state of the cells so we can restore it in the end

std::vector<state_t> states;

bool success = true;

for (const auto & ubatch : ubatches) {

// only find a suitable slot for the ubatch. don't modify the cells yet

const auto sinfo_new = find_slot(ubatch, false);

if (sinfo_new.empty()) {

success = false;

break;

}

// remember the position that we found

res.push_back(sinfo_new);

// store the old state of the cells in the recovery stack

{

state_t state = { sinfo_new, v_heads, {} };

for (uint32_t s = 0; s < sinfo_new.n_stream(); ++s) {

auto & cells = v_cells[sinfo_new.strm[s]];

state.v_cells.push_back(cells.cp(sinfo_new.idxs[s]));

}

states.push_back(std::move(state));

}

// now emplace the ubatch

apply_ubatch(sinfo_new, ubatch);

}

GGML_ASSERT(!states.empty() || !success);

// iterate backwards and restore the cells to their original state

for (auto it = states.rbegin(); it != states.rend(); ++it) {

const auto & sinfo = it->sinfo;

for (uint32_t s = 0; s < sinfo.n_stream(); ++s) {

auto & cells = v_cells[sinfo.strm[s]];

auto & head = v_heads[sinfo.strm[s]];

cells.set(sinfo.idxs[s], it->v_cells[s]);

head = it->v_heads_old[s];

}

}

if (!success) {

return {};

}

return res;

}

bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_copy_info & sc_info) {

// TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]

if (other) {

return true;

}

bool updated = false;

auto * sched = lctx->get_sched();

if (!sc_info.empty()) {

assert(n_stream > 1 && "stream copy should never happen with a single stream");

llama_synchronize(lctx);

const size_t n_copy = sc_info.ssrc.size();

for (size_t i = 0; i < n_copy; ++i) {

const auto ssrc = sc_info.ssrc[i];

const auto sdst = sc_info.sdst[i];

assert(ssrc < n_stream);

assert(sdst < n_stream);

LLAMA_LOG_DEBUG("%s: copying KV buffer: stream %d to stream %d\n", __func__, ssrc, sdst);

assert(ssrc != sdst);

for (uint32_t il = 0; il < layers.size(); ++il) {

const auto & layer = layers[il];

ggml_backend_tensor_copy(layer.k_stream[ssrc], layer.k_stream[sdst]);

if (layer.v_stream[ssrc]) {

ggml_backend_tensor_copy(layer.v_stream[ssrc], layer.v_stream[sdst]);

}

}

}

}

if (do_shift) {

if (!get_can_shift()) {

GGML_ABORT("The current KV cache / model configuration does not support K-shift");

}

LLAMA_LOG_DEBUG("%s: applying K-shift\n", __func__);

// apply K-shift if needed

if (hparams.rope_type != LLAMA_ROPE_TYPE_NONE) {

ggml_backend_sched_reset(sched);

auto * res = lctx->get_gf_res_reserve();

res->reset();

auto * gf = build_graph_shift(res, lctx);

if (!ggml_backend_sched_alloc_graph(sched, gf)) {

LLAMA_LOG_ERROR("%s: failed to allocate compute graph for K-shift\n", __func__);

return updated;

}

res->set_inputs(nullptr);

if (lctx->graph_compute(gf, false) != GGML_STATUS_SUCCESS) {

LLAMA_LOG_ERROR("%s: failed to compute K-shift\n", __func__);

return updated;

}

updated = true;

}

for (uint32_t s = 0; s < n_stream; ++s) {

auto & cells = v_cells[s];

cells.reset_shift();

}

}

return updated;

}

llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, bool cont) const {

if (debug > 0) {

for (uint32_t s = 0; s < ubatch.n_seqs_unq; ++s) {

const auto seq_id = ubatch.seq_id_unq[s];

const auto stream_id = seq_to_stream[seq_id];

const auto & cells = v_cells[stream_id];

const uint32_t head_cur = v_heads[stream_id];

LLAMA_LOG_DEBUG("%s: stream[%d], n = %5d, used = %5d, head = %5d, size = %5d, n_swa = %5d\n",

__func__, stream_id, cells.used_max_p1(), cells.get_used(), head_cur, get_size(), n_swa);

if ((debug == 2 && n_swa > 0) || debug > 2) {

std::string ss;

for (uint32_t i = 0; i < cells.size(); ++i) {

if (cells.is_empty(i)) {

ss += '.';

} else {

assert(cells.seq_count(i) >= 1);

if (cells.seq_count(i) == 1) {

ss += std::to_string(cells.seq_get(i));

} else {

ss += 'M';

}

}

if (i%256 == 255) {

ss += " *";

ss += '\n';

}

}

LLAMA_LOG_DEBUG("\n%s\n", ss.c_str());

}

if ((debug == 2 && n_swa > 0) || debug > 2) {

std::string ss;

for (uint32_t i = 0; i < cells.size(); ++i) {

std::string cur;

if (cells.is_empty(i)) {

cur = '.';

} else {

cur = std::to_string(cells.pos_get(i));

}

const int n = cur.size();

for (int j = 0; j < 5 - n; ++j) {

cur += ' ';

}

ss += cur;

if (i%256 == 255) {

ss += " *";

}

if (i%64 == 63) {

ss += '\n';

}

}

LLAMA_LOG_DEBUG("\n%s\n", ss.c_str());

}

for (int s = 0; s < LLAMA_MAX_SEQ; ++s) {

if (cells.seq_pos_min(s) < 0) {

continue;

}

LLAMA_LOG_DEBUG("%s: stream[%d] min[%d] = %5d, max[%d] = %5d\n", __func__, stream_id, s, cells.seq_pos_min(s), s, cells.seq_pos_max(s));

}

}

}

uint32_t n_tokens = ubatch.n_tokens;

uint32_t n_seqs = 1;

if (n_stream > 1) {

GGML_ASSERT(n_tokens % ubatch.n_seqs_unq == 0);

n_seqs = ubatch.n_seqs_unq;

n_tokens = n_tokens / n_seqs;

}

slot_info res = {

/*.s0 =*/ LLAMA_MAX_SEQ,

/*.s1 =*/ 0,

/*.strm =*/ { },

/*.idxs =*/ { },

};

res.resize(n_seqs);

for (uint32_t s = 0; s < n_seqs; ++s) {

const auto seq_id = ubatch.seq_id_unq[s];

if (n_stream > 1) {

GGML_ASSERT(ubatch.n_seq_id[s*n_tokens] == 1);

GGML_ASSERT(ubatch.seq_id [s*n_tokens][0] == seq_id);

}

res.s0 = std::min<uint32_t>(res.s0, seq_to_stream[seq_id]);

res.s1 = std::max<uint32_t>(res.s1, seq_to_stream[seq_id]);

res.strm[s] = seq_to_stream[seq_id];

res.idxs[s].reserve(n_tokens);

const auto & cells = v_cells[seq_to_stream[seq_id]];

uint32_t head_cur = v_heads[seq_to_stream[seq_id]];

// if we have enough unused cells before the current head ->

// better to start searching from the beginning of the cache, hoping to fill it