405 lines
13 KiB
Rust
405 lines
13 KiB
Rust
//! Comprehensive kernel validation tests for Flash Attention
|
|
//! Tests the correctness of Flash Attention kernel implementations
|
|
//!
|
|
//! NOTE: Disabled until Flash Attention API is fully implemented
|
|
|
|
#![cfg(all(feature = "cuda", feature = "disabled_tests"))]
|
|
|
|
use approx::assert_relative_eq;
|
|
use proptest::prelude::*;
|
|
use rtx_flash_attention::*;
|
|
use rtx_tensor::Tensor;
|
|
|
|
/// Test basic Flash Attention forward pass correctness
|
|
#[test]
|
|
fn test_flash_attention_forward_basic() {
|
|
let seq_len = 128;
|
|
let num_heads = 8;
|
|
let batch_size = 2;
|
|
|
|
let mut config = FlashAttentionConfig::new(num_heads, 64);
|
|
config.block_size_q = 32;
|
|
config.block_size_kv = 32;
|
|
config.causal = false;
|
|
config.softmax_scale = None;
|
|
|
|
// Create input tensors
|
|
let q = Tensor::randn(&[batch_size, num_heads, seq_len, config.head_dim]);
|
|
let k = Tensor::randn(&[batch_size, num_heads, seq_len, config.head_dim]);
|
|
let v = Tensor::randn(&[batch_size, num_heads, seq_len, config.head_dim]);
|
|
|
|
let result = flash_attention_forward(&q, &k, &v, &config).unwrap();
|
|
|
|
// Verify output shape
|
|
assert_eq!(
|
|
result.shape(),
|
|
&[batch_size, num_heads, seq_len, config.head_dim]
|
|
);
|
|
|
|
// Verify output is finite
|
|
assert!(result.all_finite().unwrap());
|
|
}
|
|
|
|
/// Test Flash Attention with causal masking
|
|
#[test]
|
|
fn test_flash_attention_causal() {
|
|
let seq_len = 64;
|
|
|
|
let mut config = FlashAttentionConfig::new(1, 64);
|
|
config.block_size_q = 32;
|
|
config.block_size_kv = 32;
|
|
config.causal = true;
|
|
config.softmax_scale = Some(0.125);
|
|
let q = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
|
|
let k = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
|
|
let v = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
|
|
|
|
let result = flash_attention_forward(&q, &k, &v, &config).unwrap();
|
|
|
|
// For causal attention, verify triangular structure
|
|
assert_eq!(result.shape(), &[1, 1, seq_len, config.head_dim]);
|
|
assert!(result.all_finite().unwrap());
|
|
}
|
|
|
|
/// Test Flash Attention backward pass
|
|
#[test]
|
|
fn test_flash_attention_backward() {
|
|
let seq_len = 32;
|
|
|
|
let mut config = FlashAttentionConfig::new(1, 32);
|
|
config.block_size_q = 16;
|
|
config.block_size_kv = 16;
|
|
config.causal = false;
|
|
config.softmax_scale = None;
|
|
let q = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
|
|
let k = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
|
|
let v = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
|
|
let grad_out = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
|
|
|
|
let (grad_q, grad_k, grad_v) =
|
|
flash_attention_backward(&grad_out, &q, &k, &v, &config).unwrap();
|
|
|
|
// Verify gradient shapes
|
|
assert_eq!(grad_q.shape(), q.shape());
|
|
assert_eq!(grad_k.shape(), k.shape());
|
|
assert_eq!(grad_v.shape(), v.shape());
|
|
|
|
// Verify gradients are finite
|
|
assert!(grad_q.all_finite().unwrap());
|
|
assert!(grad_k.all_finite().unwrap());
|
|
assert!(grad_v.all_finite().unwrap());
|
|
}
|
|
|
|
/// Test memory efficiency for large sequences
|
|
#[test]
|
|
fn test_memory_efficiency_large_sequence() {
|
|
let seq_len = 2048; // Large sequence
|
|
|
|
let mut config = FlashAttentionConfig::new(1, 64);
|
|
config.block_size_q = 64;
|
|
config.block_size_kv = 64;
|
|
config.causal = false;
|
|
config.softmax_scale = None;
|
|
let q = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
|
|
let k = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
|
|
let v = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
|
|
|
|
let result = flash_attention_forward(&q, &k, &v, &config).unwrap();
|
|
assert_eq!(result.shape(), &[1, 1, seq_len, config.head_dim]);
|
|
}
|
|
|
|
/// Test numerical equivalence with reference implementation
|
|
#[test]
|
|
fn test_numerical_equivalence() {
|
|
let seq_len = 16;
|
|
|
|
let mut config = FlashAttentionConfig::new(1, 16);
|
|
config.block_size_q = 8;
|
|
config.block_size_kv = 8;
|
|
config.causal = false;
|
|
config.softmax_scale = Some(0.25);
|
|
let q = Tensor::ones(&[1, 1, seq_len, config.head_dim]);
|
|
let k = Tensor::ones(&[1, 1, seq_len, config.head_dim]);
|
|
let v = Tensor::ones(&[1, 1, seq_len, config.head_dim]);
|
|
|
|
let flash_result = flash_attention_forward(&q, &k, &v, &config).unwrap();
|
|
let reference_result = reference_attention(&q, &k, &v, config.softmax_scale).unwrap();
|
|
|
|
// Compare results with tolerance
|
|
for i in 0..flash_result.numel() {
|
|
assert_relative_eq!(
|
|
flash_result.get_item(i).unwrap(),
|
|
reference_result.get_item(i).unwrap(),
|
|
epsilon = 1e-4
|
|
);
|
|
}
|
|
}
|
|
|
|
/// Property-based test for Flash Attention invariants
|
|
proptest! {
|
|
#[test]
|
|
fn test_flash_attention_properties(
|
|
batch_size in 1..4usize,
|
|
num_heads in 1..8usize,
|
|
seq_len in 16..128usize,
|
|
head_dim in prop::sample::select(vec![16, 32, 64]),
|
|
) {
|
|
let mut config = FlashAttentionConfig::new(num_heads, head_dim);
|
|
config.block_size_q = 32;
|
|
config.block_size_kv = 32;
|
|
config.causal = false;
|
|
config.softmax_scale = None;
|
|
|
|
let q = Tensor::randn(&[batch_size, num_heads, seq_len, head_dim]);
|
|
let k = Tensor::randn(&[batch_size, num_heads, seq_len, head_dim]);
|
|
let v = Tensor::randn(&[batch_size, num_heads, seq_len, head_dim]);
|
|
|
|
let result = flash_attention_forward(&q, &k, &v, &config).unwrap();
|
|
|
|
// Property: Output shape matches expected
|
|
prop_assert_eq!(result.shape(), &[batch_size, num_heads, seq_len, head_dim]);
|
|
|
|
// Property: Output is finite
|
|
prop_assert!(result.all_finite().unwrap());
|
|
|
|
// Property: Output norm is reasonable (not NaN/Inf)
|
|
let output_norm = result.norm().unwrap();
|
|
prop_assert!(output_norm.is_finite());
|
|
prop_assert!(output_norm > 0.0);
|
|
}
|
|
}
|
|
|
|
/// Test Flash Attention with different block sizes
|
|
#[test]
|
|
fn test_different_block_sizes() {
|
|
let block_sizes = vec![(16, 16), (32, 32), (64, 64), (128, 128)];
|
|
let seq_len = 256;
|
|
|
|
let q = Tensor::randn(&[1, 1, seq_len, 64]);
|
|
let k = Tensor::randn(&[1, 1, seq_len, 64]);
|
|
let v = Tensor::randn(&[1, 1, seq_len, 64]);
|
|
|
|
for (block_q, block_k) in block_sizes {
|
|
if block_q <= seq_len && block_k <= seq_len {
|
|
let mut config = FlashAttentionConfig::new(1, 64);
|
|
config.block_size_q = block_q;
|
|
config.block_size_kv = block_k;
|
|
config.causal = false;
|
|
config.softmax_scale = None;
|
|
|
|
let result = flash_attention_forward(&q, &k, &v, &config).unwrap();
|
|
assert_eq!(result.shape(), &[1, 1, seq_len, 64]);
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Test online softmax correctness
|
|
#[test]
|
|
fn test_online_softmax() {
|
|
use rtx_flash_attention::kernels::utils::online_softmax;
|
|
|
|
let input = vec![1.0, 2.0, 3.0, 4.0];
|
|
let result = online_softmax(&input);
|
|
|
|
// Verify softmax properties
|
|
let sum: f32 = result.iter().sum();
|
|
assert_relative_eq!(sum, 1.0, epsilon = 1e-6);
|
|
|
|
// Verify all values are positive
|
|
for &val in &result {
|
|
assert!(val > 0.0);
|
|
}
|
|
|
|
// Verify monotonicity for this input
|
|
for i in 1..result.len() {
|
|
assert!(result[i] > result[i - 1]);
|
|
}
|
|
}
|
|
|
|
/// Test Flash Attention with mixed precision
|
|
#[test]
|
|
fn test_mixed_precision() {
|
|
let seq_len = 128;
|
|
|
|
let mut config = FlashAttentionConfig::new(1, 64);
|
|
config.block_size_q = 32;
|
|
config.block_size_kv = 32;
|
|
config.causal = false;
|
|
config.softmax_scale = None;
|
|
|
|
// Test with fp16 inputs (simulated with fp32)
|
|
let q = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
|
|
let k = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
|
|
let v = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
|
|
|
|
let result = flash_attention_forward(&q, &k, &v, &config).unwrap();
|
|
assert!(result.all_finite().unwrap());
|
|
}
|
|
|
|
/// Test Flash Attention variants (edge, neuromorphic, quantum)
|
|
#[test]
|
|
fn test_flash_attention_variants() {
|
|
use rtx_flash_attention::variants::*;
|
|
|
|
let seq_len = 64;
|
|
let head_dim = 32;
|
|
|
|
let q = Tensor::randn(&[1, 1, seq_len, head_dim]);
|
|
let k = Tensor::randn(&[1, 1, seq_len, head_dim]);
|
|
let v = Tensor::randn(&[1, 1, seq_len, head_dim]);
|
|
|
|
// Test edge variant
|
|
let edge_result = edge::flash_attention_edge(&q, &k, &v).unwrap();
|
|
assert_eq!(edge_result.shape(), &[1, 1, seq_len, head_dim]);
|
|
|
|
// Test neuromorphic variant
|
|
let neuro_result = neuromorphic::flash_attention_neuromorphic(&q, &k, &v).unwrap();
|
|
assert_eq!(neuro_result.shape(), &[1, 1, seq_len, head_dim]);
|
|
|
|
// Test quantum variant
|
|
let quantum_result = quantum::flash_attention_quantum(&q, &k, &v).unwrap();
|
|
assert_eq!(quantum_result.shape(), &[1, 1, seq_len, head_dim]);
|
|
}
|
|
|
|
/// Helper function for reference attention implementation
|
|
fn reference_attention(
|
|
q: &Tensor,
|
|
k: &Tensor,
|
|
v: &Tensor,
|
|
scale: Option<f32>,
|
|
) -> Result<Tensor, Box<dyn std::error::Error>> {
|
|
// Simple reference implementation for testing
|
|
let scale = scale.unwrap_or(1.0 / (q.shape()[3] as f32).sqrt());
|
|
|
|
// QK^T
|
|
let scores = q.matmul(&k.transpose(-2, -1)?)?;
|
|
let scaled_scores = scores.mul_scalar(scale)?;
|
|
|
|
// Softmax
|
|
let attention_probs = scaled_scores.softmax(-1)?;
|
|
|
|
// Apply to values
|
|
let output = attention_probs.matmul(v)?;
|
|
Ok(output)
|
|
}
|
|
|
|
/// Benchmark comparison test
|
|
#[test]
|
|
fn test_performance_improvement() {
|
|
let seq_len = 1024;
|
|
|
|
let mut config = FlashAttentionConfig::new(8, 64);
|
|
config.block_size_q = 64;
|
|
config.block_size_kv = 64;
|
|
config.causal = false;
|
|
config.softmax_scale = None;
|
|
let q = Tensor::randn(&[1, 8, seq_len, config.head_dim]);
|
|
let k = Tensor::randn(&[1, 8, seq_len, config.head_dim]);
|
|
let v = Tensor::randn(&[1, 8, seq_len, config.head_dim]);
|
|
|
|
let start = std::time::Instant::now();
|
|
let _flash_result = flash_attention_forward(&q, &k, &v, &config).unwrap();
|
|
let flash_time = start.elapsed();
|
|
|
|
let start = std::time::Instant::now();
|
|
let _reference_result = reference_attention(&q, &k, &v, config.softmax_scale).unwrap();
|
|
let reference_time = start.elapsed();
|
|
|
|
// Flash Attention should be faster for large sequences
|
|
println!("Flash Attention time: {:?}", flash_time);
|
|
println!("Reference time: {:?}", reference_time);
|
|
|
|
// For large sequences, Flash Attention should show improvement
|
|
if seq_len > 512 {
|
|
assert!(flash_time < reference_time * 2); // Allow some margin for test environment
|
|
}
|
|
}
|
|
|
|
/// Test error handling
|
|
#[test]
|
|
fn test_error_handling() {
|
|
let mut config = FlashAttentionConfig::new(1, 64);
|
|
config.block_size_q = 32;
|
|
config.block_size_kv = 32;
|
|
config.causal = false;
|
|
config.softmax_scale = None;
|
|
|
|
// Mismatched shapes should error
|
|
let q = Tensor::randn(&[1, 1, 128, 64]);
|
|
let k = Tensor::randn(&[1, 1, 64, 64]); // Different seq_len
|
|
let v = Tensor::randn(&[1, 1, 128, 64]);
|
|
|
|
let result = flash_attention_forward(&q, &k, &v, &config);
|
|
assert!(result.is_err());
|
|
|
|
// Mismatched head dimensions should error
|
|
let q = Tensor::randn(&[1, 1, 128, 64]);
|
|
let k = Tensor::randn(&[1, 1, 128, 32]); // Different head_dim
|
|
let v = Tensor::randn(&[1, 1, 128, 64]);
|
|
|
|
let result = flash_attention_forward(&q, &k, &v, &config);
|
|
assert!(result.is_err());
|
|
}
|
|
|
|
/// Test gradient computation correctness
|
|
#[test]
|
|
fn test_gradient_correctness() {
|
|
let seq_len = 16;
|
|
|
|
let mut config = FlashAttentionConfig::new(1, 16);
|
|
config.block_size_q = 8;
|
|
config.block_size_kv = 8;
|
|
config.causal = false;
|
|
config.softmax_scale = Some(0.25);
|
|
let q = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
|
|
let k = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
|
|
let v = Tensor::randn(&[1, 1, seq_len, config.head_dim]);
|
|
|
|
// Forward pass
|
|
let output = flash_attention_forward(&q, &k, &v, &config).unwrap();
|
|
|
|
// Create gradient of loss w.r.t. output
|
|
let grad_output = Tensor::ones_like(&output);
|
|
|
|
// Backward pass
|
|
let (grad_q, grad_k, grad_v) =
|
|
flash_attention_backward(&grad_output, &q, &k, &v, &config).unwrap();
|
|
|
|
// Verify gradient shapes match input shapes
|
|
assert_eq!(grad_q.shape(), q.shape());
|
|
assert_eq!(grad_k.shape(), k.shape());
|
|
assert_eq!(grad_v.shape(), v.shape());
|
|
|
|
// Verify gradients are reasonable (not too large)
|
|
assert!(grad_q.abs().max().unwrap() < 100.0);
|
|
assert!(grad_k.abs().max().unwrap() < 100.0);
|
|
assert!(grad_v.abs().max().unwrap() < 100.0);
|
|
}
|
|
|
|
/// Test multi-head attention
|
|
#[test]
|
|
fn test_multi_head_attention() {
|
|
let seq_len = 64;
|
|
let num_heads = 12; // Large number of heads
|
|
|
|
let mut config = FlashAttentionConfig::new(num_heads, 32);
|
|
config.block_size_q = 16;
|
|
config.block_size_kv = 16;
|
|
config.causal = false;
|
|
config.softmax_scale = None;
|
|
|
|
let q = Tensor::randn(&[2, num_heads, seq_len, config.head_dim]);
|
|
let k = Tensor::randn(&[2, num_heads, seq_len, config.head_dim]);
|
|
let v = Tensor::randn(&[2, num_heads, seq_len, config.head_dim]);
|
|
|
|
let result = flash_attention_forward(&q, &k, &v, &config).unwrap();
|
|
assert_eq!(result.shape(), &[2, num_heads, seq_len, config.head_dim]);
|
|
|
|
// Verify each head produces reasonable outputs
|
|
for head in 0..num_heads {
|
|
let head_output = result.select(1, head).unwrap();
|
|
assert!(head_output.all_finite().unwrap());
|
|
assert!(head_output.norm().unwrap() > 0.0);
|
|
}
|
|
}
|