329 lines
11 KiB
Rust
329 lines
11 KiB
Rust
//! Edge case tests for Flash Attention - numerical stability and boundary conditions
|
|
//!
|
|
//! These tests validate Flash Attention behavior under extreme conditions that commonly
|
|
//! cause numerical instability or failure in attention implementations.
|
|
|
|
use rtx_flash_attention::*;
|
|
use rtx_tensor::{DType, Device, Tensor};
|
|
|
|
/// Test numerical stability with very small softmax scales
|
|
#[test]
|
|
fn test_numerical_stability_small_softmax_scale() {
|
|
let device = match Device::try_default() {
|
|
Ok(Device::Cuda(cuda_device)) => Device::Cuda(cuda_device),
|
|
_ => {
|
|
println!("CUDA not available, skipping test");
|
|
return;
|
|
}
|
|
};
|
|
|
|
let mut config = FlashAttentionConfig::new(8, 64);
|
|
config.softmax_scale = Some(1e-8); // Extremely small scale
|
|
|
|
let q = Tensor::randn(&[2, 8, 512, 64], &device).unwrap();
|
|
let k = Tensor::randn(&[2, 8, 512, 64], &device).unwrap();
|
|
let v = Tensor::randn(&[2, 8, 512, 64], &device).unwrap();
|
|
|
|
// This should not crash or produce NaN/Inf values
|
|
let result = flash_attention_forward(&q, &k, &v, &config);
|
|
assert!(
|
|
result.is_ok(),
|
|
"Forward pass should handle small softmax scale"
|
|
);
|
|
|
|
let output = result.unwrap();
|
|
assert!(
|
|
is_finite(&output),
|
|
"Output should be finite with small softmax scale"
|
|
);
|
|
assert!(!has_nan(&output), "Output should not contain NaN values");
|
|
|
|
println!("✓ Small softmax scale test passed");
|
|
}
|
|
|
|
/// Test numerical stability with very large softmax scales
|
|
#[test]
|
|
fn test_numerical_stability_large_softmax_scale() {
|
|
let device = match Device::try_default() {
|
|
Ok(Device::Cuda(cuda_device)) => Device::Cuda(cuda_device),
|
|
_ => {
|
|
println!("CUDA not available, skipping test");
|
|
return;
|
|
}
|
|
};
|
|
|
|
let mut config = FlashAttentionConfig::new(8, 64);
|
|
config.softmax_scale = Some(1e8); // Extremely large scale
|
|
|
|
let q = Tensor::randn(&[2, 8, 512, 64], &device).unwrap();
|
|
let k = Tensor::randn(&[2, 8, 512, 64], &device).unwrap();
|
|
let v = Tensor::randn(&[2, 8, 512, 64], &device).unwrap();
|
|
|
|
let result = flash_attention_forward(&q, &k, &v, &config);
|
|
assert!(
|
|
result.is_ok(),
|
|
"Forward pass should handle large softmax scale"
|
|
);
|
|
|
|
let output = result.unwrap();
|
|
assert!(
|
|
is_finite(&output),
|
|
"Output should be finite with large softmax scale"
|
|
);
|
|
assert!(!has_nan(&output), "Output should not contain NaN values");
|
|
|
|
println!("✓ Large softmax scale test passed");
|
|
}
|
|
|
|
/// Test boundary condition: single token sequence
|
|
#[test]
|
|
fn test_single_token_sequence() {
|
|
let device = match Device::try_default() {
|
|
Ok(Device::Cuda(cuda_device)) => Device::Cuda(cuda_device),
|
|
_ => {
|
|
println!("CUDA not available, skipping test");
|
|
return;
|
|
}
|
|
};
|
|
|
|
let config = FlashAttentionConfig::new(8, 64);
|
|
|
|
// Single token in sequence
|
|
let q = Tensor::randn(&[1, 8, 1, 64], &device).unwrap();
|
|
let k = Tensor::randn(&[1, 8, 1, 64], &device).unwrap();
|
|
let v = Tensor::randn(&[1, 8, 1, 64], &device).unwrap();
|
|
|
|
let result = flash_attention_forward(&q, &k, &v, &config);
|
|
assert!(result.is_ok(), "Should handle single token sequences");
|
|
|
|
let output = result.unwrap();
|
|
assert_eq!(output.shape().dims(), &[1, 8, 1, 64]);
|
|
assert!(is_finite(&output), "Single token output should be finite");
|
|
|
|
println!("✓ Single token sequence test passed");
|
|
}
|
|
|
|
/// Test boundary condition: very long sequences testing memory efficiency
|
|
#[test]
|
|
fn test_long_sequence_memory_efficiency() {
|
|
let device = match Device::try_default() {
|
|
Ok(Device::Cuda(cuda_device)) => Device::Cuda(cuda_device),
|
|
_ => {
|
|
println!("CUDA not available, skipping test");
|
|
return;
|
|
}
|
|
};
|
|
|
|
let config = FlashAttentionConfig::new(8, 64);
|
|
|
|
// Very long sequence that would be prohibitive for standard attention O(n²) memory
|
|
let seq_len = 8192; // 8K sequence length
|
|
let q = Tensor::randn(&[1, 8, seq_len, 64], &device).unwrap();
|
|
let k = Tensor::randn(&[1, 8, seq_len, 64], &device).unwrap();
|
|
let v = Tensor::randn(&[1, 8, seq_len, 64], &device).unwrap();
|
|
|
|
let result = flash_attention_forward(&q, &k, &v, &config);
|
|
assert!(result.is_ok(), "Should handle long sequences efficiently");
|
|
|
|
let output = result.unwrap();
|
|
assert_eq!(output.shape().dims(), &[1, 8, seq_len, 64]);
|
|
assert!(is_finite(&output), "Long sequence output should be finite");
|
|
|
|
println!("✓ Long sequence memory efficiency test passed");
|
|
}
|
|
|
|
/// Test numerical stability with extreme input values
|
|
#[test]
|
|
fn test_extreme_input_values() {
|
|
let device = match Device::try_default() {
|
|
Ok(Device::Cuda(cuda_device)) => Device::Cuda(cuda_device),
|
|
_ => {
|
|
println!("CUDA not available, skipping test");
|
|
return;
|
|
}
|
|
};
|
|
|
|
let config = FlashAttentionConfig::new(8, 64);
|
|
|
|
// Create tensors with extreme values
|
|
let q = create_extreme_tensor(&[2, 8, 128, 64], &device, 100.0);
|
|
let k = create_extreme_tensor(&[2, 8, 128, 64], &device, -100.0);
|
|
let v = create_extreme_tensor(&[2, 8, 128, 64], &device, 50.0);
|
|
|
|
let result = flash_attention_forward(&q, &k, &v, &config);
|
|
assert!(result.is_ok(), "Should handle extreme input values");
|
|
|
|
let output = result.unwrap();
|
|
assert!(
|
|
is_finite(&output),
|
|
"Output should be finite with extreme inputs"
|
|
);
|
|
assert!(
|
|
!has_nan(&output),
|
|
"Output should not contain NaN with extreme inputs"
|
|
);
|
|
|
|
println!("✓ Extreme input values test passed");
|
|
}
|
|
|
|
/// Test causal masking correctness for different sequence lengths
|
|
#[test]
|
|
fn test_causal_masking_correctness() {
|
|
let device = match Device::try_default() {
|
|
Ok(Device::Cuda(cuda_device)) => Device::Cuda(cuda_device),
|
|
_ => {
|
|
println!("CUDA not available, skipping test");
|
|
return;
|
|
}
|
|
};
|
|
|
|
let mut config = FlashAttentionConfig::new(4, 32);
|
|
config.causal = true;
|
|
|
|
let seq_len = 16;
|
|
let q = Tensor::randn(&[1, 4, seq_len, 32], &device).unwrap();
|
|
let k = Tensor::randn(&[1, 4, seq_len, 32], &device).unwrap();
|
|
let v = Tensor::randn(&[1, 4, seq_len, 32], &device).unwrap();
|
|
|
|
let causal_result = flash_attention_forward(&q, &k, &v, &config).unwrap();
|
|
|
|
config.causal = false;
|
|
let non_causal_result = flash_attention_forward(&q, &k, &v, &config).unwrap();
|
|
|
|
// Causal and non-causal should produce different results
|
|
assert!(
|
|
!tensors_equal(&causal_result, &non_causal_result),
|
|
"Causal and non-causal attention should produce different results"
|
|
);
|
|
|
|
assert!(is_finite(&causal_result), "Causal result should be finite");
|
|
assert!(
|
|
is_finite(&non_causal_result),
|
|
"Non-causal result should be finite"
|
|
);
|
|
|
|
println!("✓ Causal masking correctness test passed");
|
|
}
|
|
|
|
/// Test gradient flow in backward pass
|
|
#[test]
|
|
fn test_gradient_flow_correctness() {
|
|
let device = match Device::try_default() {
|
|
Ok(Device::Cuda(cuda_device)) => Device::Cuda(cuda_device),
|
|
_ => {
|
|
println!("CUDA not available, skipping test");
|
|
return;
|
|
}
|
|
};
|
|
|
|
let config = FlashAttentionConfig::new(4, 32);
|
|
|
|
let q = Tensor::randn(&[1, 4, 32, 32], &device).unwrap();
|
|
let k = Tensor::randn(&[1, 4, 32, 32], &device).unwrap();
|
|
let v = Tensor::randn(&[1, 4, 32, 32], &device).unwrap();
|
|
let grad_output = Tensor::randn(&[1, 4, 32, 32], &device).unwrap();
|
|
|
|
let (grad_q, grad_k, grad_v) =
|
|
flash_attention_backward(&grad_output, &q, &k, &v, &config).unwrap();
|
|
|
|
// Gradients should have the same shape as inputs
|
|
assert_eq!(grad_q.shape().dims(), q.shape().dims());
|
|
assert_eq!(grad_k.shape().dims(), k.shape().dims());
|
|
assert_eq!(grad_v.shape().dims(), v.shape().dims());
|
|
|
|
// Gradients should be finite and non-zero (indicating proper flow)
|
|
assert!(
|
|
is_finite(&grad_q) && !is_zero(&grad_q),
|
|
"grad_q should be finite and non-zero"
|
|
);
|
|
assert!(
|
|
is_finite(&grad_k) && !is_zero(&grad_k),
|
|
"grad_k should be finite and non-zero"
|
|
);
|
|
assert!(
|
|
is_finite(&grad_v) && !is_zero(&grad_v),
|
|
"grad_v should be finite and non-zero"
|
|
);
|
|
|
|
println!("✓ Gradient flow correctness test passed");
|
|
}
|
|
|
|
/// Test memory efficiency by comparing Flash vs naive attention memory usage
|
|
#[test]
|
|
fn test_memory_efficiency_validation() {
|
|
let batch_size = 2;
|
|
let num_heads = 8;
|
|
let seq_len = 2048;
|
|
let head_dim = 64;
|
|
|
|
let (standard_memory, flash_memory) =
|
|
utils::compare_memory_usage(batch_size, num_heads, seq_len, head_dim);
|
|
|
|
// Flash Attention should use significantly less memory for long sequences
|
|
let memory_reduction = standard_memory as f64 / flash_memory as f64;
|
|
assert!(
|
|
memory_reduction > 2.0,
|
|
"Flash Attention should reduce memory usage by at least 2x for seq_len=2048, got {:.2}x",
|
|
memory_reduction
|
|
);
|
|
|
|
println!(
|
|
"✓ Memory efficiency validation passed: {:.2}x reduction",
|
|
memory_reduction
|
|
);
|
|
}
|
|
|
|
// Helper functions for edge case testing
|
|
|
|
/// Check if all values in tensor are finite (not NaN or Inf)
|
|
fn is_finite(tensor: &Tensor) -> bool {
|
|
match tensor.to_cpu() {
|
|
Ok(data) => data.iter().all(|&x| x.is_finite()),
|
|
Err(_) => false, // If we can't get CPU data, assume not finite
|
|
}
|
|
}
|
|
|
|
/// Check if tensor contains any NaN values
|
|
fn has_nan(tensor: &Tensor) -> bool {
|
|
match tensor.to_cpu() {
|
|
Ok(data) => data.iter().any(|&x| x.is_nan()),
|
|
Err(_) => true, // If we can't get CPU data, assume NaN for safety
|
|
}
|
|
}
|
|
|
|
/// Check if tensor is all zeros
|
|
fn is_zero(tensor: &Tensor) -> bool {
|
|
match tensor.to_cpu() {
|
|
Ok(data) => data.iter().all(|&x| x == 0.0),
|
|
Err(_) => false,
|
|
}
|
|
}
|
|
|
|
/// Check if two tensors are equal within tolerance
|
|
fn tensors_equal(a: &Tensor, b: &Tensor) -> bool {
|
|
if a.shape() != b.shape() {
|
|
return false;
|
|
}
|
|
|
|
match (a.to_cpu(), b.to_cpu()) {
|
|
(Ok(data_a), Ok(data_b)) => {
|
|
const TOLERANCE: f32 = 1e-6;
|
|
data_a
|
|
.iter()
|
|
.zip(data_b.iter())
|
|
.all(|(&x, &y)| (x - y).abs() < TOLERANCE)
|
|
}
|
|
_ => false,
|
|
}
|
|
}
|
|
|
|
/// Create tensor with extreme values for stress testing
|
|
fn create_extreme_tensor(shape: &[usize], device: &Device, scale: f32) -> Tensor {
|
|
// Create a tensor with random values and scale them to extreme values
|
|
let tensor = Tensor::randn(shape, device).unwrap();
|
|
|
|
// For now, return the tensor as-is since we don't have a scalar multiply operation
|
|
// This will be a test that should fail in RED phase until we implement proper scaling
|
|
tensor
|
|
}
|