#include "llama-memory-recurrent.h"

#include "ggml-backend.h"

#include "llama-impl.h"

#include "llama-io.h"

#include "llama-batch.h"

#include "llama-model.h"

#include <algorithm>

#include <cassert>

#include <cstring>

#include <limits>

#include <map>

#include <stdexcept>

//

// llama_memory_recurrent

//

llama_memory_recurrent::llama_memory_recurrent(

const llama_model & model,

ggml_type type_r,

ggml_type type_s,

bool offload,

uint32_t mem_size,

uint32_t n_seq_max,

uint32_t n_rs_seq,

const layer_filter_cb & filter) : hparams(model.hparams), n_seq_max(n_seq_max) {

const int32_t n_layer = hparams.n_layer();

head = 0;

size = mem_size;

used = 0;

this->n_rs_seq = n_rs_seq;

rs_idx.assign(n_seq_max, 0);

cells.clear();

cells.resize(mem_size);

// 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*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();

};

r_l.resize(n_layer);

s_l.resize(n_layer);

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

if (filter && !filter(i)) {

LLAMA_LOG_DEBUG("%s: layer %3d: skipped\n", __func__, i);

continue;

}

const char * dev_name = "CPU";

ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type();

if (offload) {

auto * dev = model.dev_layer(i);

buft = ggml_backend_dev_buffer_type(dev);

dev_name = ggml_backend_dev_name(dev);

}

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

ggml_context * ctx = ctx_for_buft(buft);

if (!ctx) {

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

}

const uint32_t n_rows = mem_size * (1 + n_rs_seq);

ggml_tensor * r = ggml_new_tensor_2d(ctx, type_r, hparams.n_embd_r(), n_rows);

ggml_tensor * s = ggml_new_tensor_2d(ctx, type_s, hparams.n_embd_s(), n_rows);

ggml_format_name(r, "cache_r_l%d", i);

ggml_format_name(s, "cache_s_l%d", i);

r_l[i] = r;

s_l[i] = s;

}

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

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

ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft);

if (!buf) {

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

}

ggml_backend_buffer_clear(buf, 0);

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

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

}

{

const size_t memory_size_r = size_r_bytes();

const size_t memory_size_s = size_s_bytes();

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

(float)(memory_size_r + memory_size_s) / (1024.0f * 1024.0f), mem_size, n_layer, n_seq_max, n_rs_seq,

ggml_type_name(type_r), (float)memory_size_r / (1024.0f * 1024.0f),

ggml_type_name(type_s), (float)memory_size_s / (1024.0f * 1024.0f));

}

}

void llama_memory_recurrent::clear(bool data) {

for (int32_t i = 0; i < (int32_t) size; ++i) {

cells[i].pos = -1;

cells[i].seq_id.clear();

cells[i].src = -1;

cells[i].tail = -1;

}

head = 0;

used = 0;

if (data) {

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

ggml_backend_buffer_clear(buf.get(), 0);

}

}

std::fill(rs_idx.begin(), rs_idx.end(), 0);

}

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

uint32_t new_head = size;

if (p0 < 0) {

p0 = 0;

}

if (p1 < 0) {

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

}

const bool rm_all = p0 == 0 && p1 == std::numeric_limits<llama_pos>::max();

if (rm_all) {

if (seq_id >= 0) {

set_rs_idx(seq_id, 0);

} else {

std::fill(rs_idx.begin(), rs_idx.end(), 0);

}

}

// models like Mamba or RWKV can't have a state partially erased at the end

// of the sequence because their state isn't preserved for previous tokens

if (seq_id >= (int64_t) size) {

// could be fatal

return false;

}

