@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e05095c9e1

[Rust Frontend] Feed parser tests and benchmarks attributed input (#59048)

Signed-off-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

Bugen Zhao committed Sep 29, 2026 at 14:23 UTC e05095c9e1ae37f208c1ad70b08c17f1152878fc
7 files changed +334 -142
rust/src/chat/tests/roundtrip.rs
+27 -7
@@ -780,10 +780,29 @@ fn decoded_completion_stream(
780 assistant_stop_suffix, text
781 )
782 })?;
783 + // The engine returns token ids, so attribute the rendered text
784 + // through the production decoder, then split it into small char
785 + // chunks to exercise marker splits.
786 + let token_ids =
787 + tokenizer.encode(body, false).context("failed to encode text completion")?;
788 + let mut rest = vllm_text::DecodedText::default();
789 + for chunk in incremental_decode_chunks(
790 + tokenizer,
791 + &prompt_token_ids,
792 + &token_ids,
793 + TOKEN_COMPLETION_CHUNK_TOKENS,
794 + )? {
795 + rest.append(chunk.delta);
796 + }
797 + ensure!(
798 + rest.text == body,
799 + "text completion must roundtrip through the tokenizer\nbody: {body:?}\ndecoded: {:?}",
800 + rest.text
801 + );
802 split_by_chars(body, TEXT_COMPLETION_CHUNK_CHARS)
803 .into_iter()
804 .map(|delta| DecodedCompletionChunk {
786 - delta,
805 + delta: rest.drain_prefix(delta.len()),
806 token_ids: Vec::new(), // unused for text-level roundtrip cases
807 })
808 .collect()
@@ -836,7 +855,7 @@ fn decoded_completion_stream(
855 sampling_mask: None,
856 });
857 events.push(DecodedTextEvent::TextDelta {
839 - decoded: vllm_text::DecodedText::unattributed(chunk.delta),
858 + decoded: chunk.delta,
859 sampled: vllm_text::SampledDelta {
860 token_ids: chunk.token_ids,
861 logprobs: None,
@@ -851,7 +870,8 @@ fn decoded_completion_stream(
870
871 /// One decoded completion chunk fed into the output processor.
872 struct DecodedCompletionChunk {
854 - delta: String,
873 + /// Attributed for token-id completions, as the production decoder emits it.
874 + delta: vllm_text::DecodedText,
875 token_ids: Vec<u32>,
876 }
877
@@ -893,11 +913,11 @@ fn incremental_decode_chunks(
913 let mut decoder = tokenizer.create_decode_stream(prompt_token_ids, false, 0);
914 let mut chunks = Vec::new();
915 for chunk_token_ids in split_by_count(token_ids, chunk_size) {
896 - let mut delta = String::new();
916 + let mut delta = vllm_text::DecodedText::default();
917 for token_id in chunk_token_ids.iter().copied() {
918 decoder.push_token(token_id)?;
919 while let Some(chunk) = decoder.next_chunk() {
900 - delta.push_str(&chunk.text);
920 + delta.append(chunk);
921 }
922 }
923 chunks.push(DecodedCompletionChunk {
@@ -909,10 +929,10 @@ fn incremental_decode_chunks(
929 let (last_chunk, _) = decoder.flush(None)?;
930 if let Some(last_chunk) = last_chunk {
931 if let Some(delta) = chunks.last_mut() {
912 - delta.delta.push_str(&last_chunk.text);
932 + delta.delta.append(last_chunk);
933 } else {
934 chunks.push(DecodedCompletionChunk {
915 - delta: last_chunk.text,
935 + delta: last_chunk,
936 token_ids: Vec::new(),
937 });
938 }
rust/src/parser/Cargo.toml
+5
@@ -81,5 +81,10 @@ name = "granite4"
81 harness = false
82 required-features = ["test-util"]
83
84 +[[bench]]
85 +name = "kimi_k3"
86 +harness = false
87 +required-features = ["test-util"]
88 +
89 [lints]
90 workspace = true
rust/src/parser/benches/gemma4.rs
+33 -19
@@ -1,15 +1,17 @@
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +use std::sync::Arc;
5 use std::time::Duration;
6
7 use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
7 -use vllm_parser::tool::test_utils::{split_by_chars, test_tools};
8 -use vllm_parser::tool::{Tool, ToolParser};
8 +use vllm_parser::tool::Tool;
9 +use vllm_parser::tool::test_utils::test_tools;
10 use vllm_parser::unified::Gemma4UnifiedParser;
11 +use vllm_tokenizer::test_utils::TestTokenizer;
12
13 mod utils;
12 -use utils::{UnifiedToolParserAdapter, feed_parser};
14 +use utils::{attributed_chunks, feed_unified_parser};
15
16 const CHUNK_CHARS: usize = 7;
17 const LONG_NORMAL_TEXT_REPEATS: usize = 2048;
@@ -71,8 +73,17 @@ fn long_tool_argument_fixture() -> String {
73 )
74 }
75
74 -fn parser(tools: &[Tool]) -> Box<dyn ToolParser> {
75 - UnifiedToolParserAdapter::<Gemma4UnifiedParser>::create(tools)
76 +fn tokenizer() -> TestTokenizer {
77 + TestTokenizer::new()
78 + .with_special_token("<|tool_call>", 256)
79 + .with_special_token("<tool_call|>", 257)
80 + .with_special_token("<|\"|>", 258)
81 + .with_special_token("<|channel>", 259)
82 + .with_special_token("<channel|>", 260)
83 +}
84 +
85 +fn parser(tools: &[Tool], tokenizer: &TestTokenizer) -> Gemma4UnifiedParser {
86 + Gemma4UnifiedParser::new(tools, Arc::new(tokenizer.clone()))
87 .expect("Gemma4 unified parser should initialize")
88 }
89
@@ -85,7 +96,14 @@ fn run_stream_group(
96 expected_normal_text: &str,
97 expected_calls_len: usize,
98 ) {
88 - let chunks = split_by_chars(text, chunk_chars);
99 + let tokenizer = tokenizer();
100 + let chunks = attributed_chunks(&tokenizer, text, chunk_chars);
101 +
102 + // Check the fixture once, outside the measurement: a marker lost to
103 + // missing attribution would silently turn the whole stream into content.
104 + let summary = feed_unified_parser(&mut parser(tools, &tokenizer), &[], chunks.clone());
105 + assert_eq!(summary.normal_text, expected_normal_text, "{name}");
106 + assert_eq!(summary.calls_len, expected_calls_len, "{name}");
107
108 let mut group = c.benchmark_group(name);
109 group.sample_size(50);
@@ -94,23 +112,19 @@ fn run_stream_group(
112 group.throughput(Throughput::Bytes(text.len() as u64));
113
114 group.bench_function("reuse_parser", |b| {
97 - let mut parser = parser(tools);
98 - b.iter(|| {
99 - let result = feed_parser(&mut *parser, black_box(&chunks));
100 - debug_assert_eq!(result.0, expected_normal_text);
101 - debug_assert_eq!(result.1, expected_calls_len);
102 - black_box(result);
103 - })
115 + let mut parser = parser(tools, &tokenizer);
116 + b.iter_batched(
117 + || chunks.clone(),
118 + |chunks| black_box(feed_unified_parser(&mut parser, &[], black_box(chunks))),
119 + BatchSize::SmallInput,
120 + )
121 });
122
123 group.bench_function("create_parser", |b| {
124 b.iter_batched(
108 - || parser(tools),
109 - |mut parser| {
110 - let result = feed_parser(&mut *parser, black_box(&chunks));
111 - debug_assert_eq!(result.0, expected_normal_text);
112 - debug_assert_eq!(result.1, expected_calls_len);
113 - black_box(result);
125 + || (parser(tools, &tokenizer), chunks.clone()),
126 + |(mut parser, chunks)| {
127 + black_box(feed_unified_parser(&mut parser, &[], black_box(chunks)))
128 },
129 BatchSize::SmallInput,
130 )
rust/src/parser/benches/kimi_k3.rs new
+191
@@ -0,0 +1,191 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::sync::Arc;
5 +use std::time::Duration;
6 +
7 +use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
8 +use vllm_parser::unified::KimiK3UnifiedParser;
9 +use vllm_tokenizer::Tokenizer;
10 +use vllm_tokenizer::test_utils::TestTokenizer;
11 +
12 +mod utils;
13 +use utils::{attributed_chunks, feed_unified_parser};
14 +
15 +const CHUNK_CHARS: usize = 7;
16 +const LONG_REASONING_REPEATS: usize = 2048;
17 +const LONG_TOOL_ARGUMENT_REPEATS: usize = 256;
18 +
19 +/// The chat generation prompt for thinking mode ends inside the `think`
20 +/// channel, so the output starts with reasoning.
21 +const THINKING_PROMPT_TAIL: &str = "<|open|>think<|sep|>";
22 +
23 +fn tokenizer() -> TestTokenizer {
24 + TestTokenizer::new()
25 + .with_special_token("<|open|>", 256)
26 + .with_special_token("<|close|>", 257)
27 + .with_special_token("<|sep|>", 258)
28 + .with_special_token("<|end_of_msg|>", 259)
29 +}
30 +
31 +fn argument(key: &str, arg_type: &str, value: &str) -> String {
32 + format!(
33 + "<|open|>argument key=\"{key}\" type=\"{arg_type}\"<|sep|>{value}<|close|>argument<|sep|>"
34 + )
35 +}
36 +
37 +fn call(tool: &str, index: usize, body: &str) -> String {
38 + format!("<|open|>call tool=\"{tool}\" index=\"{index}\"<|sep|>{body}<|close|>call<|sep|>")
39 +}
40 +
41 +fn message(reasoning: &str, response: &str, calls: &str) -> String {
42 + format!(
43 + "{reasoning}<|close|>think<|sep|>\
44 + <|open|>response<|sep|>{response}<|close|>response<|sep|>\
45 + <|open|>tools<|sep|>{calls}<|close|>tools<|sep|>\
46 + <|close|>message<|sep|>"
47 + )
48 +}
49 +
50 +fn mixed_fixture() -> String {
51 + let calls = [
52 + call(
53 + "convert",
54 + 1,
55 + &[
56 + argument("whole", "number", "114.514"),
57 + argument("flag", "boolean", "true"),
58 + argument("name", "string", "demo"),
59 + argument("payload", "object", r#"{"count":42,"tags":["red","blue"]}"#),
60 + ]
61 + .concat(),
62 + ),
63 + call("update_record", 2, &argument("id", "integer", "7")),
64 + ]
65 + .concat();
66 + message(
67 + "I should inspect the data first.",
68 + "I will convert it.",
69 + &calls,
70 + )
71 +}
72 +
73 +fn long_reasoning_fixture() -> String {
74 + let line = "Ordinary reasoning text with no K3 channel markers at all.\n";
75 + message(&line.repeat(LONG_REASONING_REPEATS), "Done.", "")
76 +}
77 +
78 +fn long_tool_argument_fixture() -> String {
79 + let line = "<section><p>Literal <tag> and quotes \" inside content.</p></section>\n";
80 + let content = line.repeat(LONG_TOOL_ARGUMENT_REPEATS);
81 + let body = [
82 + argument("path", "string", "index.html"),
83 + argument("content", "string", &content),
84 + ]
85 + .concat();
86 + message(
87 + "Write the file.",
88 + "I will write the file.",
89 + &call("write_file", 1, &body),
90 + )
91 +}
92 +
93 +fn parser(tokenizer: &TestTokenizer) -> KimiK3UnifiedParser {
94 + KimiK3UnifiedParser::new(&[], Arc::new(tokenizer.clone()))
95 + .expect("Kimi K3 unified parser should initialize")
96 +}
97 +
98 +fn run_stream_group(
99 + c: &mut Criterion,
100 + name: &str,
101 + text: &str,
102 + expected_normal_text: &str,
103 + expected_reasoning_len: usize,
104 + expected_calls_len: usize,
105 +) {
106 + let tokenizer = tokenizer();
107 + let prompt_token_ids = tokenizer.encode(THINKING_PROMPT_TAIL, false).unwrap();
108 + let chunks = attributed_chunks(&tokenizer, text, CHUNK_CHARS);
109 +
110 + // Check the fixture once, outside the measurement: a marker lost to
111 + // missing attribution would silently turn the whole stream into content.
112 + let summary = feed_unified_parser(&mut parser(&tokenizer), &prompt_token_ids, chunks.clone());
113 + assert_eq!(summary.normal_text, expected_normal_text, "{name}");
114 + assert_eq!(
115 + summary.reasoning_text.len(),
116 + expected_reasoning_len,
117 + "{name}"
118 + );
119 + assert_eq!(summary.calls_len, expected_calls_len, "{name}");
120 +
121 + let mut group = c.benchmark_group(name);
122 + group.sample_size(50);
123 + group.warm_up_time(Duration::from_millis(300));
124 + group.measurement_time(Duration::from_secs(2));
125 + group.throughput(Throughput::Bytes(text.len() as u64));
126 +
127 + group.bench_function("reuse_parser", |b| {
128 + let mut parser = parser(&tokenizer);
129 + b.iter_batched(
130 + || chunks.clone(),
131 + |chunks| {
132 + black_box(feed_unified_parser(
133 + &mut parser,
134 + &prompt_token_ids,
135 + black_box(chunks),
136 + ))
137 + },
138 + BatchSize::SmallInput,
139 + )
140 + });
141 +
142 + group.bench_function("create_parser", |b| {
143 + b.iter_batched(
144 + || (parser(&tokenizer), chunks.clone()),
145 + |(mut parser, chunks)| {
146 + black_box(feed_unified_parser(
147 + &mut parser,
148 + &prompt_token_ids,
149 + black_box(chunks),
150 + ))
151 + },
152 + BatchSize::SmallInput,
153 + )
154 + });
155 +
156 + group.finish();
157 +}
158 +
159 +fn bench_kimi_k3(c: &mut Criterion) {
160 + run_stream_group(
161 + c,
162 + "kimi_k3/mixed_complex_tool_call",
163 + &mixed_fixture(),
164 + "I will convert it.",
165 + "I should inspect the data first.".len(),
166 + 2,
167 + );
168 +
169 + let long_reasoning = long_reasoning_fixture();
170 + let reasoning_len = long_reasoning.find("<|close|>think").unwrap();
171 + run_stream_group(
172 + c,
173 + "kimi_k3/long_reasoning",
174 + &long_reasoning,
175 + "Done.",
176 + reasoning_len,
177 + 0,
178 + );
179 +
180 + run_stream_group(
181 + c,
182 + "kimi_k3/long_tool_argument",
183 + &long_tool_argument_fixture(),
184 + "I will write the file.",
185 + "Write the file.".len(),
186 + 1,
187 + );
188 +}
189 +
190 +criterion_group!(benches, bench_kimi_k3);
191 +criterion_main!(benches);
rust/src/parser/benches/utils/adapter.rs deleted
-114
@@ -1,114 +0,0 @@
1 -// SPDX-License-Identifier: Apache-2.0
2 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -use std::sync::Arc;
5 -
6 -use vllm_parser::tool::{Result, Tool, ToolParser, ToolParserError, ToolParserOutput};
7 -use vllm_parser::unified::{
8 - UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput,
9 -};
10 -use vllm_tokenizer::{DecodedText, Tokenizer};
11 -
12 -/// Tokenizer stub used by unified-parser benchmarks.
13 -struct BenchTokenizer;
14 -
15 -impl Tokenizer for BenchTokenizer {
16 - fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
17 - Ok(text.chars().map(|_| u32::MAX).collect())
18 - }
19 -
20 - fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result<Vec<u32>> {
21 - self.encode(text, false)
22 - }
23 -
24 - fn decode(
25 - &self,
26 - token_ids: &[u32],
27 - _skip_special_tokens: bool,
28 - ) -> vllm_tokenizer::Result<String> {
29 - Ok("\u{FFFD}".repeat(token_ids.len()))
30 - }
31 -
32 - fn token_to_id(&self, _token: &str) -> Option<u32> {
33 - Some(u32::MAX)
34 - }
35 -
36 - fn id_to_token(&self, _id: u32) -> Option<String> {
37 - Some("\u{FFFD}".to_string())
38 - }
39 -}
40 -
41 -/// Bench-only adapter that exposes a unified parser through the tool-parser
42 -/// benchmark harness.
43 -///
44 -/// Returns error if the unified parser produces reasoning events.
45 -pub struct UnifiedToolParserAdapter<T> {
46 - inner: Box<dyn UnifiedParser>,
47 - _marker: std::marker::PhantomData<T>,
48 -}
49 -
50 -fn map_unified_error(error: UnifiedParserError) -> ToolParserError {
51 - ToolParserError::ParsingFailed {
52 - message: format!("unified parser failed: {error}"),
53 - }
54 -}
55 -
56 -fn append_unified_output(
57 - output: UnifiedParserOutput,
58 - tool_output: &mut ToolParserOutput,
59 -) -> Result<()> {
60 - for event in output.events {
61 - match event {
62 - UnifiedParserEvent::Text(text) => tool_output.push_text(text),
63 - UnifiedParserEvent::ToolCall(call) => tool_output.push_call(call),
64 - UnifiedParserEvent::Reasoning(_) => {
65 - return Err(ToolParserError::ParsingFailed {
66 - message: "unified parser emitted reasoning in tool-parser adapter".to_string(),
67 - });
68 - }
69 - }
70 - }
71 - Ok(())
72 -}
73 -
74 -impl<T: UnifiedParser> ToolParser for UnifiedToolParserAdapter<T> {
75 - fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
76 - where
77 - Self: Sized + 'static,
78 - {
79 - let inner = T::create(tools, Arc::new(BenchTokenizer)).map_err(map_unified_error)?;
80 - Ok(Box::new(Self {
81 - inner,
82 - _marker: std::marker::PhantomData,
83 - }))
84 - }
85 -
86 - fn preserve_special_tokens(&self) -> bool {
87 - self.inner.preserve_special_tokens()
88 - }
89 -
90 - fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
91 - self.inner.tool_call_id(tool_index)
92 - }
93 -
94 - fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
95 - let mut unified_output = UnifiedParserOutput::default();
96 - let result = self
97 - .inner
98 - .parse_into(DecodedText::unattributed(chunk), &mut unified_output)
99 - .map_err(map_unified_error);
100 - append_unified_output(unified_output, output)?;
101 - result
102 - }
103 -
104 - fn finish(&mut self) -> Result<ToolParserOutput> {
105 - let unified_output = self.inner.finish().map_err(map_unified_error)?;
106 - let mut output = ToolParserOutput::default();
107 - append_unified_output(unified_output, &mut output)?;
108 - Ok(output)
109 - }
110 -
111 - fn reset(&mut self) -> String {
112 - self.inner.reset()
113 - }
114 -}
rust/src/parser/benches/utils/mod.rs
+2 -2
@@ -6,12 +6,12 @@
6 #![allow(dead_code)]
7 #![allow(unused_imports)]
8
9 -mod adapter;
9 +mod unified;
10
11 -pub(super) use adapter::UnifiedToolParserAdapter;
11 use futures::FutureExt as _;
12 use openai_protocol::common::{Function as OpenAiFunction, Tool as OpenAiTool};
13 use tool_parser::traits::ToolParser as ExternalToolParser;
14 +pub(super) use unified::{attributed_chunks, feed_unified_parser};
15 use vllm_parser::tool::test_utils::collect_stream;
16 use vllm_parser::tool::{Tool, ToolParser};
17
rust/src/parser/benches/utils/unified.rs new
+76
@@ -0,0 +1,76 @@
1 +// SPDX-License-Identifier: Apache-2.0
2 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +use std::collections::BTreeSet;
5 +
6 +use vllm_parser::tool::test_utils::split_by_chars;
7 +use vllm_parser::unified::{UnifiedParser, UnifiedParserEvent, UnifiedParserOutput};
8 +use vllm_tokenizer::{DecodedText, Tokenizer};
9 +
10 +/// Split `text` into token-attributed chunks of at most `chunk_chars`
11 +/// characters, as the frontend streams them to a unified parser.
12 +///
13 +/// The whole text is encoded with special tokens and run through the
14 +/// incremental decoder once before splitting, so a marker cut by a chunk
15 +/// boundary keeps its special-token anchor. Prepare chunks outside the measured
16 +/// routine.
17 +pub fn attributed_chunks(
18 + tokenizer: &dyn Tokenizer,
19 + text: &str,
20 + chunk_chars: usize,
21 +) -> Vec<DecodedText> {
22 + let mut decoder = tokenizer.create_decode_stream(&[], false, 0);
23 + for token_id in tokenizer.encode(text, false).expect("fixture should encode") {
24 + decoder.push_token(token_id).expect("fixture token should decode");
25 + }
26 + let mut rest = decoder.flush(None).expect("decoder should flush").1;
27 + assert_eq!(
28 + rest.text, text,
29 + "fixture should round-trip through the tokenizer"
30 + );
31 +
32 + split_by_chars(text, chunk_chars)
33 + .into_iter()
34 + .map(|chunk| rest.drain_prefix(chunk.len()))
35 + .collect()
36 +}
37 +
38 +/// Event totals of one parsed stream, for benchmark sanity checks.
39 +#[derive(Debug, Default)]
40 +pub struct UnifiedStreamSummary {
41 + pub normal_text: String,
42 + pub reasoning_text: String,
43 + pub calls_len: usize,
44 +}
45 +
46 +/// Run one request stream through `parser`: initialize it from the prompt,
47 +/// feed every chunk, and finish.
48 +pub fn feed_unified_parser(
49 + parser: &mut dyn UnifiedParser,
50 + prompt_token_ids: &[u32],
51 + chunks: Vec<DecodedText>,
52 +) -> UnifiedStreamSummary {
53 + parser.initialize(prompt_token_ids).expect("parser should initialize");
54 +
55 + let mut output = UnifiedParserOutput::default();
56 + for chunk in chunks {
57 + parser.parse_into(chunk, &mut output).expect("chunk should parse");
58 + }
59 + output.append(parser.finish().expect("stream should finish"));
60 +
61 + let mut summary = UnifiedStreamSummary::default();
62 + let mut tool_indices = BTreeSet::new();
63 + for event in output.events {
64 + match event {
65 + UnifiedParserEvent::Text(text) => summary.normal_text.push_str(&text),
66 + UnifiedParserEvent::Reasoning(reasoning) => {
67 + summary.reasoning_text.push_str(&reasoning.text);
68 + }
69 + UnifiedParserEvent::ToolCall(call) => {
70 + tool_indices.insert(call.tool_index);
71 + }
72 + }
73 + }
74 + summary.calls_len = tool_indices.len();
75 + summary
76 +}