Struct FlashAttention 

pub struct FlashAttention { /* private fields */ }

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

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);

Implementations

impl FlashAttention

pub fn new(dim: usize, num_heads: usize, causal: bool) -> Self

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 forward( &self, q: &UnifiedTensor, k: &UnifiedTensor, v: &UnifiedTensor, ctx: &WgpuContext, engine: &ComputeEngine, ) -> UnifiedTensor

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_multi_head( &self, q: &UnifiedTensor, k: &UnifiedTensor, v: &UnifiedTensor, ctx: &WgpuContext, engine: &ComputeEngine, ) -> UnifiedTensor

Multi-head attention forward pass with automatic head splitting.

This method automatically splits the input into heads and applies Flash Attention to each head in parallel.

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 multi_head_forward( &self, hidden: &UnifiedTensor, ctx: &WgpuContext, engine: &ComputeEngine, ) -> UnifiedTensor

Multi-head attention with Flash Attention.

Handles Q, K, V projection and output projection internally.

pub fn memory_savings(&self, seq_len: usize) -> String

Get memory savings compared to standard attention.

pub fn backward( &self, q: &UnifiedTensor, k: &UnifiedTensor, v: &UnifiedTensor, output: &UnifiedTensor, grad_output: &UnifiedTensor, ctx: &WgpuContext, engine: &ComputeEngine, ) -> FlashAttentionGrads

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

Trait Implementations

impl Debug for FlashAttention

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.