rustytorch/crates/training/rtx-flash-attention/tests/simple_integration.rs

//! Simple integration tests for Flash Attention - focused on compilation and basic functionality
//!
//! NOTE: Disabled until Flash Attention API is fully implemented

#![cfg(all(feature = "cuda", feature = "disabled_tests"))]

use rtx_flash_attention::*;
use rtx_tensor::{DType, Device, Tensor};

/// Basic test to ensure Flash Attention forward pass compiles and runs
#[test]
fn test_basic_flash_attention_forward() {
    // Skip test if CUDA is not available
    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(1, 64);
    config.block_size_q = 32;
    config.block_size_kv = 32;
    config.causal = false;

    let q = Tensor::randn(&[1, 1, 64, 64], &device).unwrap();
    let k = Tensor::randn(&[1, 1, 64, 64], &device).unwrap();
    let v = Tensor::randn(&[1, 1, 64, 64], &device).unwrap();

    // Test forward pass
    let result = flash_attention_forward(&q, &k, &v, &config);
    assert!(result.is_ok(), "Flash attention forward should succeed");

    let output = result.unwrap();
    assert_eq!(output.shape().dims(), &[1, 1, 64, 64]);

    println!("✓ Basic Flash Attention forward test passed");
}

/// Basic test to ensure Flash Attention backward pass compiles and runs
#[test]
fn test_basic_flash_attention_backward() {
    // Skip test if CUDA is not available
    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(1, 64);
    config.block_size_q = 32;
    config.block_size_kv = 32;
    config.causal = false;

    let q = Tensor::randn(&[1, 1, 64, 64], &device).unwrap();
    let k = Tensor::randn(&[1, 1, 64, 64], &device).unwrap();
    let v = Tensor::randn(&[1, 1, 64, 64], &device).unwrap();
    let grad_output = Tensor::randn(&[1, 1, 64, 64], &device).unwrap();

    // Test backward pass
    let result = flash_attention_backward(&grad_output, &q, &k, &v, &config);
    assert!(result.is_ok(), "Flash attention backward should succeed");

    let (grad_q, grad_k, grad_v) = result.unwrap();
    assert_eq!(grad_q.shape().dims(), &[1, 1, 64, 64]);
    assert_eq!(grad_k.shape().dims(), &[1, 1, 64, 64]);
    assert_eq!(grad_v.shape().dims(), &[1, 1, 64, 64]);

    println!("✓ Basic Flash Attention backward test passed");
}

/// Test quantum variant basic compilation
#[test]
fn test_quantum_variant_compilation() {
    use rtx_flash_attention::variants::quantum::*;

    // Skip test if CUDA is not available
    let device = match Device::try_default() {
        Ok(Device::Cuda(cuda_device)) => Device::Cuda(cuda_device),
        _ => {
            println!("CUDA not available, skipping test");
            return;
        }
    };

    let q = Tensor::randn(&[1, 1, 32, 32], &device).unwrap();
    let k = Tensor::randn(&[1, 1, 32, 32], &device).unwrap();
    let v = Tensor::randn(&[1, 1, 32, 32], &device).unwrap();

    // Test that quantum function exists and can be called
    let result = flash_attention_quantum(&q, &k, &v);
    assert!(
        result.is_ok(),
        "Quantum flash attention should succeed with basic implementation"
    );

    let output = result.unwrap();
    assert_eq!(output.shape().dims(), &[1, 1, 32, 32]);

    println!("✓ Quantum Flash Attention compilation test passed");
}

/// Test neuromorphic variant compilation
#[test]
fn test_neuromorphic_variant_compilation() {
    use rtx_flash_attention::variants::neuromorphic::*;

    // Skip test if CUDA is not available
    let device = match Device::try_default() {
        Ok(Device::Cuda(cuda_device)) => Device::Cuda(cuda_device),
        _ => {
            println!("CUDA not available, skipping test");
            return;
        }
    };

    let q = Tensor::randn(&[1, 1, 32, 32], &device).unwrap();
    let k = Tensor::randn(&[1, 1, 32, 32], &device).unwrap();
    let v = Tensor::randn(&[1, 1, 32, 32], &device).unwrap();

    let config = NeuromorphicConfig {
        spike_threshold: 0.5,
        membrane_potential_decay: 0.9,
        synaptic_plasticity: true,
    };

    // Test that neuromorphic function exists but returns error (not implemented yet)
    let result = flash_attention_neuromorphic(&q, &k, &v, &config);
    assert!(
        result.is_err(),
        "Neuromorphic flash attention should return error until implemented"
    );

    println!("✓ Neuromorphic Flash Attention compilation test passed");
}

/// Test different sequence lengths
#[test]
fn test_different_sequence_lengths() {
    // Skip test if CUDA is not available
    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(1, 64);
    config.block_size_q = 32;
    config.block_size_kv = 32;
    config.causal = false;

    let seq_lengths = vec![32, 64, 128];

    for seq_len in seq_lengths {
        let q = Tensor::randn(&[1, 1, seq_len, 64], &device).unwrap();
        let k = Tensor::randn(&[1, 1, seq_len, 64], &device).unwrap();
        let v = Tensor::randn(&[1, 1, seq_len, 64], &device).unwrap();

        let result = flash_attention_forward(&q, &k, &v, &config);
        assert!(
            result.is_ok(),
            "Should work for sequence length {}",
            seq_len
        );

        let output = result.unwrap();
        assert_eq!(output.shape().dims(), &[1, 1, seq_len, 64]);
    }

    println!("✓ Different sequence lengths test passed");
}

/// Test causal attention
#[test]
fn test_causal_attention() {
    // Skip test if CUDA is not available
    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(1, 64);
    config.block_size_q = 32;
    config.block_size_kv = 32;
    config.causal = true; // Enable causal masking

    let q = Tensor::randn(&[1, 1, 64, 64], &device).unwrap();
    let k = Tensor::randn(&[1, 1, 64, 64], &device).unwrap();
    let v = Tensor::randn(&[1, 1, 64, 64], &device).unwrap();

    let result = flash_attention_forward(&q, &k, &v, &config);
    assert!(result.is_ok(), "Causal attention should work");

    let output = result.unwrap();
    assert_eq!(output.shape().dims(), &[1, 1, 64, 64]);

    println!("✓ Causal attention test passed");
}