ferrotensor_graph/

flash_attention.rs

//! Flash Attention - Memory-Efficient Attention Mechanism
//!
//! Implements FlashAttention-2 algorithm for O(N) memory instead of O(N²).
//! Uses tiled matrix multiplication with online softmax normalization.

use ferrotensor_backend::{AttentionOptions, ComputeEngine, UnifiedTensor, WgpuContext};
use std::sync::Arc;

/// Flash Attention layer with tiled computation and online softmax.
///
/// Memory efficient: O(N) instead of O(N²) for standard attention.
/// Speed: 2-4x faster than naive implementation.
///
/// # Example
/// ```no_run
/// # use ferrotensor_backend::{ComputeEngine, UnifiedTensor, WgpuContext};
/// # use ferrotensor_graph::flash_attention::{FlashAttention};
/// let ctx = pollster::block_on(WgpuContext::new());
/// let engine = ComputeEngine::new(&ctx);
/// // [batch, seq_len, hidden]
/// let q = UnifiedTensor::zeros(&ctx, &[1, 128, 768]);
/// let k = UnifiedTensor::zeros(&ctx, &[1, 128, 768]);
/// let v = UnifiedTensor::zeros(&ctx, &[1, 128, 768]);
/// let flash_attn = FlashAttention::new(768, 12, true);
/// let output = flash_attn.forward(&q, &k, &v, &ctx, &engine);
/// ```
pub struct FlashAttention {
    dim: usize,
    num_heads: usize,
    head_dim: usize,
    causal: bool,
    #[allow(dead_code)]
    scale: f32,

    // Optional: Use optimized multi-head kernel
    #[allow(dead_code)]
    use_multi_head_kernel: bool,
}

/// KV-Cache for efficient autoregressive generation
pub struct KVCache {
    /// Cached keys for each layer
    pub k_cache: Vec<Arc<UnifiedTensor>>,
    /// Cached values for each layer
    pub v_cache: Vec<Arc<UnifiedTensor>>,
    /// Maximum sequence length
    pub max_seq_len: usize,
    /// Current position in cache
    pub current_pos: usize,
}

impl FlashAttention {
    /// Create a new FlashAttention layer.
    ///
    /// # Arguments
    /// * `dim` - Total dimension (will be split across heads)
    /// * `num_heads` - Number of attention heads
    /// * `causal` - Whether to apply causal masking (for autoregressive models)
    pub fn new(dim: usize, num_heads: usize, causal: bool) -> Self {
        let head_dim = dim / num_heads;
        let scale = 1.0 / (head_dim as f32).sqrt();
        Self {
            dim,
            num_heads,
            head_dim,
            causal,
            scale,
            use_multi_head_kernel: num_heads > 1,
        }
    }

    /// `[seq, dim]`, `[batch, seq, dim]` or `[batch, seq, heads, head_dim]`
    /// viewed as `[batch, seq, num_heads, head_dim]`.
    fn heads_view(&self, t: &UnifiedTensor) -> UnifiedTensor {
        let (batch, seq) = match t.shape.len() {
            2 => (1, t.shape[0]),
            3 | 4 => (t.shape[0], t.shape[1]),
            n => panic!("FlashAttention expects a 2-D, 3-D or 4-D tensor, got {n}-D"),
        };
        assert_eq!(
            t.shape.iter().product::<usize>(),
            batch * seq * self.dim,
            "tensor {:?} does not have dim {} per token",
            t.shape,
            self.dim
        );
        t.reshaped(&[batch, seq, self.num_heads, self.head_dim])
    }

    /// Forward pass using Flash Attention algorithm.
    ///
    /// # Arguments
    /// * `q` - Query tensor `[batch, seq_q, dim]`, `[seq_q, dim]` or
    ///   `[batch, seq_q, num_heads, head_dim]`
    /// * `k` - Key tensor with the same layout and `seq_k` tokens
    /// * `v` - Value tensor with the same layout as `k`
    ///
    /// # Returns
    /// Output tensor with the same shape as `q`.
    pub fn forward(
        &self,
        q: &UnifiedTensor,
        k: &UnifiedTensor,
        v: &UnifiedTensor,
        ctx: &WgpuContext,
        engine: &ComputeEngine,
    ) -> UnifiedTensor {
        let output = UnifiedTensor::empty(ctx, &q.shape);
        engine.flash_attention_v2(
            ctx,
            &self.heads_view(q),
            &self.heads_view(k),
            &self.heads_view(v),
            &self.heads_view(&output),
            self.causal,
            0.0,
            0,
            None,
        );
        output
    }

