rustytorch/crates/training/rtx-flash-metal-attention/examples/backward_test.rs

//! Test Flash Attention backward pass
//! Run with: cargo run -p rtx-flash-metal-attention --example backward_test --release

use rtx_flash_metal_attention::{FlashAttention, FlashAttentionConfig};
use rtx_tensor::{Device, Tensor};
use std::time::Instant;

fn main() {
    println!("=== Flash Attention Backward Pass Test ===\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;
        }
    };

    // Create FlashAttention
    println!("Creating FlashAttention...");
    let config = FlashAttentionConfig::default();
    let attn = match FlashAttention::new(config) {
        Ok(a) => {
            println!("FlashAttention created successfully!");
            a
        }
        Err(e) => {
            eprintln!("Failed to create FlashAttention: {:?}", e);
            return;
        }
    };

    // Test parameters
    let batch_size = 2;
    let num_heads = 4;
    let seq_len = 64;
    let head_dim = 64;

    println!("\nTest configuration:");
    println!("  Batch size: {}", batch_size);
    println!("  Num heads: {}", num_heads);
    println!("  Sequence length: {}", seq_len);
    println!("  Head dimension: {}", head_dim);

    let shape = &[batch_size, num_heads, seq_len, head_dim];

    // Create input tensors
    println!("\nCreating tensors...");
    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");

    // Forward pass
    println!("\n--- Forward Pass ---");
    let start = Instant::now();
    let (output, lse) = match attn.forward(&q, &k, &v) {
        Ok((o, l)) => {
            println!("Forward pass successful!");
            println!("  Output shape: {:?}", o.shape());
            println!("  LSE shape: {:?}", l.shape());
            (o, l)
        }
        Err(e) => {
            eprintln!("Forward pass failed: {:?}", e);
            return;
        }
    };
    let forward_time = start.elapsed();
    println!(
        "  Forward time: {:.3}ms",
        forward_time.as_secs_f64() * 1000.0
    );

    // Create gradient output (same shape as output)
    let grad_output = Tensor::ones(shape, &metal_device).expect("grad_output tensor");

    // Backward pass
    println!("\n--- Backward Pass ---");
    let start = Instant::now();
    let (dq, dk, dv) = match attn.backward(&grad_output, &q, &k, &v, &output, &lse) {
        Ok((dq, dk, dv)) => {
            println!("Backward pass successful!");
            println!("  dQ shape: {:?}", dq.shape());
            println!("  dK shape: {:?}", dk.shape());
            println!("  dV shape: {:?}", dv.shape());
            (dq, dk, dv)
        }
        Err(e) => {
            eprintln!("Backward pass failed: {:?}", e);
            return;
        }
    };
    let backward_time = start.elapsed();
    println!(
        "  Backward time: {:.3}ms",
        backward_time.as_secs_f64() * 1000.0
    );

    // Verify gradient shapes match input shapes
    println!("\n--- Shape Verification ---");
    let q_shape = q.shape();
    let k_shape = k.shape();
    let v_shape = v.shape();

    let dq_matches = dq.shape() == q_shape;
    let dk_matches = dk.shape() == k_shape;
    let dv_matches = dv.shape() == v_shape;

    println!(
        "  dQ shape matches Q: {} ({:?} == {:?})",
        dq_matches,
        dq.shape(),
        q_shape
    );
    println!(
        "  dK shape matches K: {} ({:?} == {:?})",
        dk_matches,
        dk.shape(),
        k_shape
    );
    println!(
        "  dV shape matches V: {} ({:?} == {:?})",
        dv_matches,
        dv.shape(),
        v_shape
    );

    if dq_matches && dk_matches && dv_matches {
        println!("\n✓ All gradient shapes are correct!");
    } else {
        println!("\n✗ Shape mismatch detected!");
    }

    // Benchmark multiple iterations
    println!("\n--- Performance Benchmark ---");
    let iterations = 50;

    // Warmup
    for _ in 0..10 {
        let _ = attn.forward(&q, &k, &v);
        let _ = attn.backward(&grad_output, &q, &k, &v, &output, &lse);
    }

    // Forward benchmark
    let start = Instant::now();
    for _ in 0..iterations {
        let _ = attn.forward(&q, &k, &v);
    }
    let avg_forward = start.elapsed().as_secs_f64() * 1000.0 / iterations as f64;

    // Backward benchmark
    let start = Instant::now();
    for _ in 0..iterations {
        let _ = attn.backward(&grad_output, &q, &k, &v, &output, &lse);
    }
    let avg_backward = start.elapsed().as_secs_f64() * 1000.0 / iterations as f64;

    println!("  Average forward time: {:.4}ms", avg_forward);
    println!("  Average backward time: {:.4}ms", avg_backward);
    println!(
        "  Backward/Forward ratio: {:.2}x",
        avg_backward / avg_forward
    );

    let total_elements = (batch_size * num_heads * seq_len * head_dim) as f64;
    let forward_throughput = total_elements / (avg_forward / 1000.0) / 1e6;
    let backward_throughput = total_elements / (avg_backward / 1000.0) / 1e6;

    println!(
        "\n  Forward throughput: {:.2} M elements/s",
        forward_throughput
    );
    println!(
        "  Backward throughput: {:.2} M elements/s",
        backward_throughput
    );

    println!("\n=== Test Complete ===");
}