From 6eaf59dd74e1b8c9e05cf6e397c40f762796e291 Mon Sep 17 00:00:00 2001 From: Evandro Fonseca Junior Date: Thu, 26 Mar 2026 17:54:48 -0300 Subject: [PATCH] Revert "perf: stop generation on default turn boundaries" This reverts commit caf4654c0a2a856c8f71dd465d4a4b1d67cde5bc. --- crates/inference-core/src/session.rs | 63 ++++----------------------- crates/inference-server/src/compat.rs | 30 ++----------- 2 files changed, 12 insertions(+), 81 deletions(-) diff --git a/crates/inference-core/src/session.rs b/crates/inference-core/src/session.rs index 02dd0af..a0defd6 100644 --- a/crates/inference-core/src/session.rs +++ b/crates/inference-core/src/session.rs @@ -27,8 +27,6 @@ pub struct GenerationSession { vocab: *const llama_vocab, sampler: *mut llama_sys::llama_sampler, stop_sequences: Vec, - stop_suffix_len: usize, - generated_tail: String, // generation state n_decoded: i32, // number of tokens decoded so far (prefill + gen) @@ -165,21 +163,12 @@ impl GenerationSession { ); // create session state - let stop_suffix_len = sampler_cfg - .stop_sequences - .iter() - .map(|stop| stop.len()) - .max() - .unwrap_or(0); - Ok(Self { lib, ctx, vocab, sampler: sampler_chain, // use sampler_chain here stop_sequences: sampler_cfg.stop_sequences.clone(), - stop_suffix_len, - generated_tail: String::new(), n_decoded: n_prompt_tokens, prompt_tokens: n_prompt_tokens as u32, max_tokens: sampler_cfg.max_tokens, @@ -238,10 +227,14 @@ impl GenerationSession { let token_text = self.detokenize_token(new_token); // Check stop sequences - if self.matches_stop_sequence(&token_text) { - self.finished = true; - self.finalize_timing(); - return None; + // (simple check: if generated text so far ends with a stop seq) + // For streaming we check per-token — good enough for Phase 2. + for stop in &self.stop_sequences { + if token_text.trim_end() == stop.as_str() { + self.finished = true; + self.finalize_timing(); + return None; + } } Some(token_text) @@ -298,28 +291,6 @@ impl GenerationSession { .into_owned() } } - - fn matches_stop_sequence(&mut self, token_text: &str) -> bool { - if self.stop_sequences.is_empty() { - return false; - } - - self.generated_tail.push_str(token_text); - tail_ends_with_stop( - &mut self.generated_tail, - self.stop_suffix_len, - &self.stop_sequences, - ) - } -} - -fn tail_ends_with_stop(tail: &mut String, max_len: usize, stop_sequences: &[String]) -> bool { - if tail.len() > max_len { - let trim_from = tail.len() - max_len; - tail.drain(..trim_from); - } - - stop_sequences.iter().any(|stop| tail.ends_with(stop)) } fn decode_token_piece_with(mut call: F) -> String @@ -347,7 +318,7 @@ where #[cfg(test)] mod tests { - use super::{decode_token_piece_with, tail_ends_with_stop}; + use super::decode_token_piece_with; #[test] fn retries_when_piece_buffer_is_too_small() { @@ -393,22 +364,6 @@ mod tests { assert_eq!(decoded, " Quantum computer"); assert_eq!(calls, 2); } - - #[test] - fn matches_stop_sequence_across_multiple_tokens() { - let stops = vec!["\nAssistant:".to_owned()]; - let mut tail = String::new(); - - tail.push_str("\nAssis"); - assert!(!tail_ends_with_stop( - &mut tail, - "\nAssistant:".len(), - &stops - )); - - tail.push_str("tant:"); - assert!(tail_ends_with_stop(&mut tail, "\nAssistant:".len(), &stops)); - } } impl Drop for GenerationSession { diff --git a/crates/inference-server/src/compat.rs b/crates/inference-server/src/compat.rs index 5e1eff9..0adbc6d 100644 --- a/crates/inference-server/src/compat.rs +++ b/crates/inference-server/src/compat.rs @@ -514,17 +514,7 @@ pub fn build_sampler(req: &ChatCompletionRequest) -> inference_core::SamplerConf sampler.temperature = req.temperature.clamp(0.0, 2.0); sampler.top_p = req.top_p.clamp(0.0, 1.0); sampler.max_tokens = req.max_tokens.max(1); - sampler.stop_sequences = if req.stop.is_empty() { - vec![ - "\nUser:".to_owned(), - "\nSystem:".to_owned(), - "\nDeveloper:".to_owned(), - "\nAssistant:".to_owned(), - "\nTool:".to_owned(), - ] - } else { - req.stop.clone() - }; + sampler.stop_sequences = req.stop.clone(); sampler.repeat_penalty = if req.tools.is_empty() { 1.05 } else { 1.0 }; sampler.seed = req.seed.unwrap_or(0); sampler @@ -575,8 +565,8 @@ pub fn ensure_supported_mode(req: &ChatCompletionRequest) -> Result<()> { #[cfg(test)] mod tests { use super::{ - authorize_request, build_sampler, finalize_assistant_message, prepare_prompt, - validate_request, FinalizedAssistant, ModelCatalog, ServerConfig, StreamingMode, + authorize_request, finalize_assistant_message, prepare_prompt, validate_request, + FinalizedAssistant, ModelCatalog, ServerConfig, StreamingMode, }; use api_contracts::chat::ChatCompletionRequest; use axum::http::{HeaderMap, HeaderValue, StatusCode}; @@ -728,18 +718,4 @@ mod tests { Some("Use references and lifetimes.") ); } - - #[test] - fn adds_default_turn_stops_when_request_has_none() { - let request = parse_request( - r#"{ - "model": "alias-model", - "messages": [{"role":"user","content":"hi"}] - }"#, - ); - - let sampler = build_sampler(&request); - assert!(sampler.stop_sequences.contains(&"\nUser:".to_owned())); - assert!(sampler.stop_sequences.contains(&"\nAssistant:".to_owned())); - } }