rustytorch/crates/training/rtx-flash-metal-attention/src/lib.rs

//! # rtx-flash-metal-attention
//!
//! Metal-native Flash Attention implementation for Apple Silicon.
//!
//! This crate provides a high-performance implementation of the Flash Attention
//! algorithm using Metal compute shaders, optimized for Apple Silicon's unified
//! memory architecture.
//!
//! ## Features
//!
//! - **Memory Efficient**: Uses online softmax to avoid materializing the full
//!   attention matrix, reducing memory from O(N²) to O(N).
//! - **Apple Silicon Optimized**: Leverages Metal's unified memory and
//!   threadgroup shared memory for efficient tiling.
//! - **Causal Masking**: Built-in support for autoregressive attention masks.
//! - **Backward Pass**: Full gradient computation for training.
//!
//! ## Usage
//!
//! ```rust,ignore
//! use rtx_flash_metal_attention::{FlashAttention, FlashAttentionConfig};
//! use rtx_tensor::Tensor;
//!
//! // Create Flash Attention with default config
//! let attn = FlashAttention::new(FlashAttentionConfig::default())?;
//!
//! // Forward pass
//! let (output, lse) = attn.forward(&q, &k, &v)?;
//!
//! // Backward pass (for training)
//! let (dq, dk, dv) = attn.backward(&grad_output, &q, &k, &v, &output, &lse)?;
//! ```
//!
//! ## Tensor Layout
//!
//! All tensors use the layout `[batch, heads, sequence, head_dim]`:
//! - `batch`: Number of sequences in the batch
//! - `heads`: Number of attention heads
//! - `sequence`: Sequence length (can differ between Q and K/V)
//! - `head_dim`: Dimension per head (typically 64 or 128)

// Re-export config and error types (platform-independent)
mod config;
mod error;

pub use config::FlashAttentionConfig;
pub use error::{FlashError, FlashResult};

// macOS-only implementation
#[cfg(target_os = "macos")]
mod backward;
#[cfg(target_os = "macos")]
mod forward;
#[cfg(target_os = "macos")]
mod kernels;

#[cfg(target_os = "macos")]
pub use kernels::{DeviceCapabilities, query_device_capabilities};

#[cfg(target_os = "macos")]
use objc2::rc::Retained;
#[cfg(target_os = "macos")]
use objc2::runtime::ProtocolObject;
#[cfg(target_os = "macos")]
use objc2_metal::{MTLCommandQueue, MTLComputePipelineState, MTLDevice};
#[cfg(target_os = "macos")]
use rtx_tensor::Tensor;
#[cfg(target_os = "macos")]
use tracing::info;

/// Flash Attention implementation for Metal
///
/// This struct holds the compiled Metal pipelines and configuration
/// for executing Flash Attention operations.
#[cfg(target_os = "macos")]
pub struct FlashAttention {
    /// Metal device
    #[allow(dead_code)]
    device: Retained<ProtocolObject<dyn MTLDevice>>,
    /// Command queue for submitting work
    command_queue: Retained<ProtocolObject<dyn MTLCommandQueue>>,
    /// Forward pass compute pipeline
    forward_pipeline: Retained<ProtocolObject<dyn MTLComputePipelineState>>,
    /// Backward dQ compute pipeline
    backward_dq_pipeline: Retained<ProtocolObject<dyn MTLComputePipelineState>>,
    /// Backward dK/dV compute pipeline
    backward_dkv_pipeline: Retained<ProtocolObject<dyn MTLComputePipelineState>>,
    /// Configuration
    config: FlashAttentionConfig,
}

#[cfg(target_os = "macos")]
impl FlashAttention {
    /// Create a new Flash Attention instance with the given configuration
    ///
    /// This compiles the Metal shaders and creates compute pipelines.
    /// The compilation is done once at creation time.
    ///
    /// # Arguments
    /// * `config` - Configuration for block sizes, causal masking, etc.
    ///
    /// # Errors
    /// Returns an error if:
    /// - No Metal device is available
    /// - Shader compilation fails
    /// - Pipeline creation fails
    pub fn new(config: FlashAttentionConfig) -> FlashResult<Self> {
        info!("Creating Metal Flash Attention with config: {:?}", config);

        let device = kernels::get_default_device()?;

        let command_queue = device
            .newCommandQueue()
            .ok_or_else(|| FlashError::device("Failed to create command queue"))?;

        let (forward_pipeline, backward_dq_pipeline, backward_dkv_pipeline) =
            kernels::compile_pipelines(&device, &config)?;

        info!("Metal Flash Attention initialized successfully");

        Ok(Self {
            device,
            command_queue,
            forward_pipeline,
            backward_dq_pipeline,
            backward_dkv_pipeline,
            config,
        })
    }

