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

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
+2 -1
View File
@@ -104,7 +104,8 @@ public:
// make ubatches of equal-length sequences sets
// if sequential == true, the tokens in the ubatch will have increasing sequential sequence ids
llama_ubatch split_equal(uint32_t n_ubatch, bool sequential);
// n_keep_tail = minimum trailing tokens of a seq that must land in the same ubatch
llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail);
// sequence-set-wise split - each ubatch contains a single sequence-set
llama_ubatch split_seq(uint32_t n_ubatch);