Initial commit

crates/training/rtx-flash-attention/src/kernels/metal_test.rs

@@ -0,0 +1,204 @@
+//! Tests for Metal Flash Attention implementation
+//!
+//! Following strict TDD - tests verify Metal backend behavior
+
+#![cfg(all(test, feature = "metal", target_os = "macos"))]
+
+use super::metal::{FlashMetalKernels, MetalKernelResult};
+use crate::config::FlashAttentionConfig;
+use crate::error::FlashResult;
+use objc2_metal::MTLCreateSystemDefaultDevice;
+
+/// Test: Metal device availability on macOS
+#[test]
+fn test_metal_device_available() {
+    let device = unsafe { MTLCreateSystemDefaultDevice() };
+    assert!(device.is_some(), "Metal device must be available on macOS");
+}
+
+/// Test: FlashMetalKernels creation succeeds
+#[test]
+fn test_metal_kernels_creation() {
+    let config = FlashAttentionConfig::new(8, 64);
+    let result = FlashMetalKernels::new(&config);
+
+    assert!(result.is_ok(), "Metal kernels creation should succeed");
+
+    let kernels = result.unwrap();
+    assert_eq!(kernels.config().num_heads, 8);
+    assert_eq!(kernels.config().head_dim, 64);
+}
+
+/// Test: Device info retrieval
+#[test]
+fn test_get_device_info() {
+    let config = FlashAttentionConfig::new(8, 64);
+    let kernels = FlashMetalKernels::new(&config).unwrap();
+
+    let device_info = kernels.get_device_info();
+    assert!(device_info.is_ok());
+
+    let info_str = device_info.unwrap();
+    assert!(info_str.contains("Metal Device"));
+}
+
+/// Test: Block size optimization
+#[test]
+fn test_optimize_block_sizes() {
+    let config = FlashAttentionConfig::new(8, 64);
+    let kernels = FlashMetalKernels::new(&config).unwrap();
+
+    let test_cases = vec![
+        (512, 64, 128, 128),   // seq_len=512 -> should be clamped to 128
+        (1024, 128, 128, 128), // seq_len=1024 -> should be 128
+        (256, 64, 128, 128),   // seq_len=256 -> should be 128
+        (64, 64, 64, 64),      // seq_len=64 -> should be 64
+    ];
+
+    for (seq_len, head_dim, expected_q, expected_kv) in test_cases {
+        let (block_q, block_kv) = kernels.optimize_block_sizes(seq_len, head_dim);
+
+        // Validate block sizes are within reasonable bounds
+        assert!(block_q >= 32, "Block Q must be at least 32");
+        assert!(block_q <= 256, "Block Q must not exceed 256");
+        assert!(block_kv >= 32, "Block KV must be at least 32");
+        assert!(block_kv <= 256, "Block KV must not exceed 256");
+
+        // Check alignment to warp size (32 for Metal)
+        assert_eq!(block_q % 32, 0, "Block Q should be aligned to 32");
+        assert_eq!(block_kv % 32, 0, "Block KV should be aligned to 32");
+    }
+}
+
+/// Test: Multiple kernel instances can coexist
+#[test]
+fn test_multiple_kernel_instances() {
+    let config1 = FlashAttentionConfig::new(8, 64);
+    let config2 = FlashAttentionConfig::new(16, 128);
+
+    let kernels1 = FlashMetalKernels::new(&config1);
+    let kernels2 = FlashMetalKernels::new(&config2);
+
+    assert!(kernels1.is_ok());
+    assert!(kernels2.is_ok());
+
+    let k1 = kernels1.unwrap();
+    let k2 = kernels2.unwrap();
+
+    assert_eq!(k1.config().num_heads, 8);
+    assert_eq!(k2.config().num_heads, 16);
+    assert_eq!(k1.config().head_dim, 64);
+    assert_eq!(k2.config().head_dim, 128);
+}
+
+/// Test: Forward pass stub returns expected structure
+#[tokio::test]
+async fn test_forward_pass_stub_structure() {
+    use rtx_tensor::{Tensor, Device};
+
+    let config = FlashAttentionConfig::new(8, 64);
+    let kernels = FlashMetalKernels::new(&config).unwrap();
+
+    let device = Device::Metal(0); // Assuming Metal device 0
+    let batch_size = 2;
+    let num_heads = 8;
+    let seq_len = 128;
+    let head_dim = 64;
+
+    // Create test tensors (would be on Metal device in production)
+    let q = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
+    let k = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
+    let v = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
+    let mut output = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
+    let mut lse = Tensor::zeros(&[batch_size, num_heads, seq_len], &device).unwrap();
+
+    let result = kernels.flash_attention_forward(
+        &q, &k, &v, &mut output, &mut lse, 0.125, false
+    ).await;
+
+    // Forward pass should succeed with real execution metrics
+    assert!(result.is_ok(), "Forward pass should succeed");
+
+    let kernel_result = result.unwrap();
+    // With fixed threadgroup memory (32x32 blocks = 28KB < 32KB limit), kernel should execute
+    assert!(kernel_result.execution_time_us > 0, "Kernel should have executed (time > 0)");
+    assert!(kernel_result.memory_throughput >= 0.0, "Memory throughput should be non-negative");
+}
+
+/// Test: Backward pass execution (with 16x16 blocks = 20KB < 32KB limit)
+#[tokio::test]
+async fn test_backward_pass_execution() {
+    use rtx_tensor::{Tensor, Device};
+
+    let config = FlashAttentionConfig::new(8, 64);
+    let kernels = FlashMetalKernels::new(&config).unwrap();
+
+    let device = Device::Metal(0);
+    let batch_size = 2;
+    let num_heads = 8;
+    let seq_len = 128;
+    let head_dim = 64;
+
+    let dout = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
+    let q = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
+    let k = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
+    let v = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
+    let output = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
+    let lse = Tensor::zeros(&[batch_size, num_heads, seq_len], &device).unwrap();
+    let mut grad_q = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
+    let mut grad_k = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
+    let mut grad_v = Tensor::zeros(&[batch_size, num_heads, seq_len, head_dim], &device).unwrap();
+
+    let result = kernels.flash_attention_backward(
+        &dout, &q, &k, &v, &output, &lse,
+        &mut grad_q, &mut grad_k, &mut grad_v,
+        0.125, false
+    ).await;
+
+    // Backward pass should succeed with 16x16 blocks (20KB < 32KB limit)
+    assert!(result.is_ok(), "Backward pass should succeed with fixed block sizes");
+
+    let kernel_result = result.unwrap();
+    assert!(kernel_result.execution_time_us > 0, "Backward kernel should have executed");
+}
+
+/// Test: Configuration validation
+#[test]
+fn test_config_validation() {
+    let valid_configs = vec![
+        FlashAttentionConfig::new(8, 64),
+        FlashAttentionConfig::new(16, 128),
+        FlashAttentionConfig::new(32, 256),
+    ];
+
+    for config in valid_configs {
+        let result = FlashMetalKernels::new(&config);
+        assert!(result.is_ok(), "Valid config should create kernels successfully");
+    }
+}
+
+/// Test: Thread safety - kernels can be shared across threads
+#[test]
+fn test_thread_safety() {
+    use std::sync::Arc;
+    use std::thread;
+
+    let config = FlashAttentionConfig::new(8, 64);
+    let kernels = Arc::new(FlashMetalKernels::new(&config).unwrap());
+
+    let mut handles = vec![];
+
+    for _ in 0..4 {
+        let kernels_clone = Arc::clone(&kernels);
+        let handle = thread::spawn(move || {
+            let (block_q, block_kv) = kernels_clone.optimize_block_sizes(1024, 64);
+            assert!(block_q >= 32 && block_q <= 256);
+            assert!(block_kv >= 32 && block_kv <= 256);
+        });
+        handles.push(handle);
+    }
+
+    for handle in handles {
+        handle.join().unwrap();
+    }
+}