48 lines
1.4 KiB
Rust
48 lines
1.4 KiB
Rust
|
use std::sync::Arc;
|
||
|
|
|
||
|
|
use criterion::{black_box, criterion_group, criterion_main, Criterion};
|
||
|
|
use rand::Rng;
|
||
|
|
|
||
|
|
use text_generation_router_v3::block_allocator::Allocator;
|
||
|
|
use text_generation_router_v3::radix::RadixAllocator;
|
||
|
|
|
||
|
|
fn prefix_cache_benchmark(c: &mut Criterion) {
|
||
|
|
// let prefixes: Vec<Vec<u32>> = (0..8192)
|
||
|
|
// .chunks(256)
|
||
|
|
// .into_iter()
|
||
|
|
// .map(|c| c.collect())
|
||
|
|
// .collect();
|
||
|
|
|
||
|
|
let mut cache = RadixAllocator::new(1, 262144, None);
|
||
|
|
|
||
|
|
c.bench_function("Radix allocator", |b| {
|
||
|
|
b.iter_batched(
|
||
|
|
|| {
|
||
|
|
//prefixes
|
||
|
|
// .choose_multiple(&mut rand::thread_rng(), 5)
|
||
|
|
// .fold(Vec::new(), |mut v, s| {
|
||
|
|
// v.extend(s);
|
||
|
|
// v
|
||
|
|
// })
|
||
|
|
|
||
|
|
(0..7936)
|
||
|
|
.map(|_| rand::thread_rng().gen_range(0..1024))
|
||
|
|
.collect::<Vec<u32>>()
|
||
|
|
},
|
||
|
|
|prefill| {
|
||
|
|
let alloc = cache.allocate(
|
||
|
|
prefill.len() as u32 + 13,
|
||
|
|
Some(Arc::new(black_box(prefill))),
|
||
|
|
);
|
||
|
|
if let Some(alloc) = alloc {
|
||
|
|
cache.free(alloc.blocks.clone(), alloc.allocation_id);
|
||
|
|
}
|
||
|
|
},
|
||
|
|
criterion::BatchSize::SmallInput,
|
||
|
|
);
|
||
|
|
});
|
||
|
|
}
|
||
|
|
|
||
|
|
criterion_group!(benches, prefix_cache_benchmark);
|
||
|
|
criterion_main!(benches);
|