if (0 <= seq_id) {

int32_t & tail_id = cells[seq_id].tail;

if (tail_id >= 0) {

auto & cell = cells[tail_id];

// partial rollback via per-token snapshot index (bounded by n_rs_seq)

if (0 < p0 && p0 <= cell.pos && p1 > cell.pos) {

const llama_pos rollback = cell.pos - (p0 - 1);

if (rollback >= 1 && rollback <= (llama_pos) n_rs_seq) {

set_rs_idx(seq_id, (uint32_t) rollback);

cell.pos = p0 - 1;

return true;

}

return false;

}

// invalidate tails which will be cleared

if (p0 <= cell.pos && cell.pos < p1) {

tail_id = -1;

}

}

} else {

// seq_id is negative, then the range should include everything or nothing

if (p0 != p1 && (p0 != 0 || p1 != std::numeric_limits<llama_pos>::max())) {

//printf("[DEBUG] inside `llama_memory_recurrent::seq_rm`: `seq_id` is negative, so returning false\n");

return false;

}

}

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

if (cells[i].pos >= p0 && cells[i].pos < p1) {

if (seq_id < 0) {

cells[i].seq_id.clear();

} else if (cells[i].has_seq_id(seq_id)) {

cells[i].seq_id.erase(seq_id);

} else {

continue;

}

if (cells[i].is_empty()) {

// keep count of the number of used cells

if (cells[i].pos >= 0) {

used--;

}

cells[i].pos = -1;

cells[i].src = -1;

if (new_head == size) {

new_head = i;

}

}

}

}

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

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

head = new_head;

}

return true;

}

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

if (seq_id_src == seq_id_dst) {

return;

}

if (p0 < 0) {

p0 = 0;

}

if (p1 < 0) {

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

}

if ((uint32_t) seq_id_dst < size && (uint32_t) seq_id_src < size) {

auto & tail_src = cells[seq_id_src];

auto & tail_dst = cells[seq_id_dst];

if (tail_dst.tail >= 0) {

// clear destination seq_id if it wasn't empty

auto & cell_dst = cells[tail_dst.tail];

cell_dst.seq_id.erase(seq_id_dst);

tail_dst.tail = -1;

if (cell_dst.seq_id.empty()) {

cell_dst.pos = -1;

cell_dst.src = -1;

used -= 1;

}

}

if (tail_src.tail >= 0) {

auto & cell_src = cells[tail_src.tail];

cell_src.seq_id.insert(seq_id_dst);

tail_dst.tail = tail_src.tail;

}

}

}

void llama_memory_recurrent::seq_keep(llama_seq_id seq_id) {

uint32_t new_head = size;

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

if ((llama_seq_id) i != seq_id) {

cells[i].tail = -1;

}

if (!cells[i].has_seq_id(seq_id)) {

if (cells[i].pos >= 0) {

used--;

}

cells[i].pos = -1;

cells[i].src = -1;

cells[i].seq_id.clear();

if (new_head == size){

new_head = i;

}

} else {

cells[i].seq_id.clear();

cells[i].seq_id.insert(seq_id);

}

}

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

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

head = new_head;

}

}

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

if (shift == 0) {

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

if (p0 == p1) {

return;

}

// for Mamba-like or RWKV models, only the pos needs to be shifted

if (0 <= seq_id && seq_id < (int64_t) size) {

const int32_t tail_id = cells[seq_id].tail;

if (tail_id >= 0) {

auto & cell = cells[tail_id];

if (cell.has_seq_id(seq_id) && p0 <= cell.pos && cell.pos < p1) {

cell.pos += shift;

}

}

}

}

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

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 Mamba-like or RWKV models, only the pos needs to be changed

if (0 <= seq_id && seq_id < (int64_t) size) {

const int32_t tail_id = cells[seq_id].tail;

if (tail_id >= 0) {

auto & cell = cells[tail_id];

if (cell.has_seq_id(seq_id) && p0 <= cell.pos && cell.pos < p1) {

cell.pos /= d;

}

}

}

}

llama_pos llama_memory_recurrent::seq_pos_min(llama_seq_id seq_id) const {

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

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

if (cells[i].has_seq_id(seq_id)) {

result = std::min(result, cells[i].pos);

}

}

