haagenti-sparse 0.1.0

//! Sparse attention kernel execution

use crate::{AttentionMask, Result, SparseError};
use serde::{Deserialize, Serialize};

/// Configuration for sparse attention kernel
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KernelConfig {
    /// Hidden dimension
    pub hidden_dim: usize,
    /// Number of heads
    pub num_heads: usize,
    /// Head dimension
    pub head_dim: usize,
    /// Sequence length
    pub seq_len: usize,
    /// Use flash attention
    pub use_flash: bool,
    /// Memory format (contiguous, channels_last, etc.)
    pub memory_format: MemoryFormat,
}

/// Memory layout format
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum MemoryFormat {
    /// Standard contiguous layout [B, S, H, D]
    Contiguous,
    /// Channels last [B, H, S, D]
    ChannelsLast,
    /// Grouped format for sparse [B, G, S, D] where G = active heads
    Grouped,
}

impl Default for KernelConfig {
    fn default() -> Self {
        Self {
            hidden_dim: 2048,
            num_heads: 32,
            head_dim: 64,
            seq_len: 4096,
            use_flash: true,
            memory_format: MemoryFormat::Contiguous,
        }
    }
}

/// Manager for sparse attention kernel execution.
///
/// Handles the execution of attention computations with dynamic head sparsity,
/// including index mapping caching for efficient repeated execution with the
/// same sparsity pattern.
///
/// # Performance Optimizations
///
/// - Pre-computes index mappings for gather/scatter operations
/// - Caches mappings by pattern hash for repeated use
/// - Supports flash attention for memory efficiency
/// - Handles different memory layouts (contiguous, channels-last, grouped)
///
/// # Example
///
/// ```ignore
/// let mut kernel = SparseKernel::new(KernelConfig::default());
/// kernel.prepare(&mask, layer_idx)?;
/// let output = kernel.execute(&query, &key, &value, &mask, layer_idx)?;
/// ```
#[derive(Debug)]
pub struct SparseKernel {
    config: KernelConfig,
    /// Precomputed index mappings for each sparsity pattern
    index_cache: std::collections::HashMap<u64, IndexMapping>,
}

/// Index mapping for sparse computation
#[derive(Debug, Clone)]
struct IndexMapping {
    /// Active head indices
    active_heads: Vec<usize>,
    /// Output scatter indices (for GPU kernel scatter operation)
    #[allow(dead_code)]
    scatter_indices: Vec<usize>,
    /// Pattern hash (for cache lookup validation)
    #[allow(dead_code)]
    pattern_hash: u64,
}

impl SparseKernel {
    /// Create a new kernel with config
    pub fn new(config: KernelConfig) -> Self {
        Self {
            config,
            index_cache: std::collections::HashMap::new(),
        }
    }

    /// Prepare kernel for a specific mask
    pub fn prepare(&mut self, mask: &AttentionMask, layer: usize) -> Result<()> {
        let pattern_hash = self.compute_pattern_hash(mask, layer);

        self.index_cache.entry(pattern_hash).or_insert_with(|| {
            let active_heads = mask.active_heads(layer);
            let scatter_indices: Vec<usize> =
                active_heads.iter().enumerate().map(|(i, _)| i).collect();

            IndexMapping {
                active_heads,
                scatter_indices,
                pattern_hash,
            }
        });

        Ok(())
    }

    /// Execute sparse attention for a layer
    ///
    /// This is a simulated implementation. In practice, this would:
    /// 1. Gather only active Q, K, V heads
    /// 2. Compute attention only for active heads
    /// 3. Scatter results back to full head positions
    pub fn execute(
        &self,
        mask: &AttentionMask,
        layer: usize,
        _q: &[f32], // [batch, seq, num_heads, head_dim]
        _k: &[f32],
        _v: &[f32],
    ) -> Result<Vec<f32>> {
        let pattern_hash = self.compute_pattern_hash(mask, layer);

        let mapping = self
            .index_cache
            .get(&pattern_hash)
            .ok_or_else(|| SparseError::KernelError("Mask pattern not prepared".into()))?;

        // Simulated output
        let batch_size = 1; // Would be inferred from input
        let output_size =
            batch_size * self.config.seq_len * self.config.num_heads * self.config.head_dim;

        // In real implementation:
        // 1. Extract active heads from Q, K, V
        // 2. Compute attention: softmax(QK^T / sqrt(d)) * V
        // 3. Scatter back to full size

        let mut output = vec![0.0f32; output_size];

        // Mark active positions (simulated)
        for &head in &mapping.active_heads {
            let offset = head * self.config.head_dim;
            for d in 0..self.config.head_dim {
                if offset + d < output.len() {
                    output[offset + d] = 1.0; // Placeholder
                }
            }
        }

        Ok(output)
    }

