rustytorch/crates/training/rtx-flash-attention/METAL_SUPPORT.md

Metal GPU Support for rtx-flash-attention

Overview

This document describes the Metal GPU support implementation for the rtx-flash-attention crate, enabling Flash Attention on Apple Silicon and macOS devices.

Implementation Status

✅ Completed

1. Feature Configuration (Cargo.toml)

   • Added metal feature flag
   • Added Metal dependencies (objc2, objc2-metal, objc2-foundation, block2)
   • Feature-gated CUDA dependencies with dep: syntax
   • Properly configured for macOS target

2. CUDA Code Feature-Gating

   • All CUDA-specific code wrapped in #[cfg(feature = "cuda")]
   • Clean separation between CUDA and Metal backends
   • Files affected:
     • src/lib.rs - Feature-gated CUDA imports
     • src/core.rs - Separate impl blocks for CUDA and Metal
     • src/kernels/mod.rs - Feature-gated kernel modules
     • src/kernels/flash_forward.rs - CUDA-only
     • src/kernels/flash_backward.rs - CUDA-only
     • src/kernels/simple.rs - CUDA-only

3. Metal Kernel Implementation (src/kernels/metal.rs)

   • FlashMetalKernels struct with Metal device and command queue
   • MetalKernelResult for performance metrics
   • Forward pass stub implementation
   • Backward pass stub (returns not-implemented error)
   • Block size optimization for Metal hardware
   • Device information querying
   • Tensor validation

4. Metal Shader Source (metal/flash_attention.metal)

   • MSL (Metal Shading Language) shader structure
   • Stub implementations showing:
     • Forward kernel signature with threadgroup parameters
     • Backward kernels (dQ, dK/dV) signatures
     • Comprehensive implementation notes
   • Ready for full tiled attention implementation

5. Metal Backend in Core (src/core.rs)

   • Separate FlashAttention struct for Metal feature
   • Full FlashAttentionBackend trait implementation
   • Performance statistics tracking
   • Validation and error handling
   • Factory methods for Metal backend creation

6. Build System (build.rs)

   • Feature-gated CUDA compilation
   • Metal shader detection and path setup
   • Proper warnings for missing backends
   • Environment variable setup for shader paths

7. Test Suite (src/kernels/metal_test.rs)

   • Device availability tests
   • Kernel creation tests
   • Block size optimization validation
   • Thread safety tests
   • Forward/backward pass behavior verification
   • Multiple kernel instance tests

Key Design Decisions

Feature Flag Strategy

[features]
default = ["cuda"]
cuda = ["rtx-kernel/cuda", "rtx-tensor/cuda", "dep:cudarc"]
metal = ["rtx-kernel/metal", "rtx-tensor/metal", "dep:objc2", "dep:objc2-metal", "dep:objc2-foundation", "dep:block2"]

• Default: CUDA for backward compatibility
• Explicit dependencies: Using dep: prefix for optional dependencies
• Cascading features: Enabling rtx-kernel and rtx-tensor GPU features

Backend Abstraction

pub trait FlashAttentionBackend {
    async fn forward(...) -> FlashResult<FlashOutput>;
    async fn backward(...) -> FlashResult<FlashGradOutput>;
    fn name(&self) -> &str;
    fn supports_config(&self, config: &FlashAttentionConfig) -> bool;
    fn optimize_config(&self, config: FlashAttentionConfig) -> FlashResult<FlashAttentionConfig>;
}

• Trait-based: Clean abstraction over CUDA and Metal
• Async: Non-blocking GPU operations
• Configuration: Backend-specific optimization

Metal Hardware Considerations

