iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Rust Frontend] Make Granite4 string argument scanning incremental (#46507)

Browse Source 
Signed-off-by: reidliu41 <reid201711@gmail.com>
This commit is contained in:
Reid
2026-06-26 03:54:03 +00:00
committed by GitHub
parent 1502cf6274
commit e312c5cb25
4 changed files with 256 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/src/parser/Cargo.toml
+5
View File
@@ -74,5 +74,10 @@ name = "gemma4"
harness = false
required-features = ["test-util"]
[[bench]]
name = "granite4"
harness = false
required-features = ["test-util"]
[lints]
workspace = true
rust/src/parser/benches/granite4.rs
+75
View File
@@ -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#"<tool_call>{{"name":"f","arguments":{encoded_arguments}}}</tool_call>"#)
}
fn object_args_fixture() -> String {
format!(
r#"<tool_call>{{"name":"f","arguments":{{"data":"{}"}}}}</tool_call>"#,
"x".repeat(LONG_ARGUMENT_BYTES)
)
}
fn parser(tools: &[Tool]) -> Box<dyn ToolParser> {
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);
rust/src/parser/src/tool/json/granite4.rs
+59 -16
View File
@@ -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<JsonObjectScanState>,
args_scan: Option<Granite4ArgsScan>,
},
/// Arguments done; consume the object's closing `}` and `</tool_call>`.
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<Granite4Event> {
/// once seen whole and unescaped.
fn args_event(
input: &mut JsonToolInput<'_>,
json_scan: &mut Option<JsonObjectScanState>,
args_scan: &mut Option<Granite4ArgsScan>,
) -> ModalResult<Granite4Event> {
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<Granite4Event> {
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 `</tool_call>` end marker.
fn close_event(input: &mut JsonToolInput<'_>) -> ModalResult<Granite4Event> {
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#"<tool_call>{{"name":"f","arguments":{encoded_arguments}}}</tool_call>"#);
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 (`<tool_call> {…} </tool_call>`) and the JSON
rust/src/parser/src/utils.rs
+117 -17
View File
@@ -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<String> {
/// 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<usize> {
let text = **input;
if text.is_empty() {
return incomplete();
@@ -305,37 +319,58 @@ pub fn json_str(input: &mut Partial<&str>) -> ModalResult<String> {
));
}
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::<String>(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<String> {
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<String> {
serde_json::from_str::<String>(raw).map_err(|_| {
json_scan_error(
"JSON string",
StrContextValue::Description("valid JSON string"),
)
})
}
fn json_scan_error(label: &'static str, expected: StrContextValue) -> ErrMode<ContextError> {
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"#);