Mirrors/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggerganov/llama.cpp.git synced 2024-10-29 22:20:15 +01:00
Code Issues Packages Projects Releases Wiki Activity

llama: correct reverting of the entire batch.

Browse Source 
also updates `llama_kv_cache_find_slot`, will correctly count the number of `used` cells for recurrent models
This commit is contained in:
Michael Podvitskiy 2024-10-22 19:57:15 +02:00
parent 0c05c60d56
commit bbf27cc2b8
1 changed files with 63 additions and 57 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

120
src/llama.cpp
View File

@ -2815,22 +2815,6 @@ struct llama_kv_cache {
}
};
// saves the kv_cache state for future recovery
// used to preserve the kv_cache state before searching for a slot
struct llama_kv_slot_restorer {
struct llama_kv_cache_state {
uint32_t head = 0;
uint32_t size = 0;
uint32_t used = 0;
uint32_t n = 0;
} old_state;
std::vector<llama_kv_cell> recurrent_cells; // for recurrent models only
std::pair<uint32_t, uint32_t> slot_boundaries; // for non-recurrent models only
bool restore = false;
};
struct llama_control_vector {
std::vector<struct ggml_tensor *> tensors; // per layer
std::vector<struct ggml_context *> ctxs;
@ -3666,21 +3650,24 @@ static bool llama_kv_cache_init(
// updates the cache head
// Note: On success, it's important that cache.head points
// to the first cell of the slot.
static bool llama_kv_cache_find_slot(
struct llama_kv_cache_slot_info {
std::pair<uint32_t, uint32_t> boundaries;
bool found = false;
explicit llama_kv_cache_slot_info(bool found_) : found{found_} {}
llama_kv_cache_slot_info(uint32_t begin, uint32_t end) : boundaries{begin, end}, found{true} {}
operator bool() const { return found; }
};
static const llama_kv_cache_slot_info llama_kv_cache_slot_info_failed{false};
static struct llama_kv_cache_slot_info llama_kv_cache_find_slot(
struct llama_kv_cache & cache,
const struct llama_ubatch & batch,
struct llama_kv_slot_restorer * slot_restorer = nullptr) {
const struct llama_ubatch & batch) {
const uint32_t n_tokens = batch.n_tokens;
const uint32_t n_seqs = batch.n_seqs;
const uint32_t n_seq_tokens = batch.n_seq_tokens;
if (slot_restorer != nullptr) {
slot_restorer->old_state.head = cache.head;
slot_restorer->old_state.size = cache.size;
slot_restorer->old_state.used = cache.used;
slot_restorer->old_state.n = cache.n;
}
if (cache.recurrent) {
// For recurrent state architectures (like Mamba or RWKV),
// each cache cell can store the state for a whole sequence.
@ -3689,11 +3676,6 @@ static bool llama_kv_cache_find_slot(
// can only process batches with an equal number of new tokens in each sequence
GGML_ASSERT(batch.equal_seqs);
if (slot_restorer != nullptr) {
slot_restorer->recurrent_cells = cache.cells;
slot_restorer->restore = true;
}
int32_t min = cache.size - 1;
int32_t max = 0;
@ -3707,7 +3689,7 @@ static bool llama_kv_cache_find_slot(
// too big seq_id
// TODO: would it be possible to resize the cache instead?
LLAMA_LOG_ERROR("%s: seq_id=%d >= n_seq_max=%d Try using a bigger --parallel value\n", __func__, seq_id, cache.size);
return false;
return llama_kv_cache_slot_info_failed;
}
if (j > 0) {
llama_kv_cell & seq = cache.cells[seq_id];
@ -3842,15 +3824,17 @@ static bool llama_kv_cache_find_slot(
// allow getting the range of used cells, from head to head + n
cache.head = min;
cache.n = max - min + 1;
cache.used = std::count_if(cache.cells.begin(), cache.cells.end(),
[](const llama_kv_cell& cell){ return !cell.is_empty(); });
// sanity check
return cache.n >= n_seqs;
return llama_kv_cache_slot_info(cache.n >= n_seqs);
}
// otherwise, one cell per token.
if (n_tokens > cache.size) {
LLAMA_LOG_ERROR("%s: n_tokens=%d > cache.size=%d\n", __func__, n_tokens, cache.size);
return false;
return llama_kv_cache_slot_info_failed;
}
uint32_t n_tested = 0;
@ -3878,15 +3862,10 @@ static bool llama_kv_cache_find_slot(
if (n_tested >= cache.size) {
//LLAMA_LOG_ERROR("%s: failed to find a slot for %d tokens\n", __func__, n_tokens);
return false;
return llama_kv_cache_slot_info_failed;
}
}
if (slot_restorer != nullptr) {
slot_restorer->slot_boundaries = std::make_pair(cache.head, cache.head + n_tokens);
slot_restorer->restore = true;
}
for (uint32_t s = 0; s < n_seqs; s++) {
for (uint32_t i = 0; i < n_seq_tokens; ++i) {
uint32_t k = s*n_seq_tokens + i;
@ -3900,7 +3879,7 @@ static bool llama_kv_cache_find_slot(
cache.used += n_tokens;
return true;
return llama_kv_cache_slot_info(cache.head, cache.head + n_tokens);
}
// find how many cells are currently in use
@ -4176,22 +4155,47 @@ static uint32_t llama_kv_cache_get_padding(const struct llama_cparams & cparams)
return cparams.flash_attn ? 256u : 32u;
}
static void llama_kv_cache_slot_restore(
const struct llama_kv_slot_restorer & restorer,
struct llama_kv_cache & cache) {
if (restorer.restore) {
cache.head = restorer.old_state.head;
cache.size = restorer.old_state.size;
cache.used = restorer.old_state.used;
cache.n = restorer.old_state.n;
// saves the kv_cache state for future recovery.
// used to rollback llama_kv_cache_find_slot changes.
struct llama_kv_slot_restorer {
struct llama_kv_cache_state {
uint32_t head = 0;
uint32_t n = 0;
} old_state;
if (cache.recurrent) {
cache.cells = restorer.recurrent_cells;
} else {
llama_kv_cache_seq_rm(cache, -1, restorer.slot_boundaries.first, restorer.slot_boundaries.second + 1);
std::vector<std::pair<uint32_t, uint32_t>> slot_boundaries; // for non-recurrent models only
bool do_restore = false;
explicit llama_kv_slot_restorer(const struct llama_kv_cache & cache) {
old_state.head = cache.head;
old_state.n = cache.n;
}
void save(const struct llama_kv_cache_slot_info& slot) {
if (slot) {
do_restore = true;
if (slot.boundaries.first != slot.boundaries.second) {
slot_boundaries.push_back(slot.boundaries);
}
}
}
}
void restore(struct llama_kv_cache & cache) {
if (do_restore) {
cache.head = old_state.head;
cache.n = old_state.n;
if (cache.recurrent) { // recurrent models like Mamba or RWKV can't have a state partially erased
llama_kv_cache_seq_rm(cache, -1, -1, -1);
} else {
for (auto & slot : slot_boundaries) {
llama_kv_cache_seq_rm(cache, -1, slot.first, slot.second);
}
}
}
}
};
//
// model loading and saving
@ -17235,7 +17239,7 @@ static int llama_decode_internal(
lctx.n_queued_tokens += n_tokens_all;
auto & kv_self = lctx.kv_self;
llama_kv_slot_restorer kv_slot_restorer;
llama_kv_slot_restorer kv_slot_restorer(kv_self);
const int64_t n_embd = hparams.n_embd;
const int64_t n_vocab = hparams.n_vocab;
@ -17320,9 +17324,11 @@ static int llama_decode_internal(
kv_self.head = 0;
}
if (!llama_kv_cache_find_slot(kv_self, ubatch, &kv_slot_restorer)) {
const auto slot = llama_kv_cache_find_slot(kv_self, ubatch);
if (!slot) {
return 1;
}
kv_slot_restorer.save(slot);
if (!kv_self.recurrent) {
// a heuristic, to avoid attending the full cache if it is not yet utilized
@ -17371,7 +17377,7 @@ static int llama_decode_internal(
const auto compute_status = llama_graph_compute(lctx, gf, n_threads, threadpool);
if (compute_status != GGML_STATUS_SUCCESS) {
llama_kv_cache_slot_restore(kv_slot_restorer, kv_self);
kv_slot_restorer.restore(kv_self);
switch (compute_status) {
case GGML_STATUS_ABORTED:
return 2;