rustytorch/crates/training/rtx-flash-metal-attention/benches/metal_vs_pytorch.rs

//! Metal Flash Attention Benchmark Suite
//!
//! Compares RustyTorch Metal Flash Attention against PyTorch MPS baseline.
//! Run the PyTorch baseline with: python3 scripts/bench_pytorch_mps.py

use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
use rtx_flash_metal_attention::{FlashAttention, FlashAttentionConfig};
use rtx_tensor::{Device, Tensor};
use std::thread;
use std::time::Duration;

/// Create FlashAttention with retry logic for XPC connection issues
fn create_flash_attention_with_retry(config: FlashAttentionConfig) -> Option<FlashAttention> {
    for attempt in 0..3 {
        match FlashAttention::new(config.clone()) {
            Ok(attn) => return Some(attn),
            Err(e) => {
                let err_str = format!("{:?}", e);
                if err_str.contains("XPC") {
                    eprintln!("XPC error on attempt {}, retrying in 500ms...", attempt + 1);
                    thread::sleep(Duration::from_millis(500));
                } else {
                    eprintln!("Failed to create FlashAttention: {:?}", e);
                    return None;
                }
            }
        }
    }
    eprintln!("Failed after 3 retries");
    None
}

fn bench_metal_flash(c: &mut Criterion) {
    // Get CPU device for tensor generation (Metal RNG not implemented yet)
    let cpu_device = Device::cpu();

    // Skip if Metal not available
    let metal_device = match Device::metal(0) {
        Ok(d) => d,
        Err(e) => {
            eprintln!("Metal device not available, skipping benchmark: {:?}", e);
            return;
        }
    };

    // Create FlashAttention instance once (compiles Metal shaders)
    let attn = match create_flash_attention_with_retry(FlashAttentionConfig::default()) {
        Some(a) => a,
        None => {
            eprintln!("Could not create FlashAttention, skipping benchmark");
            return;
        }
    };

    // Test scenarios matching PyTorch benchmark
    let scenarios = vec![
        ("Latency_BS1", 1usize, 128usize, 64usize), // Inference (Robotics/Control)
        ("Throughput_BS64", 64, 128, 64),           // Training (Standard)
        ("Heavy_BS256", 256, 128, 64),              // GPU Saturation
    ];

    let mut group = c.benchmark_group("Metal_Flash_Attention");
    group.measurement_time(Duration::from_secs(10));
    group.warm_up_time(Duration::from_secs(3));

    for (name, batch_size, seq_len, head_dim) in scenarios {
        let num_heads = 8usize;
        let shape = &[batch_size, num_heads, seq_len, head_dim];

        // Generate random data on CPU first
        let q_cpu = match Tensor::randn(shape, &cpu_device) {
            Ok(t) => t,
            Err(e) => {
                eprintln!("Failed to create Q tensor on CPU: {:?}", e);
                continue;
            }
        };
        let k_cpu = match Tensor::randn(shape, &cpu_device) {
            Ok(t) => t,
            Err(e) => {
                eprintln!("Failed to create K tensor on CPU: {:?}", e);
                continue;
            }
        };
        let v_cpu = match Tensor::randn(shape, &cpu_device) {
            Ok(t) => t,
            Err(e) => {
                eprintln!("Failed to create V tensor on CPU: {:?}", e);
                continue;
            }
        };

        // Transfer to Metal device (fast on Apple Silicon unified memory)
        let q = match q_cpu.to_device(&metal_device) {
            Ok(t) => t,
            Err(e) => {
                eprintln!("Failed to transfer Q to Metal: {:?}", e);
                continue;
            }
        };
        let k = match k_cpu.to_device(&metal_device) {
            Ok(t) => t,
            Err(e) => {
                eprintln!("Failed to transfer K to Metal: {:?}", e);
                continue;
            }
        };
        let v = match v_cpu.to_device(&metal_device) {
            Ok(t) => t,
            Err(e) => {
                eprintln!("Failed to transfer V to Metal: {:?}", e);
                continue;
            }
        };

        // Throughput: total elements processed per forward pass
        let total_elements = (batch_size * num_heads * seq_len * head_dim) as u64;
        group.throughput(Throughput::Elements(total_elements));

        group.bench_function(BenchmarkId::new("RustyTorch_Metal", name), |b| {
            b.iter(|| {
                // forward() calls waitUntilCompleted() internally for synchronous timing
                let _ = attn.forward(&q, &k, &v).expect("Forward failed");
            })
        });
    }
    group.finish();
}

fn bench_metal_flash_causal(c: &mut Criterion) {
    // Get CPU device for tensor generation (Metal RNG not implemented yet)
    let cpu_device = Device::cpu();

    // Skip if Metal not available
    let metal_device = match Device::metal(0) {
        Ok(d) => d,
        Err(e) => {
            eprintln!(
                "Metal device not available, skipping causal benchmark: {:?}",
                e
            );
            return;
        }
    };

    // Create causal FlashAttention instance once
    let attn = match create_flash_attention_with_retry(FlashAttentionConfig::causal()) {
        Some(a) => a,
        None => {
            eprintln!("Could not create causal FlashAttention, skipping benchmark");
            return;
        }
    };

    let mut group = c.benchmark_group("Metal_Flash_Attention_Causal");
    group.measurement_time(Duration::from_secs(10));
    group.warm_up_time(Duration::from_secs(3));

    // Causal attention scenarios (autoregressive generation)
    let scenarios = vec![
        ("Causal_BS1", 1usize, 128usize, 64usize),
        ("Causal_BS32", 32, 256, 64),
    ];

    for (name, batch_size, seq_len, head_dim) in scenarios {
        let num_heads = 8usize;
        let shape = &[batch_size, num_heads, seq_len, head_dim];

        // Generate on CPU, transfer to Metal
        let q_cpu = match Tensor::randn(shape, &cpu_device) {
            Ok(t) => t,
            Err(e) => {
                eprintln!("Failed to create Q tensor on CPU: {:?}", e);
                continue;
            }
        };
        let k_cpu = match Tensor::randn(shape, &cpu_device) {
            Ok(t) => t,
            Err(e) => {
                eprintln!("Failed to create K tensor on CPU: {:?}", e);
                continue;
            }
        };
        let v_cpu = match Tensor::randn(shape, &cpu_device) {
            Ok(t) => t,
            Err(e) => {
                eprintln!("Failed to create V tensor on CPU: {:?}", e);
                continue;
            }
        };

        let q = match q_cpu.to_device(&metal_device) {
            Ok(t) => t,
            Err(e) => {
                eprintln!("Failed to transfer Q to Metal: {:?}", e);
                continue;
            }
        };
        let k = match k_cpu.to_device(&metal_device) {
            Ok(t) => t,
            Err(e) => {
                eprintln!("Failed to transfer K to Metal: {:?}", e);
                continue;
            }
        };
        let v = match v_cpu.to_device(&metal_device) {
            Ok(t) => t,
            Err(e) => {
                eprintln!("Failed to transfer V to Metal: {:?}", e);
                continue;
            }
        };

        let total_elements = (batch_size * num_heads * seq_len * head_dim) as u64;
        group.throughput(Throughput::Elements(total_elements));

        group.bench_function(BenchmarkId::new("RustyTorch_Metal_Causal", name), |b| {
            b.iter(|| {
                let _ = attn.forward(&q, &k, &v).expect("Forward failed");
            })
        });
    }
    group.finish();
}

criterion_group!(benches, bench_metal_flash, bench_metal_flash_causal);
criterion_main!(benches);