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

//! 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<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::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> {
        [
            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<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");
        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
    }
}