@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
rs 128 lines 3.86 KB
Raw
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 use std::time::Duration;
5
6 use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
7 use tool_parser::parsers::QwenParser as ExternalQwenParser;
8 use vllm_parser::tool::test_utils::{split_by_chars, test_tools};
9 use vllm_parser::tool::{Qwen3XmlToolParser, Tool, ToolParser};
10
11 mod utils;
12 use utils::{feed_external_parser, feed_parser, openai_tools};
13
14 const CHUNK_CHARS: usize = 7;
15 const LONG_NORMAL_TEXT_REPEATS: usize = 2048;
16
17 fn tool_call(function_name: &str, arguments: &str) -> String {
18 format!("<tool_call>\n{{\"name\":\"{function_name}\",\"arguments\":{arguments}}}\n</tool_call>")
19 }
20
21 fn mixed_fixture() -> String {
22 format!(
23 "I will check two cities before answering.\n{}{}",
24 tool_call("get_weather", r#"{"location":"Hangzhou","days":3}"#),
25 tool_call("get_weather", r#"{"location":"San Francisco","days":2}"#),
26 )
27 }
28
29 fn long_normal_text_fixture() -> String {
30 let line = "This is ordinary assistant text with no Qwen XML tool markers at all.\n";
31 line.repeat(LONG_NORMAL_TEXT_REPEATS)
32 }
33
34 fn native_parser(tools: &[Tool]) -> Box<dyn ToolParser> {
35 Qwen3XmlToolParser::create(tools).expect("Qwen XML parser should initialize")
36 }
37
38 fn run_stream_group(
39 c: &mut Criterion,
40 name: &str,
41 tools: &[Tool],
42 text: &str,
43 chunk_chars: usize,
44 expected_normal_text: &str,
45 expected_native_calls_len: usize,
46 ) {
47 let chunks = split_by_chars(text, chunk_chars);
48 let openai_tools = openai_tools(tools);
49
50 let mut group = c.benchmark_group(name);
51 group.sample_size(50);
52 group.warm_up_time(Duration::from_millis(300));
53 group.measurement_time(Duration::from_secs(2));
54 group.throughput(Throughput::Bytes(text.len() as u64));
55
56 group.bench_function("native_reuse_parser", |b| {
57 let mut parser = native_parser(tools);
58 b.iter(|| {
59 let result = feed_parser(&mut *parser, black_box(&chunks));
60 debug_assert_eq!(result.0, expected_normal_text);
61 debug_assert_eq!(result.1, expected_native_calls_len);
62 black_box(result);
63 })
64 });
65
66 group.bench_function("native_create_parser", |b| {
67 b.iter_batched(
68 || native_parser(tools),
69 |mut parser| {
70 let result = feed_parser(&mut *parser, black_box(&chunks));
71 debug_assert_eq!(result.0, expected_normal_text);
72 debug_assert_eq!(result.1, expected_native_calls_len);
73 black_box(result);
74 },
75 BatchSize::SmallInput,
76 )
77 });
78
79 group.bench_function("external_reuse_parser", |b| {
80 let mut parser = ExternalQwenParser::new();
81 b.iter(|| {
82 let result = feed_external_parser(&mut parser, &openai_tools, black_box(&chunks));
83 black_box(result);
84 })
85 });
86
87 group.bench_function("external_create_parser", |b| {
88 b.iter_batched(
89 ExternalQwenParser::new,
90 |mut parser| {
91 let result = feed_external_parser(&mut parser, &openai_tools, black_box(&chunks));
92 black_box(result);
93 },
94 BatchSize::SmallInput,
95 )
96 });
97
98 group.finish();
99 }
100
101 fn bench_qwen3_xml(c: &mut Criterion) {
102 let tools = test_tools();
103 let mixed_text = mixed_fixture();
104 let long_normal_text = long_normal_text_fixture();
105
106 run_stream_group(
107 c,
108 "qwen3_xml/mixed_text_tool_call",
109 &tools,
110 &mixed_text,
111 CHUNK_CHARS,
112 "I will check two cities before answering.\n",
113 2,
114 );
115
116 run_stream_group(
117 c,
118 "qwen3_xml/long_normal_text",
119 &tools,
120 &long_normal_text,
121 CHUNK_CHARS,
122 &long_normal_text,
123 0,
124 );
125 }
126
127 criterion_group!(benches, bench_qwen3_xml);
128 criterion_main!(benches);