#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>

//

// llama_kv_cache

//

llama_kv_cache::llama_kv_cache(

const llama_model & model,

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,

const layer_filter_cb & filter,

const layer_reuse_cb & reuse) :

model(model), hparams(model.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) {

GGML_ASSERT(kv_size % n_pad == 0);

const uint32_t n_layer_kv = hparams.n_layer_kv();

// 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_kv*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 < hparams.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;

}

// [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 < hparams.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 (model.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));

}

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) {

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) {

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) {

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) {

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) {

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 {

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 {

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);

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;

}

// remeber 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) {

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

if (head_cur > cells.get_used() + 2*n_tokens) {

head_cur = 0;

}

if (n_tokens > cells.size()) {

LLAMA_LOG_ERROR("%s: n_tokens = %d > size = %u\n", __func__, n_tokens, cells.size());

return { };

}

uint32_t n_tested = 0;

// for continuous slots, we test that all tokens in the ubatch fit, starting from the current head

// for non-continuous slots, we test the tokens one by one

const uint32_t n_test = cont ? n_tokens : 1;

while (true) {

if (head_cur + n_test > cells.size()) {

n_tested += cells.size() - head_cur;

head_cur = 0;

continue;

}

for (uint32_t i = 0; i < n_test; i++) {

const auto idx = head_cur;

head_cur++;

n_tested++;

//const llama_pos pos = ubatch.pos[i];

//const llama_seq_id seq_id = ubatch.seq_id[i][0];

// can we use this cell? either:

// - the cell is empty

// - the cell is occupied only by one sequence:

// - (disabled) mask causally, if the sequence is the same as the one we are inserting

// - mask SWA, using current max pos for that sequence in the cache

// always insert in the cell with minimum pos

bool can_use = cells.is_empty(idx);

if (!can_use && cells.seq_count(idx) == 1) {

const llama_pos pos_cell = cells.pos_get(idx);

// (disabled) causal mask

// note: it's better to purge any "future" tokens beforehand

//if (cells.seq_has(idx, seq_id)) {

// can_use = pos_cell >= pos;

//}

if (!can_use) {

const llama_seq_id seq_id_cell = cells.seq_get(idx);

// SWA mask

if (llama_hparams::is_masked_swa(n_swa, swa_type, pos_cell, cells.seq_pos_max(seq_id_cell) + 1)) {

can_use = true;

}

}

}

if (can_use) {

res.idxs[s].push_back(idx);

} else {

if (cont) {

break;

}

}

}

if (res.idxs[s].size() == n_tokens) {

break;

}

if (cont) {

res.idxs[s].clear();

}

if (n_tested >= cells.size()) {

//LLAMA_LOG_ERROR("%s: failed to find a slot for %d tokens\n", __func__, n_tokens);

return { };

}

}

// we didn't find a suitable slot - return empty result

if (res.idxs[s].size() < n_tokens) {

return { };

}

}

assert(res.s1 >= res.s0);

return res;

}

void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & ubatch) {

// keep track of the max sequence position that we would overwrite with this ubatch

// for non-SWA cache, this would be always empty

llama_seq_id seq_pos_max_rm[LLAMA_MAX_SEQ];

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

seq_pos_max_rm[s] = -1;

}

assert(ubatch.n_tokens == sinfo.n_stream()*sinfo.size());

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

for (uint32_t ii = 0; ii < sinfo.size(); ++ii) {

const uint32_t i = s*sinfo.size() + ii;

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

const auto idx = sinfo.idxs[s][ii];

if (!cells.is_empty(idx)) {

assert(cells.seq_count(idx) == 1);

const llama_seq_id seq_id = cells.seq_get(idx);

const llama_pos pos = cells.pos_get(idx);

seq_pos_max_rm[seq_id] = std::max(seq_pos_max_rm[seq_id], pos);

cells.rm(idx);

}

cells.pos_set(idx, ubatch.pos[i]);

if (ubatch.is_pos_2d()) {

llama_kv_cell_ext ext {

/*.x =*/ ubatch.pos[i + ubatch.n_tokens*2],

/*.y =*/ ubatch.pos[i + ubatch.n_tokens],

};

cells.ext_set(idx, ext);

}

for (int32_t s = 0; s < ubatch.n_seq_id[i]; s++) {

cells.seq_add(idx, ubatch.seq_id[i][s]);

}

}

}

// note: we want to preserve the invariant that all positions between [pos_min, pos_max] for each sequence

// will be present in the cache. so we have to purge any position which is less than those we would overwrite

// ref: https://github.com/ggml-org/llama.cpp/pull/13746#issuecomment-2916057092

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

if (seq_pos_max_rm[s] == -1) {

continue;

}

GGML_ASSERT(s < seq_to_stream.size());

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

if (cells.seq_pos_min(s) <= seq_pos_max_rm[s]) {

LLAMA_LOG_DEBUG("%s: purging positions [%d, %d] of sequence %d from KV cache\n",

__func__, cells.seq_pos_min(s), seq_pos_max_rm[s], s);

seq_rm(s, cells.seq_pos_min(s), seq_pos_max_rm[s] + 1);

}

}

// move the head at the end of the slot

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

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

head = sinfo.idxs[s].back() + 1;

}

}

bool llama_kv_cache::get_can_shift() const {

// Step35 uses per-layer RoPE dims; K-shift assumes a single global n_rot.

if (model.arch == LLM_ARCH_STEP35) {

return false;

}

if (hparams.n_pos_per_embd() > 1) {

return false;

}

return true;

}

uint32_t llama_kv_cache::get_size() const {

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

return cells.size();

}

uint32_t llama_kv_cache::get_n_stream() const {

return n_stream;

}

bool llama_kv_cache::get_has_shift() const {

bool result = false;

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