Initial commit

2026-03-04 00:08:42 +00:00
commit 4d88dc0584
4449 changed files with 1556714 additions and 0 deletions
@@ -0,0 +1,227 @@
+//! 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);