    /// Memory for one head's attention scores: the full `seq x seq` matrix
    /// that standard attention stores versus the `seq x head_dim` accumulator
    /// Flash Attention keeps instead (f32, per head, scores only).
    pub fn memory_savings(&self, seq_len: usize) -> String {
        let standard_memory = seq_len * seq_len * 4;
        let flash_memory = seq_len * self.head_dim * 4;
        let savings = standard_memory as f32 / flash_memory as f32;
        format!(
            "{:.1}x less memory ({} vs {} per head for seq={})",
            savings,
            format_bytes(flash_memory),
            format_bytes(standard_memory),
            seq_len
        )
    }

    /// Backward pass for Flash Attention.
    ///
    /// Computes gradients dQ, dK, dV using recomputation to avoid
    /// storing the full attention matrix.
    ///
    /// # Arguments
    /// * `q` - Query tensor from forward pass
    /// * `k` - Key tensor from forward pass
    /// * `v` - Value tensor from forward pass
    /// * `output` - Output from forward pass
    /// * `grad_output` - Gradient of the output
    ///
    /// # Returns
    /// FlashAttentionGrads containing dQ, dK, dV
    #[allow(clippy::too_many_arguments)]
    pub fn backward(
        &self,
        q: &UnifiedTensor,
        k: &UnifiedTensor,
        v: &UnifiedTensor,
        output: &UnifiedTensor,
        grad_output: &UnifiedTensor,
        ctx: &WgpuContext,
        engine: &ComputeEngine,
    ) -> FlashAttentionGrads {
        let grad_q = UnifiedTensor::empty(ctx, &q.shape);
        let grad_k = UnifiedTensor::empty(ctx, &k.shape);
        let grad_v = UnifiedTensor::empty(ctx, &v.shape);

        engine.flash_attention_v2_backward(
            ctx,
            &self.heads_view(q),
            &self.heads_view(k),
            &self.heads_view(v),
            &self.heads_view(output),
            &self.heads_view(grad_output),
            &self.heads_view(&grad_q),
            &self.heads_view(&grad_k),
            &self.heads_view(&grad_v),
            self.causal,
        );

        FlashAttentionGrads {
            grad_q,
            grad_k,
            grad_v,
        }
    }
}

/// Gradients from Flash Attention backward pass.
#[derive(Debug)]
pub struct FlashAttentionGrads {
    /// Gradient with respect to queries
    pub grad_q: UnifiedTensor,
    /// Gradient with respect to keys
    pub grad_k: UnifiedTensor,
    /// Gradient with respect to values
    pub grad_v: UnifiedTensor,
}

impl std::fmt::Debug for FlashAttention {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("FlashAttention")
            .field("dim", &self.dim)
            .field("num_heads", &self.num_heads)
            .field("head_dim", &self.head_dim)
            .field("causal", &self.causal)
            .finish()
    }
}