if (result == std::numeric_limits<llama_pos>::max()) {

result = -1;

}

return result;

}

llama_pos llama_memory_recurrent::seq_pos_max(llama_seq_id seq_id) const {

llama_pos result = -1;

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

if (cells[i].has_seq_id(seq_id)) {

result = std::max(result, cells[i].pos);

}

}

return result;

}

void llama_memory_recurrent::set_rs_idx(llama_seq_id seq_id, uint32_t idx) {

if (seq_id < 0 || (size_t) seq_id >= rs_idx.size()) {

return;

}

rs_idx[seq_id] = (idx > n_rs_seq) ? n_rs_seq : idx;

}

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

std::map<ggml_backend_buffer_type_t, size_t> ret;

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

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

}

return ret;

}

llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) {

do {

balloc.split_reset();

std::vector<llama_ubatch> ubatches;

while (true) {

llama_ubatch ubatch;

if (embd_all) {

// if all tokens are output, split by sequence

ubatch = balloc.split_seq(n_ubatch);

} else {

// TODO: non-sequential equal split can be done if using unified KV cache

// for simplicity, we always use sequential equal split for now

// [TAG_RECURRENT_ROLLBACK_SPLITS]

// the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch

// so that the rollback snapshots remain valid

ubatch = balloc.split_equal(n_ubatch, true, n_rs_seq > 0 ? n_rs_seq + 1 : 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;

}

if (!prepare(ubatches)) {

break;

}

return std::make_unique<llama_memory_recurrent_context>(this, std::move(ubatches));

} while (false);

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

}

llama_memory_context_ptr llama_memory_recurrent::init_full() {

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

}

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

GGML_UNUSED(lctx);

GGML_UNUSED(optimize);

return std::make_unique<llama_memory_recurrent_context>(LLAMA_MEMORY_STATUS_NO_UPDATE);

}

bool llama_memory_recurrent::prepare(const std::vector<llama_ubatch> & ubatches) {

// simply remember the full state because it is very small for this type of cache

// TODO: optimize

auto org_cells = cells;

auto org_used = used;

auto org_head = head;

bool success = true;

for (const auto & ubatch : ubatches) {

if (!find_slot(ubatch)) {

success = false;

break;

}

}

// restore the original state

cells = std::move(org_cells);

used = org_used;

head = org_head;

return success;

}

bool llama_memory_recurrent::find_slot(const llama_ubatch & ubatch) {

const uint32_t n_seq_tokens = ubatch.n_seq_tokens;

const uint32_t n_seqs = ubatch.n_seqs;

// 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 > used + 2*n_seqs) {

head = 0;

}

// For recurrent state architectures (like Mamba or RWKV),

// each cache cell can store the state for a whole sequence.

// A slot should be always be contiguous.

// can only process batches with an equal number of new tokens in each sequence

GGML_ASSERT(ubatch.equal_seqs());

int32_t min = size - 1;

int32_t max = 0;

// everything should fit if all seq_ids are smaller than the max

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

const uint32_t i = s*n_seq_tokens; // first token of sequence set s

const uint32_t n_seq_id = ubatch.n_seq_id[i];

for (uint32_t j = 0; j < n_seq_id; ++j) {

const llama_seq_id seq_id = ubatch.seq_id[i][j];

if (seq_id < 0 || (uint32_t) seq_id >= size) {

// too big seq_id

// TODO: would it be possible to resize the cache instead?

LLAMA_LOG_ERROR("%s: seq_id=%d >= n_seq_max=%u Try using a bigger --parallel value\n", __func__, seq_id, n_seq_max);

return false;

}

if (j > 0) {

auto & seq = cells[seq_id];

if (seq.tail >= 0) {

auto & cell = cells[seq.tail];

// clear cells from seq_ids that become shared

// (should not normally happen, but let's handle it anyway)

cell.seq_id.erase(seq_id);

seq.tail = -1;

if (cell.seq_id.empty()) {

cell.pos = -1;

cell.src = -1;

used -= 1;

}

}

}

}

}

