@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
rs 125 lines 3.98 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::{HuggingFaceTokenizer, Tokenizer};
8
9 const MODEL_ID: &str = "Qwen/Qwen3.5-0.8B";
10 const SAMPLE_TEXT: &str = "\
11 <|im_start|>system
12 You are Qwen3.5, a helpful assistant.
13 <|im_end|>
14 <|im_start|>user
15 请用中英混合总结以下需求，并给出一个简短的 JSON 示例。
16 The service should stop cleanly at EOS, avoid leaking the next template turn, and keep decode latency low.
17 Input: 4 concurrent requests, 10240 prompt tokens, 16 generated tokens.
18 <|im_end|>
19 <|im_start|>assistant
20 ";
21
22 struct BenchFixture {
23 fastokens: HuggingFaceTokenizer,
24 hf: HuggingFaceTokenizer,
25 text: String,
26 token_ids: Vec<u32>,
27 }
28
29 impl BenchFixture {
30 fn load() -> Self {
31 let path = tokenizer_json();
32 let fastokens =
33 HuggingFaceTokenizer::new_fastokens(&path).expect("load fastokens tokenizer");
34 let hf = HuggingFaceTokenizer::new_hf(&path).expect("load huggingface tokenizer");
35
36 let text = SAMPLE_TEXT.repeat(32);
37 let hf_token_ids =
38 hf.encode(text.as_str(), false).expect("encode sample text with hf tokenizer");
39 let fastokens_token_ids = fastokens
40 .encode(text.as_str(), false)
41 .expect("encode sample text with fastokens");
42 assert_eq!(fastokens_token_ids, hf_token_ids);
43
44 let hf_decoded = hf
45 .decode(hf_token_ids.as_slice(), false)
46 .expect("decode sample token ids with hf tokenizer");
47 let fastokens_decoded = fastokens
48 .decode(hf_token_ids.as_slice(), false)
49 .expect("decode sample token ids with fastokens");
50 assert_eq!(fastokens_decoded, hf_decoded);
51
52 Self {
53 fastokens,
54 hf,
55 text,
56 token_ids: hf_token_ids,
57 }
58 }
59 }
60
61 fn tokenizer_json() -> std::path::PathBuf {
62 Runtime::new().expect("build tokio runtime").block_on(async {
63 ApiBuilder::from_env()
64 .with_progress(false)
65 .build()
66 .expect("build hf-hub api")
67 .model(MODEL_ID.to_string())
68 .get("tokenizer.json")
69 .await
70 .expect("fetch tokenizer.json from hf-hub")
71 })
72 }
73
74 fn bench_encode(c: &mut Criterion) {
75 let fixture = BenchFixture::load();
76 let mut group = c.benchmark_group("tokenizer_encode");
77 group.throughput(Throughput::Bytes(fixture.text.len() as u64));
78
79 group.bench_function("fastokens", |b| {
80 b.iter(|| {
81 fixture
82 .fastokens
83 .encode(black_box(fixture.text.as_str()), black_box(false))
84 .expect("encode sample text with fastokens")
85 })
86 });
87 group.bench_function("hf_tokenizers", |b| {
88 b.iter(|| {
89 fixture
90 .hf
91 .encode(black_box(fixture.text.as_str()), black_box(false))
92 .expect("encode sample text with hf tokenizer")
93 })
94 });
95
96 group.finish();
97 }
98
99 fn bench_decode(c: &mut Criterion) {
100 let fixture = BenchFixture::load();
101 let mut group = c.benchmark_group("tokenizer_decode");
102 group.throughput(Throughput::Elements(fixture.token_ids.len() as u64));
103
104 group.bench_function("fastokens", |b| {
105 b.iter(|| {
106 fixture
107 .fastokens
108 .decode(black_box(fixture.token_ids.as_slice()), black_box(false))
109 .expect("decode sample token ids with fastokens")
110 })
111 });
112 group.bench_function("hf_tokenizers", |b| {
113 b.iter(|| {
114 fixture
115 .hf
116 .decode(black_box(fixture.token_ids.as_slice()), black_box(false))
117 .expect("decode sample token ids with hf tokenizer")
118 })
119 });
120
121 group.finish();
122 }
123
124 criterion_group!(benches, bench_encode, bench_decode);
125 criterion_main!(benches);