/// Format a byte count with a binary unit (B, KiB, MiB, GiB).
fn format_bytes(bytes: usize) -> String {
    const UNITS: [&str; 4] = ["B", "KiB", "MiB", "GiB"];
    let mut value = bytes as f64;
    let mut unit = 0;
    while value >= 1024.0 && unit < UNITS.len() - 1 {
        value /= 1024.0;
        unit += 1;
    }
    if unit == 0 {
        format!("{bytes} B")
    } else {
        format!("{value:.1} {}", UNITS[unit])
    }
}

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

    #[test]
    fn test_memory_savings() {
        let attn = FlashAttention::new(512, 8, true);
        let savings = attn.memory_savings(2048);
        assert!(savings.contains("x less memory"));
        // head_dim 64: 512*64*4 = 128 KiB vs 512*512*4 = 1 MiB.
        let attn = FlashAttention::new(512, 8, true);
        assert_eq!(
            attn.memory_savings(512),
            "8.0x less memory (128.0 KiB vs 1.0 MiB per head for seq=512)"
        );
    }

    #[test]
    fn test_format_bytes() {
        assert_eq!(format_bytes(0), "0 B");
        assert_eq!(format_bytes(1023), "1023 B");
        assert_eq!(format_bytes(1536), "1.5 KiB");
        assert_eq!(format_bytes(3 * 1024 * 1024 * 1024), "3.0 GiB");
    }

    #[test]
    fn test_creation() {
        let attn = FlashAttention::new(768, 12, true);
        assert_eq!(attn.dim, 768);
        assert_eq!(attn.num_heads, 12);
        assert_eq!(attn.head_dim, 64);
        assert!(attn.causal);
    }
}

// ============================================================================
// GROUPED QUERY ATTENTION (GQA) - Used by Llama 3, Mistral, etc.
// ============================================================================

/// Grouped Query Attention (GQA) - Efficient attention with fewer KV heads
///
/// GQA uses fewer key-value heads than query heads, reducing memory and
/// computation while maintaining quality. This is the attention mechanism
/// used by Llama 3, Mistral, and many modern LLMs.
///
/// # Memory Savings
/// - Standard MHA: `num_heads` K and V projections
/// - GQA: `num_kv_heads` K and V projections (typically 2-8x fewer)
///
/// # Example
/// ```no_run
/// # use ferrotensor_backend::{ComputeEngine, UnifiedTensor, WgpuContext};
/// # use ferrotensor_graph::flash_attention::{GroupedQueryAttention};
/// let ctx = pollster::block_on(WgpuContext::new());
/// let engine = ComputeEngine::new(&ctx);
/// // [batch, seq_len, hidden]
/// let q = UnifiedTensor::zeros(&ctx, &[1, 128, 4096]);
/// let k = UnifiedTensor::zeros(&ctx, &[1, 128, 1024]);
/// let v = UnifiedTensor::zeros(&ctx, &[1, 128, 1024]);
/// // Llama 3 8B: 32 query heads share 8 KV heads (no dropout, no mask)
/// let gqa = GroupedQueryAttention::new(4096, 32, 8, true, 0.0);
/// let output = gqa.forward(&q, &k, &v, None, &ctx, &engine);
/// ```
pub struct GroupedQueryAttention {
    /// Total hidden dimension
    pub dim: usize,
    /// Number of query heads
    pub num_heads: usize,
    /// Number of key-value heads (fewer than query heads)
    pub num_kv_heads: usize,
    /// Dimension per head
    pub head_dim: usize,
    /// Number of query heads per KV head
    pub num_groups: usize,
    /// Whether to use causal masking
    pub causal: bool,
    /// Scaling factor for attention scores
    #[allow(dead_code)]
    scale: f32,
    /// Dropout probability (0.0 = no dropout)
    pub dropout_prob: f32,
}

impl GroupedQueryAttention {
    /// Create a new GQA layer.
    ///
    /// # Arguments
    /// * `dim` - Total hidden dimension
    /// * `num_heads` - Number of query heads
    /// * `num_kv_heads` - Number of key-value heads (must divide num_heads evenly)
    /// * `causal` - Whether to apply causal masking
    /// * `dropout_prob` - Dropout probability
    ///
    /// # Panics
    /// Panics if `num_heads` is not divisible by `num_kv_heads`
    pub fn new(
        dim: usize,
        num_heads: usize,
        num_kv_heads: usize,
        causal: bool,
        dropout_prob: f32,
    ) -> Self {
        assert!(
            num_heads.is_multiple_of(num_kv_heads),
            "num_heads ({}) must be divisible by num_kv_heads ({})",
            num_heads,
            num_kv_heads
        );

        let head_dim = dim / num_heads;
        let num_groups = num_heads / num_kv_heads;
        let scale = 1.0 / (head_dim as f32).sqrt();

        Self {
            dim,
            num_heads,
            num_kv_heads,
            head_dim,
            num_groups,
            causal,
            scale,
            dropout_prob,
        }
    }

