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