From fc7096f80b14240a0a0fdecc79dc570c8620d219 Mon Sep 17 00:00:00 2001 From: Praneth Paruchuri Date: Mon, 26 Jan 2026 03:24:15 +0530 Subject: [PATCH] [model-gateway] Optimize special token search using Aho-Corasick (#17387) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- sgl-model-gateway/Cargo.toml | 5 ++ .../benches/special_token_search.rs | 53 ++++++++++++++++ sgl-model-gateway/src/tokenizer/stop.rs | 62 +++++++++++++------ 3 files changed, 100 insertions(+), 20 deletions(-) create mode 100644 sgl-model-gateway/benches/special_token_search.rs diff --git a/sgl-model-gateway/Cargo.toml b/sgl-model-gateway/Cargo.toml index 45c9f1c2a..bae768800 100644 --- a/sgl-model-gateway/Cargo.toml +++ b/sgl-model-gateway/Cargo.toml @@ -139,6 +139,7 @@ sha2 = "0.10" wasmtime = { version = "38.0", features = ["component-model", "async"] } wasmtime-wasi = "38.0" async-channel = "2.5" +aho-corasick = "1.1.4" [build-dependencies] tonic-prost-build = "0.14.2" @@ -160,6 +161,10 @@ tonic-v12 = { version = "0.12.3", package = "tonic" } serial_test = "3.0" rsa = { version = "0.9", features = ["sha2"] } +[[bench]] +name = "special_token_search" +harness = false +path = "benches/special_token_search.rs" [[bench]] name = "wasm_middleware_latency" harness = false diff --git a/sgl-model-gateway/benches/special_token_search.rs b/sgl-model-gateway/benches/special_token_search.rs new file mode 100644 index 000000000..2c57a44b0 --- /dev/null +++ b/sgl-model-gateway/benches/special_token_search.rs @@ -0,0 +1,53 @@ +use aho_corasick::AhoCorasick; +use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput}; + +fn find_special_token_boundaries_naive( + text: &str, + special_tokens: &[String], +) -> Vec<(usize, usize)> { + let mut boundaries = Vec::new(); + for token in special_tokens { + let mut start = 0; + while let Some(pos) = text[start..].find(token) { + let actual_pos = start + pos; + boundaries.push((actual_pos, actual_pos + token.len())); + start = actual_pos + token.len(); + } + } + boundaries.sort_by_key(|b| b.0); + boundaries +} + +fn find_special_token_boundaries_aho(text: &str, ac: &AhoCorasick) -> Vec<(usize, usize)> { + ac.find_iter(text) + .map(|mat| (mat.start(), mat.end())) + .collect() +} + +fn bench_token_search_comparison(c: &mut Criterion) { + let mut group = c.benchmark_group("token_boundary_search"); + let text = "User: Hello! Assistant: How can I help you today? ".repeat(1000); + + for token_count in [5, 50] { + let special_tokens: Vec = (0..token_count) + .map(|i| format!("<|stop_sequence_{}|>", i)) + .collect(); + + let ac = AhoCorasick::new(&special_tokens).unwrap(); + group.throughput(Throughput::Bytes(text.len() as u64)); + + group.bench_function(format!("naive_tokens_{}", token_count), |b| { + b.iter(|| { + find_special_token_boundaries_naive(black_box(&text), black_box(&special_tokens)) + }) + }); + + group.bench_function(format!("aho_tokens_{}", token_count), |b| { + b.iter(|| find_special_token_boundaries_aho(black_box(&text), black_box(&ac))) + }); + } + group.finish(); +} + +criterion_group!(benches, bench_token_search_comparison); +criterion_main!(benches); diff --git a/sgl-model-gateway/src/tokenizer/stop.rs b/sgl-model-gateway/src/tokenizer/stop.rs index c6f9ddc72..5d45b4cd1 100644 --- a/sgl-model-gateway/src/tokenizer/stop.rs +++ b/sgl-model-gateway/src/tokenizer/stop.rs @@ -1,5 +1,6 @@ use std::{collections::HashSet, sync::Arc}; +use aho_corasick::AhoCorasick; use anyhow::Result; use super::{ @@ -64,6 +65,8 @@ pub struct StopSequenceDecoder { /// Sequence for incremental decoding (replaces token_buffer + offsets) sequence: Sequence, config: StopSequenceConfig, + aho_corasick: Option, + visible_boundary_idx: usize, /// Buffer for partial matches (the "jail") jail_buffer: String, /// Whether we've stopped @@ -77,9 +80,31 @@ impl StopSequenceDecoder { config: StopSequenceConfig, skip_special_tokens: bool, ) -> Self { + let mut patterns: Vec = config + .stop_sequences + .iter() + .filter(|s| !s.is_empty()) + .cloned() + .collect(); + let visible_boundary_idx = patterns.len(); + patterns.extend( + config + .visible_stop_sequences + .iter() + .filter(|s| !s.is_empty()) + .cloned(), + ); + + let aho_corasick = if patterns.is_empty() { + None + } else { + Some(AhoCorasick::new(patterns).expect("Failed to build Aho-Corasick automaton")) + }; StopSequenceDecoder { sequence: Sequence::new_with_options(tokenizer, skip_special_tokens), config, + aho_corasick, + visible_boundary_idx, jail_buffer: String::new(), stopped: false, } @@ -122,28 +147,25 @@ impl StopSequenceDecoder { self.jail_buffer.push_str(&new_text); - // Check for hidden stop sequences - for stop_seq in &self.config.stop_sequences { - if let Some(pos) = self.jail_buffer.find(stop_seq) { + // Check for stop sequences + if let Some(ac) = &self.aho_corasick { + if let Some(mat) = ac.find(&self.jail_buffer) { self.stopped = true; - let output = self.jail_buffer[..pos].to_string(); - self.jail_buffer.clear(); - return Ok(if output.is_empty() { - SequenceDecoderOutput::Stopped - } else { - SequenceDecoderOutput::StoppedWithText(output) - }); - } - } + let is_visible = mat.pattern().as_usize() >= self.visible_boundary_idx; - // Check for visible stop sequences - for stop_seq in &self.config.visible_stop_sequences { - if let Some(pos) = self.jail_buffer.find(stop_seq) { - self.stopped = true; - let end_pos = pos + stop_seq.len(); - let output = self.jail_buffer[..end_pos].to_string(); - self.jail_buffer.clear(); - return Ok(SequenceDecoderOutput::StoppedWithText(output)); + if is_visible { + let output = self.jail_buffer[..mat.end()].to_string(); + self.jail_buffer.clear(); + return Ok(SequenceDecoderOutput::StoppedWithText(output)); + } else { + let output = self.jail_buffer[..mat.start()].to_string(); + self.jail_buffer.clear(); + return Ok(if output.is_empty() { + SequenceDecoderOutput::Stopped + } else { + SequenceDecoderOutput::StoppedWithText(output) + }); + } } }