    /// Forward pass with grouped query attention.
    ///
    /// # Arguments
    /// * `q` - Query tensor [batch, seq_q, num_heads, head_dim]
    /// * `k` - Key tensor [batch, seq_k, num_kv_heads, head_dim]
    /// * `v` - Value tensor [batch, seq_k, num_kv_heads, head_dim]
    ///
    /// # Returns
    /// Output tensor [batch, seq_q, num_heads, head_dim]
    pub fn forward(
        &self,
        q: &UnifiedTensor,
        k: &UnifiedTensor,
        v: &UnifiedTensor,
        mask: Option<&UnifiedTensor>,
        ctx: &WgpuContext,
        engine: &ComputeEngine,
    ) -> UnifiedTensor {
        self.forward_at(q, k, v, mask, 0, ctx, engine)
    }

    /// Like [`forward`](Self::forward), with the first query at absolute
    /// position `pos`: query `i` attends to keys `0..=pos + i` when causal.
    /// Used to decode against a KV cache, where `k`/`v` hold the cache.
    #[allow(clippy::too_many_arguments)]
    pub fn forward_at(
        &self,
        q: &UnifiedTensor,
        k: &UnifiedTensor,
        v: &UnifiedTensor,
        mask: Option<&UnifiedTensor>,
        pos: usize,
        ctx: &WgpuContext,
        engine: &ComputeEngine,
    ) -> UnifiedTensor {
        // Validation: Expect 4D tensors [batch, seq, heads, dim]
        if q.shape.len() != 4 || k.shape.len() != 4 || v.shape.len() != 4 {
            panic!("GroupedQueryAttention expects 4D tensors (batch, seq, heads, dim)");
        }

        let batch_size = q.shape[0];
        let seq_q = q.shape[1];

        let output = UnifiedTensor::empty(ctx, &[batch_size, seq_q, self.num_heads, self.head_dim]);

        // Use time-based seed for stochastic dropout
        let seed = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_nanos() as u64;

        let opts = AttentionOptions {
            causal: self.causal,
            causal_offset: pos as u32,
            dropout_prob: self.dropout_prob,
            seed,
            ..AttentionOptions::default()
        };
        engine.flash_attention_v2_with(ctx, q, k, v, &output, &opts, mask);

        output
    }

    /// Get memory savings compared to standard MHA
    pub fn memory_savings(&self) -> String {
        let mha_kv_params = 2 * self.num_heads * self.head_dim;
        let gqa_kv_params = 2 * self.num_kv_heads * self.head_dim;
        let savings = mha_kv_params as f32 / gqa_kv_params as f32;

        format!(
            "{:.1}x KV memory savings ({} vs {} KV heads)",
            savings, self.num_kv_heads, self.num_heads
        )
    }
}

// ============================================================================
// MULTI-QUERY ATTENTION (MQA) - Single KV head for maximum efficiency
// ============================================================================

/// Multi-Query Attention (MQA) - Extreme efficiency with single KV head
///
/// MQA uses a single key-value head shared across all query heads.
/// This provides maximum memory efficiency at some quality cost.
///
/// # Use Cases
/// - Inference-optimized models
/// - Memory-constrained deployments
/// - High-throughput serving
///
/// # Example
/// ```no_run
/// # use ferrotensor_backend::{ComputeEngine, UnifiedTensor, WgpuContext};
/// # use ferrotensor_graph::flash_attention::{MultiQueryAttention};
/// let ctx = pollster::block_on(WgpuContext::new());
/// let engine = ComputeEngine::new(&ctx);
/// // [batch, seq_len, hidden]
/// let q = UnifiedTensor::zeros(&ctx, &[1, 128, 768]);
/// let k = UnifiedTensor::zeros(&ctx, &[1, 128, 64]);
/// let v = UnifiedTensor::zeros(&ctx, &[1, 128, 64]);
/// // 12 query heads share one KV head
/// let mqa = MultiQueryAttention::new(768, 12, true, 0.0);
/// let output = mqa.forward(&q, &k, &v, None, &ctx, &engine);
/// ```
pub struct MultiQueryAttention {
    /// Total hidden dimension
    pub dim: usize,
    /// Number of query heads
    pub num_heads: usize,
    /// Dimension per head
    pub head_dim: usize,
    /// Whether to use causal masking
    pub causal: bool,
    /// Internal GQA with num_kv_heads = 1
    gqa: GroupedQueryAttention,
}

