//! 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, /// FlashAttention v3 - WGMMA + TMA + warp specialization (Hopper/Blackwell) FlashAttentionV3, /// 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, /// Variable-length packed-sequence attention (no padding) VarLen, /// Flash Decoding: split-K attention optimised for long-context single-token decode. /// /// Parallelises across KV-sequence chunks then combines with the online-softmax /// log-sum-exp trick. Preferred over [`SdpaBackend::FlashAttention`] when /// `seq_len_kv >= 1024` and `seq_len_q == 1` (decode phase). FlashDecode, } 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::FlashAttentionV3 => write!(f, "FlashAttentionV3"), 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"), SdpaBackend::VarLen => write!(f, "VarLen"), SdpaBackend::FlashDecode => write!(f, "FlashDecode"), } } } // ============================================================================= // 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 a specific CUDA compute capability. /// /// Call this instead of `detect_cuda()` when the SM version is known at call-site /// (e.g. after querying the driver or from a build-time constant). pub fn for_compute_capability(major: u32, minor: u32) -> Self { let is_hopper_plus = major >= 9; let is_blackwell_plus = major >= 12; Self { device_type: DeviceType::Cuda, compute_capability: Some((major, minor)), total_memory: 16 * 1024 * 1024 * 1024, // conservative 16 GB available_memory: 14 * 1024 * 1024 * 1024, has_tensor_cores: true, has_fp16: true, has_bf16: major >= 8, has_fp8: is_hopper_plus, // FP8 requires SM_90+ num_compute_units: if is_blackwell_plus { 84 } else { 108 }, memory_bandwidth_gbps: if is_blackwell_plus { 960.0 } else { 2039.0 }, supports_flash_attention: true, supports_cudnn_attention: major >= 8, } } /// Returns `true` if this hardware can run FlashAttention-3 (WGMMA + TMA). /// /// FA3 requires SM_90+ (Hopper or newer). pub fn supports_flash_v3(&self) -> bool { matches!(self.compute_capability, Some((major, _)) if major >= 9) } /// 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::FlashAttentionV3 => self.supports_flash_v3(), SdpaBackend::Math => true, // Always supported SdpaBackend::MemoryEfficient => true, SdpaBackend::CuDnn => self.supports_cudnn_attention, SdpaBackend::Metal => self.device_type == DeviceType::Metal, SdpaBackend::Cpu => true, SdpaBackend::VarLen => true, // CPU simulation always available; CUDA variant when feature is on // FlashDecode CPU reference is always available; CUDA path requires cuda feature. // The variant is only *useful* for decode-phase (seq_len_q == 1), but we // advertise it as generically supported — the selector's scoring logic // will penalise it for prefill workloads. SdpaBackend::FlashDecode => 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, /// Disabled backends pub disabled_backends: Vec, /// 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>, /// Performance history for tuning perf_history: std::sync::RwLock>, } /// 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::FlashAttentionV3 => { // FA3 shares FA2 head-dim constraints; additionally requires SM_90+ 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::FlashDecode, SdpaBackend::FlashAttentionV3, SdpaBackend::FlashAttention, SdpaBackend::CuDnn, SdpaBackend::MemoryEfficient, SdpaBackend::Math, SdpaBackend::Metal, SdpaBackend::Cpu, SdpaBackend::VarLen, ] .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::FlashAttentionV3 => { // FA3 beats FA2 across the board on Hopper/Blackwell if seq_len >= 2048 { (0.98, format!("FlashAttentionV3 (WGMMA+TMA) optimal for seq_len={seq_len}")) } else if seq_len >= 512 { (0.90, "FlashAttentionV3 excellent for medium sequences".to_string()) } else { (0.70, "FlashAttentionV3 has overhead for short sequences".to_string()) } } 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()) } SdpaBackend::VarLen => { (0.75, "VarLen eliminates padding waste for mixed-length batches".to_string()) } SdpaBackend::FlashDecode => { // Flash Decoding is purpose-built for single-token decode with long KV contexts. // It achieves ~50× over naive decode on sequences >= 8 K by parallelising across // KV chunks. For prefill (seq_len_q > 1) it degrades gracefully to standard // attention but offers no advantage. let is_decode_phase = input.seq_len_q == 1; let kv_len = input.seq_len_kv; if is_decode_phase && kv_len >= 1024 { (0.97, format!("FlashDecode optimal for decode phase with kv_len={kv_len}")) } else if is_decode_phase && kv_len >= 256 { (0.80, format!("FlashDecode good for decode phase with kv_len={kv_len}")) } else if is_decode_phase { (0.55, "FlashDecode marginal benefit for short KV in decode phase".to_string()) } else { (0.20, "FlashDecode not designed for prefill (seq_len_q > 1)".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::FlashAttentionV3 => { // ~2x over FA2 from WGMMA tile efficiency + async TMA overlap if seq_len >= 4096 { 10.0 } else if seq_len >= 2048 { 7.0 } else if seq_len >= 1024 { 5.0 } else if seq_len >= 512 { 3.5 } else { 2.0 } } 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, SdpaBackend::VarLen => 2.0, // avoids padding overhead for mixed-length batches SdpaBackend::FlashDecode => { // Split-K parallelism yields ~50× speedup on very long decode contexts. let kv_len = input.seq_len_kv; if kv_len >= 32_768 { 50.0 } else if kv_len >= 8_192 { 20.0 } else if kv_len >= 4_096 { 10.0 } else if kv_len >= 1_024 { 4.0 } else { 1.5 } } } } /// Estimate memory for backend fn estimate_memory(&self, backend: SdpaBackend, input: &AttentionInputInfo) -> usize { let base_memory = input.estimate_memory_bytes(); match backend { SdpaBackend::FlashAttentionV3 | SdpaBackend::FlashAttention => { // Both FA2 and FA3 use O(N) tiling; same memory footprint 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 } SdpaBackend::FlashDecode => { // Flash Decoding carries O(K) partial buffers where K = num_splits (≤64). // For a single decode token: partial_out[num_heads, splits, head_dim] + // partial_max/sum[num_heads, splits]. This is negligible vs. the KV cache. let num_splits = crate::kernels::FlashDecodeKernel::num_splits_for_seq_len( input.seq_len_kv, input.head_dim, ); let partial_bytes = input.num_heads * num_splits * (input.head_dim + 2) * 4; let qkvo = input.batch_size * input.num_heads * input.seq_len_kv * input.head_dim * 4 * 2; // K + V qkvo + partial_bytes } _ => 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 = 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"); assert_eq!(format!("{}", SdpaBackend::FlashDecode), "FlashDecode"); } #[test] fn test_selector_decode_phase_long_context() { // Flash Decoding should be selected for single-token decode with long KV. let hw = HardwareCapabilities::detect_cuda(0); let selector = SdpaBackendSelector::with_hardware(SdpaConfig::default(), hw); let input = AttentionInputInfo { batch_size: 1, num_heads: 32, seq_len_q: 1, // single decode token seq_len_kv: 8192, // long KV context head_dim: 128, dtype: AttentionDType::Float16, is_causal: true, has_mask: false, dropout: 0.0, }; let rec = selector.select(&input); assert_eq!( rec.backend, SdpaBackend::FlashDecode, "expected FlashDecode for decode phase with kv_len=8192, got {:?}", rec.backend ); assert!( rec.expected_speedup >= 4.0, "FlashDecode should report >=4× speedup for long context, got {}", rec.expected_speedup ); } #[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 } }