FlashAttention Implementation in PyTorch

This implementation builds FlashAttention-style multi-head attention from scratch. It includes a standard scaled dot-product attention module, a tiled/blocked variant that processes the sequence in chunks to reduce peak memory, and a test function covering shapes and variable sequence lengths.

import torch
import torch.nn as nn
import torch.nn.functional as F

class FlashAttention(nn.Module):
    def __init__(self, dim, num_heads=8, dropout=0.1, max_seq_length=1024):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scaling = self.head_dim ** -0.5
        
        self.max_seq_length = max_seq_length
        self.dropout = dropout

        # Linear projections
        self.q_proj = nn.Linear(dim, dim)
        self.k_proj = nn.Linear(dim, dim)
        self.v_proj = nn.Linear(dim, dim)
        self.out_proj = nn.Linear(dim, dim)

    def forward(self, x, mask=None):
        batch_size, seq_length, dim = x.shape

        # Project and reshape queries, keys, and values
        q = self.q_proj(x).reshape(batch_size, seq_length, self.num_heads, self.head_dim)
        k = self.k_proj(x).reshape(batch_size, seq_length, self.num_heads, self.head_dim)
        v = self.v_proj(x).reshape(batch_size, seq_length, self.num_heads, self.head_dim)
        
        q = q.transpose(1, 2)  # (batch_size, num_heads, seq_length, head_dim)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # Scaled dot-product attention 
        attn_weights = torch.matmul(q, k.transpose(-2, -1)) * self.scaling
        
        if mask is not None:
            attn_weights = attn_weights.masked_fill(mask == 0, float('-inf'))
        
        attn_weights = F.softmax(attn_weights, dim=-1)
        attn_weights = F.dropout(attn_weights, p=self.dropout, training=self.training)

        attn_output = torch.matmul(attn_weights, v)
        
        attn_output = attn_output.transpose(1, 2).reshape(batch_size, seq_length, dim)
        attn_output = self.out_proj(attn_output)

        return attn_output
    
class FlashAttentionWithTiling(FlashAttention):
    """
    tiling for better memory efficiency.
    This splits the input sequence into blocks and processes them in chunks.
    """
    def __init__(self, dim, num_heads=8, dropout=0.1, max_seq_length=1024, block_size=64):
        super().__init__(dim, num_heads, dropout, max_seq_length)
        self.block_size = block_size

    def forward(self, x, mask=None):
        batch_size, seq_length, dim = x.shape

        q = self.q_proj(x).reshape(batch_size, seq_length, self.num_heads, self.head_dim)
        k = self.k_proj(x).reshape(batch_size, seq_length, self.num_heads, self.head_dim)
        v = self.v_proj(x).reshape(batch_size, seq_length, self.num_heads, self.head_dim)
        
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        output = torch.zeros_like(q)

        # Process in blocks
        for i in range(0, seq_length, self.block_size):
            q_block = q[:, :, i:min(i + self.block_size, seq_length), :]

            # Initialize block attention
            block_weights = torch.zeros(batch_size, self.num_heads, 
                                     q_block.size(2), seq_length, 
                                     device=q.device)
            
            for j in range(0, seq_length, self.block_size):
                k_block = k[:, :, j:min(j + self.block_size, seq_length), :]
                v_block = v[:, :, j:min(j + self.block_size, seq_length), :]
                
                scores = torch.matmul(q_block, k_block.transpose(-2, -1)) * self.scaling
                
                if mask is not None:
                    block_mask = mask[:, :, i:i + self.block_size, j:j + self.block_size]
                    scores = scores.masked_fill(block_mask == 0, float('-inf'))
                
                block_weights[:, :, :, j:j + k_block.size(2)] = scores

            block_weights = F.softmax(block_weights, dim=-1)
            block_weights = F.dropout(block_weights, p=self.dropout, training=self.training)
            
            block_output = torch.zeros_like(q_block)

            # Compute block output by processing each key/value block
            for j in range(0, seq_length, self.block_size):
                v_block = v[:, :, j:min(j + self.block_size, seq_length), :]
                block_weights_slice = block_weights[:, :, :, j:j + v_block.size(2)]
                block_output += torch.matmul(block_weights_slice, v_block)

            output[:, :, i:i + q_block.size(2), :] = block_output

        output = output.transpose(1, 2).reshape(batch_size, seq_length, dim)
        output = self.out_proj(output)

        return output
    
def test_flash_attention():
    batch_size = 2
    seq_length = 128
    dim = 256
    num_heads = 8

    flash_attn = FlashAttention(dim=dim, num_heads=num_heads)
    flash_attn_tiling = FlashAttentionWithTiling(dim=dim, num_heads=num_heads, block_size=32)
    
    x = torch.randn(batch_size, seq_length, dim)
    
    mask = torch.ones(batch_size, 1, seq_length, seq_length)
    
    output1 = flash_attn(x, mask)
    
    output2 = flash_attn_tiling(x, mask)
    
    assert output1.shape == (batch_size, seq_length, dim)
    assert output2.shape == (batch_size, seq_length, dim)

    # Verify that outputs are different
    assert not torch.allclose(output1, output2, atol=1e-4)
    
    print("Basic shape tests passed!")

    # Test with different sequence lengths
    x_short = torch.randn(batch_size, seq_length // 2, dim)
    mask_short = torch.ones(batch_size, 1, seq_length // 2, seq_length // 2)
    
    output_short1 = flash_attn(x_short, mask_short)
    output_short2 = flash_attn_tiling(x_short, mask_short)
    
    assert output_short1.shape == (batch_size, seq_length // 2, dim)
    assert output_short2.shape == (batch_size, seq_length // 2, dim)
    
    print("Variable sequence length tests passed!")

    return "All tests passed successfully!"

if __name__ == "__main__":
    test_flash_attention()