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

//! Test Flash Attention with float16 (half precision)
//! Run with: cargo run -p rtx-flash-metal-attention --example f16_test --release

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

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

    // Test parameters
    let batch_size = 2;
    let num_heads = 4;
    let seq_len = 128;
    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];

    // ===== Test Float32 (f32) =====
    println!("\n--- Float32 (f32) Mode ---");
    let config_f32 = FlashAttentionConfig::default();
    println!(
        "  Config: block_q={}, block_kv={}, use_f16={}",
        config_f32.block_q, config_f32.block_kv, config_f32.use_f16
    );

    let attn_f32 = match FlashAttention::new(config_f32) {
        Ok(a) => a,
        Err(e) => {
            eprintln!("Failed to create f32 FlashAttention: {:?}", e);
            return;
        }
    };

    let q_f32 = Tensor::ones_typed(shape, DType::F32, &metal_device).expect("Q f32");
    let k_f32 = Tensor::ones_typed(shape, DType::F32, &metal_device).expect("K f32");
    let v_f32 = Tensor::ones_typed(shape, DType::F32, &metal_device).expect("V f32");

    println!("  Q dtype: {:?}", q_f32.dtype());

    // Warmup
    for _ in 0..5 {
        let _ = attn_f32.forward(&q_f32, &k_f32, &v_f32);
    }

    // Benchmark f32
    let iterations = 50;
    let start = Instant::now();
    for _ in 0..iterations {
        let _ = attn_f32.forward(&q_f32, &k_f32, &v_f32);
    }
    let f32_time = start.elapsed().as_secs_f64() * 1000.0 / iterations as f64;
    println!("  Average time: {:.4} ms", f32_time);

    // ===== Test Float16 (f16) =====
    println!("\n--- Float16 (f16) Mode ---");
    let config_f16 = FlashAttentionConfig::f16();
    println!(
        "  Config: block_q={}, block_kv={}, use_f16={}",
        config_f16.block_q, config_f16.block_kv, config_f16.use_f16
    );

    let attn_f16 = match FlashAttention::new(config_f16) {
        Ok(a) => a,
        Err(e) => {
            eprintln!("Failed to create f16 FlashAttention: {:?}", e);
            return;
        }
    };

    // Now using ones_typed since fill() supports F16 on Metal
    let q_f16 = Tensor::ones_typed(shape, DType::F16, &metal_device).expect("Q f16");
    let k_f16 = Tensor::ones_typed(shape, DType::F16, &metal_device).expect("K f16");
    let v_f16 = Tensor::ones_typed(shape, DType::F16, &metal_device).expect("V f16");

    println!("  Q dtype: {:?}", q_f16.dtype());

    // Warmup
    for _ in 0..5 {
        let _ = attn_f16.forward(&q_f16, &k_f16, &v_f16);
    }

    // Benchmark f16
    let start = Instant::now();
    for _ in 0..iterations {
        let _ = attn_f16.forward(&q_f16, &k_f16, &v_f16);
    }
    let f16_time = start.elapsed().as_secs_f64() * 1000.0 / iterations as f64;
    println!("  Average time: {:.4} ms", f16_time);

    // ===== Test Float16 Causal =====
    println!("\n--- Float16 Causal Mode ---");
    let config_f16_causal = FlashAttentionConfig::f16_causal();
    println!(
        "  Config: block_q={}, block_kv={}, use_f16={}, causal={}",
        config_f16_causal.block_q,
        config_f16_causal.block_kv,
        config_f16_causal.use_f16,
        config_f16_causal.causal
    );

    let attn_f16_causal = match FlashAttention::new(config_f16_causal) {
        Ok(a) => a,
        Err(e) => {
            eprintln!("Failed to create f16 causal FlashAttention: {:?}", e);
            return;
        }
    };

    // Warmup
    for _ in 0..5 {
        let _ = attn_f16_causal.forward(&q_f16, &k_f16, &v_f16);
    }

    // Benchmark f16 causal
    let start = Instant::now();
    for _ in 0..iterations {
        let _ = attn_f16_causal.forward(&q_f16, &k_f16, &v_f16);
    }
    let f16_causal_time = start.elapsed().as_secs_f64() * 1000.0 / iterations as f64;
    println!("  Average time: {:.4} ms", f16_causal_time);

    // ===== Summary =====
    println!("\n--- Performance Summary ---");
    println!("  f32 forward:       {:.4} ms", f32_time);
    println!(
        "  f16 forward:       {:.4} ms ({})",
        f16_time,
        if f16_time < f32_time {
            format!("{:.1}x faster", f32_time / f16_time)
        } else {
            format!("{:.1}x slower", f16_time / f32_time)
        }
    );
    println!("  f16 causal:        {:.4} ms", f16_causal_time);

    // Memory savings
    let f32_bytes = (batch_size * num_heads * seq_len * head_dim * 4) as f64 / 1024.0 / 1024.0;
    let f16_bytes = (batch_size * num_heads * seq_len * head_dim * 2) as f64 / 1024.0 / 1024.0;
    println!("\n--- Memory Usage (Q+K+V) ---");
    println!("  f32: {:.2} MB", f32_bytes * 3.0);
    println!(
        "  f16: {:.2} MB ({:.0}% reduction)",
        f16_bytes * 3.0,
        (1.0 - f16_bytes / f32_bytes) * 100.0
    );

    // ===== Test F16 Backward Pass =====
    println!("\n--- Float16 Backward Pass Test ---");

    // Forward pass to get output and LSE
    let (output_f16, lse_f16) = match attn_f16.forward(&q_f16, &k_f16, &v_f16) {
        Ok((o, l)) => {
            println!("  Forward pass successful");
            println!("    Output dtype: {:?}", o.dtype());
            println!("    LSE dtype: {:?}", l.dtype());
            (o, l)
        }
        Err(e) => {
            eprintln!("  Forward pass failed: {:?}", e);
            println!("\n=== Test Complete ===");
            return;
        }
    };

    // Create grad_output (same shape as output)
    let grad_output_f16 =
        Tensor::zeros_typed(shape, DType::F16, &metal_device).expect("grad_output f16");

    // Backward pass
    match attn_f16.backward(
        &grad_output_f16,
        &q_f16,
        &k_f16,
        &v_f16,
        &output_f16,
        &lse_f16,
    ) {
        Ok((dq, dk, dv)) => {
            println!("  Backward pass successful!");
            println!("    dQ dtype: {:?}, shape: {:?}", dq.dtype(), dq.shape());
            println!("    dK dtype: {:?}, shape: {:?}", dk.dtype(), dk.shape());
            println!("    dV dtype: {:?}, shape: {:?}", dv.dtype(), dv.shape());

            // Verify shapes match
            let shapes_match = dq.shape() == q_f16.shape()
                && dk.shape() == k_f16.shape()
                && dv.shape() == v_f16.shape();
            println!("    Shapes match inputs: {}", shapes_match);
        }
        Err(e) => {
            eprintln!("  Backward pass failed: {:?}", e);
        }
    }

    // Benchmark f16 backward
    println!("\n--- F16 Backward Performance ---");
    let (output_bench, lse_bench) = attn_f16.forward(&q_f16, &k_f16, &v_f16).expect("forward");

    // Warmup
    for _ in 0..5 {
        let _ = attn_f16.backward(
            &grad_output_f16,
            &q_f16,
            &k_f16,
            &v_f16,
            &output_bench,
            &lse_bench,
        );
    }

    let start = Instant::now();
    for _ in 0..iterations {
        let _ = attn_f16.backward(
            &grad_output_f16,
            &q_f16,
            &k_f16,
            &v_f16,
            &output_bench,
            &lse_bench,
        );
    }
    let f16_backward_time = start.elapsed().as_secs_f64() * 1000.0 / iterations as f64;
    println!("  f16 backward: {:.4} ms", f16_backward_time);
    println!(
        "  Backward/Forward ratio: {:.2}x",
        f16_backward_time / f16_time
    );

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