Fix batch_pos

Signed-off-by: Adrien Gallouët <angt@huggingface.co>
This commit is contained in:
Adrien Gallouët 2025-01-31 18:20:45 +00:00
parent e88a527fcf
commit f38c34aeb7
No known key found for this signature in database

View File

@ -226,7 +226,7 @@ impl Llamacpp {
pos: bindings::llama_pos, pos: bindings::llama_pos,
seq_id: bindings::llama_seq_id, seq_id: bindings::llama_seq_id,
logits: bool, logits: bool,
) { ) -> usize {
let n = self.batch.n_tokens as usize; let n = self.batch.n_tokens as usize;
unsafe { unsafe {
*self.batch.token.add(n) = token; *self.batch.token.add(n) = token;
@ -236,6 +236,7 @@ impl Llamacpp {
*self.batch.logits.add(n) = logits as i8; *self.batch.logits.add(n) = logits as i8;
} }
self.batch.n_tokens += 1; self.batch.n_tokens += 1;
n
} }
// useless ? // useless ?
@ -578,8 +579,7 @@ impl LlamacppBackend {
for seq in seqs.iter_mut() { for seq in seqs.iter_mut() {
if seq.running { if seq.running {
llamacpp.batch_push(seq.token, seq.pos, seq.id as _, true); seq.batch_pos = llamacpp.batch_push(seq.token, seq.pos, seq.id as _, true);
seq.batch_pos = 0;
seq.pos += 1; seq.pos += 1;
} }
} }