@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
rs 213 lines 6.44 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::Glm4MoeParser as ExternalGlm4MoeParser;
8 use vllm_parser::tool::test_utils::{split_by_chars, test_tools};
9 use vllm_parser::tool::{Glm45MoeToolParser, Glm47MoeToolParser, Tool, ToolParser};
10
11 mod utils;
12 use utils::{feed_external_parser, feed_parser, openai_tools};
13
14 const GLM45_PARSER_NAME: &str = "glm45";
15 const GLM47_PARSER_NAME: &str = "glm47";
16 const CHUNK_CHARS: usize = 7;
17 const LONG_NORMAL_TEXT_REPEATS: usize = 2048;
18
19 fn glm45_mixed_fixture() -> String {
20 concat!(
21 "I will check two cities before answering.\n",
22 "<tool_call>get_weather\n",
23 "<arg_key>city</arg_key>\n",
24 "<arg_value>Hangzhou</arg_value>\n",
25 "<arg_key>date</arg_key>\n",
26 "<arg_value>2026-05-07</arg_value>\n",
27 "<arg_key>unit</arg_key>\n",
28 "<arg_value>celsius</arg_value>\n",
29 "<arg_key>days</arg_key>\n",
30 "<arg_value>3</arg_value>\n",
31 "</tool_call>\n",
32 "<tool_call>get_weather\n",
33 "<arg_key>city</arg_key>\n",
34 "<arg_value>San Francisco</arg_value>\n",
35 "<arg_key>date</arg_key>\n",
36 "<arg_value>2026-05-07</arg_value>\n",
37 "<arg_key>unit</arg_key>\n",
38 "<arg_value>fahrenheit</arg_value>\n",
39 "<arg_key>days</arg_key>\n",
40 "<arg_value>2</arg_value>\n",
41 "</tool_call>",
42 )
43 .to_string()
44 }
45
46 fn glm47_mixed_fixture() -> String {
47 concat!(
48 "I will check two cities before answering.\n",
49 "<tool_call>get_weather",
50 "<arg_key>city</arg_key>",
51 "<arg_value>Hangzhou</arg_value>",
52 "<arg_key>date</arg_key>",
53 "<arg_value>2026-05-07</arg_value>",
54 "<arg_key>unit</arg_key>",
55 "<arg_value>celsius</arg_value>",
56 "<arg_key>days</arg_key>",
57 "<arg_value>3</arg_value>",
58 "</tool_call>",
59 "<tool_call>get_weather",
60 "<arg_key>city</arg_key>",
61 "<arg_value>San Francisco</arg_value>",
62 "<arg_key>date</arg_key>",
63 "<arg_value>2026-05-07</arg_value>",
64 "<arg_key>unit</arg_key>",
65 "<arg_value>fahrenheit</arg_value>",
66 "<arg_key>days</arg_key>",
67 "<arg_value>2</arg_value>",
68 "</tool_call>",
69 )
70 .to_string()
71 }
72
73 fn long_normal_text_fixture() -> String {
74 let line = "This is ordinary assistant text with no GLM MoE tool markers at all.\n";
75 line.repeat(LONG_NORMAL_TEXT_REPEATS)
76 }
77
78 fn native_parser(name: &str, tools: &[Tool]) -> Box<dyn ToolParser> {
79 match name {
80 GLM45_PARSER_NAME => Glm45MoeToolParser::create(tools),
81 GLM47_PARSER_NAME => Glm47MoeToolParser::create(tools),
82 _ => unreachable!("unexpected GLM parser name"),
83 }
84 .expect("GLM MoE parser should initialize")
85 }
86
87 fn external_parser(name: &str) -> ExternalGlm4MoeParser {
88 match name {
89 GLM45_PARSER_NAME => ExternalGlm4MoeParser::glm45(),
90 GLM47_PARSER_NAME => ExternalGlm4MoeParser::glm47(),
91 _ => unreachable!("unexpected GLM parser name"),
92 }
93 }
94
95 fn run_stream_group(
96 c: &mut Criterion,
97 name: &str,
98 parser_name: &str,
99 tools: &[Tool],
100 text: &str,
101 chunk_chars: usize,
102 expected_normal_text: &str,
103 expected_native_calls_len: usize,
104 ) {
105 let chunks = split_by_chars(text, chunk_chars);
106 let openai_tools = openai_tools(tools);
107
108 let mut group = c.benchmark_group(name);
109 group.sample_size(50);
110 group.warm_up_time(Duration::from_millis(300));
111 group.measurement_time(Duration::from_secs(2));
112 group.throughput(Throughput::Bytes(text.len() as u64));
113
114 group.bench_function("native_reuse_parser", |b| {
115 let mut parser = native_parser(parser_name, tools);
116 b.iter(|| {
117 let result = feed_parser(&mut *parser, black_box(&chunks));
118 debug_assert_eq!(result.0, expected_normal_text);
119 debug_assert_eq!(result.1, expected_native_calls_len);
120 black_box(result);
121 })
122 });
123
124 group.bench_function("native_create_parser", |b| {
125 b.iter_batched(
126 || native_parser(parser_name, tools),
127 |mut parser| {
128 let result = feed_parser(&mut *parser, black_box(&chunks));
129 debug_assert_eq!(result.0, expected_normal_text);
130 debug_assert_eq!(result.1, expected_native_calls_len);
131 black_box(result);
132 },
133 BatchSize::SmallInput,
134 )
135 });
136
137 group.bench_function("external_reuse_parser", |b| {
138 let mut parser = external_parser(parser_name);
139 b.iter(|| {
140 let result = feed_external_parser(&mut parser, &openai_tools, black_box(&chunks));
141 debug_assert_eq!(result.0, expected_normal_text);
142 black_box(result);
143 })
144 });
145
146 group.bench_function("external_create_parser", |b| {
147 b.iter_batched(
148 || external_parser(parser_name),
149 |mut parser| {
150 let result = feed_external_parser(&mut parser, &openai_tools, black_box(&chunks));
151 debug_assert_eq!(result.0, expected_normal_text);
152 black_box(result);
153 },
154 BatchSize::SmallInput,
155 )
156 });
157
158 group.finish();
159 }
160
161 fn bench_glm45_moe(c: &mut Criterion) {
162 let tools = test_tools();
163 let glm45_mixed_text = glm45_mixed_fixture();
164 let glm47_mixed_text = glm47_mixed_fixture();
165 let long_normal_text = long_normal_text_fixture();
166
167 run_stream_group(
168 c,
169 "glm45/mixed_text_tool_call",
170 GLM45_PARSER_NAME,
171 &tools,
172 &glm45_mixed_text,
173 CHUNK_CHARS,
174 "I will check two cities before answering.\n",
175 2,
176 );
177
178 run_stream_group(
179 c,
180 "glm47/mixed_text_tool_call",
181 GLM47_PARSER_NAME,
182 &tools,
183 &glm47_mixed_text,
184 CHUNK_CHARS,
185 "I will check two cities before answering.\n",
186 2,
187 );
188
189 run_stream_group(
190 c,
191 "glm45/long_normal_text",
192 GLM45_PARSER_NAME,
193 &tools,
194 &long_normal_text,
195 CHUNK_CHARS,
196 &long_normal_text,
197 0,
198 );
199
200 run_stream_group(
201 c,
202 "glm47/long_normal_text",
203 GLM47_PARSER_NAME,
204 &tools,
205 &long_normal_text,
206 CHUNK_CHARS,
207 &long_normal_text,
208 0,
209 );
210 }
211
212 criterion_group!(benches, bench_glm45_moe);
213 criterion_main!(benches);