perf: stop generation on default turn boundaries

crates/inference-core/src/session.rs

@@ -27,6 +27,8 @@ pub struct GenerationSession {
     vocab: *const llama_vocab,
     sampler: *mut llama_sys::llama_sampler,
     stop_sequences: Vec<String>,
+    stop_suffix_len: usize,
+    generated_tail: String,
 
     // generation state
     n_decoded: i32, // number of tokens decoded so far (prefill + gen)
@@ -163,12 +165,21 @@ 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,
@@ -227,14 +238,10 @@ impl GenerationSession {
             let token_text = self.detokenize_token(new_token);
 
             // Check stop sequences
-            // (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;
-                }
+            if self.matches_stop_sequence(&token_text) {
+                self.finished = true;
+                self.finalize_timing();
+                return None;
             }
 
             Some(token_text)
@@ -291,6 +298,28 @@ 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<F>(mut call: F) -> String
@@ -318,7 +347,7 @@ where
 
 #[cfg(test)]
 mod tests {
-    use super::decode_token_piece_with;
+    use super::{decode_token_piece_with, tail_ends_with_stop};
 
     #[test]
     fn retries_when_piece_buffer_is_too_small() {
@@ -364,6 +393,22 @@ 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 {