nicolas
/
ik_llama_opt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Discard very first compute graph for recurrent models (#1393)

This commit is contained in:
Kawrakow 2026-03-10 09:41:47 +01:00 committed by GitHub
parent f90b4c2f27
commit cda15bf175
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 14 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
src/llama.cpp
View File

@ -3514,8 +3514,12 @@ static int llama_decode_internal(
printf("graph_compute(...): %d us\n", int(tim2-tim1)); printf("graph_compute(...): %d us\n", int(tim2-tim1));
#endif #endif
bool reset_previous = false;
// update the kv ring buffer // update the kv ring buffer
{ {
if (llama_model_has_recurrent(&lctx.model) && kv_self.head == 0) {
reset_previous = true;
}
kv_self.head += n_tokens; kv_self.head += n_tokens;
// Ensure kv cache head points to a valid index. // Ensure kv cache head points to a valid index.
@ -3607,6 +3611,12 @@ static int llama_decode_internal(
} }
n_outputs_prev += lctx.n_outputs; n_outputs_prev += lctx.n_outputs;
cur_token += n_tokens; cur_token += n_tokens;
if (reset_previous) {
// We need to discard this graph. Otherwise, iwith CUDA graphs enabled, the graph will get resused and this will reset the
// recurrent state for each new token. This is probably not very relevant in practice because we basically never run TG with
// empty context, but for the sake of correctness let's just do it.
lctx.prev.reset();
}
} }
// set to total number of outputs in the batch, for use in llama_get_logits_ith // set to total number of outputs in the batch, for use in llama_get_logits_ith