perf: stop generation on default turn boundaries

This commit is contained in:
Evandro Fonseca Junior committed 2026-03-26 17:52:59 -03:00
1 parent 1bd3365a31
commit 653f5b5bc1
2 files changed
+81 -12

No files matched your search

+54 -9
View File
@@ -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 {
+27 -3
View File
@@ -514,7 +514,17 @@ 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 = req.stop.clone();
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.repeat_penalty = if req.tools.is_empty() { 1.05 } else { 1.0 };
sampler.seed = req.seed.unwrap_or(0);
sampler
@@ -565,8 +575,8 @@ pub fn ensure_supported_mode(req: &ChatCompletionRequest) -> Result<()> {
#[cfg(test)]
mod tests {
use super::{
authorize_request, finalize_assistant_message, prepare_prompt, validate_request,
FinalizedAssistant, ModelCatalog, ServerConfig, StreamingMode,
authorize_request, build_sampler, finalize_assistant_message, prepare_prompt,
validate_request, FinalizedAssistant, ModelCatalog, ServerConfig, StreamingMode,
};
use api_contracts::chat::ChatCompletionRequest;
use axum::http::{HeaderMap, HeaderValue, StatusCode};
@@ -718,4 +728,18 @@ 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()));
}
}