llama_cpp_bindings/
llama_batch.rs

1use std::marker::PhantomData;
2
3use llama_cpp_bindings_sys::{
4    llama_batch, llama_batch_free, llama_batch_init, llama_pos, llama_seq_id,
5};
6
7use crate::batch_add_error::BatchAddError;
8use crate::sampled_token::SampledToken;
9use crate::token::LlamaToken;
10
11fn checked_n_tokens_plus_one_as_usize(n_tokens: i32) -> Result<usize, BatchAddError> {
12    let incremented = n_tokens.checked_add(1).ok_or_else(|| {
13        BatchAddError::IntegerOverflow(format!("n_tokens + 1 overflows i32: {n_tokens}"))
14    })?;
15
16    usize::try_from(incremented).map_err(|convert_error| {
17        BatchAddError::IntegerOverflow(format!("cannot fit n_tokens into a usize: {convert_error}"))
18    })
19}
20
21fn checked_i32_as_usize(value: i32, description: &str) -> Result<usize, BatchAddError> {
22    usize::try_from(value).map_err(|convert_error| {
23        BatchAddError::IntegerOverflow(format!(
24            "cannot fit {description} into a usize: {convert_error}"
25        ))
26    })
27}
28
29fn checked_usize_as_llama_seq_id(
30    value: usize,
31    description: &str,
32) -> Result<llama_seq_id, BatchAddError> {
33    llama_seq_id::try_from(value).map_err(|convert_error| {
34        BatchAddError::IntegerOverflow(format!(
35            "cannot fit {description} into a llama_seq_id: {convert_error}"
36        ))
37    })
38}
39
40fn checked_usize_as_i32(value: usize, description: &str) -> Result<i32, BatchAddError> {
41    i32::try_from(value).map_err(|convert_error| {
42        BatchAddError::IntegerOverflow(format!(
43            "cannot fit {description} into a i32: {convert_error}"
44        ))
45    })
46}
47
48fn checked_usize_as_llama_pos(value: usize, description: &str) -> Result<llama_pos, BatchAddError> {
49    llama_pos::try_from(value).map_err(|convert_error| {
50        BatchAddError::IntegerOverflow(format!(
51            "cannot fit {description} into a llama_pos: {convert_error}"
52        ))
53    })
54}
55
56#[derive(Debug)]
57pub struct LlamaBatch<'tokens> {
58    allocated: usize,
59    pub initialized_logits: Vec<i32>,
60    pub llama_batch: llama_batch,
61    phantom: PhantomData<&'tokens [LlamaToken]>,
62}
63
64impl<'tokens> LlamaBatch<'tokens> {
65    pub fn clear(&mut self) {
66        self.llama_batch.n_tokens = 0;
67        self.initialized_logits.clear();
68    }
69
70    /// # Errors
71    ///
72    /// Returns an error if there is insufficient space in the buffer or if integer conversions fail.
73    pub fn add(
74        &mut self,
75        sampled_token: &SampledToken,
76        pos: llama_pos,
77        seq_ids: &[i32],
78        logits: bool,
79    ) -> Result<(), BatchAddError> {
80        let (SampledToken::Content(LlamaToken(id))
81        | SampledToken::Reasoning(LlamaToken(id))
82        | SampledToken::ToolCall(LlamaToken(id))
83        | SampledToken::Undeterminable(LlamaToken(id))) = *sampled_token;
84        let required = checked_n_tokens_plus_one_as_usize(self.n_tokens())?;
85
86        if self.allocated < required {
87            return Err(BatchAddError::InsufficientSpace(self.allocated));
88        }
89
90        let offset = self.llama_batch.n_tokens;
91        let offset_usize = checked_i32_as_usize(offset, "n_tokens")?;
92        let n_seq_id = checked_usize_as_llama_seq_id(seq_ids.len(), "seq_ids.len()")?;
93
94        unsafe {
95            self.llama_batch.token.add(offset_usize).write(id);
96            self.llama_batch.pos.add(offset_usize).write(pos);
97            self.llama_batch.n_seq_id.add(offset_usize).write(n_seq_id);
98            for (seq_index, seq_id) in seq_ids.iter().enumerate() {
99                let tmp = *self.llama_batch.seq_id.add(offset_usize);
100                tmp.add(seq_index).write(*seq_id);
101            }
102            self.llama_batch
103                .logits
104                .add(offset_usize)
105                .write(i8::from(logits));
106        }
107
108        if logits {
109            self.initialized_logits.push(offset);
110        }
111
112        self.llama_batch.n_tokens += 1;
113
114        Ok(())
115    }
116
117    /// # Errors
118    ///
119    /// Returns an error if there is insufficient space in the buffer or if integer conversions fail.
120    pub fn add_sequence(
121        &mut self,
122        tokens: &[LlamaToken],
123        seq_id: i32,
124        logits_all: bool,
125    ) -> Result<(), BatchAddError> {
126        let last_index = checked_usize_as_llama_pos(tokens.len().saturating_sub(1), "n_tokens")?;
127
128        for (position, token) in (0..).zip(tokens.iter()) {
129            self.add(
130                &SampledToken::Content(*token),
131                position,
132                &[seq_id],
133                logits_all || position == last_index,
134            )?;
135        }
136
137        Ok(())
138    }
139
140    /// # Errors
141    ///
142    /// Returns an error if `n_tokens` exceeds `i32::MAX`.
143    pub fn new(n_tokens: usize, n_seq_max: i32) -> Result<Self, BatchAddError> {
144        let n_tokens_i32 = checked_usize_as_i32(n_tokens, "n_tokens")?;
145        let batch = unsafe { llama_batch_init(n_tokens_i32, 0, n_seq_max) };
146
147        Ok(LlamaBatch {
148            allocated: n_tokens,
149            initialized_logits: vec![],
150            llama_batch: batch,
151            phantom: PhantomData,
152        })
153    }
154
155    /// # Errors
156    ///
157    /// Returns an error if the provided token buffer is empty or if integer conversions fail.
158    pub fn get_one(tokens: &'tokens [LlamaToken]) -> Result<Self, BatchAddError> {
159        if tokens.is_empty() {
160            return Err(BatchAddError::EmptyBuffer);
161        }
162
163        let token_count = checked_usize_as_i32(tokens.len(), "token count")?;
164
165        let batch = unsafe {
166            let ptr = tokens.as_ptr().cast::<i32>().cast_mut();
167            llama_cpp_bindings_sys::llama_batch_get_one(ptr, token_count)
168        };
169
170        let last_token_index = checked_usize_as_i32(tokens.len() - 1, "last token index")?;
171
172        Ok(Self {
173            allocated: 0,
174            initialized_logits: vec![last_token_index],
175            llama_batch: batch,
176            phantom: PhantomData,
177        })
178    }
179
180    #[must_use]
181    pub const fn n_tokens(&self) -> i32 {
182        self.llama_batch.n_tokens
183    }
184}
185
186impl Drop for LlamaBatch<'_> {
187    fn drop(&mut self) {
188        unsafe {
189            if self.allocated > 0 {
190                llama_batch_free(self.llama_batch);
191            }
192        }
193    }
194}
195
196#[cfg(test)]
197mod tests {
198    use crate::sampled_token::SampledToken;
199    use crate::token::LlamaToken;
200
201    use super::{
202        BatchAddError, LlamaBatch, checked_i32_as_usize, checked_n_tokens_plus_one_as_usize,
203        checked_usize_as_i32, checked_usize_as_llama_pos, checked_usize_as_llama_seq_id,
204    };
205
206    #[test]
207    fn new_creates_empty_batch() {
208        let batch = LlamaBatch::new(16, 1).unwrap();
209
210        assert_eq!(batch.n_tokens(), 0);
211        assert!(batch.initialized_logits.is_empty());
212    }
213
214    #[test]
215    fn clear_resets_batch() {
216        let mut batch = LlamaBatch::new(16, 1).unwrap();
217        batch
218            .add(&SampledToken::Content(LlamaToken::new(1)), 0, &[0], true)
219            .unwrap();
220        assert_eq!(batch.n_tokens(), 1);
221
222        batch.clear();
223
224        assert_eq!(batch.n_tokens(), 0);
225        assert!(batch.initialized_logits.is_empty());
226    }
227
228    #[test]
229    fn add_increments_token_count() {
230        let mut batch = LlamaBatch::new(16, 1).unwrap();
231
232        batch
233            .add(&SampledToken::Content(LlamaToken::new(1)), 0, &[0], false)
234            .unwrap();
235        assert_eq!(batch.n_tokens(), 1);
236
237        batch
238            .add(&SampledToken::Content(LlamaToken::new(2)), 1, &[0], false)
239            .unwrap();
240        assert_eq!(batch.n_tokens(), 2);
241    }
242
243    #[test]
244    fn add_tracks_logits() {
245        let mut batch = LlamaBatch::new(16, 1).unwrap();
246
247        batch
248            .add(&SampledToken::Content(LlamaToken::new(1)), 0, &[0], false)
249            .unwrap();
250        assert!(batch.initialized_logits.is_empty());
251
252        batch
253            .add(&SampledToken::Content(LlamaToken::new(2)), 1, &[0], true)
254            .unwrap();
255        assert_eq!(batch.initialized_logits, vec![1]);
256    }
257
258    #[test]
259    fn add_returns_insufficient_space_when_full() {
260        let mut batch = LlamaBatch::new(1, 1).unwrap();
261        batch
262            .add(&SampledToken::Content(LlamaToken::new(1)), 0, &[0], false)
263            .unwrap();
264
265        let result = batch.add(&SampledToken::Content(LlamaToken::new(2)), 1, &[0], false);
266
267        assert_eq!(result, Err(BatchAddError::InsufficientSpace(1)));
268    }
269
270    #[test]
271    fn add_accepts_reasoning_sampled_token_variant() {
272        let mut batch = LlamaBatch::new(4, 1).unwrap();
273
274        batch
275            .add(&SampledToken::Reasoning(LlamaToken::new(11)), 0, &[0], true)
276            .unwrap();
277
278        assert_eq!(batch.n_tokens(), 1);
279    }
280
281    #[test]
282    fn add_accepts_tool_call_sampled_token_variant() {
283        let mut batch = LlamaBatch::new(4, 1).unwrap();
284
285        batch
286            .add(&SampledToken::ToolCall(LlamaToken::new(22)), 0, &[0], true)
287            .unwrap();
288
289        assert_eq!(batch.n_tokens(), 1);
290    }
291
292    #[test]
293    fn add_accepts_undeterminable_sampled_token_variant() {
294        let mut batch = LlamaBatch::new(4, 1).unwrap();
295
296        batch
297            .add(
298                &SampledToken::Undeterminable(LlamaToken::new(33)),
299                0,
300                &[0],
301                false,
302            )
303            .unwrap();
304
305        assert_eq!(batch.n_tokens(), 1);
306    }
307
308    #[test]
309    fn add_sequence_adds_all_tokens() {
310        let mut batch = LlamaBatch::new(16, 1).unwrap();
311        let tokens = vec![
312            LlamaToken::new(10),
313            LlamaToken::new(20),
314            LlamaToken::new(30),
315        ];
316
317        batch.add_sequence(&tokens, 0, false).unwrap();
318
319        assert_eq!(batch.n_tokens(), 3);
320    }
321
322    #[test]
323    fn add_sequence_sets_logits_on_last_token() {
324        let mut batch = LlamaBatch::new(16, 1).unwrap();
325        let tokens = vec![
326            LlamaToken::new(10),
327            LlamaToken::new(20),
328            LlamaToken::new(30),
329        ];
330
331        batch.add_sequence(&tokens, 0, false).unwrap();
332
333        assert_eq!(batch.initialized_logits, vec![2]);
334    }
335
336    #[test]
337    fn add_sequence_insufficient_space() {
338        let mut batch = LlamaBatch::new(2, 1).unwrap();
339        let tokens = vec![
340            LlamaToken::new(10),
341            LlamaToken::new(20),
342            LlamaToken::new(30),
343        ];
344
345        let result = batch.add_sequence(&tokens, 0, false);
346
347        assert!(result.is_err());
348    }
349
350    #[test]
351    fn add_sequence_fails_mid_loop_when_batch_fills() {
352        let mut batch = LlamaBatch::new(2, 1).unwrap();
353        batch
354            .add(&SampledToken::Content(LlamaToken::new(1)), 0, &[0], false)
355            .unwrap();
356
357        let tokens = vec![LlamaToken::new(10), LlamaToken::new(20)];
358        let result = batch.add_sequence(&tokens, 0, false);
359
360        assert!(result.is_err());
361    }
362
363    #[test]
364    fn get_one_with_valid_tokens() {
365        let tokens = vec![LlamaToken::new(1), LlamaToken::new(2)];
366        let batch = LlamaBatch::get_one(&tokens).expect("test: get_one should succeed");
367
368        assert_eq!(batch.n_tokens(), 2);
369        assert_eq!(batch.initialized_logits, vec![1]);
370    }
371
372    #[test]
373    fn get_one_empty_slice_returns_error() {
374        let tokens: Vec<LlamaToken> = vec![];
375        let result = LlamaBatch::get_one(&tokens);
376
377        assert_eq!(result.unwrap_err(), BatchAddError::EmptyBuffer);
378    }
379
380    #[test]
381    fn get_one_single_token() {
382        let tokens = vec![LlamaToken::new(42)];
383        let batch = LlamaBatch::get_one(&tokens).expect("test: get_one should succeed");
384
385        assert_eq!(batch.n_tokens(), 1);
386        assert_eq!(batch.initialized_logits, vec![0]);
387    }
388
389    #[test]
390    fn add_with_logits_false_retains_only_previous_logits() {
391        let mut batch = LlamaBatch::new(16, 1).unwrap();
392
393        batch
394            .add(&SampledToken::Content(LlamaToken::new(1)), 0, &[0], true)
395            .unwrap();
396        assert_eq!(batch.initialized_logits, vec![0]);
397
398        batch
399            .add(&SampledToken::Content(LlamaToken::new(2)), 0, &[0], false)
400            .unwrap();
401        assert_eq!(batch.initialized_logits, vec![0]);
402    }
403
404    #[test]
405    fn add_sequence_with_logits_all_marks_every_token() -> Result<(), BatchAddError> {
406        let mut batch = LlamaBatch::new(16, 1)?;
407        let tokens = vec![
408            LlamaToken::new(10),
409            LlamaToken::new(20),
410            LlamaToken::new(30),
411        ];
412
413        batch.add_sequence(&tokens, 0, true)?;
414
415        assert_eq!(batch.n_tokens(), 3);
416        assert_eq!(batch.initialized_logits, vec![0, 1, 2]);
417
418        Ok(())
419    }
420
421    #[test]
422    fn add_with_multiple_seq_ids() -> Result<(), BatchAddError> {
423        let mut batch = LlamaBatch::new(16, 4)?;
424
425        batch.add(
426            &SampledToken::Content(LlamaToken::new(1)),
427            0,
428            &[0, 1, 2],
429            true,
430        )?;
431
432        assert_eq!(batch.n_tokens(), 1);
433        assert_eq!(batch.initialized_logits, vec![0]);
434
435        Ok(())
436    }
437
438    #[test]
439    fn drop_does_not_free_get_one_batch() {
440        let tokens = vec![LlamaToken::new(1), LlamaToken::new(2)];
441        let batch = LlamaBatch::get_one(&tokens).expect("test: get_one should succeed");
442
443        assert_eq!(batch.allocated, 0);
444        drop(batch);
445    }
446
447    #[test]
448    fn checked_n_tokens_plus_one_as_usize_succeeds_for_zero() {
449        let result = checked_n_tokens_plus_one_as_usize(0);
450
451        assert_eq!(result, Ok(1));
452    }
453
454    #[test]
455    fn checked_n_tokens_plus_one_as_usize_fails_for_negative() {
456        let result = checked_n_tokens_plus_one_as_usize(-2);
457
458        assert_eq!(
459            std::mem::discriminant(&result.unwrap_err()),
460            std::mem::discriminant(&BatchAddError::IntegerOverflow(String::new())),
461        );
462    }
463
464    #[test]
465    fn checked_n_tokens_plus_one_as_usize_fails_for_i32_max() {
466        let result = checked_n_tokens_plus_one_as_usize(i32::MAX);
467
468        assert_eq!(
469            std::mem::discriminant(&result.unwrap_err()),
470            std::mem::discriminant(&BatchAddError::IntegerOverflow(String::new())),
471        );
472    }
473
474    #[test]
475    fn checked_i32_as_usize_succeeds_for_zero() {
476        let result = checked_i32_as_usize(0, "test_value");
477
478        assert_eq!(result, Ok(0));
479    }
480
481    #[test]
482    fn checked_i32_as_usize_fails_for_negative() {
483        let result = checked_i32_as_usize(i32::MIN, "test_value");
484
485        assert_eq!(
486            std::mem::discriminant(&result.unwrap_err()),
487            std::mem::discriminant(&BatchAddError::IntegerOverflow(String::new())),
488        );
489    }
490
491    #[test]
492    fn checked_usize_as_llama_seq_id_succeeds_for_zero() {
493        let result = checked_usize_as_llama_seq_id(0, "test_value");
494
495        assert_eq!(result, Ok(0));
496    }
497
498    #[test]
499    fn checked_usize_as_llama_seq_id_fails_for_overflow() {
500        let result = checked_usize_as_llama_seq_id(usize::MAX, "test_value");
501
502        assert_eq!(
503            std::mem::discriminant(&result.unwrap_err()),
504            std::mem::discriminant(&BatchAddError::IntegerOverflow(String::new())),
505        );
506    }
507
508    #[test]
509    fn checked_usize_as_i32_succeeds_for_zero() {
510        let result = checked_usize_as_i32(0, "test_value");
511
512        assert_eq!(result, Ok(0));
513    }
514
515    #[test]
516    fn checked_usize_as_i32_fails_for_overflow() {
517        let result = checked_usize_as_i32(usize::MAX, "test_value");
518
519        assert_eq!(
520            std::mem::discriminant(&result.unwrap_err()),
521            std::mem::discriminant(&BatchAddError::IntegerOverflow(String::new())),
522        );
523    }
524
525    #[test]
526    fn checked_usize_as_llama_pos_succeeds_for_zero() {
527        let result = checked_usize_as_llama_pos(0, "test_value");
528
529        assert_eq!(result, Ok(0));
530    }
531
532    #[test]
533    fn checked_usize_as_llama_pos_fails_for_overflow() {
534        let result = checked_usize_as_llama_pos(usize::MAX, "test_value");
535
536        assert_eq!(
537            std::mem::discriminant(&result.unwrap_err()),
538            std::mem::discriminant(&BatchAddError::IntegerOverflow(String::new())),
539        );
540    }
541
542    #[test]
543    fn new_fails_for_oversized_n_tokens() {
544        let result = LlamaBatch::new(usize::MAX, 1);
545
546        assert_eq!(
547            std::mem::discriminant(&result.unwrap_err()),
548            std::mem::discriminant(&BatchAddError::IntegerOverflow(String::new())),
549        );
550    }
551
552    #[test]
553    fn add_fails_when_required_token_count_overflows_i32() {
554        let mut batch = LlamaBatch::new(16, 1).unwrap();
555        batch.llama_batch.n_tokens = i32::MAX;
556
557        let result = batch.add(&SampledToken::Content(LlamaToken::new(1)), 0, &[0], false);
558
559        assert_eq!(
560            std::mem::discriminant(&result.unwrap_err()),
561            std::mem::discriminant(&BatchAddError::IntegerOverflow(String::new())),
562        );
563    }
564
565    #[test]
566    fn add_fails_when_existing_offset_is_negative() {
567        let mut batch = LlamaBatch::new(16, 1).unwrap();
568        batch.llama_batch.n_tokens = -1;
569
570        let result = batch.add(&SampledToken::Content(LlamaToken::new(1)), 0, &[0], false);
571
572        assert_eq!(
573            std::mem::discriminant(&result.unwrap_err()),
574            std::mem::discriminant(&BatchAddError::IntegerOverflow(String::new())),
575        );
576    }
577}
llama_cpp_bindings/llama_batch.rs

llama_cpp_bindings/
llama_batch.rs
