From cac9227a6af947b52dc510ce89ba2d1b1419062c Mon Sep 17 00:00:00 2001 From: Evandro Fonseca Junior Date: Thu, 26 Mar 2026 17:38:51 -0300 Subject: [PATCH] Revert "perf: avoid heap allocs during token detokenization" This reverts commit 5269f43e522ed03bfd888723c2ba492c51cf41b1. --- crates/inference-core/src/session.rs | 40 +++++++++++----------------- 1 file changed, 16 insertions(+), 24 deletions(-) diff --git a/crates/inference-core/src/session.rs b/crates/inference-core/src/session.rs index f6a9566..a0defd6 100644 --- a/crates/inference-core/src/session.rs +++ b/crates/inference-core/src/session.rs @@ -297,30 +297,22 @@ fn decode_token_piece_with(mut call: F) -> String where F: FnMut(*mut std::ffi::c_char, i32) -> i32, { - 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]; + let mut capacity = 64usize; loop { + let mut buffer = vec![0u8; capacity]; let written = call( - heap_buffer.as_mut_ptr() as *mut std::ffi::c_char, - heap_buffer.len() as i32, + buffer.as_mut_ptr() as *mut std::ffi::c_char, + capacity as i32, ); if written < 0 { - heap_buffer.resize((-written) as usize, 0); + capacity = (-written) as usize; continue; } - return String::from_utf8_lossy(&heap_buffer[..written as usize]).into_owned(); + let written = written as usize; + return String::from_utf8_lossy(&buffer[..written]).into_owned(); } } @@ -329,14 +321,17 @@ mod tests { use super::decode_token_piece_with; #[test] - fn uses_stack_buffer_when_piece_fits() { + fn retries_when_piece_buffer_is_too_small() { 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()); } @@ -355,21 +350,18 @@ mod tests { calls += 1; if calls == 1 { - return -300; + return -18; } - let bytes = b" Quantum computer rendered from a resized heap buffer"; - assert!(len >= bytes.len() as i32); + assert!(len >= 18); + let bytes = b" Quantum computer"; unsafe { std::ptr::copy_nonoverlapping(bytes.as_ptr(), buf as *mut u8, bytes.len()); } bytes.len() as i32 }); - assert_eq!( - decoded, - " Quantum computer rendered from a resized heap buffer" - ); + assert_eq!(decoded, " Quantum computer"); assert_eq!(calls, 2); } }