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 memory_savings(&self, seq_len: usize) -> String

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 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.