@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
rs 116 lines 3.77 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 vllm_parser::tool::test_utils::{split_by_chars, test_tools};
8 use vllm_parser::tool::{DeepSeekV32ToolParser, Tool, ToolParser};
9
10 mod utils;
11 use utils::feed_parser;
12
13 const CHUNK_CHARS: usize = 7;
14 const LONG_NORMAL_TEXT_REPEATS: usize = 2048;
15
16 fn mixed_fixture() -> String {
17 concat!(
18 "I will check two cities before answering.\n",
19 "<｜DSML｜function_calls>\n",
20 "<｜DSML｜invoke name=\"get_weather\">\n",
21 "<｜DSML｜parameter name=\"location\" string=\"true\">Hangzhou</｜DSML｜parameter>\n",
22 "<｜DSML｜parameter name=\"date\" string=\"true\">2026-04-28</｜DSML｜parameter>\n",
23 "<｜DSML｜parameter name=\"unit\" string=\"true\">celsius</｜DSML｜parameter>\n",
24 "<｜DSML｜parameter name=\"days\" string=\"false\">3</｜DSML｜parameter>\n",
25 "</｜DSML｜invoke>\n",
26 "<｜DSML｜invoke name=\"get_weather\">\n",
27 "<｜DSML｜parameter name=\"location\" string=\"true\">San Francisco</｜DSML｜parameter>\n",
28 "<｜DSML｜parameter name=\"date\" string=\"true\">2026-04-28</｜DSML｜parameter>\n",
29 "<｜DSML｜parameter name=\"unit\" string=\"true\">fahrenheit</｜DSML｜parameter>\n",
30 "<｜DSML｜parameter name=\"days\" string=\"false\">2</｜DSML｜parameter>\n",
31 "</｜DSML｜invoke>\n",
32 "</｜DSML｜function_calls>",
33 )
34 .to_string()
35 }
36
37 fn long_normal_text_fixture() -> String {
38 let line = "This is ordinary assistant text with no DSML tool markers at all.\n";
39 line.repeat(LONG_NORMAL_TEXT_REPEATS)
40 }
41
42 fn parser(tools: &[Tool]) -> Box<dyn ToolParser> {
43 DeepSeekV32ToolParser::create(tools).expect("DeepSeek V3.2 parser should initialize")
44 }
45
46 fn run_stream_group(
47 c: &mut Criterion,
48 name: &str,
49 tools: &[Tool],
50 text: &str,
51 chunk_chars: usize,
52 expected_normal_text: &str,
53 expected_calls_len: usize,
54 ) {
55 let chunks = split_by_chars(text, chunk_chars);
56
57 let mut group = c.benchmark_group(name);
58 group.sample_size(50);
59 group.warm_up_time(Duration::from_millis(300));
60 group.measurement_time(Duration::from_secs(2));
61 group.throughput(Throughput::Bytes(text.len() as u64));
62
63 group.bench_function("reuse_parser", |b| {
64 let mut parser = parser(tools);
65 b.iter(|| {
66 let result = feed_parser(&mut *parser, black_box(&chunks));
67 debug_assert_eq!(result.0, expected_normal_text);
68 debug_assert_eq!(result.1, expected_calls_len);
69 black_box(result);
70 })
71 });
72
73 group.bench_function("create_parser", |b| {
74 b.iter_batched(
75 || parser(tools),
76 |mut parser| {
77 let result = feed_parser(&mut *parser, black_box(&chunks));
78 debug_assert_eq!(result.0, expected_normal_text);
79 debug_assert_eq!(result.1, expected_calls_len);
80 black_box(result);
81 },
82 BatchSize::SmallInput,
83 )
84 });
85
86 group.finish();
87 }
88
89 fn bench_deepseek_v32(c: &mut Criterion) {
90 let tools = test_tools();
91 let mixed_text = mixed_fixture();
92 let long_normal_text = long_normal_text_fixture();
93
94 run_stream_group(
95 c,
96 "deepseek_v32/mixed_text_tool_call",
97 &tools,
98 &mixed_text,
99 CHUNK_CHARS,
100 "I will check two cities before answering.\n",
101 2,
102 );
103
104 run_stream_group(
105 c,
106 "deepseek_v32/long_normal_text",
107 &tools,
108 &long_normal_text,
109 CHUNK_CHARS,
110 &long_normal_text,
111 0,
112 );
113 }
114
115 criterion_group!(benches, bench_deepseek_v32);
116 criterion_main!(benches);