#ifndef NDEBUG

{

std::vector<int32_t> tails_verif;

tails_verif.assign(size, -1);

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

auto & cell = cells[i];

for (llama_seq_id seq_id : cell.seq_id) {

if (tails_verif[seq_id] != -1) {

LLAMA_LOG_ERROR("%s: duplicate tail for seq_id %d in cell %d and %d\n", __func__, seq_id, i, tails_verif[seq_id]);

}

tails_verif[seq_id] = i;

}

}

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

if (tails_verif[i] != cells[i].tail) {

LLAMA_LOG_ERROR("%s: wrong tail for seq_id %d, (%d instead of %d)\n", __func__, i, cells[i].tail, tails_verif[i]);

}

}

}

#endif

// find next empty cell

uint32_t next_empty_cell = head;

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

if (next_empty_cell >= size) { next_empty_cell -= size; }

auto & cell = cells[next_empty_cell];

if (cell.is_empty()) { break; }

next_empty_cell += 1;

}

// find usable cell range

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

const uint32_t i = s*n_seq_tokens;

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

auto & seq_meta = cells[seq_id];

bool has_cell = false;

if (seq_meta.tail >= 0) {

auto & cell = cells[seq_meta.tail];

GGML_ASSERT(cell.has_seq_id(seq_id));

// does this seq_id "own" the cell?

if (cell.seq_id.size() == 1) { has_cell = true; }

}

if (!has_cell) {

auto & empty_cell = cells[next_empty_cell];

GGML_ASSERT(empty_cell.is_empty());

// copy old tail into the empty cell

if (seq_meta.tail >= 0) {

auto & orig_cell = cells[seq_meta.tail];

empty_cell.pos = orig_cell.pos;

empty_cell.src = orig_cell.src;

orig_cell.seq_id.erase(seq_id);

empty_cell.seq_id.insert(seq_id); // will be overwritten

GGML_ASSERT(!orig_cell.is_empty()); // has at least one remaining seq_id

}

seq_meta.tail = next_empty_cell;

// find next empty cell

if (s + 1 < n_seqs) {

for (uint32_t j = 0; j < size; ++j) {

next_empty_cell += 1;

if (next_empty_cell >= size) { next_empty_cell -= size; }

auto & cell = cells[next_empty_cell];

if (cell.is_empty()) { break; }

}

}

}

if (min > seq_meta.tail) { min = seq_meta.tail; }

if (max < seq_meta.tail) { max = seq_meta.tail; }

}

// gather and re-order

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

const uint32_t i = s*n_seq_tokens;

const int32_t dst_id = s + min;

const int32_t src_id = cells[ubatch.seq_id[i][0]].tail;

if (dst_id != src_id) {

auto & dst_cell = cells[dst_id];

auto & src_cell = cells[src_id];

std::swap(dst_cell.pos, src_cell.pos);

std::swap(dst_cell.src, src_cell.src);

std::swap(dst_cell.seq_id, src_cell.seq_id);

// swap tails

for (uint32_t j = 0; j < size; ++j) {

int32_t & tail = cells[j].tail;

if (tail == src_id) {

tail = dst_id;

} else if (tail == dst_id) {

tail = src_id;

}

}

}

}

// update the pos of the used seqs

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

const uint32_t i = s*n_seq_tokens;

const llama_pos last_pos = ubatch.pos[i + n_seq_tokens - 1];

const int32_t cell_id = s + min;

auto & cell = cells[cell_id];

