From e312c5cb25427e76fc3830ab14e7b6bc0963a55c Mon Sep 17 00:00:00 2001 From: Reid <61492567+reidliu41@users.noreply.github.com> Date: Fri, 26 Jun 2026 11:54:03 +0800 Subject: [PATCH] [Rust Frontend] Make Granite4 string argument scanning incremental (#46507) Signed-off-by: reidliu41 --- rust/src/parser/Cargo.toml | 5 + rust/src/parser/benches/granite4.rs | 75 ++++++++++++ rust/src/parser/src/tool/json/granite4.rs | 75 +++++++++--- rust/src/parser/src/utils.rs | 134 +++++++++++++++++++--- 4 files changed, 256 insertions(+), 33 deletions(-) create mode 100644 rust/src/parser/benches/granite4.rs diff --git a/rust/src/parser/Cargo.toml b/rust/src/parser/Cargo.toml index 67c74bb560..09c3d6b5dc 100644 --- a/rust/src/parser/Cargo.toml +++ b/rust/src/parser/Cargo.toml @@ -74,5 +74,10 @@ name = "gemma4" harness = false required-features = ["test-util"] +[[bench]] +name = "granite4" +harness = false +required-features = ["test-util"] + [lints] workspace = true diff --git a/rust/src/parser/benches/granite4.rs b/rust/src/parser/benches/granite4.rs new file mode 100644 index 0000000000..17ef867160 --- /dev/null +++ b/rust/src/parser/benches/granite4.rs @@ -0,0 +1,75 @@ +use std::time::Duration; + +use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main}; +use vllm_parser::tool::test_utils::{split_by_chars, test_tools}; +use vllm_parser::tool::{Granite4ToolParser, Tool, ToolParser}; + +mod utils; +use utils::feed_parser; + +const CHUNK_CHARS: usize = 7; +const LONG_ARGUMENT_BYTES: usize = 64 * 1024; + +fn string_args_fixture() -> String { + let arguments = format!(r#"{{"data":"{}"}}"#, "x".repeat(LONG_ARGUMENT_BYTES)); + let encoded_arguments = serde_json::to_string(&arguments).unwrap(); + format!(r#"{{"name":"f","arguments":{encoded_arguments}}}"#) +} + +fn object_args_fixture() -> String { + format!( + r#"{{"name":"f","arguments":{{"data":"{}"}}}}"#, + "x".repeat(LONG_ARGUMENT_BYTES) + ) +} + +fn parser(tools: &[Tool]) -> Box { + Granite4ToolParser::create(tools).expect("Granite4 parser should initialize") +} + +fn run_stream_group(c: &mut Criterion, name: &str, tools: &[Tool], text: &str) { + let chunks = split_by_chars(text, CHUNK_CHARS); + + let mut group = c.benchmark_group(name); + group.sample_size(50); + group.warm_up_time(Duration::from_millis(300)); + group.measurement_time(Duration::from_secs(2)); + group.throughput(Throughput::Bytes(text.len() as u64)); + + group.bench_function("reuse_parser", |b| { + let mut parser = parser(tools); + b.iter(|| { + let result = feed_parser(&mut *parser, black_box(&chunks)); + debug_assert_eq!(result.0, ""); + debug_assert_eq!(result.1, 1); + black_box(result); + }) + }); + + group.bench_function("create_parser", |b| { + b.iter_batched( + || parser(tools), + |mut parser| { + let result = feed_parser(&mut *parser, black_box(&chunks)); + debug_assert_eq!(result.0, ""); + debug_assert_eq!(result.1, 1); + black_box(result); + }, + BatchSize::SmallInput, + ) + }); + + group.finish(); +} + +fn bench_granite4(c: &mut Criterion) { + let tools = test_tools(); + let string_args = string_args_fixture(); + let object_args = object_args_fixture(); + + run_stream_group(c, "granite4/long_string_arguments", &tools, &string_args); + run_stream_group(c, "granite4/long_object_arguments", &tools, &object_args); +} + +criterion_group!(benches, bench_granite4); +criterion_main!(benches); diff --git a/rust/src/parser/src/tool/json/granite4.rs b/rust/src/parser/src/tool/json/granite4.rs index fe1cf19022..4989578011 100644 --- a/rust/src/parser/src/tool/json/granite4.rs +++ b/rust/src/parser/src/tool/json/granite4.rs @@ -9,7 +9,8 @@ use super::{ tool_call_header_event, }; use crate::tool::utils::{ - JsonObjectScanState, json_str, parse_buffered_event, safe_text_len, take_json_object, + JsonObjectScanState, JsonStringScanState, decode_json_str, parse_buffered_event, safe_text_len, + take_json_object, take_json_string, }; use crate::tool::{Result, Tool, ToolCallDelta, ToolParser, ToolParserOutput}; @@ -22,14 +23,20 @@ enum Granite4Mode { Header, /// Parsing the arguments value: /// `None` until the first byte decides object vs string; - /// `Some` while streaming an object value. + /// `Some` while streaming the selected value shape. Args { - json_scan: Option, + args_scan: Option, }, /// Arguments done; consume the object's closing `}` and ``. Close, } +#[derive(Debug, Clone, PartialEq, Eq)] +enum Granite4ArgsScan { + Object(JsonObjectScanState), + String(JsonStringScanState), +} + #[derive(Debug, Clone, PartialEq, Eq)] enum Granite4Event { Text { @@ -92,7 +99,7 @@ impl Granite4ToolParser { let tool_index = self.emitted_tool_count; self.emitted_tool_count += 1; self.active_tool_index = Some(tool_index); - self.mode = Granite4Mode::Args { json_scan: None }; + self.mode = Granite4Mode::Args { args_scan: None }; output.push_call(ToolCallDelta { tool_index, name: Some(function_name), @@ -187,7 +194,7 @@ fn parse_next_granite4_event( match mode { Granite4Mode::Text => text_event(input), Granite4Mode::Header => header_event(input), - Granite4Mode::Args { json_scan } => args_event(input, json_scan), + Granite4Mode::Args { args_scan } => args_event(input, args_scan), Granite4Mode::Close => close_event(input), } } @@ -237,14 +244,19 @@ fn header_event(input: &mut JsonToolInput<'_>) -> ModalResult { /// once seen whole and unescaped. fn args_event( input: &mut JsonToolInput<'_>, - json_scan: &mut Option, + args_scan: &mut Option, ) -> ModalResult { - if let Some(scan) = json_scan { - let len = take_json_object(input, scan)?; - return Ok(Granite4Event::ObjectArgsDelta { - len, - complete: scan.complete(), - }); + if let Some(scan) = args_scan { + return match scan { + Granite4ArgsScan::Object(scan) => { + let len = take_json_object(input, scan)?; + Ok(Granite4Event::ObjectArgsDelta { + len, + complete: scan.complete(), + }) + } + Granite4ArgsScan::String(scan) => string_args_event(input, scan), + }; } match peek(any).parse_next(input)? { @@ -252,12 +264,16 @@ fn args_event( let mut scan = JsonObjectScanState::default(); let len = take_json_object(input, &mut scan)?; let complete = scan.complete(); - *json_scan = Some(scan); + *args_scan = Some(Granite4ArgsScan::Object(scan)); Ok(Granite4Event::ObjectArgsDelta { len, complete }) } - '"' => Ok(Granite4Event::StringArgs { - decoded: json_str(input)?, - }), + '"' => { + *args_scan = Some(Granite4ArgsScan::String(JsonStringScanState::default())); + let Some(Granite4ArgsScan::String(scan)) = args_scan else { + unreachable!("Granite4 string scan state was just initialized") + }; + string_args_event(input, scan) + } _ => { let mut error = ContextError::new(); error.push(StrContext::Label("Granite4 arguments")); @@ -266,6 +282,17 @@ fn args_event( } } +fn string_args_event( + input: &mut JsonToolInput<'_>, + scan: &mut JsonStringScanState, +) -> ModalResult { + let text = **input; + let len = take_json_string(input, scan)?; + Ok(Granite4Event::StringArgs { + decoded: decode_json_str(&text[..len])?, + }) +} + /// Parse the tool-call object's closing `}` and the `` end marker. fn close_event(input: &mut JsonToolInput<'_>) -> ModalResult { seq!(_: ws0, _: literal("}"), _: ws0, _: literal(TOOL_CALL_END)) @@ -420,6 +447,22 @@ mod tests { assert_eq!(output.calls()[0].arguments, r#"{"a":1}"#); } + #[test] + fn granite4_long_string_args_stream_without_reparse() { + let arguments = format!(r#"{{"data":"{}"}}"#, "x".repeat(64 * 1024)); + let encoded_arguments = serde_json::to_string(&arguments).unwrap(); + let input = + format!(r#"{{"name":"f","arguments":{encoded_arguments}}}"#); + let chunks = split_by_chars(&input, 7); + let mut parser = Granite4ToolParser::new(&test_tools()); + + let output = collect_stream(&mut parser, &chunks); + + assert_eq!(output.calls().len(), 1); + assert_eq!(output.calls()[0].name.as_deref(), Some("f")); + assert_eq!(output.calls()[0].arguments, arguments); + } + #[test] fn granite4_streaming_handles_marker_and_json_whitespace() { // Granite spaces the markers (` {…} `) and the JSON diff --git a/rust/src/parser/src/utils.rs b/rust/src/parser/src/utils.rs index 7025521539..bd8f6f48e9 100644 --- a/rust/src/parser/src/utils.rs +++ b/rust/src/parser/src/utils.rs @@ -290,8 +290,22 @@ pub fn take_json_object( Ok(text.len()) } -/// Parse a JSON string literal. -pub fn json_str(input: &mut Partial<&str>) -> ModalResult { +/// Streaming lexical state for a JSON string literal. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct JsonStringScanState { + scanned_len: usize, + escape: bool, +} + +/// Parse a raw JSON string literal, resuming from the last scanned byte. +/// +/// The returned length covers the quoted JSON string. This only scans for the +/// string boundary; callers that need the decoded value should pass the raw +/// slice to [`decode_json_str`]. +pub fn take_json_string( + input: &mut Partial<&str>, + state: &mut JsonStringScanState, +) -> ModalResult { let text = **input; if text.is_empty() { return incomplete(); @@ -305,37 +319,58 @@ pub fn json_str(input: &mut Partial<&str>) -> ModalResult { )); } - let mut escape = false; - let mut index = 1; + let mut index = if state.scanned_len == 0 { + 1 + } else if state.scanned_len <= bytes.len() { + state.scanned_len + } else { + return incomplete(); + }; + while index < bytes.len() { let byte = bytes[index]; index += 1; - if escape { - escape = false; + if state.escape { + state.escape = false; continue; } match byte { - b'\\' => escape = true, + b'\\' => state.escape = true, b'"' => { - let raw = &text[..index]; - let value = serde_json::from_str::(raw).map_err(|_| { - json_scan_error( - "JSON string", - StrContextValue::Description("valid JSON string"), - ) - })?; input.next_slice(index); - return Ok(value); + return Ok(index); } _ => {} } } + state.scanned_len = text.len(); incomplete() } +/// Parse a JSON string literal. +pub fn json_str(input: &mut Partial<&str>) -> ModalResult { + let text = **input; + let checkpoint = input.checkpoint(); + let mut state = JsonStringScanState::default(); + let len = take_json_string(input, &mut state)?; + decode_json_str(&text[..len]).inspect_err(|_| { + input.reset(&checkpoint); + }) +} + +/// Decode a complete JSON string literal. +pub fn decode_json_str(raw: &str) -> ModalResult { + serde_json::from_str::(raw).map_err(|_| { + json_scan_error( + "JSON string", + StrContextValue::Description("valid JSON string"), + ) + }) +} + fn json_scan_error(label: &'static str, expected: StrContextValue) -> ErrMode { let mut error = ContextError::new(); error.push(StrContext::Label(label)); @@ -388,8 +423,8 @@ mod tests { use winnow::stream::{Offset, Partial, Stream}; use super::{ - JsonObjectScanState, MarkerScanState, json_str, partial_prefix_len, safe_text_len, - safe_text_len_mul, take_json_object, take_until_marker, + JsonObjectScanState, JsonStringScanState, MarkerScanState, json_str, partial_prefix_len, + safe_text_len, safe_text_len_mul, take_json_object, take_json_string, take_until_marker, }; #[test] @@ -714,6 +749,71 @@ mod tests { .assert_eq(&error.to_string()); } + #[test] + fn take_json_string_consumes_complete_string() { + let mut state = JsonStringScanState::default(); + let mut input = Partial::new(r#""say_\"hi\u0021" rest"#); + let checkpoint = input.checkpoint(); + + let len = take_json_string(&mut input, &mut state).unwrap(); + + assert_eq!(len, r#""say_\"hi\u0021""#.len()); + assert_eq!(input.offset_from(&checkpoint), len); + assert_eq!(*input, " rest"); + } + + #[test] + fn take_json_string_resumes_after_incomplete_input() { + let mut state = JsonStringScanState::default(); + let mut input = Partial::new(r#""{\"data\":\"partial"#); + let checkpoint = input.checkpoint(); + + let error = take_json_string(&mut input, &mut state).unwrap_err(); + + assert!(matches!(error, ErrMode::Incomplete(_))); + assert_eq!(input.offset_from(&checkpoint), 0); + assert_eq!(state.scanned_len, r#""{\"data\":\"partial"#.len()); + + let mut input = Partial::new(r#""{\"data\":\"partial string\"}" tail"#); + let len = take_json_string(&mut input, &mut state).unwrap(); + + assert_eq!(len, r#""{\"data\":\"partial string\"}""#.len()); + assert_eq!(*input, " tail"); + } + + #[test] + fn take_json_string_tracks_escape_across_chunks() { + let mut state = JsonStringScanState::default(); + let mut input = Partial::new(r#""abc\"#); + + let error = take_json_string(&mut input, &mut state).unwrap_err(); + + assert!(matches!(error, ErrMode::Incomplete(_))); + assert!(state.escape); + + let mut input = Partial::new(r#""abc\"def" tail"#); + let len = take_json_string(&mut input, &mut state).unwrap(); + + assert_eq!(len, r#""abc\"def""#.len()); + assert_eq!(*input, " tail"); + } + + #[test] + fn take_json_string_rejects_non_string_start() { + let mut state = JsonStringScanState::default(); + let mut input = Partial::new("42"); + + let error = take_json_string(&mut input, &mut state).unwrap_err(); + + let ErrMode::Cut(error) = error else { + panic!("expected cut error"); + }; + expect![[r#" + invalid JSON string + expected `"`"#]] + .assert_eq(&error.to_string()); + } + #[test] fn json_str_decodes_escaped_content() { let mut input = Partial::new(r#""say_\"hi\u0021" rest"#);