462 lines
15 KiB
Rust
462 lines
15 KiB
Rust
//! Benchmarks for Metal-accelerated Mamba (State Space Model)
|
|
//!
|
|
//! Run with: cargo bench --bench metal_mamba_bench --features metal
|
|
|
|
use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main};
|
|
use rtx_tensor::{DType, Device, Tensor};
|
|
|
|
/// Benchmark configuration parameters for Mamba
|
|
struct MambaBenchConfig {
|
|
batch_size: usize,
|
|
seq_len: usize,
|
|
d_model: usize,
|
|
d_state: usize,
|
|
d_conv: usize,
|
|
expand: usize,
|
|
}
|
|
|
|
impl MambaBenchConfig {
|
|
fn small() -> Self {
|
|
Self {
|
|
batch_size: 1,
|
|
seq_len: 128,
|
|
d_model: 256,
|
|
d_state: 8,
|
|
d_conv: 4,
|
|
expand: 2,
|
|
}
|
|
}
|
|
|
|
fn medium() -> Self {
|
|
Self {
|
|
batch_size: 4,
|
|
seq_len: 512,
|
|
d_model: 768,
|
|
d_state: 16,
|
|
d_conv: 4,
|
|
expand: 2,
|
|
}
|
|
}
|
|
|
|
fn large() -> Self {
|
|
Self {
|
|
batch_size: 8,
|
|
seq_len: 2048,
|
|
d_model: 1024,
|
|
d_state: 16,
|
|
d_conv: 4,
|
|
expand: 2,
|
|
}
|
|
}
|
|
|
|
fn d_inner(&self) -> usize {
|
|
self.d_model * self.expand
|
|
}
|
|
|
|
fn total_tokens(&self) -> usize {
|
|
self.batch_size * self.seq_len
|
|
}
|
|
|
|
fn name(&self) -> String {
|
|
format!(
|
|
"b{}s{}d{}n{}",
|
|
self.batch_size, self.seq_len, self.d_model, self.d_state
|
|
)
|
|
}
|
|
}
|
|
|
|
/// Benchmark input projection (d_model -> 2 * d_inner)
|
|
fn bench_input_projection(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("mamba_input_projection");
|
|
|
|
for config in [MambaBenchConfig::small(), MambaBenchConfig::medium()] {
|
|
let device = Device::cpu();
|
|
let d_inner = config.d_inner();
|
|
let total_tokens = config.total_tokens();
|
|
|
|
let input = Tensor::randn(&[total_tokens, config.d_model], DType::F32, &device).unwrap();
|
|
|
|
let in_proj = Tensor::randn(&[2 * d_inner, config.d_model], DType::F32, &device).unwrap();
|
|
|
|
let flops = 2 * total_tokens * config.d_model * 2 * d_inner;
|
|
group.throughput(Throughput::Elements(flops as u64));
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("cpu", config.name()),
|
|
&(&input, &in_proj),
|
|
|b, (inp, proj)| {
|
|
b.iter(|| {
|
|
let proj_t = proj.t().unwrap();
|
|
black_box(inp.matmul(&proj_t).unwrap())
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark causal 1D convolution (critical for Mamba)
|
|
fn bench_causal_conv1d(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("mamba_causal_conv1d");
|
|
group.sample_size(30);
|
|
|
|
for config in [MambaBenchConfig::small(), MambaBenchConfig::medium()] {
|
|
let device = Device::cpu();
|
|
let d_inner = config.d_inner();
|
|
|
|
// Input: [batch, seq, d_inner]
|
|
let input = Tensor::randn(
|
|
&[config.batch_size, config.seq_len, d_inner],
|
|
DType::F32,
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
// Convolution weight: [d_inner, 1, kernel_size]
|
|
let conv_weight = Tensor::randn(&[d_inner, 1, config.d_conv], DType::F32, &device).unwrap();
|
|
|
|
// FLOPs: batch * seq * d_inner * kernel_size
|
|
let flops = config.batch_size * config.seq_len * d_inner * config.d_conv;
|
|
group.throughput(Throughput::Elements(flops as u64));
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("cpu_naive", config.name()),
|
|
&(&input, &conv_weight),
|
|
|b, (inp, weight)| {
|
|
b.iter(|| {
|
|
// Naive implementation for benchmarking
|
|
let inp_data = inp.to_cpu().unwrap();
|
|
let weight_data = weight.to_cpu().unwrap();
|
|
let mut output = vec![0.0f32; config.batch_size * config.seq_len * d_inner];
|
|
|
|
for batch in 0..config.batch_size {
|
|
for t in 0..config.seq_len {
|
|
for d in 0..d_inner {
|
|
let mut sum = 0.0f32;
|
|
for k in 0..config.d_conv {
|
|
let t_src = t as i64 - k as i64;
|
|
if t_src >= 0 {
|
|
let idx = batch * config.seq_len * d_inner
|
|
+ t_src as usize * d_inner
|
|
+ d;
|
|
let w_idx = d * config.d_conv + (config.d_conv - 1 - k);
|
|
sum += inp_data[idx] * weight_data[w_idx];
|
|
}
|
|
}
|
|
output[batch * config.seq_len * d_inner + t * d_inner + d] = sum;
|
|
}
|
|
}
|
|
}
|
|
black_box(output)
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark selective scan (the core SSM operation)
|
|
fn bench_selective_scan(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("mamba_selective_scan");
|
|
group.sample_size(20);
|
|
|
|
for config in [MambaBenchConfig::small()] {
|
|
let device = Device::cpu();
|
|
let d_inner = config.d_inner();
|
|
|
|
// Input u: [batch, seq, d_inner]
|
|
let u = Tensor::randn(
|
|
&[config.batch_size, config.seq_len, d_inner],
|
|
DType::F32,
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
// Delta (time step): [batch, seq, d_inner]
|
|
let delta = Tensor::randn(
|
|
&[config.batch_size, config.seq_len, d_inner],
|
|
DType::F32,
|
|
&device,
|
|
)
|
|
.unwrap()
|
|
.abs()
|
|
.unwrap();
|
|
|
|
// A: [d_inner, d_state]
|
|
let A = Tensor::randn(&[d_inner, config.d_state], DType::F32, &device)
|
|
.unwrap()
|
|
.neg()
|
|
.unwrap()
|
|
.exp()
|
|
.unwrap();
|
|
|
|
// B: [batch, seq, d_state]
|
|
let B = Tensor::randn(
|
|
&[config.batch_size, config.seq_len, config.d_state],
|
|
DType::F32,
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
// C: [batch, seq, d_state]
|
|
let C = Tensor::randn(
|
|
&[config.batch_size, config.seq_len, config.d_state],
|
|
DType::F32,
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
// Operations: batch * seq * d_inner * d_state * (state update + output)
|
|
let ops = config.batch_size * config.seq_len * d_inner * config.d_state * 4;
|
|
group.throughput(Throughput::Elements(ops as u64));
|
|
|
|
group.bench_function(BenchmarkId::new("cpu_sequential", config.name()), |b| {
|
|
b.iter(|| {
|
|
let u_data = u.to_cpu().unwrap();
|
|
let delta_data = delta.to_cpu().unwrap();
|
|
let A_data = A.to_cpu().unwrap();
|
|
let B_data = B.to_cpu().unwrap();
|
|
let C_data = C.to_cpu().unwrap();
|
|
|
|
let mut output = vec![0.0f32; config.batch_size * config.seq_len * d_inner];
|
|
let mut state = vec![0.0f32; config.batch_size * d_inner * config.d_state];
|
|
|
|
for batch in 0..config.batch_size {
|
|
for t in 0..config.seq_len {
|
|
for d in 0..d_inner {
|
|
let u_val = u_data[batch * config.seq_len * d_inner + t * d_inner + d];
|
|
let delta_val =
|
|
delta_data[batch * config.seq_len * d_inner + t * d_inner + d];
|
|
|
|
let mut y = 0.0f32;
|
|
for n in 0..config.d_state {
|
|
let state_idx =
|
|
batch * d_inner * config.d_state + d * config.d_state + n;
|
|
let A_val = A_data[d * config.d_state + n];
|
|
let B_val = B_data[batch * config.seq_len * config.d_state
|
|
+ t * config.d_state
|
|
+ n];
|
|
let C_val = C_data[batch * config.seq_len * config.d_state
|
|
+ t * config.d_state
|
|
+ n];
|
|
|
|
let deltaA = (delta_val * A_val).exp();
|
|
let deltaB = delta_val * B_val;
|
|
|
|
state[state_idx] = deltaA * state[state_idx] + deltaB * u_val;
|
|
y += C_val * state[state_idx];
|
|
}
|
|
|
|
output[batch * config.seq_len * d_inner + t * d_inner + d] = y;
|
|
}
|
|
}
|
|
}
|
|
|
|
black_box(output)
|
|
});
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark SiLU activation (used throughout Mamba)
|
|
fn bench_silu_activation(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("mamba_silu");
|
|
|
|
for size in [65536, 262144, 1048576] {
|
|
let device = Device::cpu();
|
|
let tensor = Tensor::randn(&[size], DType::F32, &device).unwrap();
|
|
|
|
group.throughput(Throughput::Elements(size as u64));
|
|
group.bench_with_input(BenchmarkId::new("cpu", size), &tensor, |b, t| {
|
|
b.iter(|| black_box(t.swish().unwrap()));
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark output projection (d_inner -> d_model)
|
|
fn bench_output_projection(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("mamba_output_projection");
|
|
|
|
for config in [MambaBenchConfig::small(), MambaBenchConfig::medium()] {
|
|
let device = Device::cpu();
|
|
let d_inner = config.d_inner();
|
|
let total_tokens = config.total_tokens();
|
|
|
|
let gated = Tensor::randn(&[total_tokens, d_inner], DType::F32, &device).unwrap();
|
|
|
|
let out_proj = Tensor::randn(&[config.d_model, d_inner], DType::F32, &device).unwrap();
|
|
|
|
let flops = 2 * total_tokens * d_inner * config.d_model;
|
|
group.throughput(Throughput::Elements(flops as u64));
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("cpu", config.name()),
|
|
&(&gated, &out_proj),
|
|
|b, (g, proj)| {
|
|
b.iter(|| {
|
|
let proj_t = proj.t().unwrap();
|
|
black_box(g.matmul(&proj_t).unwrap())
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark state initialization
|
|
fn bench_state_init(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("mamba_state_init");
|
|
|
|
for config in [
|
|
MambaBenchConfig::small(),
|
|
MambaBenchConfig::medium(),
|
|
MambaBenchConfig::large(),
|
|
] {
|
|
let device = Device::cpu();
|
|
let d_inner = config.d_inner();
|
|
let state_size = config.batch_size * d_inner * config.d_state;
|
|
|
|
group.throughput(Throughput::Elements(state_size as u64));
|
|
group.bench_function(BenchmarkId::new("zeros", config.name()), |b| {
|
|
b.iter(|| {
|
|
black_box(
|
|
Tensor::zeros(&[config.batch_size, d_inner, config.d_state], &device).unwrap(),
|
|
)
|
|
});
|
|
});
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Benchmark full Mamba forward simulation
|
|
fn bench_mamba_forward_simulation(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("mamba_forward_simulation");
|
|
group.sample_size(10);
|
|
|
|
let config = MambaBenchConfig::small();
|
|
let device = Device::cpu();
|
|
let d_inner = config.d_inner();
|
|
let total_tokens = config.total_tokens();
|
|
|
|
// Input
|
|
let input = Tensor::randn(
|
|
&[config.batch_size, config.seq_len, config.d_model],
|
|
DType::F32,
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
// Weights
|
|
let in_proj = Tensor::randn(&[2 * d_inner, config.d_model], DType::F32, &device).unwrap();
|
|
let conv_weight = Tensor::randn(&[d_inner, 1, config.d_conv], DType::F32, &device).unwrap();
|
|
let out_proj = Tensor::randn(&[config.d_model, d_inner], DType::F32, &device).unwrap();
|
|
|
|
group.throughput(Throughput::Elements(total_tokens as u64));
|
|
|
|
group.bench_function(BenchmarkId::new("cpu_simulation", config.name()), |b| {
|
|
b.iter(|| {
|
|
// Input projection
|
|
let flat = input.view(&[total_tokens, config.d_model]).unwrap();
|
|
let in_proj_t = in_proj.t().unwrap();
|
|
let projected = flat.matmul(&in_proj_t).unwrap();
|
|
|
|
// Split x and z
|
|
let x = projected.narrow(1, 0, d_inner).unwrap();
|
|
let z = projected.narrow(1, d_inner, d_inner).unwrap();
|
|
|
|
// Apply SiLU to both
|
|
let x_act = x.swish().unwrap();
|
|
let z_act = z.swish().unwrap();
|
|
|
|
// Gating
|
|
let gated = x_act.mul(&z_act).unwrap();
|
|
|
|
// Output projection
|
|
let out_proj_t = out_proj.t().unwrap();
|
|
let output = gated.matmul(&out_proj_t).unwrap();
|
|
|
|
// Reshape
|
|
let output = output
|
|
.view(&[config.batch_size, config.seq_len, config.d_model])
|
|
.unwrap();
|
|
|
|
// Residual
|
|
black_box(output.add(&input).unwrap())
|
|
});
|
|
});
|
|
|
|
group.finish();
|
|
}
|
|
|
|
/// Compare Mamba vs Attention complexity scaling
|
|
fn bench_complexity_scaling(c: &mut Criterion) {
|
|
let mut group = c.benchmark_group("mamba_vs_attention_scaling");
|
|
group.sample_size(10);
|
|
|
|
let d_model = 256;
|
|
let device = Device::cpu();
|
|
|
|
for seq_len in [128, 256, 512] {
|
|
let batch_size = 2;
|
|
let total_tokens = batch_size * seq_len;
|
|
|
|
// Mamba-like: O(n) - just two projections and element-wise ops
|
|
let input = Tensor::randn(&[total_tokens, d_model], DType::F32, &device).unwrap();
|
|
let weight = Tensor::randn(&[d_model, d_model], DType::F32, &device).unwrap();
|
|
|
|
group.throughput(Throughput::Elements(seq_len as u64));
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("mamba_linear", seq_len),
|
|
&(&input, &weight),
|
|
|b, (inp, w)| {
|
|
b.iter(|| {
|
|
let w_t = w.t().unwrap();
|
|
let proj = inp.matmul(&w_t).unwrap();
|
|
black_box(proj.swish().unwrap())
|
|
});
|
|
},
|
|
);
|
|
|
|
// Attention-like: O(n^2) - Q * K^T
|
|
let q = Tensor::randn(&[batch_size, seq_len, d_model], DType::F32, &device).unwrap();
|
|
let k = Tensor::randn(&[batch_size, seq_len, d_model], DType::F32, &device).unwrap();
|
|
|
|
group.bench_with_input(
|
|
BenchmarkId::new("attention_quadratic", seq_len),
|
|
&(&q, &k),
|
|
|b, (q_tensor, k_tensor)| {
|
|
b.iter(|| {
|
|
// Q * K^T: [batch, seq, d] * [batch, d, seq] -> [batch, seq, seq]
|
|
let k_t = k_tensor.transpose(1, 2).unwrap();
|
|
black_box(q_tensor.matmul(&k_t).unwrap())
|
|
});
|
|
},
|
|
);
|
|
}
|
|
|
|
group.finish();
|
|
}
|
|
|
|
criterion_group!(
|
|
benches,
|
|
bench_input_projection,
|
|
bench_causal_conv1d,
|
|
bench_selective_scan,
|
|
bench_silu_activation,
|
|
bench_output_projection,
|
|
bench_state_init,
|
|
bench_mamba_forward_simulation,
|
|
bench_complexity_scaling,
|
|
);
|
|
|
|
criterion_main!(benches);
|