if (cell.pos >= 0 && last_pos != cell.pos + (llama_pos) n_seq_tokens) {

// What should happen when the pos backtracks or skips a value?

// Clearing the state mid-batch would require special-casing which isn't done.

LLAMA_LOG_WARN("%s: non-consecutive token position %d after %d for sequence %d with %u new tokens\n",

__func__, last_pos, cell.pos, ubatch.seq_id[i][0], n_seq_tokens);

}

cell.pos = last_pos;

cell.seq_id.clear();

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

const llama_seq_id seq_id = ubatch.seq_id[i][j];

cell.seq_id.insert(seq_id);

cells[seq_id].tail = cell_id;

}

}

// Find first cell without src refs, to use as the zero-ed state

{

// TODO: bake-in src refcounts in the cell metadata

std::vector<int32_t> refcounts(size, 0);

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

const int32_t src = cells[i].src;

if (src >= 0) {

refcounts[src] += 1;

}

}

rs_z = -1;

for (int i = min; i <= max; ++i) {

if (refcounts[i] == 0) {

rs_z = i;

break;

}

}

for (int i = min; i <= max; ++i) {

if (cells[i].src < 0) {

GGML_ASSERT(rs_z >= 0);

cells[i].src0 = rs_z;

} else {

// Stage the source ids for all used cells to allow correct seq_* behavior

// and still make these values available when setting the inputs

cells[i].src0 = cells[i].src;

}

cells[i].src = i; // avoid moving or clearing twice

}

}

// allow getting the range of used cells, from head to head + n

head = min;

n = max - min + 1;

used = std::count_if(cells.begin(), cells.end(),

[](const mem_cell & cell){ return !cell.is_empty(); });

// sanity check

return n >= n_seqs;

}

bool llama_memory_recurrent::get_can_shift() const {

// shifting the pos is trivial for recurrent models

return true;

}

size_t llama_memory_recurrent::total_size() const {

size_t size = 0;

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

size += ggml_backend_buffer_get_size(buf.get());

}

return size;

}

size_t llama_memory_recurrent::size_r_bytes() const {

size_t size_r_bytes = 0;

for (const auto & r : r_l) {

if (r != nullptr) {

size_r_bytes += ggml_nbytes(r);

}

}

return size_r_bytes;

}

size_t llama_memory_recurrent::size_s_bytes() const {

size_t size_s_bytes = 0;

for (const auto & s : s_l) {

if (s != nullptr) {

size_s_bytes += ggml_nbytes(s);

}

}

return size_s_bytes;

}

void llama_memory_recurrent::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {

GGML_UNUSED(flags);

std::vector<std::pair<uint32_t, uint32_t>> cell_ranges; // ranges, from inclusive, to exclusive

std::vector<std::pair<uint32_t, uint32_t>> cell_ranges_data; // logical source row ranges

uint32_t cell_count = 0;

// Count the number of cells with the specified seq_id

// Find all the ranges of cells with this seq id (or all, when -1)

uint32_t cell_range_begin = size;

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

const auto & cell = cells[i];

if ((seq_id == -1 && !cell.is_empty()) || cell.has_seq_id(seq_id)) {

++cell_count;

uint32_t rs_idx_cur = 0;

if (n_rs_seq != 0) {

if (seq_id != -1) {

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

rs_idx_cur = rs_idx[seq_id];

} else {

bool has_rs_idx = false;

for (const llama_seq_id cell_seq_id : cell.seq_id) {

GGML_ASSERT(cell_seq_id >= 0 && (size_t) cell_seq_id < rs_idx.size());

const uint32_t seq_rs_idx = rs_idx[cell_seq_id];

if (!has_rs_idx) {

rs_idx_cur = seq_rs_idx;

has_rs_idx = true;

} else if (rs_idx_cur != seq_rs_idx) {

GGML_ABORT("cannot write shared recurrent state with different rollback indices");

}

}

}

}

const uint32_t cell_id = rs_idx_cur * size + (cell.src >= 0 ? cell.src : (int32_t) i);

if (cell_ranges_data.empty() || cell_ranges_data.back().second != cell_id) {

cell_ranges_data.emplace_back(cell_id, cell_id + 1);

} else {

cell_ranges_data.back().second++;

}

if (cell_range_begin == size) {

cell_range_begin = i;

}

} else {

if (cell_range_begin != size) {

cell_ranges.emplace_back(cell_range_begin, i);

cell_range_begin = size;

}

}

}