• Threadgroup memory: 32KB typical (vs CUDA's 160KB on RTX 5090)
• Max threads: 1024 per threadgroup
• Block sizes: Optimized for 32-256 threads
• No explicit tensor cores: Metal abstracts this

Usage

Building with Metal Support

# Metal only (macOS)
cargo build -p rtx-flash-attention --no-default-features --features metal

# Both CUDA and Metal
cargo build -p rtx-flash-attention --features "cuda,metal"

# Default (CUDA only)
cargo build -p rtx-flash-attention

Runtime Detection

use rtx_flash_attention::core::FlashAttentionFactory;

// Check backend availability
let cuda_available = FlashAttentionFactory::cuda_available();
let metal_available = FlashAttentionFactory::metal_available();

// Get device capabilities
let caps = FlashAttentionFactory::device_capabilities()?;
println!("GPU: {:?}, Memory: {} GB", caps.compute_capability, caps.total_memory / 1024 / 1024 / 1024);

Creating Flash Attention with Metal

use rtx_flash_attention::{FlashAttention, FlashAttentionConfig};

// Create configuration
let config = FlashAttentionConfig::new(8, 64);

// Create Flash Attention instance (automatically selects Metal on macOS with metal feature)
let flash = FlashAttention::new(config)?;

// Run forward pass
let output = flash.forward(&q, &k, &v, false, 0.125).await?;

Current Limitations

Known Issues

1. rtx-runtime Dependency: The crate currently cannot compile with --no-default-features --features metal because rtx-runtime has hardcoded CUDA dependencies that are not properly feature-gated. This needs to be fixed in rtx-runtime first.

2. Stub Implementation: The Metal kernels are stubs that:

   • Return zero tensors for forward pass
   • Return not-implemented error for backward pass
   • Don't perform actual attention computation

3. No Tensor Integration: Tests use rtx-tensor but Metal device support in rtx-tensor needs verification

Next Steps for Full Implementation

1. Complete Metal Shader Implementation

The shader stubs in metal/flash_attention.metal need full implementation:

kernel void flash_attention_forward(...) {
    // 1. Load Q tile into threadgroup memory
    threadgroup half Q_tile[BLOCK_SIZE_Q][HEAD_DIM];

    // 2. Loop over K/V tiles
    for (uint kv_block = 0; kv_block < num_kv_blocks; ++kv_block) {
        // Load K, V tiles

        // Compute attention scores: S = Q @ K^T

        // Online softmax with running max/sum
        float m_new = simd_max(m_old, row_max(scores));
        float l_new = l_old * exp(m_old - m_new) + row_sum(exp(scores - m_new));

        // Accumulate output: O = softmax(S) @ V
        // Apply causal masking if needed
    }

    // Write final output and LSE
}

Key Metal APIs to use:

• threadgroup memory for tiles
• simd_max() / simd_sum() for reductions
• simd_shuffle_down() for warp-level communication
• threadgroup_barrier() for synchronization

2. Kernel Loading and Compilation

Update FlashMetalKernels to actually compile and load shaders:

impl FlashMetalKernels {
    pub fn new(config: &FlashAttentionConfig) -> FlashResult<Self> {
        let device = unsafe { MTLCreateSystemDefaultDevice() }?;
        let queue = device.newCommandQueue()?;

        // Compile Metal shader
        let shader_source = include_str!(env!("METAL_SHADER_PATH"));
        let library = device.newLibraryWithSource_options_error(shader_source, None)?;

        // Create pipeline states
        let forward_function = library.newFunctionWithName(ns_string!("flash_attention_forward"))?;
        let forward_pipeline = device.newComputePipelineStateWithFunction_error(forward_function)?;

        Ok(Self { device, queue, forward_pipeline, ... })
    }
}

3. Buffer Management

Implement Metal buffer creation and data transfer:

pub async fn flash_attention_forward(
    &self,
    q: &Tensor,
    k: &Tensor,
    v: &Tensor,
    output: &mut Tensor,
    lse: &mut Tensor,
    softmax_scale: f32,
    causal: bool,
) -> FlashResult<MetalKernelResult> {
    // Get Metal buffers from tensors
    let q_buffer = q.metal_buffer()?;
    let k_buffer = k.metal_buffer()?;
    let v_buffer = v.metal_buffer()?;
    let output_buffer = output.metal_buffer_mut()?;
    let lse_buffer = lse.metal_buffer_mut()?;

    // Create command buffer
    let command_buffer = self.command_queue.commandBuffer()?;
    let encoder = command_buffer.computeCommandEncoder()?;

    // Set pipeline and buffers
    encoder.setComputePipelineState(&self.forward_pipeline);
    encoder.setBuffer_offset_atIndex(&q_buffer, 0, 0);
    encoder.setBuffer_offset_atIndex(&k_buffer, 0, 1);
    encoder.setBuffer_offset_atIndex(&v_buffer, 0, 2);
    encoder.setBuffer_offset_atIndex(&output_buffer, 0, 3);
    encoder.setBuffer_offset_atIndex(&lse_buffer, 0, 4);

    // Dispatch
    let grid_size = MTLSize::new(num_blocks, num_heads, batch_size);
    let threadgroup_size = MTLSize::new(threads_per_group, 1, 1);
    encoder.dispatchThreadgroups_threadsPerThreadgroup(grid_size, threadgroup_size);

    encoder.endEncoding();
    command_buffer.commit();
    command_buffer.waitUntilCompleted();

    Ok(MetalKernelResult { ... })
}

4. Performance Optimization

Metal-specific optimizations:

• Use simdgroup functions for warp-level operations
• Optimize threadgroup memory layout for bank-conflict-free access
• Use Metal Performance Shaders (MPS) for matrix operations if beneficial
• Profile with Xcode Instruments Metal profiler
• Consider using MTLGPUFamily capabilities detection

5. Fix rtx-runtime Dependencies

The rtx-runtime crate needs:

• Feature-gate all CUDA backend code with #[cfg(feature = "cuda")]
• Implement Metal stream and kernel launch abstractions
• Update Stream::raw_stream() to support Metal command queues
• Feature-gate kernel launch system

6. Integration Tests

Create end-to-end tests:

#[tokio::test]
#[cfg(all(feature = "metal", target_os = "macos"))]
async fn test_metal_flash_attention_correctness() {
    let config = FlashAttentionConfig::new(8, 64);
    let flash = FlashAttention::new(config)?;

    // Create random inputs
    let device = Device::Metal(0);
    let q = Tensor::randn(&[2, 8, 512, 64], &device)?;
    let k = Tensor::randn(&[2, 8, 512, 64], &device)?;
    let v = Tensor::randn(&[2, 8, 512, 64], &device)?;

    // Run flash attention
    let output = flash.forward(&q, &k, &v, false, 0.125).await?;

    // Compare with naive attention
    let expected = naive_attention(&q, &k, &v, 0.125)?;

    assert_tensors_close(&output.output, &expected, 1e-3)?;
}

Architecture Diagram

rtx-flash-attention
├── Cargo.toml (features: cuda, metal)
├── build.rs (feature-gated compilation)
├── metal/
│   └── flash_attention.metal (MSL shaders)
├── cuda/ (feature = "cuda")
│   ├── flash_attention_forward.cu
│   └── flash_attention_backward.cu
└── src/
    ├── lib.rs (public API)
    ├── core.rs (backend trait + impls)
    ├── config.rs (configuration)
    ├── error.rs (error types)
    └── kernels/
        ├── mod.rs (feature-gated exports)
        ├── metal.rs (feature = "metal")
        ├── flash_forward.rs (feature = "cuda")
        └── flash_backward.rs (feature = "cuda")

Testing

Run Metal Tests (macOS only)

# All Metal tests
cargo test -p rtx-flash-attention --features metal -- --test-threads=1

# Specific test
cargo test -p rtx-flash-attention --features metal test_metal_device_available

# With output
cargo test -p rtx-flash-attention --features metal -- --nocapture

Expected Test Results (Current Stub Implementation)

• ✅ Device availability tests pass
• ✅ Kernel creation tests pass
• ✅ Block size optimization tests pass
• ⚠️ Forward pass returns zeros (stub)
• ❌ Backward pass returns not-implemented error

References

Metal Documentation

• Metal Shading Language Specification
• Metal Best Practices Guide
• objc2-metal crate

Flash Attention

• Flash Attention Paper - Dao et al., 2022
• Flash Attention 2 - Dao, 2023

Similar Implementations

• MLX Flash Attention - Apple's ML framework with Metal Flash Attention
• PyTorch Metal Backend

Contributing

When implementing the full Metal backend:

1. Follow TDD: Write tests first, then implementation
2. No Stubs: All functions must have real implementations
3. Performance First: Optimize for Metal hardware characteristics
4. Correctness: Ensure numerical accuracy matches CUDA implementation
5. Documentation: Document all Metal-specific optimizations

License

Same as rtx-flash-attention: MIT OR Apache-2.0