@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
rs 124 lines 4.25 KB
Raw
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 use criterion::{Criterion, Throughput, black_box, criterion_group, criterion_main};
5 use hf_hub::api::tokio::ApiBuilder;
6 use tokio::runtime::Runtime;
7 use vllm_tokenizer::{TiktokenTokenizer, Tokenizer};
8
9 const MODEL_ID: &str = "moonshotai/Kimi-K2.5";
10 const SAMPLE_TEXT: &str = "\
11 <think>
12 I'm sure it's fine, but I can't say I'd trust that it's what we'd ship.
13 </think>
14 请用中英混合总结以下需求，并保留 tool-call marker:
15 <|tool_calls_section_begin|>{\"name\":\"summarize\",\"arguments\":{\"style\":\"brief\"}}<|tool_calls_section_end|>
16 The service should stop cleanly at EOS, avoid leaking the next template turn, and keep decode latency low.
17 ";
18
19 struct BenchFixture {
20 riptoken: TiktokenTokenizer,
21 tiktoken_rs: TiktokenTokenizer,
22 text: String,
23 token_ids: Vec<u32>,
24 }
25
26 impl BenchFixture {
27 fn load() -> Self {
28 let path = tiktoken_model();
29 let riptoken = TiktokenTokenizer::new_riptoken(&path).expect("load riptoken tokenizer");
30 let tiktoken_rs =
31 TiktokenTokenizer::new_tiktoken_rs(&path).expect("load tiktoken-rs tokenizer");
32
33 let text = SAMPLE_TEXT.repeat(32);
34 let riptoken_token_ids =
35 riptoken.encode(text.as_str(), false).expect("encode sample text with riptoken");
36 let tiktoken_rs_token_ids = tiktoken_rs
37 .encode(text.as_str(), false)
38 .expect("encode sample text with tiktoken-rs");
39 assert_eq!(riptoken_token_ids, tiktoken_rs_token_ids);
40
41 let riptoken_decoded = riptoken
42 .decode(riptoken_token_ids.as_slice(), false)
43 .expect("decode sample token ids with riptoken");
44 let tiktoken_rs_decoded = tiktoken_rs
45 .decode(riptoken_token_ids.as_slice(), false)
46 .expect("decode sample token ids with tiktoken-rs");
47 assert_eq!(riptoken_decoded, tiktoken_rs_decoded);
48
49 Self {
50 riptoken,
51 tiktoken_rs,
52 text,
53 token_ids: riptoken_token_ids,
54 }
55 }
56 }
57
58 fn tiktoken_model() -> std::path::PathBuf {
59 Runtime::new().expect("build tokio runtime").block_on(async {
60 let repo = ApiBuilder::from_env()
61 .with_progress(false)
62 .build()
63 .expect("build hf-hub api")
64 .model(MODEL_ID.to_string());
65 repo.get("config.json").await.expect("fetch config.json from hf-hub");
66 repo.get("tokenizer_config.json")
67 .await
68 .expect("fetch tokenizer_config.json from hf-hub");
69 repo.get("tiktoken.model").await.expect("fetch tiktoken.model from hf-hub")
70 })
71 }
72
73 fn bench_encode(c: &mut Criterion) {
74 let fixture = BenchFixture::load();
75 let mut group = c.benchmark_group("tiktoken_encode");
76 group.throughput(Throughput::Bytes(fixture.text.len() as u64));
77
78 group.bench_function("riptoken", |b| {
79 b.iter(|| {
80 fixture
81 .riptoken
82 .encode(black_box(fixture.text.as_str()), black_box(false))
83 .expect("encode sample text with riptoken")
84 })
85 });
86 group.bench_function("tiktoken_rs", |b| {
87 b.iter(|| {
88 fixture
89 .tiktoken_rs
90 .encode(black_box(fixture.text.as_str()), black_box(false))
91 .expect("encode sample text with tiktoken-rs")
92 })
93 });
94
95 group.finish();
96 }
97
98 fn bench_decode(c: &mut Criterion) {
99 let fixture = BenchFixture::load();
100 let mut group = c.benchmark_group("tiktoken_decode");
101 group.throughput(Throughput::Elements(fixture.token_ids.len() as u64));
102
103 group.bench_function("riptoken", |b| {
104 b.iter(|| {
105 fixture
106 .riptoken
107 .decode(black_box(fixture.token_ids.as_slice()), black_box(false))
108 .expect("decode sample token ids with riptoken")
109 })
110 });
111 group.bench_function("tiktoken_rs", |b| {
112 b.iter(|| {
113 fixture
114 .tiktoken_rs
115 .decode(black_box(fixture.token_ids.as_slice()), black_box(false))
116 .expect("decode sample token ids with tiktoken-rs")
117 })
118 });
119
120 group.finish();
121 }
122
123 criterion_group!(benches, bench_encode, bench_decode);
124 criterion_main!(benches);