if (cell_range_begin != size) {

cell_ranges.emplace_back(cell_range_begin, size);

}

if ((flags & LLAMA_STATE_SEQ_FLAGS_ON_DEVICE) && cell_ranges.size() > 1) {

GGML_ABORT("cannot save/load multiple ranges of cells to/from device memory\n");

}

// DEBUG CHECK: Sum of cell counts in ranges should equal the total cell count

uint32_t cell_count_check = 0;

for (const auto & range : cell_ranges) {

cell_count_check += range.second - range.first;

}

GGML_ASSERT(cell_count == cell_count_check);

cell_count_check = 0;

for (const auto & range : cell_ranges_data) {

cell_count_check += range.second - range.first;

}

GGML_ASSERT(cell_count == cell_count_check);

io.write(&cell_count, sizeof(cell_count));

state_write_meta(io, cell_ranges, seq_id);

state_write_data(io, cell_ranges_data);

}

void llama_memory_recurrent::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {

GGML_UNUSED(flags);

uint32_t cell_count;

io.read(&cell_count, sizeof(cell_count));

bool res = true;

res = res && state_read_meta(io, cell_count, seq_id);

res = res && state_read_data(io, cell_count);

if (!res) {

if (seq_id == -1) {

clear(true);

} else {

seq_rm(seq_id, -1, -1);

}

throw std::runtime_error("failed to restore kv cache");

}

if (n_rs_seq != 0) {

if (seq_id == -1) {

std::fill(rs_idx.begin(), rs_idx.end(), 0);

} else {

set_rs_idx(seq_id, 0);

}

}

}

void llama_memory_recurrent::state_write_meta(llama_io_write_i & io, const std::vector<std::pair<uint32_t, uint32_t>> & cell_ranges, llama_seq_id seq_id) const {

for (const auto & range : cell_ranges) {

for (uint32_t i = range.first; i < range.second; ++i) {

const auto & cell = cells[i];

const llama_pos pos = cell.pos;

const uint32_t n_seq_id = seq_id == -1 ? cell.seq_id.size() : 0;

io.write(&pos, sizeof(pos));

io.write(&n_seq_id, sizeof(n_seq_id));

if (n_seq_id) {

for (auto seq_id : cell.seq_id) {

io.write(&seq_id, sizeof(seq_id));

}

}

}

}

}

void llama_memory_recurrent::state_write_data(llama_io_write_i & io, const std::vector<std::pair<uint32_t, uint32_t>> & cell_ranges) const {

const uint32_t s_trans = 0;

const uint32_t n_layer = hparams.n_layer();

io.write(&s_trans, sizeof(s_trans));

io.write(&n_layer, sizeof(n_layer));

// Iterate and write all the R tensors first, each row is a cell

// Get whole range at a time

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

// skip null layers (read_data will handle this by checking "r_l" and "s_l" for null)

if (r_l[il] == nullptr) continue;

// Write R tensor type

const int32_t r_type_i = (int32_t)r_l[il]->type;

io.write(&r_type_i, sizeof(r_type_i));

// Write row size of R tensor

const uint64_t r_size_row = ggml_row_size(r_l[il]->type, hparams.n_embd_r());

io.write(&r_size_row, sizeof(r_size_row));

// Write each logical cell row range. With pending recurrent rollback,

// the logical current state may live in a rollback snapshot plane.

for (const auto & range : cell_ranges) {

const size_t range_size = range.second - range.first;

const size_t buf_size = range_size * r_size_row;

io.write_tensor(r_l[il], range.first * r_size_row, buf_size);

}

}

