Initial commit

crates/training/rtx-flash-attention/src/backend_selector.rs

@@ -0,0 +1,912 @@
+//! SDPA Backend Auto-Selection for RustyTorch++
+//!
+//! Automatically selects the optimal Scaled Dot-Product Attention backend based on:
+//! - Hardware capabilities (GPU architecture, Tensor Cores, memory)
+//! - Input characteristics (sequence length, head dimension, batch size)
+//! - User preferences and constraints
+//!
+//! # Backends
+//!
+//! 1. **FlashAttention**: Optimal for long sequences, uses tiling for O(1) memory
+//! 2. **Math**: Standard attention, best for small sequences or debugging
+//! 3. **Memory-Efficient**: Chunked attention, trades compute for memory
+//! 4. **CuDNN**: NVIDIA's optimized implementation (when available)
+//!
+//! # Usage
+//!
+//! ```rust,ignore
+//! use rtx_flash_attention::backend_selector::{SdpaBackendSelector, SdpaBackend};
+//!
+//! let selector = SdpaBackendSelector::new(SdpaConfig::default())?;
+//!
+//! // Automatic selection
+//! let backend = selector.select(&query, &key, &value)?;
+//!
+//! // Or get recommendation
+//! let recommendation = selector.recommend(seq_len, head_dim, batch_size)?;
+//! println!("Recommended backend: {:?}", recommendation.backend);
+//! println!("Expected speedup: {:.2}x", recommendation.expected_speedup);
+//! ```
+
+use std::collections::HashMap;
+
+// =============================================================================
+// Backends
+// =============================================================================
+
+/// Available SDPA backends
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum SdpaBackend {
+    /// FlashAttention v2 - optimal for long sequences
+    FlashAttention,
+    /// Standard mathematical attention - simple, debuggable
+    Math,
+    /// Memory-efficient chunked attention
+    MemoryEfficient,
+    /// NVIDIA cuDNN attention (Ampere+)
+    CuDnn,
+    /// Custom Metal implementation for Apple Silicon
+    Metal,
+    /// Fallback CPU implementation
+    Cpu,
+}
+
+impl std::fmt::Display for SdpaBackend {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            SdpaBackend::FlashAttention => write!(f, "FlashAttention"),
+            SdpaBackend::Math => write!(f, "Math"),
+            SdpaBackend::MemoryEfficient => write!(f, "MemoryEfficient"),
+            SdpaBackend::CuDnn => write!(f, "cuDNN"),
+            SdpaBackend::Metal => write!(f, "Metal"),
+            SdpaBackend::Cpu => write!(f, "CPU"),
+        }
+    }
+}
+
+// =============================================================================
+// Hardware Detection
+// =============================================================================
+
+/// Detected hardware capabilities
+#[derive(Debug, Clone)]
+pub struct HardwareCapabilities {
+    /// Device type
+    pub device_type: DeviceType,
+    /// GPU compute capability (e.g., 8.0 for A100)
+    pub compute_capability: Option<(u32, u32)>,
+    /// Total device memory (bytes)
+    pub total_memory: usize,
+    /// Available device memory (bytes)
+    pub available_memory: usize,
+    /// Has Tensor Cores
+    pub has_tensor_cores: bool,
+    /// Has FP16 support
+    pub has_fp16: bool,
+    /// Has BF16 support
+    pub has_bf16: bool,
+    /// Has FP8 support
+    pub has_fp8: bool,
+    /// Number of SMs (CUDA) or compute units
+    pub num_compute_units: u32,
+    /// Memory bandwidth (GB/s)
+    pub memory_bandwidth_gbps: f32,
+    /// Supports FlashAttention
+    pub supports_flash_attention: bool,
+    /// Supports cuDNN attention
+    pub supports_cudnn_attention: bool,
+}
+
+/// Device type
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum DeviceType {
+    /// CPU
+    Cpu,
+    /// NVIDIA CUDA GPU
+    Cuda,
+    /// AMD ROCm GPU
+    Rocm,
+    /// Apple Metal
+    Metal,
+    /// Intel XPU
+    Xpu,
+}
+
+impl HardwareCapabilities {
+    /// Detect capabilities for CPU
+    pub fn detect_cpu() -> Self {
+        Self {
+            device_type: DeviceType::Cpu,
+            compute_capability: None,
+            total_memory: 0,
+            available_memory: 0,
+            has_tensor_cores: false,
+            has_fp16: true,
+            has_bf16: false,
+            has_fp8: false,
+            num_compute_units: std::thread::available_parallelism().map(std::num::NonZero::get).unwrap_or(4) as u32,
+            memory_bandwidth_gbps: 50.0, // Typical DDR5
+            supports_flash_attention: false,
+            supports_cudnn_attention: false,
+        }
+    }
+
+    /// Detect capabilities for CUDA device
+    pub fn detect_cuda(device_id: usize) -> Self {
+        // In production, this would query CUDA runtime
+        // For now, simulate A100 capabilities
+        Self {
+            device_type: DeviceType::Cuda,
+            compute_capability: Some((8, 0)), // A100
+            total_memory: 80 * 1024 * 1024 * 1024, // 80GB
+            available_memory: 70 * 1024 * 1024 * 1024,
+            has_tensor_cores: true,
+            has_fp16: true,
+            has_bf16: true,
+            has_fp8: false, // Only Hopper+
+            num_compute_units: 108, // A100 SMs
+            memory_bandwidth_gbps: 2039.0, // A100 HBM2e
+            supports_flash_attention: true,
+            supports_cudnn_attention: true,
+        }
+    }
+
+    /// Detect capabilities for Metal device
+    pub fn detect_metal() -> Self {
+        Self {
+            device_type: DeviceType::Metal,
+            compute_capability: None,
+            total_memory: 32 * 1024 * 1024 * 1024, // 32GB unified
+            available_memory: 28 * 1024 * 1024 * 1024,
+            has_tensor_cores: false,
+            has_fp16: true,
+            has_bf16: false,
+            has_fp8: false,
+            num_compute_units: 40, // M2 Max GPU cores
+            memory_bandwidth_gbps: 400.0, // M2 Max
+            supports_flash_attention: true, // Custom MSL implementation
+            supports_cudnn_attention: false,
+        }
+    }
+
+    /// Check if hardware supports a specific backend
+    pub fn supports_backend(&self, backend: SdpaBackend) -> bool {
+        match backend {
+            SdpaBackend::FlashAttention => self.supports_flash_attention,
+            SdpaBackend::Math => true, // Always supported
+            SdpaBackend::MemoryEfficient => true,
+            SdpaBackend::CuDnn => self.supports_cudnn_attention,
+            SdpaBackend::Metal => self.device_type == DeviceType::Metal,
+            SdpaBackend::Cpu => true,
+        }
+    }
+}
+
+// =============================================================================
+// Input Characteristics
+// =============================================================================
+
+/// Characteristics of the attention input
+#[derive(Debug, Clone)]
+pub struct AttentionInputInfo {
+    /// Batch size
+    pub batch_size: usize,
+    /// Number of attention heads
+    pub num_heads: usize,
+    /// Sequence length (query)
+    pub seq_len_q: usize,
+    /// Sequence length (key/value)
+    pub seq_len_kv: usize,
+    /// Head dimension
+    pub head_dim: usize,
+    /// Data type
+    pub dtype: AttentionDType,
+    /// Is causal attention
+    pub is_causal: bool,
+    /// Has attention mask
+    pub has_mask: bool,
+    /// Dropout probability
+    pub dropout: f32,
+}
+
+/// Data types for attention
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum AttentionDType {
+    Float32,
+    Float16,
+    BFloat16,
+    Float8E4M3,
+    Float8E5M2,
+}
+
+impl AttentionInputInfo {
+    /// Estimate memory required for this attention operation
+    pub fn estimate_memory_bytes(&self) -> usize {
+        let bytes_per_element = match self.dtype {
+            AttentionDType::Float32 => 4,
+            AttentionDType::Float16 | AttentionDType::BFloat16 => 2,
+            AttentionDType::Float8E4M3 | AttentionDType::Float8E5M2 => 1,
+        };
+
+        // Q, K, V, Output, Attention weights
+        let qkvo_size = self.batch_size * self.num_heads * self.seq_len_q * self.head_dim * bytes_per_element * 4;
+
+        // Attention scores (for math backend)
+        let attn_size = self.batch_size * self.num_heads * self.seq_len_q * self.seq_len_kv * bytes_per_element;
+
+        qkvo_size + attn_size
+    }
+
+    /// Estimate FLOPs for this attention operation
+    pub fn estimate_flops(&self) -> usize {
+        // QK^T: 2 * B * H * N * N * D
+        let qk_flops = 2 * self.batch_size * self.num_heads * self.seq_len_q * self.seq_len_kv * self.head_dim;
+
+        // Softmax: ~5 * B * H * N * N
+        let softmax_flops = 5 * self.batch_size * self.num_heads * self.seq_len_q * self.seq_len_kv;
+
+        // AttnV: 2 * B * H * N * N * D
+        let av_flops = 2 * self.batch_size * self.num_heads * self.seq_len_q * self.seq_len_kv * self.head_dim;
+
+        qk_flops + softmax_flops + av_flops
+    }
+}
+
+// =============================================================================
+// Selection Configuration
+// =============================================================================
+
+/// Configuration for backend selection
+#[derive(Debug, Clone)]
+pub struct SdpaConfig {
+    /// Preferred backend (if supported)
+    pub preferred_backend: Option<SdpaBackend>,
+    /// Disabled backends
+    pub disabled_backends: Vec<SdpaBackend>,
+    /// Enable auto-tuning
+    pub enable_autotuning: bool,
+    /// Memory budget (bytes, 0 = no limit)
+    pub memory_budget: usize,
+    /// Optimize for latency vs throughput
+    pub optimize_for: OptimizeFor,
+    /// Minimum sequence length for FlashAttention
+    pub flash_min_seq_len: usize,
+    /// Enable deterministic mode
+    pub deterministic: bool,
+    /// Enable debug mode (use Math backend)
+    pub debug_mode: bool,
+}
+
+/// Optimization target
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum OptimizeFor {
+    /// Minimize latency (single request)
+    Latency,
+    /// Maximize throughput (batch processing)
+    Throughput,
+    /// Balance both
+    Balanced,
+}
+
+impl Default for SdpaConfig {
+    fn default() -> Self {
+        Self {
+            preferred_backend: None,
+            disabled_backends: Vec::new(),
+            enable_autotuning: true,
+            memory_budget: 0,
+            optimize_for: OptimizeFor::Balanced,
+            flash_min_seq_len: 128,
+            deterministic: false,
+            debug_mode: false,
+        }
+    }
+}
+
+// =============================================================================
+// Backend Recommendation
+// =============================================================================
+
+/// Recommendation for which backend to use
+#[derive(Debug, Clone)]
+pub struct BackendRecommendation {
+    /// Recommended backend
+    pub backend: SdpaBackend,
+    /// Confidence score (0.0 to 1.0)
+    pub confidence: f32,
+    /// Expected speedup vs Math backend
+    pub expected_speedup: f32,
+    /// Expected memory usage (bytes)
+    pub expected_memory: usize,
+    /// Reason for recommendation
+    pub reason: String,
+    /// Alternative backends ranked
+    pub alternatives: Vec<(SdpaBackend, f32)>,
+}
+
+// =============================================================================
+// Backend Selector
+// =============================================================================
+
+/// SDPA Backend Selector
+pub struct SdpaBackendSelector {
+    /// Configuration
+    config: SdpaConfig,
+    /// Detected hardware capabilities
+    hardware: HardwareCapabilities,
+    /// Cached recommendations
+    cache: std::sync::RwLock<HashMap<CacheKey, BackendRecommendation>>,
+    /// Performance history for tuning
+    perf_history: std::sync::RwLock<Vec<PerformanceRecord>>,
+}
+
+/// Cache key for recommendations
+#[derive(Debug, Clone, Hash, PartialEq, Eq)]
+struct CacheKey {
+    batch_size: usize,
+    num_heads: usize,
+    seq_len_q: usize,
+    seq_len_kv: usize,
+    head_dim: usize,
+    is_causal: bool,
+}
+
+/// Performance record for auto-tuning
+#[derive(Debug, Clone)]
+struct PerformanceRecord {
+    key: CacheKey,
+    backend: SdpaBackend,
+    latency_us: f64,
+    memory_bytes: usize,
+}
+
+impl SdpaBackendSelector {
+    /// Create a new selector with configuration
+    pub fn new(config: SdpaConfig) -> Self {
+        // Auto-detect hardware
+        let hardware = if cfg!(feature = "cuda") {
+            HardwareCapabilities::detect_cuda(0)
+        } else if cfg!(target_os = "macos") {
+            HardwareCapabilities::detect_metal()
+        } else {
+            HardwareCapabilities::detect_cpu()
+        };
+
+        Self {
+            config,
+            hardware,
+            cache: std::sync::RwLock::new(HashMap::new()),
+            perf_history: std::sync::RwLock::new(Vec::new()),
+        }
+    }
+
+    /// Create selector with specific hardware capabilities
+    pub fn with_hardware(config: SdpaConfig, hardware: HardwareCapabilities) -> Self {
+        Self {
+            config,
+            hardware,
+            cache: std::sync::RwLock::new(HashMap::new()),
+            perf_history: std::sync::RwLock::new(Vec::new()),
+        }
+    }
+
+    /// Get hardware capabilities
+    pub fn hardware(&self) -> &HardwareCapabilities {
+        &self.hardware
+    }
+
+    /// Select optimal backend for given input
+    pub fn select(&self, input: &AttentionInputInfo) -> BackendRecommendation {
+        // Check debug mode
+        if self.config.debug_mode {
+            return BackendRecommendation {
+                backend: SdpaBackend::Math,
+                confidence: 1.0,
+                expected_speedup: 1.0,
+                expected_memory: input.estimate_memory_bytes(),
+                reason: "Debug mode enabled".to_string(),
+                alternatives: vec![],
+            };
+        }
+
+        // Check preferred backend
+        if let Some(preferred) = self.config.preferred_backend {
+            if self.is_backend_available(preferred, input) {
+                return BackendRecommendation {
+                    backend: preferred,
+                    confidence: 1.0,
+                    expected_speedup: self.estimate_speedup(preferred, input),
+                    expected_memory: self.estimate_memory(preferred, input),
+                    reason: "User preferred backend".to_string(),
+                    alternatives: self.get_alternatives(input),
+                };
+            }
+        }
+
+        // Check cache
+        let cache_key = CacheKey {
+            batch_size: input.batch_size,
+            num_heads: input.num_heads,
+            seq_len_q: input.seq_len_q,
+            seq_len_kv: input.seq_len_kv,
+            head_dim: input.head_dim,
+            is_causal: input.is_causal,
+        };
+
+        {
+            let cache = self.cache.read().unwrap();
+            if let Some(rec) = cache.get(&cache_key) {
+                return rec.clone();
+            }
+        }
+
+        // Compute recommendation
+        let recommendation = self.compute_recommendation(input);
+
+        // Cache it
+        {
+            let mut cache = self.cache.write().unwrap();
+            cache.insert(cache_key, recommendation.clone());
+        }
+
+        recommendation
+    }
+
+    /// Check if a backend is available for the given input
+    fn is_backend_available(&self, backend: SdpaBackend, input: &AttentionInputInfo) -> bool {
+        // Check if disabled
+        if self.config.disabled_backends.contains(&backend) {
+            return false;
+        }
+
+        // Check hardware support
+        if !self.hardware.supports_backend(backend) {
+            return false;
+        }
+
+        // Backend-specific checks
+        match backend {
+            SdpaBackend::FlashAttention => {
+                // FlashAttention has constraints
+                input.seq_len_q >= self.config.flash_min_seq_len
+                    && input.head_dim <= 256
+                    && (input.head_dim == 32 || input.head_dim == 64 || input.head_dim == 128 || input.head_dim == 256)
+            }
+            SdpaBackend::CuDnn => {
+                // cuDNN constraints
+                input.head_dim <= 128 && !input.has_mask
+            }
+            _ => true,
+        }
+    }
+
+    /// Compute recommendation for input
+    fn compute_recommendation(&self, input: &AttentionInputInfo) -> BackendRecommendation {
+        let available_backends = self.get_available_backends(input);
+
+        if available_backends.is_empty() {
+            return BackendRecommendation {
+                backend: SdpaBackend::Math,
+                confidence: 0.5,
+                expected_speedup: 1.0,
+                expected_memory: input.estimate_memory_bytes(),
+                reason: "Fallback to Math (no optimized backend available)".to_string(),
+                alternatives: vec![],
+            };
+        }
+
+        // Score each backend
+        let mut scores: Vec<(SdpaBackend, f32, String)> = available_backends
+            .iter()
+            .map(|&b| {
+                let (score, reason) = self.score_backend(b, input);
+                (b, score, reason)
+            })
+            .collect();
+
+        // Sort by score (descending)
+        scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
+
+        let (best_backend, best_score, reason) = scores.first().cloned().unwrap();
+
+        let alternatives: Vec<(SdpaBackend, f32)> = scores
+            .iter()
+            .skip(1)
+            .map(|(b, s, _)| (*b, *s))
+            .collect();
+
+        BackendRecommendation {
+            backend: best_backend,
+            confidence: best_score,
+            expected_speedup: self.estimate_speedup(best_backend, input),
+            expected_memory: self.estimate_memory(best_backend, input),
+            reason,
+            alternatives,
+        }
+    }
+
+    /// Get available backends for input
+    fn get_available_backends(&self, input: &AttentionInputInfo) -> Vec<SdpaBackend> {
+        [
+            SdpaBackend::FlashAttention,
+            SdpaBackend::CuDnn,
+            SdpaBackend::MemoryEfficient,
+            SdpaBackend::Math,
+            SdpaBackend::Metal,
+            SdpaBackend::Cpu,
+        ]
+        .iter()
+        .filter(|&&b| self.is_backend_available(b, input))
+        .copied()
+        .collect()
+    }
+
+    /// Score a backend for given input
+    fn score_backend(&self, backend: SdpaBackend, input: &AttentionInputInfo) -> (f32, String) {
+        let seq_len = input.seq_len_q.max(input.seq_len_kv);
+        let memory_required = self.estimate_memory(backend, input);
+
+        // Check memory budget
+        if self.config.memory_budget > 0 && memory_required > self.config.memory_budget {
+            return (0.0, "Exceeds memory budget".to_string());
+        }
+
+        match backend {
+            SdpaBackend::FlashAttention => {
+                if seq_len >= 2048 {
+                    (0.95, format!("FlashAttention optimal for seq_len={}", seq_len))
+                } else if seq_len >= 512 {
+                    (0.85, "FlashAttention good for medium sequences".to_string())
+                } else {
+                    (0.6, "FlashAttention has overhead for short sequences".to_string())
+                }
+            }
+            SdpaBackend::CuDnn => {
+                if seq_len <= 512 && input.head_dim <= 64 {
+                    (0.9, "cuDNN optimal for small attention".to_string())
+                } else {
+                    (0.7, "cuDNN usable but not optimal".to_string())
+                }
+            }
+            SdpaBackend::MemoryEfficient => {
+                if memory_required > self.hardware.available_memory / 2 {
+                    (0.85, "Memory-efficient prevents OOM".to_string())
+                } else {
+                    (0.5, "Memory-efficient not needed".to_string())
+                }
+            }
+            SdpaBackend::Math => {
+                if seq_len <= 128 {
+                    (0.8, "Math backend efficient for tiny sequences".to_string())
+                } else {
+                    (0.3, "Math backend as fallback".to_string())
+                }
+            }
+            SdpaBackend::Metal => {
+                if self.hardware.device_type == DeviceType::Metal {
+                    (0.9, "Native Metal implementation".to_string())
+                } else {
+                    (0.0, "Metal not available".to_string())
+                }
+            }
+            SdpaBackend::Cpu => {
+                (0.1, "CPU fallback".to_string())
+            }
+        }
+    }
+
+    /// Estimate speedup for backend
+    fn estimate_speedup(&self, backend: SdpaBackend, input: &AttentionInputInfo) -> f32 {
+        let seq_len = input.seq_len_q.max(input.seq_len_kv);
+
+        match backend {
+            SdpaBackend::FlashAttention => {
+                if seq_len >= 4096 { 5.0 }
+                else if seq_len >= 2048 { 3.5 }
+                else if seq_len >= 1024 { 2.5 }
+                else if seq_len >= 512 { 1.8 }
+                else { 1.2 }
+            }
+            SdpaBackend::CuDnn => {
+                if seq_len <= 512 { 2.0 }
+                else { 1.5 }
+            }
+            SdpaBackend::MemoryEfficient => 1.3,
+            SdpaBackend::Metal => 2.5,
+            SdpaBackend::Math => 1.0,
+            SdpaBackend::Cpu => 0.1,
+        }
+    }
+
+    /// Estimate memory for backend
+    fn estimate_memory(&self, backend: SdpaBackend, input: &AttentionInputInfo) -> usize {
+        let base_memory = input.estimate_memory_bytes();
+
+        match backend {
+            SdpaBackend::FlashAttention => {
+                // FlashAttention uses O(N) instead of O(N^2) for attention matrix
+                let qkvo_size = base_memory / 2; // Q, K, V, O only
+                let softmax_lse = input.batch_size * input.num_heads * input.seq_len_q * 4;
+                qkvo_size + softmax_lse
+            }
+            SdpaBackend::MemoryEfficient => {
+                // Chunked, uses fraction of full attention matrix
+                base_memory / 4
+            }
+            _ => base_memory,
+        }
+    }
+
+    /// Get alternative backends ranked
+    fn get_alternatives(&self, input: &AttentionInputInfo) -> Vec<(SdpaBackend, f32)> {
+        self.get_available_backends(input)
+            .into_iter()
+            .map(|b| (b, self.score_backend(b, input).0))
+            .collect()
+    }
+
+    /// Record performance for auto-tuning
+    pub fn record_performance(
+        &self,
+        input: &AttentionInputInfo,
+        backend: SdpaBackend,
+        latency_us: f64,
+        memory_bytes: usize,
+    ) {
+        if !self.config.enable_autotuning {
+            return;
+        }
+
+        let record = PerformanceRecord {
+            key: CacheKey {
+                batch_size: input.batch_size,
+                num_heads: input.num_heads,
+                seq_len_q: input.seq_len_q,
+                seq_len_kv: input.seq_len_kv,
+                head_dim: input.head_dim,
+                is_causal: input.is_causal,
+            },
+            backend,
+            latency_us,
+            memory_bytes,
+        };
+
+        let mut history = self.perf_history.write().unwrap();
+        history.push(record);
+
+        // Keep last 1000 records
+        if history.len() > 1000 {
+            history.remove(0);
+        }
+    }
+
+    /// Clear caches
+    pub fn clear_cache(&self) {
+        self.cache.write().unwrap().clear();
+        self.perf_history.write().unwrap().clear();
+    }
+}
+
+// =============================================================================
+// Global Selector
+// =============================================================================
+
+static GLOBAL_SELECTOR: std::sync::OnceLock<SdpaBackendSelector> = std::sync::OnceLock::new();
+
+fn init_global_selector() -> &'static SdpaBackendSelector {
+    GLOBAL_SELECTOR.get_or_init(|| SdpaBackendSelector::new(SdpaConfig::default()))
+}
+
+/// Get the global SDPA backend selector
+pub fn get_selector() -> &'static SdpaBackendSelector {
+    init_global_selector()
+}
+
+/// Convenience function to select backend
+pub fn select_backend(input: &AttentionInputInfo) -> BackendRecommendation {
+    get_selector().select(input)
+}
+
+// =============================================================================
+// Tests
+// =============================================================================
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_sdpa_backend_display() {
+        assert_eq!(format!("{}", SdpaBackend::FlashAttention), "FlashAttention");
+        assert_eq!(format!("{}", SdpaBackend::Math), "Math");
+    }
+
+    #[test]
+    fn test_hardware_capabilities_cpu() {
+        let hw = HardwareCapabilities::detect_cpu();
+        assert_eq!(hw.device_type, DeviceType::Cpu);
+        assert!(!hw.supports_flash_attention);
+        assert!(hw.supports_backend(SdpaBackend::Math));
+    }
+
+    #[test]
+    fn test_hardware_capabilities_cuda() {
+        let hw = HardwareCapabilities::detect_cuda(0);
+        assert_eq!(hw.device_type, DeviceType::Cuda);
+        assert!(hw.has_tensor_cores);
+        assert!(hw.supports_flash_attention);
+    }
+
+    #[test]
+    fn test_attention_input_info_memory() {
+        let input = AttentionInputInfo {
+            batch_size: 4,
+            num_heads: 32,
+            seq_len_q: 1024,
+            seq_len_kv: 1024,
+            head_dim: 64,
+            dtype: AttentionDType::Float16,
+            is_causal: true,
+            has_mask: false,
+            dropout: 0.0,
+        };
+
+        let memory = input.estimate_memory_bytes();
+        assert!(memory > 0);
+    }
+
+    #[test]
+    fn test_sdpa_config_default() {
+        let config = SdpaConfig::default();
+        assert!(config.preferred_backend.is_none());
+        assert!(config.enable_autotuning);
+        assert!(!config.deterministic);
+    }
+
+    #[test]
+    fn test_selector_debug_mode() {
+        let config = SdpaConfig {
+            debug_mode: true,
+            ..Default::default()
+        };
+        let selector = SdpaBackendSelector::new(config);
+
+        let input = AttentionInputInfo {
+            batch_size: 4,
+            num_heads: 32,
+            seq_len_q: 1024,
+            seq_len_kv: 1024,
+            head_dim: 64,
+            dtype: AttentionDType::Float16,
+            is_causal: true,
+            has_mask: false,
+            dropout: 0.0,
+        };
+
+        let rec = selector.select(&input);
+        assert_eq!(rec.backend, SdpaBackend::Math);
+    }
+
+    #[test]
+    fn test_selector_long_sequence() {
+        let hw = HardwareCapabilities::detect_cuda(0);
+        let selector = SdpaBackendSelector::with_hardware(SdpaConfig::default(), hw);
+
+        let input = AttentionInputInfo {
+            batch_size: 4,
+            num_heads: 32,
+            seq_len_q: 8192,
+            seq_len_kv: 8192,
+            head_dim: 64,
+            dtype: AttentionDType::Float16,
+            is_causal: true,
+            has_mask: false,
+            dropout: 0.0,
+        };
+
+        let rec = selector.select(&input);
+        assert_eq!(rec.backend, SdpaBackend::FlashAttention);
+        assert!(rec.expected_speedup > 2.0);
+    }
+
+    #[test]
+    fn test_selector_short_sequence() {
+        let hw = HardwareCapabilities::detect_cuda(0);
+        let selector = SdpaBackendSelector::with_hardware(SdpaConfig::default(), hw);
+
+        let input = AttentionInputInfo {
+            batch_size: 1,
+            num_heads: 8,
+            seq_len_q: 64,
+            seq_len_kv: 64,
+            head_dim: 64,
+            dtype: AttentionDType::Float16,
+            is_causal: false,
+            has_mask: false,
+            dropout: 0.0,
+        };
+
+        let rec = selector.select(&input);
+        // Should prefer cuDNN or Math for short sequences
+        assert!(rec.backend == SdpaBackend::CuDnn || rec.backend == SdpaBackend::Math);
+    }
+
+    #[test]
+    fn test_selector_preferred_backend() {
+        let config = SdpaConfig {
+            preferred_backend: Some(SdpaBackend::MemoryEfficient),
+            ..Default::default()
+        };
+        let hw = HardwareCapabilities::detect_cuda(0);
+        let selector = SdpaBackendSelector::with_hardware(config, hw);
+
+        let input = AttentionInputInfo {
+            batch_size: 4,
+            num_heads: 32,
+            seq_len_q: 2048,
+            seq_len_kv: 2048,
+            head_dim: 64,
+            dtype: AttentionDType::Float16,
+            is_causal: true,
+            has_mask: false,
+            dropout: 0.0,
+        };
+
+        let rec = selector.select(&input);
+        assert_eq!(rec.backend, SdpaBackend::MemoryEfficient);
+    }
+
+    #[test]
+    fn test_backend_recommendation() {
+        let rec = BackendRecommendation {
+            backend: SdpaBackend::FlashAttention,
+            confidence: 0.95,
+            expected_speedup: 3.5,
+            expected_memory: 1024 * 1024,
+            reason: "Test".to_string(),
+            alternatives: vec![(SdpaBackend::CuDnn, 0.7)],
+        };
+
+        assert_eq!(rec.backend, SdpaBackend::FlashAttention);
+        assert!(rec.confidence > 0.9);
+    }
+
+    #[test]
+    fn test_estimate_flops() {
+        let input = AttentionInputInfo {
+            batch_size: 4,
+            num_heads: 32,
+            seq_len_q: 1024,
+            seq_len_kv: 1024,
+            head_dim: 64,
+            dtype: AttentionDType::Float16,
+            is_causal: true,
+            has_mask: false,
+            dropout: 0.0,
+        };
+
+        let flops = input.estimate_flops();
+        assert!(flops > 0);
+    }
+
+    #[test]
+    fn test_record_performance() {
+        let selector = SdpaBackendSelector::new(SdpaConfig::default());
+
+        let input = AttentionInputInfo {
+            batch_size: 4,
+            num_heads: 32,
+            seq_len_q: 1024,
+            seq_len_kv: 1024,
+            head_dim: 64,
+            dtype: AttentionDType::Float16,
+            is_causal: true,
+            has_mask: false,
+            dropout: 0.0,
+        };
+
+        selector.record_performance(&input, SdpaBackend::FlashAttention, 100.0, 1024);
+
+        // No panic = success
+    }
+}