impl MultiQueryAttention {
    /// Create a new MQA layer.
    ///
    /// # Arguments
    /// * `dim` - Total hidden dimension
    /// * `num_heads` - Number of query heads
    /// * `causal` - Whether to apply causal masking
    pub fn new(dim: usize, num_heads: usize, causal: bool, dropout_prob: f32) -> Self {
        let head_dim = dim / num_heads;
        let gqa = GroupedQueryAttention::new(dim, num_heads, 1, causal, dropout_prob);

        Self {
            dim,
            num_heads,
            head_dim,
            causal,
            gqa,
        }
    }

    /// Forward pass with multi-query attention.
    ///
    /// # Arguments
    /// * `q` - Query tensor [batch, seq_q, num_heads, head_dim]
    /// * `k` - Key tensor [batch, seq_k, 1, head_dim]
    /// * `v` - Value tensor [batch, seq_k, 1, head_dim]
    ///
    /// # Returns
    /// Output tensor [batch, seq_q, num_heads, head_dim]
    pub fn forward(
        &self,
        q: &UnifiedTensor,
        k: &UnifiedTensor,
        v: &UnifiedTensor,
        mask: Option<&UnifiedTensor>,
        ctx: &WgpuContext,
        engine: &ComputeEngine,
    ) -> UnifiedTensor {
        // Delegate to GQA with num_kv_heads = 1
        self.gqa.forward(q, k, v, mask, ctx, engine)
    }

    /// Get memory savings compared to MHA
    pub fn memory_savings(&self) -> String {
        format!(
            "{}x KV memory savings (1 vs {} KV heads)",
            self.num_heads, self.num_heads
        )
    }
}

// ============================================================================
// SLIDING WINDOW ATTENTION - For long sequences
// ============================================================================

/// Sliding Window Attention for efficient long-context processing
///
/// Limits attention to a local window around each position, enabling
/// linear memory scaling for very long sequences.
///
/// # Use Cases
/// - Long document processing (100K+ tokens)
/// - Streaming input processing
/// - Memory-efficient inference
///
/// # Example
/// ```no_run
/// # use ferrotensor_backend::{ComputeEngine, UnifiedTensor, WgpuContext};
/// # use ferrotensor_graph::flash_attention::{SlidingWindowAttention};
/// let ctx = pollster::block_on(WgpuContext::new());
/// let engine = ComputeEngine::new(&ctx);
/// // [batch, seq_len, hidden]
/// let q = UnifiedTensor::zeros(&ctx, &[1, 128, 768]);
/// let k = UnifiedTensor::zeros(&ctx, &[1, 128, 768]);
/// let v = UnifiedTensor::zeros(&ctx, &[1, 128, 768]);
/// let swa = SlidingWindowAttention::new(768, 12, 4096, true);
/// let output = swa.forward(&q, &k, &v, &ctx, &engine);
/// ```
pub struct SlidingWindowAttention {
    /// Total hidden dimension
    pub dim: usize,
    /// Number of attention heads
    pub num_heads: usize,
    /// Dimension per head
    pub head_dim: usize,
    /// Window size for local attention
    pub window_size: usize,
    /// Whether to use causal masking
    pub causal: bool,
    #[allow(dead_code)]
    scale: f32,
}

