527 lines
17 KiB
Rust
527 lines
17 KiB
Rust
//! Comprehensive performance validation tests for Flash Attention
|
|
//!
|
|
//! Validates the 5-8x speedup claims over Flash Attention 2 and other baselines
|
|
//!
|
|
//! NOTE: Disabled until Flash Attention API is fully implemented
|
|
|
|
#![cfg(all(feature = "cuda", feature = "disabled_tests"))]
|
|
|
|
use rtx_flash_attention::{
|
|
FlashAttention, FlashAttentionConfig, FlashAttentionFactory,
|
|
error::{FlashError, FlashResult},
|
|
utils::{compare_memory_usage, naive_attention},
|
|
};
|
|
use rtx_tensor::{DType, Device, Tensor};
|
|
use std::time::{Duration, Instant};
|
|
use tokio::test;
|
|
|
|
/// Performance test configuration
|
|
struct PerfTestConfig {
|
|
name: String,
|
|
batch_size: usize,
|
|
num_heads: usize,
|
|
seq_len: usize,
|
|
head_dim: usize,
|
|
causal: bool,
|
|
num_warmup_runs: usize,
|
|
num_benchmark_runs: usize,
|
|
}
|
|
|
|
impl PerfTestConfig {
|
|
fn new(
|
|
name: &str,
|
|
batch_size: usize,
|
|
num_heads: usize,
|
|
seq_len: usize,
|
|
head_dim: usize,
|
|
) -> Self {
|
|
Self {
|
|
name: name.to_string(),
|
|
batch_size,
|
|
num_heads,
|
|
seq_len,
|
|
head_dim,
|
|
causal: false,
|
|
num_warmup_runs: 3,
|
|
num_benchmark_runs: 10,
|
|
}
|
|
}
|
|
|
|
fn with_causal(mut self, causal: bool) -> Self {
|
|
self.causal = causal;
|
|
self
|
|
}
|
|
}
|
|
|
|
/// Benchmark results for comparison
|
|
#[derive(Debug, Clone)]
|
|
struct BenchmarkResult {
|
|
name: String,
|
|
avg_time_us: f64,
|
|
min_time_us: u64,
|
|
max_time_us: u64,
|
|
std_dev_us: f64,
|
|
memory_usage_bytes: usize,
|
|
throughput_gb_s: f64,
|
|
}
|
|
|
|
impl BenchmarkResult {
|
|
fn new(name: String, times_us: Vec<u64>, memory_usage: usize) -> Self {
|
|
let avg_time_us = times_us.iter().sum::<u64>() as f64 / times_us.len() as f64;
|
|
let min_time_us = *times_us.iter().min().unwrap();
|
|
let max_time_us = *times_us.iter().max().unwrap();
|
|
|
|
let variance = times_us
|
|
.iter()
|
|
.map(|&t| (t as f64 - avg_time_us).powi(2))
|
|
.sum::<f64>()
|
|
/ times_us.len() as f64;
|
|
let std_dev_us = variance.sqrt();
|
|
|
|
// Calculate throughput (simplified)
|
|
let avg_time_s = avg_time_us / 1_000_000.0;
|
|
let throughput_gb_s = (memory_usage as f64) / (avg_time_s * 1_000_000_000.0);
|
|
|
|
Self {
|
|
name,
|
|
avg_time_us,
|
|
min_time_us,
|
|
max_time_us,
|
|
std_dev_us,
|
|
memory_usage_bytes: memory_usage,
|
|
throughput_gb_s,
|
|
}
|
|
}
|
|
|
|
fn speedup_vs(&self, baseline: &BenchmarkResult) -> f64 {
|
|
baseline.avg_time_us / self.avg_time_us
|
|
}
|
|
}
|
|
|
|
/// Generate test tensors for benchmarking
|
|
fn create_test_tensors(config: &PerfTestConfig) -> FlashResult<(Tensor, Tensor, Tensor)> {
|
|
let device = Device::Cuda(0);
|
|
let shape = [
|
|
config.batch_size,
|
|
config.num_heads,
|
|
config.seq_len,
|
|
config.head_dim,
|
|
];
|
|
|
|
let q = Tensor::randn(&shape, DType::F16, &device)
|
|
.map_err(|e| FlashError::tensor(format!("Failed to create Q tensor: {}", e)))?;
|
|
let k = Tensor::randn(&shape, DType::F16, &device)
|
|
.map_err(|e| FlashError::tensor(format!("Failed to create K tensor: {}", e)))?;
|
|
let v = Tensor::randn(&shape, DType::F16, &device)
|
|
.map_err(|e| FlashError::tensor(format!("Failed to create V tensor: {}", e)))?;
|
|
|
|
Ok((q, k, v))
|
|
}
|
|
|
|
/// Benchmark Flash Attention implementation
|
|
async fn benchmark_flash_attention(
|
|
config: &PerfTestConfig,
|
|
q: &Tensor,
|
|
k: &Tensor,
|
|
v: &Tensor,
|
|
) -> FlashResult<BenchmarkResult> {
|
|
let flash_config = FlashAttentionConfig::for_inference(config.num_heads, config.head_dim);
|
|
let flash = FlashAttention::new(flash_config)?;
|
|
|
|
let softmax_scale = 1.0 / (config.head_dim as f32).sqrt();
|
|
|
|
// Warmup runs
|
|
for _ in 0..config.num_warmup_runs {
|
|
let _ = flash.forward(q, k, v, config.causal, softmax_scale).await?;
|
|
}
|
|
|
|
// Benchmark runs
|
|
let mut times = Vec::new();
|
|
for _ in 0..config.num_benchmark_runs {
|
|
let start = Instant::now();
|
|
let result = flash.forward(q, k, v, config.causal, softmax_scale).await?;
|
|
let elapsed = start.elapsed();
|
|
times.push(elapsed.as_micros() as u64);
|
|
|
|
// Use the result to prevent optimization
|
|
std::hint::black_box(&result);
|
|
}
|
|
|
|
let memory_usage = flash
|
|
.config()
|
|
.estimate_memory_usage(config.batch_size, config.seq_len);
|
|
|
|
Ok(BenchmarkResult::new(
|
|
format!("Flash Attention ({})", config.name),
|
|
times,
|
|
memory_usage,
|
|
))
|
|
}
|
|
|
|
/// Benchmark naive attention for comparison
|
|
async fn benchmark_naive_attention(
|
|
config: &PerfTestConfig,
|
|
q: &Tensor,
|
|
k: &Tensor,
|
|
v: &Tensor,
|
|
) -> FlashResult<BenchmarkResult> {
|
|
let softmax_scale = 1.0 / (config.head_dim as f32).sqrt();
|
|
|
|
// Warmup runs
|
|
for _ in 0..config.num_warmup_runs {
|
|
let _ = naive_attention(q, k, v, config.causal, softmax_scale)?;
|
|
}
|
|
|
|
// Benchmark runs
|
|
let mut times = Vec::new();
|
|
for _ in 0..config.num_benchmark_runs {
|
|
let start = Instant::now();
|
|
let result = naive_attention(q, k, v, config.causal, softmax_scale)?;
|
|
let elapsed = start.elapsed();
|
|
times.push(elapsed.as_micros() as u64);
|
|
|
|
// Use the result to prevent optimization
|
|
std::hint::black_box(&result);
|
|
}
|
|
|
|
// Calculate O(n²) memory usage for standard attention
|
|
let (standard_memory, _) = compare_memory_usage(
|
|
config.batch_size,
|
|
config.num_heads,
|
|
config.seq_len,
|
|
config.head_dim,
|
|
);
|
|
|
|
Ok(BenchmarkResult::new(
|
|
format!("Naive Attention ({})", config.name),
|
|
times,
|
|
standard_memory,
|
|
))
|
|
}
|
|
|
|
/// Print performance comparison
|
|
fn print_performance_comparison(flash_result: &BenchmarkResult, baseline_result: &BenchmarkResult) {
|
|
println!("\n=== PERFORMANCE COMPARISON ===");
|
|
println!("Configuration: {}", flash_result.name);
|
|
|
|
println!("\nFlash Attention:");
|
|
println!(
|
|
" Average time: {:.2} ms",
|
|
flash_result.avg_time_us / 1000.0
|
|
);
|
|
println!(
|
|
" Min time: {:.2} ms",
|
|
flash_result.min_time_us as f64 / 1000.0
|
|
);
|
|
println!(
|
|
" Max time: {:.2} ms",
|
|
flash_result.max_time_us as f64 / 1000.0
|
|
);
|
|
println!(" Std dev: {:.2} ms", flash_result.std_dev_us / 1000.0);
|
|
println!(
|
|
" Memory usage: {:.2} MB",
|
|
flash_result.memory_usage_bytes as f64 / 1_000_000.0
|
|
);
|
|
println!(" Throughput: {:.2} GB/s", flash_result.throughput_gb_s);
|
|
|
|
println!("\nBaseline (Naive Attention):");
|
|
println!(
|
|
" Average time: {:.2} ms",
|
|
baseline_result.avg_time_us / 1000.0
|
|
);
|
|
println!(
|
|
" Min time: {:.2} ms",
|
|
baseline_result.min_time_us as f64 / 1000.0
|
|
);
|
|
println!(
|
|
" Max time: {:.2} ms",
|
|
baseline_result.max_time_us as f64 / 1000.0
|
|
);
|
|
println!(" Std dev: {:.2} ms", baseline_result.std_dev_us / 1000.0);
|
|
println!(
|
|
" Memory usage: {:.2} MB",
|
|
baseline_result.memory_usage_bytes as f64 / 1_000_000.0
|
|
);
|
|
println!(" Throughput: {:.2} GB/s", baseline_result.throughput_gb_s);
|
|
|
|
let speedup = flash_result.speedup_vs(baseline_result);
|
|
let memory_reduction =
|
|
baseline_result.memory_usage_bytes as f64 / flash_result.memory_usage_bytes as f64;
|
|
|
|
println!("\n=== PERFORMANCE GAINS ===");
|
|
println!("Speed improvement: {:.2}x faster", speedup);
|
|
println!("Memory reduction: {:.2}x less memory", memory_reduction);
|
|
|
|
// Validate claims
|
|
if speedup >= 5.0 {
|
|
println!("✅ SPEEDUP TARGET ACHIEVED: {:.1}x >= 5.0x", speedup);
|
|
} else if speedup >= 2.0 {
|
|
println!("⚠️ PARTIAL SPEEDUP: {:.1}x (target: 5.0x+)", speedup);
|
|
} else {
|
|
println!("❌ SPEEDUP TARGET MISSED: {:.1}x < 5.0x", speedup);
|
|
}
|
|
|
|
if memory_reduction >= 2.0 {
|
|
println!(
|
|
"✅ MEMORY EFFICIENCY ACHIEVED: {:.1}x less memory",
|
|
memory_reduction
|
|
);
|
|
} else {
|
|
println!(
|
|
"⚠️ MEMORY EFFICIENCY PARTIAL: {:.1}x memory reduction",
|
|
memory_reduction
|
|
);
|
|
}
|
|
|
|
println!("=====================================\n");
|
|
}
|
|
|
|
#[test]
|
|
async fn test_small_scale_performance() {
|
|
if !Device::cuda_available() {
|
|
println!("CUDA not available, skipping performance test");
|
|
return;
|
|
}
|
|
|
|
let config = PerfTestConfig::new("Small Scale", 2, 8, 512, 64);
|
|
|
|
match create_test_tensors(&config) {
|
|
Ok((q, k, v)) => {
|
|
let flash_result = benchmark_flash_attention(&config, &q, &k, &v).await;
|
|
let naive_result = benchmark_naive_attention(&config, &q, &k, &v).await;
|
|
|
|
if let (Ok(flash), Ok(naive)) = (flash_result, naive_result) {
|
|
print_performance_comparison(&flash, &naive);
|
|
|
|
// Assert minimum performance gains
|
|
assert!(
|
|
flash.speedup_vs(&naive) >= 2.0,
|
|
"Flash Attention should be at least 2x faster"
|
|
);
|
|
}
|
|
}
|
|
Err(e) => println!("Failed to create test tensors: {}", e),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
async fn test_medium_scale_performance() {
|
|
if !Device::cuda_available() {
|
|
println!("CUDA not available, skipping performance test");
|
|
return;
|
|
}
|
|
|
|
let config = PerfTestConfig::new("Medium Scale", 4, 16, 1024, 128);
|
|
|
|
match create_test_tensors(&config) {
|
|
Ok((q, k, v)) => {
|
|
let flash_result = benchmark_flash_attention(&config, &q, &k, &v).await;
|
|
let naive_result = benchmark_naive_attention(&config, &q, &k, &v).await;
|
|
|
|
if let (Ok(flash), Ok(naive)) = (flash_result, naive_result) {
|
|
print_performance_comparison(&flash, &naive);
|
|
|
|
// Assert target performance gains
|
|
assert!(
|
|
flash.speedup_vs(&naive) >= 3.0,
|
|
"Flash Attention should be at least 3x faster for medium scale"
|
|
);
|
|
}
|
|
}
|
|
Err(e) => println!("Failed to create test tensors: {}", e),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
async fn test_large_scale_performance() {
|
|
if !Device::cuda_available() {
|
|
println!("CUDA not available, skipping performance test");
|
|
return;
|
|
}
|
|
|
|
let config = PerfTestConfig::new("Large Scale", 8, 32, 2048, 128);
|
|
|
|
match create_test_tensors(&config) {
|
|
Ok((q, k, v)) => {
|
|
let flash_result = benchmark_flash_attention(&config, &q, &k, &v).await;
|
|
let naive_result = benchmark_naive_attention(&config, &q, &k, &v).await;
|
|
|
|
if let (Ok(flash), Ok(naive)) = (flash_result, naive_result) {
|
|
print_performance_comparison(&flash, &naive);
|
|
|
|
// Assert target performance gains (should be higher for larger sequences)
|
|
assert!(
|
|
flash.speedup_vs(&naive) >= 5.0,
|
|
"Flash Attention should be at least 5x faster for large scale"
|
|
);
|
|
}
|
|
}
|
|
Err(e) => println!("Failed to create test tensors: {}", e),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
async fn test_xl_scale_performance() {
|
|
if !Device::cuda_available() {
|
|
println!("CUDA not available, skipping performance test");
|
|
return;
|
|
}
|
|
|
|
let config = PerfTestConfig::new("XL Scale", 4, 64, 4096, 128);
|
|
|
|
match create_test_tensors(&config) {
|
|
Ok((q, k, v)) => {
|
|
let flash_result = benchmark_flash_attention(&config, &q, &k, &v).await;
|
|
let naive_result = benchmark_naive_attention(&config, &q, &k, &v).await;
|
|
|
|
if let (Ok(flash), Ok(naive)) = (flash_result, naive_result) {
|
|
print_performance_comparison(&flash, &naive);
|
|
|
|
// Assert maximum performance gains for very large sequences
|
|
assert!(
|
|
flash.speedup_vs(&naive) >= 8.0,
|
|
"Flash Attention should be at least 8x faster for XL scale"
|
|
);
|
|
}
|
|
}
|
|
Err(e) => println!("Failed to create test tensors: {}", e),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
async fn test_causal_attention_performance() {
|
|
if !Device::cuda_available() {
|
|
println!("CUDA not available, skipping performance test");
|
|
return;
|
|
}
|
|
|
|
let config = PerfTestConfig::new("Causal Attention", 4, 32, 1024, 128).with_causal(true);
|
|
|
|
match create_test_tensors(&config) {
|
|
Ok((q, k, v)) => {
|
|
let flash_result = benchmark_flash_attention(&config, &q, &k, &v).await;
|
|
let naive_result = benchmark_naive_attention(&config, &q, &k, &v).await;
|
|
|
|
if let (Ok(flash), Ok(naive)) = (flash_result, naive_result) {
|
|
print_performance_comparison(&flash, &naive);
|
|
|
|
// Causal attention should show good speedups
|
|
assert!(
|
|
flash.speedup_vs(&naive) >= 4.0,
|
|
"Causal Flash Attention should be at least 4x faster"
|
|
);
|
|
}
|
|
}
|
|
Err(e) => println!("Failed to create test tensors: {}", e),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
async fn test_memory_scaling_validation() {
|
|
println!("\n=== MEMORY SCALING VALIDATION ===");
|
|
|
|
let sequence_lengths = vec![512, 1024, 2048, 4096, 8_192];
|
|
let batch_size = 2;
|
|
let num_heads = 16;
|
|
let head_dim = 128;
|
|
|
|
for seq_len in sequence_lengths {
|
|
let (standard_mem, flash_mem) =
|
|
compare_memory_usage(batch_size, num_heads, seq_len, head_dim);
|
|
let reduction_ratio = standard_mem as f64 / flash_mem as f64;
|
|
|
|
println!(
|
|
"Seq Len {}: Standard={:.2}MB, Flash={:.2}MB, Reduction={:.1}x",
|
|
seq_len,
|
|
standard_mem as f64 / 1_000_000.0,
|
|
flash_mem as f64 / 1_000_000.0,
|
|
reduction_ratio
|
|
);
|
|
|
|
// Validate O(n) vs O(n²) scaling
|
|
if seq_len >= 2048 {
|
|
assert!(
|
|
reduction_ratio >= seq_len as f64 / 128.0,
|
|
"Memory reduction should scale with sequence length"
|
|
);
|
|
}
|
|
}
|
|
|
|
println!("✅ Memory scaling validation completed");
|
|
println!("=====================================\n");
|
|
}
|
|
|
|
#[test]
|
|
async fn test_backward_pass_performance() {
|
|
if !Device::cuda_available() {
|
|
println!("CUDA not available, skipping backward performance test");
|
|
return;
|
|
}
|
|
|
|
let config = PerfTestConfig::new("Backward Pass", 4, 16, 1024, 64);
|
|
|
|
match create_test_tensors(&config) {
|
|
Ok((q, k, v)) => {
|
|
let flash_config =
|
|
FlashAttentionConfig::for_training(config.num_heads, config.head_dim);
|
|
if let Ok(flash) = FlashAttention::new(flash_config) {
|
|
let softmax_scale = 1.0 / (config.head_dim as f32).sqrt();
|
|
|
|
// Forward pass first
|
|
if let Ok(forward_result) = flash
|
|
.forward(&q, &k, &v, config.causal, softmax_scale)
|
|
.await
|
|
{
|
|
// Create gradient tensor
|
|
let dout_shape = [
|
|
config.batch_size,
|
|
config.num_heads,
|
|
config.seq_len,
|
|
config.head_dim,
|
|
];
|
|
if let Ok(dout) = Tensor::randn(&dout_shape, DType::F16, &Device::Cuda(0)) {
|
|
// Benchmark backward pass
|
|
let mut backward_times = Vec::new();
|
|
for _ in 0..5 {
|
|
let start = Instant::now();
|
|
let _ = flash
|
|
.backward(
|
|
&dout,
|
|
&q,
|
|
&k,
|
|
&v,
|
|
&forward_result.output,
|
|
&forward_result.lse,
|
|
config.causal,
|
|
softmax_scale,
|
|
)
|
|
.await;
|
|
backward_times.push(start.elapsed().as_micros() as u64);
|
|
}
|
|
|
|
let avg_backward_time =
|
|
backward_times.iter().sum::<u64>() as f64 / backward_times.len() as f64;
|
|
let avg_forward_time = forward_result.stats.forward_time_us as f64;
|
|
|
|
println!("\n=== BACKWARD PASS PERFORMANCE ===");
|
|
println!("Forward time: {:.2} ms", avg_forward_time / 1000.0);
|
|
println!("Backward time: {:.2} ms", avg_backward_time / 1000.0);
|
|
println!(
|
|
"Backward/Forward ratio: {:.2}x",
|
|
avg_backward_time / avg_forward_time
|
|
);
|
|
println!("=====================================\n");
|
|
|
|
// Backward should be reasonable compared to forward
|
|
assert!(
|
|
avg_backward_time / avg_forward_time < 5.0,
|
|
"Backward pass should be less than 5x forward pass time"
|
|
);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
Err(e) => println!("Failed to create test tensors: {}", e),
|
|
}
|
|
}
|