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

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

Browse Source 
Signed-off-by: Bugen Zhao <i@bugenzhao.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Bugen Zhao
2026-09-29 06:23:54 +00:00
committed by GitHub
co-authored by Claude Opus 5.5
parent 35d6fb3187
commit e05095c9e1
7 changed files with 334 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/src/chat/tests/roundtrip.rs
+27 -7
View File
@@ -780,10 +780,29 @@ fn decoded_completion_stream(
assistant_stop_suffix, text
)
})?;
// The engine returns token ids, so attribute the rendered text
// through the production decoder, then split it into small char
// chunks to exercise marker splits.
let token_ids =
tokenizer.encode(body, false).context("failed to encode text completion")?;
let mut rest = vllm_text::DecodedText::default();
for chunk in incremental_decode_chunks(
tokenizer,
&prompt_token_ids,
&token_ids,
TOKEN_COMPLETION_CHUNK_TOKENS,
)? {
rest.append(chunk.delta);
}
ensure!(
rest.text == body,
"text completion must roundtrip through the tokenizer\nbody: {body:?}\ndecoded: {:?}",
rest.text
);
split_by_chars(body, TEXT_COMPLETION_CHUNK_CHARS)
.into_iter()
.map(|delta| DecodedCompletionChunk {
delta,
delta: rest.drain_prefix(delta.len()),
token_ids: Vec::new(), // unused for text-level roundtrip cases
})
.collect()
@@ -836,7 +855,7 @@ fn decoded_completion_stream(
sampling_mask: None,
});
events.push(DecodedTextEvent::TextDelta {
decoded: vllm_text::DecodedText::unattributed(chunk.delta),
decoded: chunk.delta,
sampled: vllm_text::SampledDelta {
token_ids: chunk.token_ids,
logprobs: None,
@@ -851,7 +870,8 @@ fn decoded_completion_stream(
/// One decoded completion chunk fed into the output processor.
struct DecodedCompletionChunk {
delta: String,
/// Attributed for token-id completions, as the production decoder emits it.
delta: vllm_text::DecodedText,
token_ids: Vec<u32>,
}
@@ -893,11 +913,11 @@ fn incremental_decode_chunks(
let mut decoder = tokenizer.create_decode_stream(prompt_token_ids, false, 0);
let mut chunks = Vec::new();
for chunk_token_ids in split_by_count(token_ids, chunk_size) {
let mut delta = String::new();
let mut delta = vllm_text::DecodedText::default();
for token_id in chunk_token_ids.iter().copied() {
decoder.push_token(token_id)?;
while let Some(chunk) = decoder.next_chunk() {
delta.push_str(&chunk.text);
delta.append(chunk);
}
}
chunks.push(DecodedCompletionChunk {
@@ -909,10 +929,10 @@ fn incremental_decode_chunks(
let (last_chunk, _) = decoder.flush(None)?;
if let Some(last_chunk) = last_chunk {
if let Some(delta) = chunks.last_mut() {
delta.delta.push_str(&last_chunk.text);
delta.delta.append(last_chunk);
} else {
chunks.push(DecodedCompletionChunk {
delta: last_chunk.text,
delta: last_chunk,
token_ids: Vec::new(),
});
}
rust/src/parser/Cargo.toml
+5
View File
@@ -81,5 +81,10 @@ name = "granite4"
harness = false
required-features = ["test-util"]
[[bench]]
name = "kimi_k3"
harness = false
required-features = ["test-util"]
[lints]
workspace = true
rust/src/parser/benches/gemma4.rs
+33 -19
View File
@@ -1,15 +1,17 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use std::sync::Arc;
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::{Tool, ToolParser};
use vllm_parser::tool::Tool;
use vllm_parser::tool::test_utils::test_tools;
use vllm_parser::unified::Gemma4UnifiedParser;
use vllm_tokenizer::test_utils::TestTokenizer;
mod utils;
use utils::{UnifiedToolParserAdapter, feed_parser};
use utils::{attributed_chunks, feed_unified_parser};
const CHUNK_CHARS: usize = 7;
const LONG_NORMAL_TEXT_REPEATS: usize = 2048;
@@ -71,8 +73,17 @@ fn long_tool_argument_fixture() -> String {
)
}
fn parser(tools: &[Tool]) -> Box<dyn ToolParser> {
UnifiedToolParserAdapter::<Gemma4UnifiedParser>::create(tools)
fn tokenizer() -> TestTokenizer {
TestTokenizer::new()
.with_special_token("<|tool_call>", 256)
.with_special_token("<tool_call|>", 257)
.with_special_token("<|\"|>", 258)
.with_special_token("<|channel>", 259)
.with_special_token("<channel|>", 260)
}
fn parser(tools: &[Tool], tokenizer: &TestTokenizer) -> Gemma4UnifiedParser {
Gemma4UnifiedParser::new(tools, Arc::new(tokenizer.clone()))
.expect("Gemma4 unified parser should initialize")
}
@@ -85,7 +96,14 @@ fn run_stream_group(
expected_normal_text: &str,
expected_calls_len: usize,
) {
let chunks = split_by_chars(text, chunk_chars);
let tokenizer = tokenizer();
let chunks = attributed_chunks(&tokenizer, text, chunk_chars);
// Check the fixture once, outside the measurement: a marker lost to
// missing attribution would silently turn the whole stream into content.
let summary = feed_unified_parser(&mut parser(tools, &tokenizer), &[], chunks.clone());
assert_eq!(summary.normal_text, expected_normal_text, "{name}");
assert_eq!(summary.calls_len, expected_calls_len, "{name}");
let mut group = c.benchmark_group(name);
group.sample_size(50);
@@ -94,23 +112,19 @@ fn run_stream_group(
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, expected_normal_text);
debug_assert_eq!(result.1, expected_calls_len);
black_box(result);
})
let mut parser = parser(tools, &tokenizer);
b.iter_batched(
|| chunks.clone(),
|chunks| black_box(feed_unified_parser(&mut parser, &[], black_box(chunks))),
BatchSize::SmallInput,
)
});
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, expected_normal_text);
debug_assert_eq!(result.1, expected_calls_len);
black_box(result);
|| (parser(tools, &tokenizer), chunks.clone()),
|(mut parser, chunks)| {
black_box(feed_unified_parser(&mut parser, &[], black_box(chunks)))
},
BatchSize::SmallInput,
)
rust/src/parser/benches/kimi_k3.rs
+191
View File
@@ -0,0 +1,191 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use std::sync::Arc;
use std::time::Duration;
use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
use vllm_parser::unified::KimiK3UnifiedParser;
use vllm_tokenizer::Tokenizer;
use vllm_tokenizer::test_utils::TestTokenizer;
mod utils;
use utils::{attributed_chunks, feed_unified_parser};
const CHUNK_CHARS: usize = 7;
const LONG_REASONING_REPEATS: usize = 2048;
const LONG_TOOL_ARGUMENT_REPEATS: usize = 256;
/// The chat generation prompt for thinking mode ends inside the `think`
/// channel, so the output starts with reasoning.
const THINKING_PROMPT_TAIL: &str = "<|open|>think<|sep|>";
fn tokenizer() -> TestTokenizer {
TestTokenizer::new()
.with_special_token("<|open|>", 256)
.with_special_token("<|close|>", 257)
.with_special_token("<|sep|>", 258)
.with_special_token("<|end_of_msg|>", 259)
}
fn argument(key: &str, arg_type: &str, value: &str) -> String {
format!(
"<|open|>argument key=\"{key}\" type=\"{arg_type}\"<|sep|>{value}<|close|>argument<|sep|>"
)
}
fn call(tool: &str, index: usize, body: &str) -> String {
format!("<|open|>call tool=\"{tool}\" index=\"{index}\"<|sep|>{body}<|close|>call<|sep|>")
}
fn message(reasoning: &str, response: &str, calls: &str) -> String {
format!(
"{reasoning}<|close|>think<|sep|>\
<|open|>response<|sep|>{response}<|close|>response<|sep|>\
<|open|>tools<|sep|>{calls}<|close|>tools<|sep|>\
<|close|>message<|sep|>"
)
}
fn mixed_fixture() -> String {
let calls = [
call(
"convert",
1,
&[
argument("whole", "number", "114.514"),
argument("flag", "boolean", "true"),
argument("name", "string", "demo"),
argument("payload", "object", r#"{"count":42,"tags":["red","blue"]}"#),
]
.concat(),
),
call("update_record", 2, &argument("id", "integer", "7")),
]
.concat();
message(
"I should inspect the data first.",
"I will convert it.",
&calls,
)
}
fn long_reasoning_fixture() -> String {
let line = "Ordinary reasoning text with no K3 channel markers at all.\n";
message(&line.repeat(LONG_REASONING_REPEATS), "Done.", "")
}
fn long_tool_argument_fixture() -> String {
let line = "<section><p>Literal <tag> and quotes \" inside content.</p></section>\n";
let content = line.repeat(LONG_TOOL_ARGUMENT_REPEATS);
let body = [
argument("path", "string", "index.html"),
argument("content", "string", &content),
]
.concat();
message(
"Write the file.",
"I will write the file.",
&call("write_file", 1, &body),
)
}
fn parser(tokenizer: &TestTokenizer) -> KimiK3UnifiedParser {
KimiK3UnifiedParser::new(&[], Arc::new(tokenizer.clone()))
.expect("Kimi K3 unified parser should initialize")
}
fn run_stream_group(
c: &mut Criterion,
name: &str,
text: &str,
expected_normal_text: &str,
expected_reasoning_len: usize,
expected_calls_len: usize,
) {
let tokenizer = tokenizer();
let prompt_token_ids = tokenizer.encode(THINKING_PROMPT_TAIL, false).unwrap();
let chunks = attributed_chunks(&tokenizer, text, CHUNK_CHARS);
// Check the fixture once, outside the measurement: a marker lost to
// missing attribution would silently turn the whole stream into content.
let summary = feed_unified_parser(&mut parser(&tokenizer), &prompt_token_ids, chunks.clone());
assert_eq!(summary.normal_text, expected_normal_text, "{name}");
assert_eq!(
summary.reasoning_text.len(),
expected_reasoning_len,
"{name}"
);
assert_eq!(summary.calls_len, expected_calls_len, "{name}");
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(&tokenizer);
b.iter_batched(
|| chunks.clone(),
|chunks| {
black_box(feed_unified_parser(
&mut parser,
&prompt_token_ids,
black_box(chunks),
))
},
BatchSize::SmallInput,
)
});
group.bench_function("create_parser", |b| {
b.iter_batched(
|| (parser(&tokenizer), chunks.clone()),
|(mut parser, chunks)| {
black_box(feed_unified_parser(
&mut parser,
&prompt_token_ids,
black_box(chunks),
))
},
BatchSize::SmallInput,
)
});
group.finish();
}
fn bench_kimi_k3(c: &mut Criterion) {
run_stream_group(
c,
"kimi_k3/mixed_complex_tool_call",
&mixed_fixture(),
"I will convert it.",
"I should inspect the data first.".len(),
2,
);
let long_reasoning = long_reasoning_fixture();
let reasoning_len = long_reasoning.find("<|close|>think").unwrap();
run_stream_group(
c,
"kimi_k3/long_reasoning",
&long_reasoning,
"Done.",
reasoning_len,
0,
);
run_stream_group(
c,
"kimi_k3/long_tool_argument",
&long_tool_argument_fixture(),
"I will write the file.",
"Write the file.".len(),
1,
);
}
criterion_group!(benches, bench_kimi_k3);
criterion_main!(benches);
rust/src/parser/benches/utils/adapter.rs
-114
View File
@@ -1,114 +0,0 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use std::sync::Arc;
use vllm_parser::tool::{Result, Tool, ToolParser, ToolParserError, ToolParserOutput};
use vllm_parser::unified::{
UnifiedParser, UnifiedParserError, UnifiedParserEvent, UnifiedParserOutput,
};
use vllm_tokenizer::{DecodedText, Tokenizer};
/// Tokenizer stub used by unified-parser benchmarks.
struct BenchTokenizer;
impl Tokenizer for BenchTokenizer {
fn encode(&self, text: &str, _add_special_tokens: bool) -> vllm_tokenizer::Result<Vec<u32>> {
Ok(text.chars().map(|_| u32::MAX).collect())
}
fn encode_ordinary(&self, text: &str) -> vllm_tokenizer::Result<Vec<u32>> {
self.encode(text, false)
}
fn decode(
&self,
token_ids: &[u32],
_skip_special_tokens: bool,
) -> vllm_tokenizer::Result<String> {
Ok("\u{FFFD}".repeat(token_ids.len()))
}
fn token_to_id(&self, _token: &str) -> Option<u32> {
Some(u32::MAX)
}
fn id_to_token(&self, _id: u32) -> Option<String> {
Some("\u{FFFD}".to_string())
}
}
/// Bench-only adapter that exposes a unified parser through the tool-parser
/// benchmark harness.
///
/// Returns error if the unified parser produces reasoning events.
pub struct UnifiedToolParserAdapter<T> {
inner: Box<dyn UnifiedParser>,
_marker: std::marker::PhantomData<T>,
}
fn map_unified_error(error: UnifiedParserError) -> ToolParserError {
ToolParserError::ParsingFailed {
message: format!("unified parser failed: {error}"),
}
}
fn append_unified_output(
output: UnifiedParserOutput,
tool_output: &mut ToolParserOutput,
) -> Result<()> {
for event in output.events {
match event {
UnifiedParserEvent::Text(text) => tool_output.push_text(text),
UnifiedParserEvent::ToolCall(call) => tool_output.push_call(call),
UnifiedParserEvent::Reasoning(_) => {
return Err(ToolParserError::ParsingFailed {
message: "unified parser emitted reasoning in tool-parser adapter".to_string(),
});
}
}
}
Ok(())
}
impl<T: UnifiedParser> ToolParser for UnifiedToolParserAdapter<T> {
fn create(tools: &[Tool]) -> Result<Box<dyn ToolParser>>
where
Self: Sized + 'static,
{
let inner = T::create(tools, Arc::new(BenchTokenizer)).map_err(map_unified_error)?;
Ok(Box::new(Self {
inner,
_marker: std::marker::PhantomData,
}))
}
fn preserve_special_tokens(&self) -> bool {
self.inner.preserve_special_tokens()
}
fn tool_call_id(&self, tool_index: usize) -> Option<&str> {
self.inner.tool_call_id(tool_index)
}
fn parse_into(&mut self, chunk: &str, output: &mut ToolParserOutput) -> Result<()> {
let mut unified_output = UnifiedParserOutput::default();
let result = self
.inner
.parse_into(DecodedText::unattributed(chunk), &mut unified_output)
.map_err(map_unified_error);
append_unified_output(unified_output, output)?;
result
}
fn finish(&mut self) -> Result<ToolParserOutput> {
let unified_output = self.inner.finish().map_err(map_unified_error)?;
let mut output = ToolParserOutput::default();
append_unified_output(unified_output, &mut output)?;
Ok(output)
}
fn reset(&mut self) -> String {
self.inner.reset()
}
}
rust/src/parser/benches/utils/mod.rs
+2 -2
View File
@@ -6,12 +6,12 @@
#![allow(dead_code)]
#![allow(unused_imports)]
mod adapter;
mod unified;
pub(super) use adapter::UnifiedToolParserAdapter;
use futures::FutureExt as _;
use openai_protocol::common::{Function as OpenAiFunction, Tool as OpenAiTool};
use tool_parser::traits::ToolParser as ExternalToolParser;
pub(super) use unified::{attributed_chunks, feed_unified_parser};
use vllm_parser::tool::test_utils::collect_stream;
use vllm_parser::tool::{Tool, ToolParser};
rust/src/parser/benches/utils/unified.rs
+76
View File
@@ -0,0 +1,76 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
use std::collections::BTreeSet;
use vllm_parser::tool::test_utils::split_by_chars;
use vllm_parser::unified::{UnifiedParser, UnifiedParserEvent, UnifiedParserOutput};
use vllm_tokenizer::{DecodedText, Tokenizer};
/// Split `text` into token-attributed chunks of at most `chunk_chars`
/// characters, as the frontend streams them to a unified parser.
///
/// The whole text is encoded with special tokens and run through the
/// incremental decoder once before splitting, so a marker cut by a chunk
/// boundary keeps its special-token anchor. Prepare chunks outside the measured
/// routine.
pub fn attributed_chunks(
tokenizer: &dyn Tokenizer,
text: &str,
chunk_chars: usize,
) -> Vec<DecodedText> {
let mut decoder = tokenizer.create_decode_stream(&[], false, 0);
for token_id in tokenizer.encode(text, false).expect("fixture should encode") {
decoder.push_token(token_id).expect("fixture token should decode");
}
let mut rest = decoder.flush(None).expect("decoder should flush").1;
assert_eq!(
rest.text, text,
"fixture should round-trip through the tokenizer"
);
split_by_chars(text, chunk_chars)
.into_iter()
.map(|chunk| rest.drain_prefix(chunk.len()))
.collect()
}
/// Event totals of one parsed stream, for benchmark sanity checks.
#[derive(Debug, Default)]
pub struct UnifiedStreamSummary {
pub normal_text: String,
pub reasoning_text: String,
pub calls_len: usize,
}
/// Run one request stream through `parser`: initialize it from the prompt,
/// feed every chunk, and finish.
pub fn feed_unified_parser(
parser: &mut dyn UnifiedParser,
prompt_token_ids: &[u32],
chunks: Vec<DecodedText>,
) -> UnifiedStreamSummary {
parser.initialize(prompt_token_ids).expect("parser should initialize");
let mut output = UnifiedParserOutput::default();
for chunk in chunks {
parser.parse_into(chunk, &mut output).expect("chunk should parse");
}
output.append(parser.finish().expect("stream should finish"));
let mut summary = UnifiedStreamSummary::default();
let mut tool_indices = BTreeSet::new();
for event in output.events {
match event {
UnifiedParserEvent::Text(text) => summary.normal_text.push_str(&text),
UnifiedParserEvent::Reasoning(reasoning) => {
summary.reasoning_text.push_str(&reasoning.text);
}
UnifiedParserEvent::ToolCall(call) => {
tool_indices.insert(call.tool_index);
}
}
}
summary.calls_len = tool_indices.len();
summary
}