perf: avoid heap allocs during token detokenization

crates/inference-core/src/session.rs

@@ -297,22 +297,30 @@ fn decode_token_piece_with<F>(mut call: F) -> String
 where
     F: FnMut(*mut std::ffi::c_char, i32) -> i32,
 {
-    let mut capacity = 64usize;
+    let mut stack_buffer = [0u8; 256];
+    let written = call(
+        stack_buffer.as_mut_ptr() as *mut std::ffi::c_char,
+        stack_buffer.len() as i32,
+    );
+
+    if written >= 0 {
+        return String::from_utf8_lossy(&stack_buffer[..written as usize]).into_owned();
+    }
+
+    let mut heap_buffer = vec![0u8; (-written) as usize];
 
     loop {
-        let mut buffer = vec![0u8; capacity];
         let written = call(
-            buffer.as_mut_ptr() as *mut std::ffi::c_char,
-            capacity as i32,
+            heap_buffer.as_mut_ptr() as *mut std::ffi::c_char,
+            heap_buffer.len() as i32,
         );
 
         if written < 0 {
-            capacity = (-written) as usize;
+            heap_buffer.resize((-written) as usize, 0);
             continue;
         }
 
-        let written = written as usize;
-        return String::from_utf8_lossy(&buffer[..written]).into_owned();
+        return String::from_utf8_lossy(&heap_buffer[..written as usize]).into_owned();
     }
 }
 
@@ -321,17 +329,14 @@ mod tests {
     use super::decode_token_piece_with;
 
     #[test]
-    fn retries_when_piece_buffer_is_too_small() {
+    fn uses_stack_buffer_when_piece_fits() {
         let mut calls = 0;
 
         let decoded = decode_token_piece_with(|buf, len| {
             calls += 1;
 
-            if len < 17 {
-                return -17;
-            }
-
             let bytes = b" Quantum computer";
+            assert!(len >= bytes.len() as i32);
             unsafe {
                 std::ptr::copy_nonoverlapping(bytes.as_ptr(), buf as *mut u8, bytes.len());
             }
@@ -350,18 +355,21 @@ mod tests {
             calls += 1;
 
             if calls == 1 {
-                return -18;
+                return -300;
             }
 
-            assert!(len >= 18);
-            let bytes = b" Quantum computer";
+            let bytes = b" Quantum computer rendered from a resized heap buffer";
+            assert!(len >= bytes.len() as i32);
             unsafe {
                 std::ptr::copy_nonoverlapping(bytes.as_ptr(), buf as *mut u8, bytes.len());
             }
             bytes.len() as i32
         });
 
-        assert_eq!(decoded, " Quantum computer");
+        assert_eq!(
+            decoded,
+            " Quantum computer rendered from a resized heap buffer"
+        );
         assert_eq!(calls, 2);
     }
 }