[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

@@ -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

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<String> = (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);

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<AhoCorasick>,
+    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<String> = 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)
+                    });
+                }
             }
         }