| 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::tool::Tool; |
| 9 | use vllm_parser::tool::test_utils::test_tools; |
| 10 | use vllm_parser::unified::Gemma4UnifiedParser; |
| 11 | use vllm_tokenizer::test_utils::TestTokenizer; |
| 12 | |
| 13 | mod utils; |
| 14 | use utils::{attributed_chunks, feed_unified_parser}; |
| 15 | |
| 16 | const CHUNK_CHARS: usize = 7; |
| 17 | const LONG_NORMAL_TEXT_REPEATS: usize = 2048; |
| 18 | const LONG_TOOL_ARGUMENT_REPEATS: usize = 256; |
| 19 | |
| 20 | fn mixed_fixture() -> String { |
| 21 | concat!( |
| 22 | "I will inspect the data before answering.\n", |
| 23 | "<|tool_call>", |
| 24 | "call:convert{", |
| 25 | "whole:114.514,", |
| 26 | "flag:true,", |
| 27 | "empty:<|\"|><|\"|>,", |
| 28 | "payload:{", |
| 29 | "name:<|\"|>demo<|\"|>,", |
| 30 | "count:42,", |
| 31 | "enabled:false,", |
| 32 | "missing:null,", |
| 33 | "nested:{level:2,label:<|\"|>deep<|\"|>},", |
| 34 | "tags:[<|\"|>red<|\"|>,<|\"|>blue<|\"|>,3,true,null,{kind:<|\"|>leaf<|\"|>}]", |
| 35 | "},", |
| 36 | "items:[", |
| 37 | "<|\"|>alpha<|\"|>,", |
| 38 | "{key:<|\"|>value<|\"|>,score:0.75},", |
| 39 | "[1,2,3]", |
| 40 | "]", |
| 41 | "}", |
| 42 | "<tool_call|>", |
| 43 | "<|tool_call>", |
| 44 | "call:update_record{", |
| 45 | "data:{id:7,active:true,notes:[<|\"|>keep<|\"|>,<|\"|>review<|\"|>]}", |
| 46 | "}", |
| 47 | "<tool_call|>", |
| 48 | " Finished.", |
| 49 | ) |
| 50 | .to_string() |
| 51 | } |
| 52 | |
| 53 | fn long_normal_text_fixture() -> String { |
| 54 | let line = "This is ordinary assistant text with no Gemma4 tool markers at all.\n"; |
| 55 | line.repeat(LONG_NORMAL_TEXT_REPEATS) |
| 56 | } |
| 57 | |
| 58 | fn long_tool_argument_fixture() -> String { |
| 59 | let line = |
| 60 | "<section><p>Literal } and <tool_call|> marker-shaped text inside content.</p></section>\n"; |
| 61 | format!( |
| 62 | concat!( |
| 63 | "I will write the file.\n", |
| 64 | "<|tool_call>", |
| 65 | "call:write_file{{", |
| 66 | "path:<|\"|>index.html<|\"|>,", |
| 67 | "content:<|\"|>{}<|\"|>", |
| 68 | "}}", |
| 69 | "<tool_call|>", |
| 70 | "Done." |
| 71 | ), |
| 72 | line.repeat(LONG_TOOL_ARGUMENT_REPEATS) |
| 73 | ) |
| 74 | } |
| 75 | |
| 76 | fn tokenizer() -> TestTokenizer { |
| 77 | TestTokenizer::new() |
| 78 | .with_special_token("<|tool_call>", 256) |
| 79 | .with_special_token("<tool_call|>", 257) |
| 80 | .with_special_token("<|\"|>", 258) |
| 81 | .with_special_token("<|channel>", 259) |
| 82 | .with_special_token("<channel|>", 260) |
| 83 | } |
| 84 | |
| 85 | fn parser(tools: &[Tool], tokenizer: &TestTokenizer) -> Gemma4UnifiedParser { |
| 86 | Gemma4UnifiedParser::new(tools, Arc::new(tokenizer.clone())) |
| 87 | .expect("Gemma4 unified parser should initialize") |
| 88 | } |
| 89 | |
| 90 | fn run_stream_group( |
| 91 | c: &mut Criterion, |
| 92 | name: &str, |
| 93 | tools: &[Tool], |
| 94 | text: &str, |
| 95 | chunk_chars: usize, |
| 96 | expected_normal_text: &str, |
| 97 | expected_calls_len: usize, |
| 98 | ) { |
| 99 | let tokenizer = tokenizer(); |
| 100 | let chunks = attributed_chunks(&tokenizer, text, chunk_chars); |
| 101 | |
| 102 | // Check the fixture once, outside the measurement: a marker lost to |
| 103 | // missing attribution would silently turn the whole stream into content. |
| 104 | let summary = feed_unified_parser(&mut parser(tools, &tokenizer), &[], chunks.clone()); |
| 105 | assert_eq!(summary.normal_text, expected_normal_text, "{name}"); |
| 106 | assert_eq!(summary.calls_len, expected_calls_len, "{name}"); |
| 107 | |
| 108 | let mut group = c.benchmark_group(name); |
| 109 | group.sample_size(50); |
| 110 | group.warm_up_time(Duration::from_millis(300)); |
| 111 | group.measurement_time(Duration::from_secs(2)); |
| 112 | group.throughput(Throughput::Bytes(text.len() as u64)); |
| 113 | |
| 114 | group.bench_function("reuse_parser", |b| { |
| 115 | let mut parser = parser(tools, &tokenizer); |
| 116 | b.iter_batched( |
| 117 | || chunks.clone(), |
| 118 | |chunks| black_box(feed_unified_parser(&mut parser, &[], black_box(chunks))), |
| 119 | BatchSize::SmallInput, |
| 120 | ) |
| 121 | }); |
| 122 | |
| 123 | group.bench_function("create_parser", |b| { |
| 124 | b.iter_batched( |
| 125 | || (parser(tools, &tokenizer), chunks.clone()), |
| 126 | |(mut parser, chunks)| { |
| 127 | black_box(feed_unified_parser(&mut parser, &[], black_box(chunks))) |
| 128 | }, |
| 129 | BatchSize::SmallInput, |
| 130 | ) |
| 131 | }); |
| 132 | |
| 133 | group.finish(); |
| 134 | } |
| 135 | |
| 136 | fn bench_gemma4(c: &mut Criterion) { |
| 137 | let tools = test_tools(); |
| 138 | let mixed_text = mixed_fixture(); |
| 139 | let long_tool_argument = long_tool_argument_fixture(); |
| 140 | let long_normal_text = long_normal_text_fixture(); |
| 141 | |
| 142 | run_stream_group( |
| 143 | c, |
| 144 | "gemma4/mixed_complex_tool_call", |
| 145 | &tools, |
| 146 | &mixed_text, |
| 147 | CHUNK_CHARS, |
| 148 | "I will inspect the data before answering.\n Finished.", |
| 149 | 2, |
| 150 | ); |
| 151 | |
| 152 | run_stream_group( |
| 153 | c, |
| 154 | "gemma4/long_tool_argument", |
| 155 | &tools, |
| 156 | &long_tool_argument, |
| 157 | CHUNK_CHARS, |
| 158 | "I will write the file.\nDone.", |
| 159 | 1, |
| 160 | ); |
| 161 | |
| 162 | run_stream_group( |
| 163 | c, |
| 164 | "gemma4/long_normal_text", |
| 165 | &tools, |
| 166 | &long_normal_text, |
| 167 | CHUNK_CHARS, |
| 168 | &long_normal_text, |
| 169 | 0, |
| 170 | ); |
| 171 | } |
| 172 | |
| 173 | criterion_group!(benches, bench_gemma4); |
| 174 | criterion_main!(benches); |