@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
rs 191 lines 5.65 KB
Raw
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 use std::sync::Arc;
5 use std::time::Duration;
6
7 use criterion::{BatchSize, Criterion, Throughput, black_box, criterion_group, criterion_main};
8 use vllm_parser::unified::KimiK3UnifiedParser;
9 use vllm_tokenizer::Tokenizer;
10 use vllm_tokenizer::test_utils::TestTokenizer;
11
12 mod utils;
13 use utils::{attributed_chunks, feed_unified_parser};
14
15 const CHUNK_CHARS: usize = 7;
16 const LONG_REASONING_REPEATS: usize = 2048;
17 const LONG_TOOL_ARGUMENT_REPEATS: usize = 256;
18
19 /// The chat generation prompt for thinking mode ends inside the `think`
20 /// channel, so the output starts with reasoning.
21 const THINKING_PROMPT_TAIL: &str = "<|open|>think<|sep|>";
22
23 fn tokenizer() -> TestTokenizer {
24 TestTokenizer::new()
25 .with_special_token("<|open|>", 256)
26 .with_special_token("<|close|>", 257)
27 .with_special_token("<|sep|>", 258)
28 .with_special_token("<|end_of_msg|>", 259)
29 }
30
31 fn argument(key: &str, arg_type: &str, value: &str) -> String {
32 format!(
33 "<|open|>argument key=\"{key}\" type=\"{arg_type}\"<|sep|>{value}<|close|>argument<|sep|>"
34 )
35 }
36
37 fn call(tool: &str, index: usize, body: &str) -> String {
38 format!("<|open|>call tool=\"{tool}\" index=\"{index}\"<|sep|>{body}<|close|>call<|sep|>")
39 }
40
41 fn message(reasoning: &str, response: &str, calls: &str) -> String {
42 format!(
43 "{reasoning}<|close|>think<|sep|>\
44 <|open|>response<|sep|>{response}<|close|>response<|sep|>\
45 <|open|>tools<|sep|>{calls}<|close|>tools<|sep|>\
46 <|close|>message<|sep|>"
47 )
48 }
49
50 fn mixed_fixture() -> String {
51 let calls = [
52 call(
53 "convert",
54 1,
55 &[
56 argument("whole", "number", "114.514"),
57 argument("flag", "boolean", "true"),
58 argument("name", "string", "demo"),
59 argument("payload", "object", r#"{"count":42,"tags":["red","blue"]}"#),
60 ]
61 .concat(),
62 ),
63 call("update_record", 2, &argument("id", "integer", "7")),
64 ]
65 .concat();
66 message(
67 "I should inspect the data first.",
68 "I will convert it.",
69 &calls,
70 )
71 }
72
73 fn long_reasoning_fixture() -> String {
74 let line = "Ordinary reasoning text with no K3 channel markers at all.\n";
75 message(&line.repeat(LONG_REASONING_REPEATS), "Done.", "")
76 }
77
78 fn long_tool_argument_fixture() -> String {
79 let line = "<section><p>Literal <tag> and quotes \" inside content.</p></section>\n";
80 let content = line.repeat(LONG_TOOL_ARGUMENT_REPEATS);
81 let body = [
82 argument("path", "string", "index.html"),
83 argument("content", "string", &content),
84 ]
85 .concat();
86 message(
87 "Write the file.",
88 "I will write the file.",
89 &call("write_file", 1, &body),
90 )
91 }
92
93 fn parser(tokenizer: &TestTokenizer) -> KimiK3UnifiedParser {
94 KimiK3UnifiedParser::new(&[], Arc::new(tokenizer.clone()))
95 .expect("Kimi K3 unified parser should initialize")
96 }
97
98 fn run_stream_group(
99 c: &mut Criterion,
100 name: &str,
101 text: &str,
102 expected_normal_text: &str,
103 expected_reasoning_len: usize,
104 expected_calls_len: usize,
105 ) {
106 let tokenizer = tokenizer();
107 let prompt_token_ids = tokenizer.encode(THINKING_PROMPT_TAIL, false).unwrap();
108 let chunks = attributed_chunks(&tokenizer, text, CHUNK_CHARS);
109
110 // Check the fixture once, outside the measurement: a marker lost to
111 // missing attribution would silently turn the whole stream into content.
112 let summary = feed_unified_parser(&mut parser(&tokenizer), &prompt_token_ids, chunks.clone());
113 assert_eq!(summary.normal_text, expected_normal_text, "{name}");
114 assert_eq!(
115 summary.reasoning_text.len(),
116 expected_reasoning_len,
117 "{name}"
118 );
119 assert_eq!(summary.calls_len, expected_calls_len, "{name}");
120
121 let mut group = c.benchmark_group(name);
122 group.sample_size(50);
123 group.warm_up_time(Duration::from_millis(300));
124 group.measurement_time(Duration::from_secs(2));
125 group.throughput(Throughput::Bytes(text.len() as u64));
126
127 group.bench_function("reuse_parser", |b| {
128 let mut parser = parser(&tokenizer);
129 b.iter_batched(
130 || chunks.clone(),
131 |chunks| {
132 black_box(feed_unified_parser(
133 &mut parser,
134 &prompt_token_ids,
135 black_box(chunks),
136 ))
137 },
138 BatchSize::SmallInput,
139 )
140 });
141
142 group.bench_function("create_parser", |b| {
143 b.iter_batched(
144 || (parser(&tokenizer), chunks.clone()),
145 |(mut parser, chunks)| {
146 black_box(feed_unified_parser(
147 &mut parser,
148 &prompt_token_ids,
149 black_box(chunks),
150 ))
151 },
152 BatchSize::SmallInput,
153 )
154 });
155
156 group.finish();
157 }
158
159 fn bench_kimi_k3(c: &mut Criterion) {
160 run_stream_group(
161 c,
162 "kimi_k3/mixed_complex_tool_call",
163 &mixed_fixture(),
164 "I will convert it.",
165 "I should inspect the data first.".len(),
166 2,
167 );
168
169 let long_reasoning = long_reasoning_fixture();
170 let reasoning_len = long_reasoning.find("<|close|>think").unwrap();
171 run_stream_group(
172 c,
173 "kimi_k3/long_reasoning",
174 &long_reasoning,
175 "Done.",
176 reasoning_len,
177 0,
178 );
179
180 run_stream_group(
181 c,
182 "kimi_k3/long_tool_argument",
183 &long_tool_argument_fixture(),
184 "I will write the file.",
185 "Write the file.".len(),
186 1,
187 );
188 }
189
190 criterion_group!(benches, bench_kimi_k3);
191 criterion_main!(benches);