if (!s_trans) {

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

// skip null layers (read_data will handle this by checking "r_l" and "s_l" for null)

if (s_l[il] == nullptr) continue;

// Write S tensor type

const int32_t s_type_i = (int32_t)s_l[il]->type;

io.write(&s_type_i, sizeof(s_type_i));

// Write row size of S tensor

const uint64_t s_size_row = ggml_row_size(s_l[il]->type, hparams.n_embd_s());

io.write(&s_size_row, sizeof(s_size_row));

// Write each logical cell row range. With pending recurrent rollback,

// the logical current state may live in a rollback snapshot plane.

for (const auto & range : cell_ranges) {

const size_t range_size = range.second - range.first;

const size_t buf_size = range_size * s_size_row;

io.write_tensor(s_l[il], range.first * s_size_row, buf_size);

}

}

} else {

// When S tensor is transposed, we also need the element size and get the element ranges from each row

const uint32_t mem_size = size;

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

// skip null layers (read_data will handle this by checking "r_l" and "s_l" for null)

if (s_l[il] == nullptr) continue;

const uint32_t n_embd_s = hparams.n_embd_s();

// Write S tensor type

const int32_t s_type_i = (int32_t)s_l[il]->type;

io.write(&s_type_i, sizeof(s_type_i));

// Write element size

const uint32_t s_size_el = ggml_type_size(s_l[il]->type);

io.write(&s_size_el, sizeof(s_size_el));

// Write GQA embedding size

io.write(&n_embd_s, sizeof(n_embd_s));

// For each row, we get the element values of each logical cell

for (uint32_t j = 0; j < n_embd_s; ++j) {

for (const auto & range : cell_ranges) {

const size_t range_size = range.second - range.first;

const size_t src_offset = (range.first + j * mem_size) * s_size_el;

const size_t buf_size = range_size * s_size_el;

io.write_tensor(s_l[il], src_offset, buf_size);

}

}

}

}

}

bool llama_memory_recurrent::state_read_meta(llama_io_read_i & io, uint32_t cell_count, llama_seq_id dest_seq_id) {

if (dest_seq_id != -1) {

// single sequence

seq_rm(dest_seq_id, -1, -1);

if (cell_count == 0) {

return true;

}

llama_batch_allocr balloc(hparams.n_pos_per_embd());

llama_ubatch ubatch = balloc.ubatch_reserve(cell_count, 1);

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

llama_pos pos;

uint32_t n_seq_id;

io.read(&pos, sizeof(pos));

io.read(&n_seq_id, sizeof(n_seq_id));

if (n_seq_id != 0) {

LLAMA_LOG_ERROR("%s: invalid seq_id-agnostic kv cell\n", __func__);

return false;

}

ubatch.pos[i] = pos;

}

ubatch.n_seq_id[0] = 1;

ubatch.seq_id[0] = &dest_seq_id;

if (!find_slot(ubatch)) {

LLAMA_LOG_ERROR("%s: failed to find available cells in kv cache\n", __func__);

return false;

}

// DEBUG CHECK: kv.head should be our first cell, kv.head + cell_count - 1 should be our last cell (verify seq_id and pos values)

// Assume that this is one contiguous block of cells

GGML_ASSERT(head + cell_count <= size);

GGML_ASSERT(cells[head].pos == ubatch.pos[0]);

GGML_ASSERT(cells[head + cell_count - 1].pos == ubatch.pos[cell_count - 1]);

GGML_ASSERT(cells[head].has_seq_id(dest_seq_id));

GGML_ASSERT(cells[head + cell_count - 1].has_seq_id(dest_seq_id));

} else {

// whole KV cache restore

if (cell_count > size) {

LLAMA_LOG_ERROR("%s: not enough cells in kv cache\n", __func__);

return false;

}

clear(true);

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

auto & cell = cells[i];

llama_pos pos;