    /// Create a Flash Attention instance with default configuration
    pub fn default_config() -> FlashResult<Self> {
        Self::new(FlashAttentionConfig::default())
    }

    /// Create a Flash Attention instance for causal (autoregressive) attention
    pub fn causal() -> FlashResult<Self> {
        Self::new(FlashAttentionConfig::causal())
    }

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

    /// Execute the forward pass
    ///
    /// Computes: `O = softmax(Q @ K^T * scale) @ V`
    ///
    /// # Arguments
    /// * `q` - Query tensor `[batch, heads, seq_q, head_dim]`
    /// * `k` - Key tensor `[batch, heads, seq_kv, head_dim]`
    /// * `v` - Value tensor `[batch, heads, seq_kv, head_dim]`
    ///
    /// # Returns
    /// * `output` - Attention output `[batch, heads, seq_q, head_dim]`
    /// * `lse` - Log-sum-exp values `[batch, heads, seq_q]` (needed for backward)
    ///
    /// # Errors
    /// Returns an error if:
    /// - Tensor shapes are invalid
    /// - Tensors are not on a Metal device
    /// - Kernel execution fails
    pub fn forward(&self, q: &Tensor, k: &Tensor, v: &Tensor) -> FlashResult<(Tensor, Tensor)> {
        forward::flash_attention_forward(self, q, k, v)
    }

    /// Execute the backward pass
    ///
    /// Computes gradients dQ, dK, dV for the attention operation.
    ///
    /// # Arguments
    /// * `grad_output` - Gradient of loss w.r.t. output `[batch, heads, seq_q, head_dim]`
    /// * `q` - Query tensor from forward pass
    /// * `k` - Key tensor from forward pass
    /// * `v` - Value tensor from forward pass
    /// * `output` - Output from forward pass
    /// * `lse` - Log-sum-exp from forward pass
    ///
    /// # Returns
    /// * `(dQ, dK, dV)` - Gradients with same shapes as Q, K, V
    ///
    /// # Errors
    /// Returns an error if:
    /// - Tensor shapes are invalid or mismatched
    /// - Tensors are not on a Metal device
    /// - Kernel execution fails
    pub fn backward(
        &self,
        grad_output: &Tensor,
        q: &Tensor,
        k: &Tensor,
        v: &Tensor,
        output: &Tensor,
        lse: &Tensor,
    ) -> FlashResult<(Tensor, Tensor, Tensor)> {
        backward::flash_attention_backward(self, grad_output, q, k, v, output, lse)
    }
}

// Stub implementation for non-macOS platforms
#[cfg(not(target_os = "macos"))]
pub struct FlashAttention {
    config: FlashAttentionConfig,
}

#[cfg(not(target_os = "macos"))]
impl FlashAttention {
    pub fn new(config: FlashAttentionConfig) -> FlashResult<Self> {
        Ok(Self { config })
    }

    pub fn default_config() -> FlashResult<Self> {
        Self::new(FlashAttentionConfig::default())
    }

    pub fn causal() -> FlashResult<Self> {
        Self::new(FlashAttentionConfig::causal())
    }

    pub fn config(&self) -> &FlashAttentionConfig {
        &self.config
    }
}

/// Check if Metal Flash Attention is available on this system
pub fn is_available() -> bool {
    #[cfg(target_os = "macos")]
    {
        kernels::get_default_device().is_ok()
    }
    #[cfg(not(target_os = "macos"))]
    {
        false
    }
}

/// Device capabilities (stub for non-macOS)
#[cfg(not(target_os = "macos"))]
#[derive(Debug, Clone)]
pub struct DeviceCapabilities {
    pub name: String,
    pub max_threadgroup_memory: usize,
    pub max_threads_per_threadgroup: usize,
}

#[cfg(not(target_os = "macos"))]
pub fn query_device_capabilities() -> FlashResult<DeviceCapabilities> {
    Err(FlashError::not_available(
        "Metal not available on this platform",
    ))
}

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

    #[test]
    fn test_is_available() {
        // On macOS, Metal should be available
        #[cfg(target_os = "macos")]
        assert!(is_available());

        // On non-macOS, Metal is not available
        #[cfg(not(target_os = "macos"))]
        assert!(!is_available());
    }

    #[test]
    fn test_config_creation() {
        let config = FlashAttentionConfig::default();
        assert_eq!(config.block_q, 16); // Optimized for 32KB threadgroup memory
        assert_eq!(config.block_kv, 16);
        assert_eq!(config.max_head_dim, 64);
        assert!(!config.causal);

        let causal_config = FlashAttentionConfig::causal();
        assert!(causal_config.causal);

        let large_config = FlashAttentionConfig::large_head_dim();
        assert_eq!(large_config.max_head_dim, 128);
        assert_eq!(large_config.block_q, 8); // Smaller blocks for larger head_dim
    }
}