A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 13:30 UTC
[Rust Frontend] Reduce Gemma4 tool parser args scan complexity (#43850)
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Bugen Zhao committed
May 28, 2026 at 22:52 UTC
61a1e30473c78412463ae32fcb95c405bfee9126
2 files changed
+254
-53
rust/src/tool-parser/benches/gemma4.rs
+30
@@ -9,6 +9,7 @@ use utils::feed_parser;
9
10
const CHUNK_CHARS: usize = 7;
11
const LONG_NORMAL_TEXT_REPEATS: usize = 2048;
12
+const LONG_TOOL_ARGUMENT_REPEATS: usize = 256;
13
14
fn mixed_fixture() -> String {
15
concat!(
@@ -48,6 +49,24 @@ fn long_normal_text_fixture() -> String {
49
line.repeat(LONG_NORMAL_TEXT_REPEATS)
50
}
51
52
+fn long_tool_argument_fixture() -> String {
53
+ let line =
54
+ "<section><p>Literal } and <tool_call|> marker-shaped text inside content.</p></section>\n";
55
+ format!(
56
+ concat!(
57
+ "I will write the file.\n",
58
+ "<|tool_call>",
59
+ "call:write_file{{",
60
+ "path:<|\"|>index.html<|\"|>,",
61
+ "content:<|\"|>{}<|\"|>",
62
+ "}}",
63
+ "<tool_call|>",
64
+ "Done."
65
+ ),
66
+ line.repeat(LONG_TOOL_ARGUMENT_REPEATS)
67
+ )
68
+}
69
+
70
fn parser(tools: &[Tool]) -> Box<dyn ToolParser> {
71
Gemma4ToolParser::create(tools).expect("Gemma4 parser should initialize")
72
}
@@ -98,6 +117,7 @@ fn run_stream_group(
117
fn bench_gemma4(c: &mut Criterion) {
118
let tools = test_tools();
119
let mixed_text = mixed_fixture();
120
+ let long_tool_argument = long_tool_argument_fixture();
121
let long_normal_text = long_normal_text_fixture();
122
123
run_stream_group(
@@ -110,6 +130,16 @@ fn bench_gemma4(c: &mut Criterion) {
130
2,
131
);
132
133
+ run_stream_group(
134
+ c,
135
+ "gemma4/long_tool_argument",
136
+ &tools,
137
+ &long_tool_argument,
138
+ CHUNK_CHARS,
139
+ "I will write the file.\nDone.",
140
+ 1,
141
+ );
142
+
143
run_stream_group(
144
c,
145
"gemma4/long_normal_text",
rust/src/tool-parser/src/gemma4.rs
+224
-53
@@ -1,12 +1,12 @@
1
use serde_json::{Map, Number, Value};
2
use winnow::ascii::multispace0 as ws0;
3
-use winnow::combinator::{alt, delimited, opt, separated, seq, terminated};
3
+use winnow::combinator::{alt, delimited, eof, opt, separated, seq, terminated};
4
use winnow::error::{ContextError, ErrMode, ModalResult};
5
use winnow::prelude::*;
6
-use winnow::stream::Partial;
6
+use winnow::stream::{Partial, Stream};
7
use winnow::token::{literal, take_till, take_until};
8
9
-use super::utils::{parse_buffered_event, safe_text_len};
9
+use super::utils::{incomplete, parse_buffered_event, partial_prefix_len, safe_text_len};
10
use super::{Result, ToolCallDelta, ToolParser, ToolParserOutput};
11
use crate::Tool;
12
@@ -19,12 +19,26 @@ type Gemma4Input<'i> = Partial<&'i str>;
19
20
#[derive(Debug, Clone, PartialEq)]
21
enum Gemma4Event {
22
- Text {
23
- len: usize,
24
- },
22
+ Text { len: usize },
23
+ ToolCallStart,
24
+ ToolCallHeader { name: String },
25
+ ToolCall { args: Map<String, Value> },
26
+}
27
+
28
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
29
+struct Gemma4ArgsScanState {
30
+ scanned_len: usize,
31
+ in_string: bool,
32
+}
33
+
34
+#[derive(Debug, Clone, Default, PartialEq, Eq)]
35
+enum Gemma4Mode {
36
+ #[default]
37
+ Text,
38
+ Header,
39
ToolCall {
40
name: String,
27
- args: Map<String, Value>,
41
+ args_scan: Gemma4ArgsScanState,
42
},
43
}
44
@@ -40,6 +54,7 @@ enum Gemma4Event {
54
/// Arguments are emitted only after a full Gemma4 tool call is parsed.
55
pub struct Gemma4ToolParser {
56
buffer: String,
57
+ mode: Gemma4Mode,
58
emitted_tool_count: usize,
59
}
60
@@ -47,6 +62,7 @@ impl Gemma4ToolParser {
62
fn new(_tools: &[Tool]) -> Self {
63
Self {
64
buffer: String::new(),
65
+ mode: Gemma4Mode::default(),
66
emitted_tool_count: 0,
67
}
68
}
@@ -56,7 +72,20 @@ impl Gemma4ToolParser {
72
Gemma4Event::Text { len: consumed_len } => {
73
output.normal_text.push_str(&self.buffer[..consumed_len]);
74
}
59
- Gemma4Event::ToolCall { name, args } => {
75
+ Gemma4Event::ToolCallStart => self.mode = Gemma4Mode::Header,
76
+ Gemma4Event::ToolCallHeader { name } => {
77
+ self.mode = Gemma4Mode::ToolCall {
78
+ name,
79
+ args_scan: Gemma4ArgsScanState::default(),
80
+ };
81
+ }
82
+ Gemma4Event::ToolCall { args } => {
83
+ let mode = std::mem::replace(&mut self.mode, Gemma4Mode::Text);
84
+ let Gemma4Mode::ToolCall { name, .. } = mode else {
85
+ return Err(parsing_failed!(
86
+ "Gemma4 arguments without an active tool call"
87
+ ));
88
+ };
89
let arguments = serde_json::to_string(&args)
90
.map_err(|error| parsing_failed!("failed to serialize arguments: {}", error))?;
91
@@ -72,8 +101,24 @@ impl Gemma4ToolParser {
101
}
102
103
fn reset(&mut self) -> String {
104
+ let raw = match std::mem::replace(&mut self.mode, Gemma4Mode::Text) {
105
+ Gemma4Mode::Text => std::mem::take(&mut self.buffer),
106
+ Gemma4Mode::Header => {
107
+ format!("{}{}", TOOL_CALL_START, std::mem::take(&mut self.buffer))
108
+ }
109
+ Gemma4Mode::ToolCall { name, .. } => {
110
+ format!(
111
+ "{}{}{}{{{}",
112
+ TOOL_CALL_START,
113
+ CALL_PREFIX,
114
+ name,
115
+ std::mem::take(&mut self.buffer)
116
+ )
117
+ }
118
+ };
119
+ self.mode = Gemma4Mode::Text;
120
self.emitted_tool_count = 0;
76
- std::mem::take(&mut self.buffer)
121
+ raw
122
}
123
}
124
@@ -92,9 +137,11 @@ impl ToolParser for Gemma4ToolParser {
137
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
138
self.buffer.push_str(chunk);
139
95
- while let Some((event, consumed_len)) =
96
- parse_buffered_event(&self.buffer, parse_next_gemma4_event)?
97
- {
140
+ while let Some((event, consumed_len)) = {
141
+ parse_buffered_event(&self.buffer, |input| {
142
+ parse_next_gemma4_event(input, &mut self.mode)
143
+ })?
144
+ } {
145
self.apply_event(event, output)?;
146
self.buffer.drain(..consumed_len);
147
}
@@ -105,11 +152,11 @@ impl ToolParser for Gemma4ToolParser {
152
fn finish(&mut self) -> Result<ToolParserOutput> {
153
let mut output = ToolParserOutput::default();
154
108
- if !self.buffer.is_empty() {
109
- if self.buffer.starts_with(TOOL_CALL_START) {
155
+ match &self.mode {
156
+ Gemma4Mode::Text => output.normal_text.push_str(&self.buffer),
157
+ Gemma4Mode::Header | Gemma4Mode::ToolCall { .. } => {
158
return Err(parsing_failed!("incomplete Gemma4 tool call"));
159
}
112
- output.normal_text.push_str(&self.buffer);
160
}
161
162
let _ = self.reset();
@@ -122,25 +169,50 @@ impl ToolParser for Gemma4ToolParser {
169
}
170
171
/// Parse one Gemma4 event from buffered streaming input.
125
-fn parse_next_gemma4_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
126
- alt((tool_call_event, safe_text_event)).parse_next(input)
172
+fn parse_next_gemma4_event(
173
+ input: &mut Gemma4Input<'_>,
174
+ mode: &mut Gemma4Mode,
175
+) -> ModalResult<Gemma4Event> {
176
+ match mode {
177
+ Gemma4Mode::Text => parse_text_event(input),
178
+ Gemma4Mode::Header => tool_call_header_event(input),
179
+ Gemma4Mode::ToolCall { args_scan, .. } => tool_call_args_event(input, args_scan),
180
+ }
181
+}
182
+
183
+/// Parse a Gemma4 text-mode event.
184
+fn parse_text_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
185
+ alt((tool_call_start_event, safe_text_event)).parse_next(input)
186
+}
187
+
188
+/// Parse a Gemma4 tool-call start marker.
189
+fn tool_call_start_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
190
+ literal(TOOL_CALL_START).value(Gemma4Event::ToolCallStart).parse_next(input)
191
}
192
129
-/// Parse a complete Gemma4 tool call.
130
-// TODO: incremental parsing arguments to reduce scanning from O(n^2) to O(n).
131
-fn tool_call_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
132
- let (name, args) = seq!(
133
- _: literal(TOOL_CALL_START),
193
+/// Parse a Gemma4 tool-call header.
194
+fn tool_call_header_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
195
+ let (name,) = seq!(
196
_: literal(CALL_PREFIX),
197
gemma4_tool_name,
198
_: literal("{"),
137
- gemma4_args,
138
- _: literal("}"),
139
- _: literal(TOOL_CALL_END),
199
)
200
.parse_next(input)?;
201
+ Ok(Gemma4Event::ToolCallHeader { name })
202
+}
203
+
204
+/// Parse complete Gemma4 tool-call arguments.
205
+fn tool_call_args_event(
206
+ input: &mut Gemma4Input<'_>,
207
+ args_scan: &mut Gemma4ArgsScanState,
208
+) -> ModalResult<Gemma4Event> {
209
+ let raw_args = gemma4_raw_args_until_tool_call_end(input, args_scan)?;
210
+ let Some(args_input) = raw_args.strip_suffix('}') else {
211
+ return Err(ErrMode::Cut(ContextError::new()));
212
+ };
213
+ let args = parse_gemma4_args(args_input)?;
214
143
- Ok(Gemma4Event::ToolCall { name, args })
215
+ Ok(Gemma4Event::ToolCall { args })
216
}
217
218
/// Parse a Gemma4 tool name.
@@ -157,8 +229,75 @@ fn safe_text_event(input: &mut Gemma4Input<'_>) -> ModalResult<Gemma4Event> {
229
safe_text_len(input, TOOL_CALL_START).map(|len| Gemma4Event::Text { len })
230
}
231
232
+/// Parse raw Gemma4 arguments through the first end marker outside a Gemma string.
233
+fn gemma4_raw_args_until_tool_call_end<'i>(
234
+ input: &mut Gemma4Input<'i>,
235
+ state: &mut Gemma4ArgsScanState,
236
+) -> ModalResult<&'i str> {
237
+ let text = **input;
238
+ if state.scanned_len > text.len() {
239
+ return incomplete();
240
+ }
241
+
242
+ loop {
243
+ let rest = &text[state.scanned_len..];
244
+ if state.in_string {
245
+ let Some(string_delim) = rest.find(STRING_DELIM) else {
246
+ state.scanned_len = safe_scan_len(text, state.scanned_len, &[STRING_DELIM]);
247
+ return incomplete();
248
+ };
249
+
250
+ state.scanned_len += string_delim + STRING_DELIM.len();
251
+ state.in_string = false;
252
+ continue;
253
+ }
254
+
255
+ let next_string_delim = rest.find(STRING_DELIM);
256
+ let next_tool_call_end = rest.find(TOOL_CALL_END);
257
+ match (next_string_delim, next_tool_call_end) {
258
+ (Some(string_delim), Some(tool_call_end)) if tool_call_end < string_delim => {
259
+ let end = state.scanned_len + tool_call_end;
260
+ state.scanned_len = end + TOOL_CALL_END.len();
261
+ input.next_slice(state.scanned_len);
262
+ return Ok(&text[..end]);
263
+ }
264
+ (Some(string_delim), _) => {
265
+ state.scanned_len += string_delim + STRING_DELIM.len();
266
+ state.in_string = true;
267
+ }
268
+ (None, Some(tool_call_end)) => {
269
+ let end = state.scanned_len + tool_call_end;
270
+ state.scanned_len = end + TOOL_CALL_END.len();
271
+ input.next_slice(state.scanned_len);
272
+ return Ok(&text[..end]);
273
+ }
274
+ (None, None) => {
275
+ state.scanned_len =
276
+ safe_scan_len(text, state.scanned_len, &[STRING_DELIM, TOOL_CALL_END]);
277
+ return incomplete();
278
+ }
279
+ }
280
+ }
281
+}
282
+
283
+/// Return the scan length while holding back a split marker prefix.
284
+fn safe_scan_len(text: &str, start: usize, markers: &[&str]) -> usize {
285
+ let max_partial = markers
286
+ .iter()
287
+ .map(|marker| partial_prefix_len(&text[start..], marker))
288
+ .max()
289
+ .unwrap_or(0);
290
+ text.len() - max_partial
291
+}
292
+
293
+/// Parse complete Gemma4 custom key-value arguments.
294
+fn parse_gemma4_args(args: &str) -> ModalResult<Map<String, Value>> {
295
+ let mut input = args;
296
+ terminated(gemma4_args, eof).parse_next(&mut input)
297
+}
298
+
299
/// Parse Gemma4's custom key-value argument object content.
161
-fn gemma4_args(input: &mut Gemma4Input<'_>) -> ModalResult<Map<String, Value>> {
300
+fn gemma4_args(input: &mut &str) -> ModalResult<Map<String, Value>> {
301
let pairs: Vec<(String, Value)> = delimited(
302
ws0,
303
terminated(
@@ -172,7 +311,7 @@ fn gemma4_args(input: &mut Gemma4Input<'_>) -> ModalResult<Map<String, Value>> {
311
}
312
313
/// Parse a Gemma4 key-value pair.
175
-fn gemma4_pair(input: &mut Gemma4Input<'_>) -> ModalResult<(String, Value)> {
314
+fn gemma4_pair(input: &mut &str) -> ModalResult<(String, Value)> {
315
let (key, value) = seq!(
316
_: ws0,
317
gemma4_key,
@@ -186,7 +325,7 @@ fn gemma4_pair(input: &mut Gemma4Input<'_>) -> ModalResult<(String, Value)> {
325
}
326
327
/// Parse a Gemma4 bare key.
189
-fn gemma4_key(input: &mut Gemma4Input<'_>) -> ModalResult<String> {
328
+fn gemma4_key(input: &mut &str) -> ModalResult<String> {
329
let key = take_till(1.., |char: char| char == ':').parse_next(input)?.trim();
330
if key.is_empty() {
331
return Err(ErrMode::Cut(ContextError::new()));
@@ -195,7 +334,7 @@ fn gemma4_key(input: &mut Gemma4Input<'_>) -> ModalResult<String> {
334
}
335
336
/// Parse a Gemma4 value.
198
-fn gemma4_value(input: &mut Gemma4Input<'_>) -> ModalResult<Value> {
337
+fn gemma4_value(input: &mut &str) -> ModalResult<Value> {
338
alt((
339
gemma4_string.map(|value: &str| Value::String(value.to_string())),
340
gemma4_object.map(Value::Object),
@@ -206,7 +345,7 @@ fn gemma4_value(input: &mut Gemma4Input<'_>) -> ModalResult<Value> {
345
}
346
347
/// Parse a Gemma4 string delimited by `<|"|>`.
209
-fn gemma4_string<'i>(input: &mut Gemma4Input<'i>) -> ModalResult<&'i str> {
348
+fn gemma4_string<'i>(input: &mut &'i str) -> ModalResult<&'i str> {
349
delimited(
350
literal(STRING_DELIM),
351
take_until(0.., STRING_DELIM),
@@ -216,17 +355,17 @@ fn gemma4_string<'i>(input: &mut Gemma4Input<'i>) -> ModalResult<&'i str> {
355
}
356
357
/// Parse a nested Gemma4 object.
219
-fn gemma4_object(input: &mut Gemma4Input<'_>) -> ModalResult<Map<String, Value>> {
358
+fn gemma4_object(input: &mut &str) -> ModalResult<Map<String, Value>> {
359
delimited(literal("{"), gemma4_args, literal("}")).parse_next(input)
360
}
361
362
/// Parse a Gemma4 array value.
224
-fn gemma4_array_value(input: &mut Gemma4Input<'_>) -> ModalResult<Vec<Value>> {
363
+fn gemma4_array_value(input: &mut &str) -> ModalResult<Vec<Value>> {
364
delimited(literal("["), gemma4_array_content, literal("]")).parse_next(input)
365
}
366
367
/// Parse Gemma4 array content.
229
-fn gemma4_array_content(input: &mut Gemma4Input<'_>) -> ModalResult<Vec<Value>> {
368
+fn gemma4_array_content(input: &mut &str) -> ModalResult<Vec<Value>> {
369
delimited(
370
ws0,
371
terminated(
@@ -239,14 +378,14 @@ fn gemma4_array_content(input: &mut Gemma4Input<'_>) -> ModalResult<Vec<Value>>
378
}
379
380
/// Parse a Gemma4 bare scalar.
242
-fn gemma4_bare_value(input: &mut Gemma4Input<'_>) -> ModalResult<Value> {
381
+fn gemma4_bare_value(input: &mut &str) -> ModalResult<Value> {
382
take_till(1.., |char: char| matches!(char, ',' | '}' | ']'))
383
.map(parse_gemma4_scalar)
384
.parse_next(input)
385
}
386
387
/// Parse a Gemma4 comma separator.
249
-fn comma_separator(input: &mut Gemma4Input<'_>) -> ModalResult<()> {
388
+fn comma_separator(input: &mut &str) -> ModalResult<()> {
389
delimited(ws0, literal(","), ws0).void().parse_next(input)
390
}
391
@@ -284,29 +423,15 @@ mod tests {
423
use winnow::combinator::{eof, terminated};
424
use winnow::error::ErrMode;
425
use winnow::prelude::*;
287
- use winnow::stream::Partial;
426
427
use super::{
290
- Gemma4ToolParser, ToolCallDelta, ToolParser, ToolParserOutput, gemma4_args,
291
- gemma4_array_content,
428
+ Gemma4ToolParser, ToolCallDelta, ToolParser, ToolParserOutput, gemma4_array_content,
429
+ parse_gemma4_args,
430
};
431
use crate::{Tool, ToolParserTestExt as _};
432
295
- fn parse_gemma4_args(args: &str) -> super::Result<serde_json::Map<String, Value>> {
296
- let mut input = Partial::new(args);
297
- let _ = input.complete();
298
- match terminated(gemma4_args, eof).parse_next(&mut input) {
299
- Ok(value) => Ok(value),
300
- Err(ErrMode::Incomplete(_)) => Err(parsing_failed!("incomplete Gemma4 arguments")),
301
- Err(ErrMode::Backtrack(error) | ErrMode::Cut(error)) => {
302
- Err(parsing_failed!("{}", error))
303
- }
304
- }
305
- }
306
-
433
fn parse_gemma4_array(array: &str) -> super::Result<Vec<Value>> {
308
- let mut input = Partial::new(array);
309
- let _ = input.complete();
434
+ let mut input = array;
435
match terminated(gemma4_array_content, eof).parse_next(&mut input) {
436
Ok(value) => Ok(value),
437
Err(ErrMode::Incomplete(_)) => Err(parsing_failed!("incomplete Gemma4 array")),
@@ -560,6 +685,21 @@ mod tests {
685
assert!(!first_call(&output).arguments.contains("<|"));
686
}
687
688
+ #[test]
689
+ fn gemma4_streaming_handles_split_tool_call_end_marker() {
690
+ let output = collect_stream(&[
691
+ "<|tool_call>",
692
+ "call:get_weather{location:<|\"|>Paris<|\"|>}<tool",
693
+ "_call|>",
694
+ ]);
695
+
696
+ assert_eq!(first_call(&output).name.as_deref(), Some("get_weather"));
697
+ assert_eq!(
698
+ serde_json::from_str::<Value>(&first_call(&output).arguments).unwrap(),
699
+ json!({ "location": "Paris" })
700
+ );
701
+ }
702
+
703
#[test]
704
fn gemma4_streaming_handles_end_marker_literal_inside_string() {
705
let output = collect_stream(&[
@@ -652,4 +792,35 @@ mod tests {
792
793
assert!(error.to_report_string().contains("incomplete Gemma4 tool call"));
794
}
795
+
796
+ #[test]
797
+ fn gemma4_reset_preserves_internally_buffered_arguments() {
798
+ let mut parser = Gemma4ToolParser::new(&test_tools());
799
+ for chunk in [
800
+ "<|tool_call>",
801
+ "call:write_file{",
802
+ "content:<|\"|>hello ",
803
+ "world<|\"|>",
804
+ ] {
805
+ parser.parse_chunk(chunk).unwrap();
806
+ }
807
+
808
+ let raw = parser.reset();
809
+
810
+ assert_eq!(
811
+ raw,
812
+ "<|tool_call>call:write_file{content:<|\"|>hello world<|\"|>"
813
+ );
814
+ }
815
+
816
+ #[test]
817
+ fn gemma4_reset_preserves_completed_arguments_after_parse_error() {
818
+ let mut parser = Gemma4ToolParser::new(&test_tools());
819
+ let input = "<|tool_call>call:set{broken}<tool_call|>";
820
+
821
+ let _error = parser.parse_chunk(input).unwrap_err();
822
+ let raw = parser.reset();
823
+
824
+ assert_eq!(raw, input);
825
+ }
826
}