99 lines
2.9 KiB
Rust
99 lines
2.9 KiB
Rust
use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
|
|
use ruvector_sona::{SonaConfig, SonaEngine};
|
|
|
|
fn trajectory_benchmark(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("trajectory");
|
|
|
|
for dim in [64, 128, 256, 512].iter() {
|
|
let engine = SonaEngine::with_config(SonaConfig {
|
|
hidden_dim: *dim,
|
|
embedding_dim: *dim,
|
|
..Default::default()
|
|
});
|
|
|
|
group.bench_with_input(BenchmarkId::new("single", dim), dim, |b, &dim| {
|
|
b.iter(|| {
|
|
let mut builder = engine.begin_trajectory(vec![0.1; dim]);
|
|
builder.add_step(vec![0.5; dim], vec![], 0.8);
|
|
builder.add_step(vec![0.6; dim], vec![], 0.9);
|
|
engine.end_trajectory(builder, black_box(0.85));
|
|
});
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn lora_application_benchmark(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("lora");
|
|
|
|
for dim in [64, 128, 256, 512].iter() {
|
|
let engine = SonaEngine::with_config(SonaConfig {
|
|
hidden_dim: *dim,
|
|
embedding_dim: *dim,
|
|
..Default::default()
|
|
});
|
|
|
|
// Warmup with some trajectories
|
|
for _ in 0..10 {
|
|
let mut builder = engine.begin_trajectory(vec![0.1; *dim]);
|
|
builder.add_step(vec![0.5; *dim], vec![], 0.8);
|
|
engine.end_trajectory(builder, 0.85);
|
|
}
|
|
engine.flush();
|
|
|
|
group.bench_with_input(BenchmarkId::new("micro", dim), dim, |b, &dim| {
|
|
let input = vec![1.0; dim];
|
|
let mut output = vec![0.0; dim];
|
|
b.iter(|| {
|
|
engine.apply_micro_lora(black_box(&input), black_box(&mut output));
|
|
});
|
|
});
|
|
|
|
group.bench_with_input(BenchmarkId::new("base", dim), dim, |b, &dim| {
|
|
let input = vec![1.0; dim];
|
|
let mut output = vec![0.0; dim];
|
|
b.iter(|| {
|
|
engine.apply_base_lora(0, black_box(&input), black_box(&mut output));
|
|
});
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn background_learning_benchmark(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("learning");
|
|
group.sample_size(10); // Fewer samples for expensive operation
|
|
|
|
let engine = SonaEngine::with_config(SonaConfig {
|
|
hidden_dim: 256,
|
|
embedding_dim: 256,
|
|
..Default::default()
|
|
});
|
|
|
|
// Prepare 100 trajectories
|
|
for _ in 0..100 {
|
|
let mut builder = engine.begin_trajectory(vec![0.1; 256]);
|
|
builder.add_step(vec![0.5; 256], vec![], 0.8);
|
|
builder.add_step(vec![0.6; 256], vec![], 0.9);
|
|
engine.end_trajectory(builder, 0.85);
|
|
}
|
|
|
|
group.bench_function("force_learn", |b| {
|
|
b.iter(|| {
|
|
black_box(engine.force_learn());
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
criterion_group!(
|
|
benches,
|
|
trajectory_benchmark,
|
|
lora_application_benchmark,
|
|
background_learning_benchmark
|
|
);
|
|
criterion_main!(benches);
|