215 lines
7.1 KiB
Rust
215 lines
7.1 KiB
Rust
//! Flash Attention Benchmark - Metal vs CPU with Causal Masking
|
|
//! Run with: cargo run -p rtx-flash-metal-attention --example simple_bench --release
|
|
|
|
use rtx_flash_metal_attention::{FlashAttention, FlashAttentionConfig};
|
|
use rtx_tensor::{Device, Tensor};
|
|
use std::time::Instant;
|
|
|
|
/// Naive CPU attention implementation for comparison
|
|
/// O = softmax(Q @ K^T / sqrt(d)) @ V
|
|
fn cpu_naive_attention(
|
|
batch_size: usize,
|
|
num_heads: usize,
|
|
seq_len: usize,
|
|
head_dim: usize,
|
|
causal: bool,
|
|
) -> f64 {
|
|
let scale = 1.0 / (head_dim as f32).sqrt();
|
|
|
|
// Simulate Q, K, V as all ones (same as Metal benchmark)
|
|
let q = vec![1.0f32; batch_size * num_heads * seq_len * head_dim];
|
|
let k = vec![1.0f32; batch_size * num_heads * seq_len * head_dim];
|
|
let v = vec![1.0f32; batch_size * num_heads * seq_len * head_dim];
|
|
let mut output = vec![0.0f32; batch_size * num_heads * seq_len * head_dim];
|
|
|
|
// For each batch and head
|
|
for b in 0..batch_size {
|
|
for h in 0..num_heads {
|
|
let base = (b * num_heads + h) * seq_len * head_dim;
|
|
|
|
// For each query position
|
|
for i in 0..seq_len {
|
|
// Compute attention scores for row i
|
|
let mut scores = vec![f32::NEG_INFINITY; seq_len];
|
|
let mut max_score = f32::NEG_INFINITY;
|
|
|
|
// Apply causal mask: only attend to positions <= i
|
|
let end_j = if causal { i + 1 } else { seq_len };
|
|
|
|
for j in 0..end_j {
|
|
let mut dot = 0.0f32;
|
|
for d in 0..head_dim {
|
|
dot += q[base + i * head_dim + d] * k[base + j * head_dim + d];
|
|
}
|
|
scores[j] = dot * scale;
|
|
max_score = max_score.max(scores[j]);
|
|
}
|
|
|
|
// Softmax (only over valid positions)
|
|
let mut sum_exp = 0.0f32;
|
|
for j in 0..end_j {
|
|
scores[j] = (scores[j] - max_score).exp();
|
|
sum_exp += scores[j];
|
|
}
|
|
for j in 0..end_j {
|
|
scores[j] /= sum_exp;
|
|
}
|
|
|
|
// Weighted sum of V
|
|
for d in 0..head_dim {
|
|
let mut acc = 0.0f32;
|
|
for j in 0..end_j {
|
|
acc += scores[j] * v[base + j * head_dim + d];
|
|
}
|
|
output[base + i * head_dim + d] = acc;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Return sum to prevent optimization
|
|
output.iter().map(|x| *x as f64).sum()
|
|
}
|
|
|
|
fn bench_attention(
|
|
metal_device: &Device,
|
|
config: FlashAttentionConfig,
|
|
scenarios: &[(&str, usize, usize, usize)],
|
|
causal: bool,
|
|
) -> Vec<f64> {
|
|
let mode = if causal { "Causal" } else { "Full" };
|
|
|
|
let attn = match FlashAttention::new(config) {
|
|
Ok(a) => a,
|
|
Err(e) => {
|
|
eprintln!("Failed to create FlashAttention: {:?}", e);
|
|
return vec![];
|
|
}
|
|
};
|
|
|
|
println!("\n{:=^80}", format!(" Metal Flash Attention ({}) ", mode));
|
|
println!(
|
|
"{:<20} {:>15} {:>15} {:>15}",
|
|
"Scenario", "Avg Time (ms)", "Throughput", "Elements"
|
|
);
|
|
println!("{:-<80}", "");
|
|
|
|
let mut metal_times = Vec::new();
|
|
|
|
for (name, batch_size, seq_len, head_dim) in scenarios {
|
|
let num_heads = 8;
|
|
let shape = &[*batch_size, num_heads, *seq_len, *head_dim];
|
|
|
|
let q = Tensor::ones(shape, metal_device).expect("Q tensor");
|
|
let k = Tensor::ones(shape, metal_device).expect("K tensor");
|
|
let v = Tensor::ones(shape, metal_device).expect("V tensor");
|
|
|
|
// Warmup
|
|
for _ in 0..10 {
|
|
let _ = attn.forward(&q, &k, &v).expect("warmup forward");
|
|
}
|
|
|
|
// Benchmark
|
|
let iterations = 100;
|
|
let start = Instant::now();
|
|
|
|
for _ in 0..iterations {
|
|
let _ = attn.forward(&q, &k, &v).expect("forward");
|
|
}
|
|
|
|
let elapsed = start.elapsed();
|
|
let avg_ms = elapsed.as_secs_f64() * 1000.0 / iterations as f64;
|
|
let total_elements = (batch_size * num_heads * seq_len * head_dim) as f64;
|
|
let throughput = total_elements / (avg_ms / 1000.0) / 1e6;
|
|
|
|
metal_times.push(avg_ms);
|
|
println!(
|
|
"{:<20} {:>12.4} ms {:>12.2} M/s {:>12.0}",
|
|
name, avg_ms, throughput, total_elements
|
|
);
|
|
}
|
|
|
|
println!("\n{:=^80}", format!(" CPU Naive Attention ({}) ", mode));
|
|
println!(
|
|
"{:<20} {:>15} {:>15} {:>15}",
|
|
"Scenario", "Avg Time (ms)", "Throughput", "Speedup"
|
|
);
|
|
println!("{:-<80}", "");
|
|
|
|
for (idx, (name, batch_size, seq_len, head_dim)) in scenarios.iter().enumerate() {
|
|
let num_heads = 8;
|
|
let total_elements = (batch_size * num_heads * seq_len * head_dim) as f64;
|
|
|
|
// Warmup
|
|
for _ in 0..2 {
|
|
let _ = cpu_naive_attention(*batch_size, num_heads, *seq_len, *head_dim, causal);
|
|
}
|
|
|
|
// Fewer iterations for CPU (it's slower)
|
|
let iterations = 10;
|
|
let start = Instant::now();
|
|
|
|
for _ in 0..iterations {
|
|
let _ = cpu_naive_attention(*batch_size, num_heads, *seq_len, *head_dim, causal);
|
|
}
|
|
|
|
let elapsed = start.elapsed();
|
|
let avg_ms = elapsed.as_secs_f64() * 1000.0 / iterations as f64;
|
|
let throughput = total_elements / (avg_ms / 1000.0) / 1e6;
|
|
let speedup = avg_ms / metal_times[idx];
|
|
|
|
println!(
|
|
"{:<20} {:>12.4} ms {:>12.2} M/s {:>12.1}x",
|
|
name, avg_ms, throughput, speedup
|
|
);
|
|
}
|
|
|
|
metal_times
|
|
}
|
|
|
|
fn main() {
|
|
println!("=== Flash Attention: Metal vs CPU Benchmark ===\n");
|
|
|
|
// Get Metal device
|
|
let metal_device = match Device::metal(0) {
|
|
Ok(d) => {
|
|
println!("Metal device: Metal(0)");
|
|
d
|
|
}
|
|
Err(e) => {
|
|
eprintln!("Metal device not available: {:?}", e);
|
|
return;
|
|
}
|
|
};
|
|
|
|
// Test scenarios (head_dim must be <= max_head_dim which is 64)
|
|
let scenarios = vec![
|
|
("Latency_BS1", 1usize, 128usize, 64usize),
|
|
("Throughput_BS32", 32, 128, 64),
|
|
("Heavy_BS128", 128, 128, 64),
|
|
];
|
|
|
|
// Benchmark full attention (non-causal)
|
|
println!("\n{:#^80}", " FULL ATTENTION (Non-Causal) ");
|
|
let config_full = FlashAttentionConfig::default();
|
|
let _full_times = bench_attention(&metal_device, config_full, &scenarios, false);
|
|
|
|
// Benchmark causal attention
|
|
println!("\n\n{:#^80}", " CAUSAL ATTENTION (Autoregressive) ");
|
|
let config_causal = FlashAttentionConfig::causal();
|
|
let _causal_times = bench_attention(&metal_device, config_causal, &scenarios, true);
|
|
|
|
// Benchmark with larger head dimension (128)
|
|
println!("\n\n{:#^80}", " LARGE HEAD DIM (128) ");
|
|
let large_scenarios = vec![
|
|
("BS1_HD128", 1usize, 128usize, 128usize),
|
|
("BS32_HD128", 32, 128, 128),
|
|
("BS64_HD128", 64, 128, 128),
|
|
];
|
|
let config_large = FlashAttentionConfig::large_head_dim();
|
|
let _large_times = bench_attention(&metal_device, config_large, &large_scenarios, false);
|
|
|
|
println!("\n{:-<80}", "");
|
|
println!("\nDone!");
|
|
}
|