//! Benchmarks for Metal-accelerated Mixture of Experts //! //! Run with: cargo bench --bench metal_moe_bench --features metal use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main}; use rtx_tensor::{DType, Device, Tensor}; // Note: Import from rtx_transformers when available // use rtx_transformers::layers::{MetalMoE, MetalMoEConfig}; /// Benchmark configuration parameters struct BenchConfig { batch_size: usize, seq_len: usize, hidden_dim: usize, expert_hidden_dim: usize, num_experts: usize, top_k: usize, } impl BenchConfig { fn small() -> Self { Self { batch_size: 1, seq_len: 128, hidden_dim: 256, expert_hidden_dim: 512, num_experts: 4, top_k: 2, } } fn medium() -> Self { Self { batch_size: 4, seq_len: 512, hidden_dim: 768, expert_hidden_dim: 3072, num_experts: 8, top_k: 2, } } fn large() -> Self { Self { batch_size: 8, seq_len: 1024, hidden_dim: 1024, expert_hidden_dim: 4096, num_experts: 16, top_k: 2, } } fn total_tokens(&self) -> usize { self.batch_size * self.seq_len } fn name(&self) -> String { format!( "b{}s{}h{}e{}k{}", self.batch_size, self.seq_len, self.hidden_dim, self.num_experts, self.top_k ) } } /// Benchmark softmax routing (simulates router gating) fn bench_softmax_routing(c: &mut Criterion) { let mut group = c.benchmark_group("moe_softmax_routing"); for config in [ BenchConfig::small(), BenchConfig::medium(), BenchConfig::large(), ] { let device = Device::cpu(); let total_tokens = config.total_tokens(); // Create gate logits let gate_logits = Tensor::randn(&[total_tokens, config.num_experts], DType::F32, &device).unwrap(); group.throughput(Throughput::Elements(total_tokens as u64)); group.bench_with_input( BenchmarkId::new("cpu", config.name()), &gate_logits, |b, logits| { b.iter(|| black_box(logits.softmax(-1).unwrap())); }, ); } group.finish(); } /// Benchmark top-k selection fn bench_topk_selection(c: &mut Criterion) { let mut group = c.benchmark_group("moe_topk_selection"); for config in [BenchConfig::small(), BenchConfig::medium()] { let device = Device::cpu(); let total_tokens = config.total_tokens(); // Create gate probabilities let gate_probs = Tensor::randn(&[total_tokens, config.num_experts], DType::F32, &device) .unwrap() .softmax(-1) .unwrap(); group.throughput(Throughput::Elements(total_tokens as u64)); group.bench_with_input( BenchmarkId::new("cpu", config.name()), &gate_probs, |b, probs| { b.iter(|| black_box(probs.topk(config.top_k, -1, true, true).unwrap())); }, ); } group.finish(); } /// Benchmark matrix multiplication (simulates expert forward) fn bench_expert_matmul(c: &mut Criterion) { let mut group = c.benchmark_group("moe_expert_matmul"); for config in [BenchConfig::small(), BenchConfig::medium()] { let device = Device::cpu(); // Simulate tokens going through one expert let tokens_per_expert = config.total_tokens() / config.num_experts; let input = Tensor::randn(&[tokens_per_expert, config.hidden_dim], DType::F32, &device).unwrap(); // Up projection weight let up_weight = Tensor::randn( &[config.expert_hidden_dim, config.hidden_dim], DType::F32, &device, ) .unwrap(); // Down projection weight let down_weight = Tensor::randn( &[config.hidden_dim, config.expert_hidden_dim], DType::F32, &device, ) .unwrap(); let flops = 2 * tokens_per_expert * config.hidden_dim * config.expert_hidden_dim * 2; group.throughput(Throughput::Elements(flops as u64)); group.bench_with_input( BenchmarkId::new("cpu", config.name()), &(&input, &up_weight, &down_weight), |b, (inp, up, down)| { b.iter(|| { let up_t = up.t().unwrap(); let hidden = inp.matmul(&up_t).unwrap(); let activated = hidden.swish().unwrap(); let down_t = down.t().unwrap(); black_box(activated.matmul(&down_t).unwrap()) }); }, ); } group.finish(); } /// Benchmark full MoE forward pass simulation fn bench_moe_forward_simulation(c: &mut Criterion) { let mut group = c.benchmark_group("moe_forward_simulation"); group.sample_size(20); // Reduce sample size for slower benchmarks for config in [BenchConfig::small()] { let device = Device::cpu(); let total_tokens = config.total_tokens(); // Input tensor let input = Tensor::randn( &[config.batch_size, config.seq_len, config.hidden_dim], DType::F32, &device, ) .unwrap(); // Gate weight let gate_weight = Tensor::randn( &[config.num_experts, config.hidden_dim], DType::F32, &device, ) .unwrap(); // Expert weights (simplified - just one set) let up_weight = Tensor::randn( &[config.expert_hidden_dim, config.hidden_dim], DType::F32, &device, ) .unwrap(); let down_weight = Tensor::randn( &[config.hidden_dim, config.expert_hidden_dim], DType::F32, &device, ) .unwrap(); group.throughput(Throughput::Elements(total_tokens as u64)); group.bench_function(BenchmarkId::new("cpu_simulation", config.name()), |b| { b.iter(|| { // Flatten input let flat = input.view(&[total_tokens, config.hidden_dim]).unwrap(); // Compute gate logits let gate_t = gate_weight.t().unwrap(); let logits = flat.matmul(&gate_t).unwrap(); // Softmax let probs = logits.softmax(-1).unwrap(); // Top-k (simulated with max for simplicity) let _top_vals = probs.max().unwrap(); // Expert forward (simplified - just one pass) let up_t = up_weight.t().unwrap(); let hidden = flat.matmul(&up_t).unwrap(); let activated = hidden.swish().unwrap(); let down_t = down_weight.t().unwrap(); let output = activated.matmul(&down_t).unwrap(); black_box(output) }); }); } group.finish(); } /// Benchmark memory bandwidth (critical for MoE dispatch) fn bench_memory_bandwidth(c: &mut Criterion) { let mut group = c.benchmark_group("moe_memory_bandwidth"); for size in [1024, 4096, 16384, 65536] { let device = Device::cpu(); let src = Tensor::randn(&[size, 768], DType::F32, &device).unwrap(); let bytes = size * 768 * 4; // f32 = 4 bytes group.throughput(Throughput::Bytes(bytes as u64)); group.bench_with_input(BenchmarkId::new("clone", size), &src, |b, tensor| { b.iter(|| black_box(tensor.to_cpu().unwrap())); }); } group.finish(); } criterion_group!( benches, bench_softmax_routing, bench_topk_selection, bench_expert_matmul, bench_moe_forward_simulation, bench_memory_bandwidth, ); criterion_main!(benches);