Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

llama-batch: add n_keep_tail in split_equal for recurrent models (#25278)

Browse Source
This commit is contained in:
Aman Gupta
2026-07-08 15:55:19 +08:00
committed by GitHub
parent f296fdfbed
commit 230ea9d214
10 changed files with 102 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-memory-recurrent.cpp
+6 -9
View File
@@ -416,15 +416,12 @@ llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr &
// if all tokens are output, split by sequence
ubatch = balloc.split_seq(n_ubatch);
} else {
if (n_rs_seq > 0) {
// [TAG_RECURRENT_ROLLBACK_SPLITS]
// TODO: recurrent state rollback does not support equal splits
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
ubatch = balloc.split_equal(n_ubatch, true);
}
// 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) {