234 lines
7.8 KiB
Rust
234 lines
7.8 KiB
Rust
//! Benchmarks for text generation algorithms
|
|
|
|
use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main};
|
|
use rtx_nlg::*;
|
|
use std::sync::Arc;
|
|
|
|
fn bench_beam_search(c: &mut Criterion) {
|
|
let rt = tokio::runtime::Runtime::new().unwrap();
|
|
|
|
let mut group = c.benchmark_group("beam_search");
|
|
group.throughput(Throughput::Elements(1));
|
|
|
|
for num_beams in [1, 4, 8, 16].iter() {
|
|
group.bench_with_input(
|
|
BenchmarkId::new("beams", num_beams),
|
|
num_beams,
|
|
|b, &num_beams| {
|
|
b.to_async(&rt).iter(|| async {
|
|
let config = GenerationConfig::beam_search()
|
|
.num_beams(num_beams)
|
|
.max_length(50);
|
|
|
|
let mut generator = TextGenerator::new(config).unwrap();
|
|
let output = generator.generate(black_box("The future of AI")).unwrap();
|
|
black_box(output);
|
|
});
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_nucleus_sampling(c: &mut Criterion) {
|
|
let rt = tokio::runtime::Runtime::new().unwrap();
|
|
|
|
let mut group = c.benchmark_group("nucleus_sampling");
|
|
group.throughput(Throughput::Elements(1));
|
|
|
|
for top_p in [0.7, 0.8, 0.9, 0.95].iter() {
|
|
group.bench_with_input(
|
|
BenchmarkId::new("top_p", (top_p * 100.0) as u32),
|
|
top_p,
|
|
|b, &top_p| {
|
|
b.to_async(&rt).iter(|| async {
|
|
let config = GenerationConfig::nucleus_sampling()
|
|
.top_p(top_p)
|
|
.max_length(50)
|
|
.seed(42);
|
|
|
|
let mut generator = TextGenerator::new(config).unwrap();
|
|
let output = generator.generate(black_box("Tell me a story")).unwrap();
|
|
black_box(output);
|
|
});
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_topk_sampling(c: &mut Criterion) {
|
|
let rt = tokio::runtime::Runtime::new().unwrap();
|
|
|
|
let mut group = c.benchmark_group("topk_sampling");
|
|
group.throughput(Throughput::Elements(1));
|
|
|
|
for top_k in [10, 25, 50, 100].iter() {
|
|
group.bench_with_input(BenchmarkId::new("top_k", top_k), top_k, |b, &top_k| {
|
|
b.to_async(&rt).iter(|| async {
|
|
let config = GenerationConfig::topk_sampling()
|
|
.top_k(top_k)
|
|
.max_length(50)
|
|
.seed(42);
|
|
|
|
let mut generator = TextGenerator::new(config).unwrap();
|
|
let output = generator.generate(black_box("Describe a sunset")).unwrap();
|
|
black_box(output);
|
|
});
|
|
});
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_batch_generation(c: &mut Criterion) {
|
|
let rt = tokio::runtime::Runtime::new().unwrap();
|
|
|
|
let mut group = c.benchmark_group("batch_generation");
|
|
|
|
for batch_size in [1, 4, 8, 16].iter() {
|
|
group.throughput(Throughput::Elements(*batch_size as u64));
|
|
group.bench_with_input(
|
|
BenchmarkId::new("batch_size", batch_size),
|
|
batch_size,
|
|
|b, &batch_size| {
|
|
b.to_async(&rt).iter(|| async {
|
|
let model = Arc::new(MockModelInterface::new("/tmp/mock").unwrap());
|
|
let config = GenerationConfig::default();
|
|
let generator =
|
|
serving::BatchGenerator::with_model(model, config, batch_size).unwrap();
|
|
|
|
let prompts: Vec<String> = (0..batch_size)
|
|
.map(|i| format!("Prompt number {}", i))
|
|
.collect();
|
|
|
|
let results = generator.generate_batch(prompts).await.unwrap();
|
|
black_box(results);
|
|
});
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_quality_checking(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("quality_checking");
|
|
group.throughput(Throughput::Elements(1));
|
|
|
|
let checker = quality::QualityChecker::new()
|
|
.add_filter(Box::new(quality::RepetitionPenalty::default()))
|
|
.add_filter(Box::new(quality::DiversityPromoter::default()))
|
|
.add_filter(Box::new(quality::ToxicityFilter::default()));
|
|
|
|
let test_texts = [
|
|
"This is a short text.",
|
|
"This is a longer text with more words and varied vocabulary to test the quality checking system performance.",
|
|
"This is an even longer text that contains multiple sentences with different structures. It includes various words and phrases to provide a comprehensive test of the quality checking algorithms. The text should be diverse enough to trigger different quality metrics and filters.",
|
|
];
|
|
|
|
for (i, text) in test_texts.iter().enumerate() {
|
|
group.bench_with_input(
|
|
BenchmarkId::new("text_length", text.split_whitespace().count()),
|
|
text,
|
|
|b, text| {
|
|
b.iter(|| {
|
|
let report = checker.check_quality(black_box(text)).unwrap();
|
|
black_box(report);
|
|
});
|
|
},
|
|
);
|
|
}
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_kv_cache(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("kv_cache");
|
|
group.throughput(Throughput::Elements(1));
|
|
|
|
let mut cache = optimization::KVCache::new(1024).unwrap();
|
|
|
|
// Pre-populate cache
|
|
for i in 0..100 {
|
|
let key = rtx_tensor::Tensor::zeros(
|
|
&[1, 10, 64],
|
|
rtx_tensor::DType::F32,
|
|
&rtx_tensor::Device::cuda(0).unwrap_or(Device::default()),
|
|
)
|
|
.unwrap();
|
|
let value = rtx_tensor::Tensor::zeros(
|
|
&[1, 10, 64],
|
|
rtx_tensor::DType::F32,
|
|
&rtx_tensor::Device::cuda(0).unwrap_or(Device::default()),
|
|
)
|
|
.unwrap();
|
|
cache
|
|
.insert(format!("key_{}", i), vec![key], vec![value])
|
|
.unwrap();
|
|
}
|
|
|
|
group.bench_function("cache_lookup", |b| {
|
|
b.iter(|| {
|
|
let key = format!("key_{}", black_box(42));
|
|
let result = cache.get(&key);
|
|
black_box(result);
|
|
});
|
|
});
|
|
|
|
group.bench_function("cache_insert", |b| {
|
|
b.iter(|| {
|
|
let key = format!("new_key_{}", black_box(rand::random::<u32>()));
|
|
let tensor_key = rtx_tensor::Tensor::zeros(
|
|
&[1, 10, 64],
|
|
rtx_tensor::DType::F32,
|
|
&rtx_tensor::Device::cuda(0).unwrap_or(Device::default()),
|
|
)
|
|
.unwrap();
|
|
let tensor_value = rtx_tensor::Tensor::zeros(
|
|
&[1, 10, 64],
|
|
rtx_tensor::DType::F32,
|
|
&rtx_tensor::Device::cuda(0).unwrap_or(Device::default()),
|
|
)
|
|
.unwrap();
|
|
cache
|
|
.insert(key, vec![tensor_key], vec![tensor_value])
|
|
.unwrap();
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_speculative_decoding(c: &mut Criterion) {
|
|
let rt = tokio::runtime::Runtime::new().unwrap();
|
|
|
|
let mut group = c.benchmark_group("speculative_decoding");
|
|
group.throughput(Throughput::Elements(1));
|
|
|
|
let decoder = optimization::SpeculativeDecoder::new(4).unwrap();
|
|
let target_model = MockModelInterface::new("/tmp/mock").unwrap();
|
|
let input_ids = rtx_tensor::Tensor::from_slice(&[1u32, 2, 3, 4, 5], &[1, 5]).unwrap();
|
|
|
|
group.bench_function("speculative_step", |b| {
|
|
b.iter(|| {
|
|
let result = decoder
|
|
.speculative_step(black_box(&target_model), black_box(&input_ids), None)
|
|
.unwrap();
|
|
black_box(result);
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
criterion_group!(
|
|
generation_benches,
|
|
bench_beam_search,
|
|
bench_nucleus_sampling,
|
|
bench_topk_sampling,
|
|
bench_batch_generation,
|
|
bench_quality_checking,
|
|
bench_kv_cache,
|
|
bench_speculative_decoding
|
|
);
|
|
|
|
criterion_main!(generation_benches);
|