    /// Compute hash for mask pattern at a layer
    fn compute_pattern_hash(&self, mask: &AttentionMask, layer: usize) -> u64 {
        use std::hash::{Hash, Hasher};
        let mut hasher = std::collections::hash_map::DefaultHasher::new();

        layer.hash(&mut hasher);
        for head in 0..mask.num_heads {
            mask.is_active(layer, head).hash(&mut hasher);
        }

        hasher.finish()
    }

    /// Estimate compute savings for a mask
    pub fn estimate_savings(&self, mask: &AttentionMask) -> ComputeEstimate {
        let total_heads = mask.num_heads * mask.num_layers;
        let active_heads: usize = (0..mask.num_layers).map(|l| mask.active_count(l)).sum();

        let compute_ratio = active_heads as f32 / total_heads as f32;

        // Memory savings from not loading inactive weights
        let memory_ratio = compute_ratio * 0.9 + 0.1; // Some overhead

        // Attention compute is quadratic in heads for multi-head attention
        let attention_ratio = compute_ratio; // Linear for independent heads

        ComputeEstimate {
            compute_ratio,
            memory_ratio,
            attention_ratio,
            estimated_speedup: 1.0 / compute_ratio,
            active_heads,
            total_heads,
        }
    }

    /// Get current configuration
    pub fn config(&self) -> &KernelConfig {
        &self.config
    }

    /// Clear cached index mappings
    pub fn clear_cache(&mut self) {
        self.index_cache.clear();
    }
}

/// Estimate of compute and memory savings
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ComputeEstimate {
    /// Fraction of compute used (1.0 = full, 0.5 = half)
    pub compute_ratio: f32,
    /// Fraction of memory bandwidth used
    pub memory_ratio: f32,
    /// Fraction of attention compute
    pub attention_ratio: f32,
    /// Estimated speedup (e.g., 2.0 = 2x faster)
    pub estimated_speedup: f32,
    /// Number of active heads
    pub active_heads: usize,
    /// Total heads across all layers
    pub total_heads: usize,
}

/// Kernel statistics
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct KernelStats {
    /// Total executions
    pub executions: u64,
    /// Cache hits
    pub cache_hits: u64,
    /// Average sparsity
    pub avg_sparsity: f32,
    /// Total compute saved (estimated GFLOPs)
    pub compute_saved: f64,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_kernel_prepare() {
        let mut kernel = SparseKernel::new(KernelConfig::default());
        let mask = AttentionMask::random(32, 10, 0.5);

        kernel.prepare(&mask, 0).unwrap();
        kernel.prepare(&mask, 5).unwrap();

        // Cache should have entries
        assert!(!kernel.index_cache.is_empty());
    }

    #[test]
    fn test_compute_estimate() {
        let kernel = SparseKernel::new(KernelConfig::default());
        let mask = AttentionMask::random(32, 10, 0.5);

        let estimate = kernel.estimate_savings(&mask);

        // With 50% sparsity, should see roughly 50% compute
        assert!(estimate.compute_ratio > 0.4 && estimate.compute_ratio < 0.7);
        assert!(estimate.estimated_speedup > 1.4 && estimate.estimated_speedup < 2.5);
    }

    #[test]
    fn test_execute() {
        let mut kernel = SparseKernel::new(KernelConfig {
            num_heads: 8,
            head_dim: 64,
            seq_len: 16,
            ..Default::default()
        });

        let mask = AttentionMask::random(8, 4, 0.5);
        kernel.prepare(&mask, 0).unwrap();

        // Create dummy inputs
        let q = vec![0.0f32; 16 * 8 * 64];
        let k = vec![0.0f32; 16 * 8 * 64];
        let v = vec![0.0f32; 16 * 8 * 64];

        let output = kernel.execute(&mask, 0, &q, &k, &v).unwrap();
        assert!(!output.is_empty());
    }
}