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