impl SlidingWindowAttention {
    /// Create a new sliding window attention layer.
    ///
    /// # Arguments
    /// * `dim` - Total hidden dimension
    /// * `num_heads` - Number of attention heads
    /// * `window_size` - Maximum distance for attention (tokens only attend within this window)
    /// * `causal` - Whether to apply causal masking
    pub fn new(dim: usize, num_heads: usize, window_size: usize, causal: bool) -> Self {
        let head_dim = dim / num_heads;
        let scale = 1.0 / (head_dim as f32).sqrt();

        Self {
            dim,
            num_heads,
            head_dim,
            window_size,
            causal,
            scale,
        }
    }

    /// Forward pass with sliding window attention.
    ///
    /// Each position can only attend to positions within `window_size` distance.
    ///
    /// # Arguments
    /// * `q` - Query tensor [batch, seq_q, dim]
    /// * `k` - Key tensor [batch, seq_k, dim]
    /// * `v` - Value tensor [batch, seq_k, dim]
    ///
    /// # Returns
    /// Output tensor [batch, seq_q, dim]
    pub fn forward(
        &self,
        q: &UnifiedTensor,
        k: &UnifiedTensor,
        v: &UnifiedTensor,
        ctx: &WgpuContext,
        engine: &ComputeEngine,
    ) -> UnifiedTensor {
        let view = |t: &UnifiedTensor| {
            t.reshaped(&[t.shape[0], t.shape[1], self.num_heads, self.head_dim])
        };
        let output = UnifiedTensor::empty(ctx, &q.shape);
        let opts = AttentionOptions {
            causal: self.causal,
            window: self.window_size as u32,
            ..AttentionOptions::default()
        };
        engine.flash_attention_v2_with(
            ctx,
            &view(q),
            &view(k),
            &view(v),
            &view(&output),
            &opts,
            None,
        );
        output
    }

    /// Get memory savings compared to full attention
    pub fn memory_savings(&self, seq_len: usize) -> String {
        let full_attention = seq_len * seq_len;
        let window_attention = seq_len * self.window_size.min(seq_len);
        let savings = full_attention as f32 / window_attention as f32;

        format!(
            "{:.1}x attention savings (window={}, seq={})",
            savings, self.window_size, seq_len
        )
    }
}

// ============================================================================
// ATTENTION MASK - Custom masking patterns
// ============================================================================

/// Attention mask for custom attention patterns
#[derive(Debug, Clone)]
pub enum AttentionMask {
    /// No mask (full attention)
    None,
    /// Causal mask (lower triangular)
    Causal,
    /// Padding mask to ignore padded tokens
    Padding(Vec<usize>), // lengths of non-padded sequences
    /// Sliding window mask
    SlidingWindow(usize),
    /// Custom mask tensor
    Custom(Arc<UnifiedTensor>),
}

impl AttentionMask {
    /// Create a causal mask
    pub fn causal() -> Self {
        Self::Causal
    }

    /// Create a padding mask from sequence lengths
    pub fn from_lengths(lengths: Vec<usize>) -> Self {
        Self::Padding(lengths)
    }

    /// Create a sliding window mask
    pub fn sliding_window(window_size: usize) -> Self {
        Self::SlidingWindow(window_size)
    }
}

impl std::fmt::Debug for GroupedQueryAttention {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("GroupedQueryAttention")
            .field("dim", &self.dim)
            .field("num_heads", &self.num_heads)
            .field("num_kv_heads", &self.num_kv_heads)
            .field("num_groups", &self.num_groups)
            .field("head_dim", &self.head_dim)
            .field("causal", &self.causal)
            .field("dropout", &self.dropout_prob)
            .finish()
    }
}

impl std::fmt::Debug for MultiQueryAttention {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("MultiQueryAttention")
            .field("dim", &self.dim)
            .field("num_heads", &self.num_heads)
            .field("head_dim", &self.head_dim)
            .field("causal", &self.causal)
            .finish()
    }
}

impl std::fmt::Debug for SlidingWindowAttention {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SlidingWindowAttention")
            .field("dim", &self.dim)
            .field("num_heads", &self.num_heads)
            .field("head_dim", &self.head_dim)
            .field("window_size", &self.window_size)
            .field("causal", &self.causal)
            .finish()
    }
}