@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
rs 78 lines 2.52 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::{Granite4ToolParser, Tool, ToolParser};
9
10 mod utils;
11 use utils::feed_parser;
12
13 const CHUNK_CHARS: usize = 7;
14 const LONG_ARGUMENT_BYTES: usize = 64 * 1024;
15
16 fn string_args_fixture() -> String {
17 let arguments = format!(r#"{{"data":"{}"}}"#, "x".repeat(LONG_ARGUMENT_BYTES));
18 let encoded_arguments = serde_json::to_string(&arguments).unwrap();
19 format!(r#"<tool_call>{{"name":"f","arguments":{encoded_arguments}}}</tool_call>"#)
20 }
21
22 fn object_args_fixture() -> String {
23 format!(
24 r#"<tool_call>{{"name":"f","arguments":{{"data":"{}"}}}}</tool_call>"#,
25 "x".repeat(LONG_ARGUMENT_BYTES)
26 )
27 }
28
29 fn parser(tools: &[Tool]) -> Box<dyn ToolParser> {
30 Granite4ToolParser::create(tools).expect("Granite4 parser should initialize")
31 }
32
33 fn run_stream_group(c: &mut Criterion, name: &str, tools: &[Tool], text: &str) {
34 let chunks = split_by_chars(text, CHUNK_CHARS);
35
36 let mut group = c.benchmark_group(name);
37 group.sample_size(50);
38 group.warm_up_time(Duration::from_millis(300));
39 group.measurement_time(Duration::from_secs(2));
40 group.throughput(Throughput::Bytes(text.len() as u64));
41
42 group.bench_function("reuse_parser", |b| {
43 let mut parser = parser(tools);
44 b.iter(|| {
45 let result = feed_parser(&mut *parser, black_box(&chunks));
46 debug_assert_eq!(result.0, "");
47 debug_assert_eq!(result.1, 1);
48 black_box(result);
49 })
50 });
51
52 group.bench_function("create_parser", |b| {
53 b.iter_batched(
54 || parser(tools),
55 |mut parser| {
56 let result = feed_parser(&mut *parser, black_box(&chunks));
57 debug_assert_eq!(result.0, "");
58 debug_assert_eq!(result.1, 1);
59 black_box(result);
60 },
61 BatchSize::SmallInput,
62 )
63 });
64
65 group.finish();
66 }
67
68 fn bench_granite4(c: &mut Criterion) {
69 let tools = test_tools();
70 let string_args = string_args_fixture();
71 let object_args = object_args_fixture();
72
73 run_stream_group(c, "granite4/long_string_arguments", &tools, &string_args);
74 run_stream_group(c, "granite4/long_object_arguments", &tools, &object_args);
75 }
76
77 criterion_group!(benches, bench_granite4);
78 criterion_main!(benches);