rustytorch/crates/training/rtx-transformers/benches/metal_moe_bench.rs

//! 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);