@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
rs 133 lines 4.05 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::DeepSeekParser as ExternalDeepSeekParser;
8 use vllm_parser::tool::test_utils::{split_by_chars, test_tools};
9 use vllm_parser::tool::{DeepSeekV3ToolParser, 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 mixed_fixture() -> String {
18 concat!(
19 "I will check two cities before answering.\n",
20 "<｜tool▁calls▁begin｜>",
21 "<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n",
22 "```json\n",
23 "{\"location\":\"Hangzhou\",\"days\":3}",
24 "\n```<｜tool▁call▁end｜>",
25 "<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_weather\n",
26 "```json\n",
27 "{\"location\":\"San Francisco\",\"days\":2}",
28 "\n```<｜tool▁call▁end｜>",
29 "<｜tool▁calls▁end｜>",
30 )
31 .to_string()
32 }
33
34 fn long_normal_text_fixture() -> String {
35 let line = "This is ordinary assistant text with no DeepSeek V3 tool markers at all.\n";
36 line.repeat(LONG_NORMAL_TEXT_REPEATS)
37 }
38
39 fn native_parser(tools: &[Tool]) -> Box<dyn ToolParser> {
40 DeepSeekV3ToolParser::create(tools).expect("DeepSeek V3 parser should initialize")
41 }
42
43 fn run_stream_group(
44 c: &mut Criterion,
45 name: &str,
46 tools: &[Tool],
47 text: &str,
48 chunk_chars: usize,
49 expected_normal_text: &str,
50 expected_native_calls_len: usize,
51 ) {
52 let chunks = split_by_chars(text, chunk_chars);
53 let openai_tools = openai_tools(tools);
54
55 let mut group = c.benchmark_group(name);
56 group.sample_size(50);
57 group.warm_up_time(Duration::from_millis(300));
58 group.measurement_time(Duration::from_secs(2));
59 group.throughput(Throughput::Bytes(text.len() as u64));
60
61 group.bench_function("native_reuse_parser", |b| {
62 let mut parser = native_parser(tools);
63 b.iter(|| {
64 let result = feed_parser(&mut *parser, black_box(&chunks));
65 debug_assert_eq!(result.0, expected_normal_text);
66 debug_assert_eq!(result.1, expected_native_calls_len);
67 black_box(result);
68 })
69 });
70
71 group.bench_function("native_create_parser", |b| {
72 b.iter_batched(
73 || native_parser(tools),
74 |mut parser| {
75 let result = feed_parser(&mut *parser, black_box(&chunks));
76 debug_assert_eq!(result.0, expected_normal_text);
77 debug_assert_eq!(result.1, expected_native_calls_len);
78 black_box(result);
79 },
80 BatchSize::SmallInput,
81 )
82 });
83
84 group.bench_function("external_reuse_parser", |b| {
85 let mut parser = ExternalDeepSeekParser::new();
86 b.iter(|| {
87 let result = feed_external_parser(&mut parser, &openai_tools, black_box(&chunks));
88 black_box(result);
89 })
90 });
91
92 group.bench_function("external_create_parser", |b| {
93 b.iter_batched(
94 ExternalDeepSeekParser::new,
95 |mut parser| {
96 let result = feed_external_parser(&mut parser, &openai_tools, black_box(&chunks));
97 black_box(result);
98 },
99 BatchSize::SmallInput,
100 )
101 });
102
103 group.finish();
104 }
105
106 fn bench_deepseek_v3(c: &mut Criterion) {
107 let tools = test_tools();
108 let mixed_text = mixed_fixture();
109 let long_normal_text = long_normal_text_fixture();
110
111 run_stream_group(
112 c,
113 "deepseek_v3/mixed_text_tool_call",
114 &tools,
115 &mixed_text,
116 CHUNK_CHARS,
117 "I will check two cities before answering.\n",
118 2,
119 );
120
121 run_stream_group(
122 c,
123 "deepseek_v3/long_normal_text",
124 &tools,
125 &long_normal_text,
126 CHUNK_CHARS,
127 &long_normal_text,
128 0,
129 );
130 }
131
132 criterion_group!(benches, bench_deepseek_v3);
133 criterion_main!(benches);