FlashAttention Implementation › Using FlashAttention in PyTorch

import torch
import torch.nn.functional as F


def demonstrate_flash_attention_pytorch():
    """
    Show how to use FlashAttention through PyTorch's scaled_dot_product_attention.
    """
    # Check if CUDA is available
    device = "cuda" if torch.cuda.is_available() else "cpu"

    # Create sample tensors
    batch_size = 2
    num_heads = 8
    seq_len = 1024
    d_head = 64

    # Shape: (batch, heads, seq_len, d_head)
    q = torch.randn(
        batch_size,
        num_heads,
        seq_len,
        d_head,
        device=device,
        dtype=torch.float16,
    )
    k = torch.randn(
        batch_size,
        num_heads,
        seq_len,
        d_head,
        device=device,
        dtype=torch.float16,
    )
    v = torch.randn(
        batch_size,
        num_heads,
        seq_len,
        d_head,
        device=device,
        dtype=torch.float16,
    )

    # Use scaled_dot_product_attention - automatically uses FlashAttention when possible
    # This is the recommended API for PyTorch 2.0+
    output = F.scaled_dot_product_attention(q, k, v)

    return output.shape, device


if torch.cuda.is_available():
    shape, device = demonstrate_flash_attention_pytorch()
    result = f"Output shape: {shape}, Device: {device}"
else:
    result = "CUDA not available - FlashAttention requires GPU"

# Code showing how to control attention backend
# This demonstrates the API, actual execution requires CUDA

example_code = """
import torch
import torch.nn.functional as F
from torch.backends.cuda import sdp_kernel

# Create inputs
q = torch.randn(2, 8, 1024, 64, device="cuda", dtype=torch.float16)
k = torch.randn(2, 8, 1024, 64, device="cuda", dtype=torch.float16)
v = torch.randn(2, 8, 1024, 64, device="cuda", dtype=torch.float16)

# Force FlashAttention only
with sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False):
    output = F.scaled_dot_product_attention(q, k, v)

# Force memory-efficient attention (xFormers-style)
with sdp_kernel(enable_flash=False, enable_math=False, enable_mem_efficient=True):
    output = F.scaled_dot_product_attention(q, k, v)

# Allow PyTorch to choose best implementation
with sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True):
    output = F.scaled_dot_product_attention(q, k, v)
"""

def flash_attention_requirements():
    """
    Document FlashAttention requirements and when PyTorch will fall back to alternatives.
    """
    requirements = {
        "Data type": {
            "supported": ["float16", "bfloat16"],
            "note": "float32 falls back to math kernel",
        },
        "Device": {
            "supported": ["CUDA (compute capability >= 8.0)"],
            "note": "A100, H100, RTX 30xx/40xx",
        },
        "Head dimension": {
            "supported": ["8, 16, 32, 64, 128, 256"],
            "note": "Must be one of these specific values",
        },
        "Sequence length": {
            "supported": ["Any length"],
            "note": "Longer sequences benefit more from FlashAttention",
        },
        "Attention mask": {
            "supported": ["Causal mask, no mask"],
            "note": "Arbitrary masks may fall back to other kernels",
        },
        "Dropout": {
            "supported": ["Yes, during training"],
            "note": "Implemented within the kernel",
        },
    }
    return requirements


reqs = flash_attention_requirements()

# Installation and direct usage of flash-attn library
flash_attn_example = """
# Install: pip install flash-attn --no-build-isolation

from flash_attn import flash_attn_func, flash_attn_qkvpacked_func

# Method 1: Separate Q, K, V tensors
# Shapes: (batch, seqlen, nheads, headdim)
q = torch.randn(2, 1024, 8, 64, device="cuda", dtype=torch.float16)
k = torch.randn(2, 1024, 8, 64, device="cuda", dtype=torch.float16)
v = torch.randn(2, 1024, 8, 64, device="cuda", dtype=torch.float16)

output = flash_attn_func(q, k, v, causal=True)

# Method 2: Packed QKV tensor (more memory efficient)
# Shape: (batch, seqlen, 3, nheads, headdim)
qkv = torch.randn(2, 1024, 3, 8, 64, device="cuda", dtype=torch.float16)
output = flash_attn_qkvpacked_func(qkv, causal=True)

# With dropout (training only)
output = flash_attn_func(q, k, v, dropout_p=0.1, causal=True)

# With softmax scaling
output = flash_attn_func(q, k, v, softmax_scale=1.0/8.0, causal=True)
"""

class FlashAttentionLayer(torch.nn.Module):
    """
    Multi-head attention using FlashAttention through PyTorch's API.
    """

    def __init__(self, d_model: int, num_heads: int, dropout: float = 0.0):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_head = d_model // num_heads
        self.dropout = dropout

        # Linear projections
        self.q_proj = torch.nn.Linear(d_model, d_model)
        self.k_proj = torch.nn.Linear(d_model, d_model)
        self.v_proj = torch.nn.Linear(d_model, d_model)
        self.o_proj = torch.nn.Linear(d_model, d_model)

    def forward(self, x: torch.Tensor, causal: bool = True) -> torch.Tensor:
        batch_size, seq_len, _ = x.shape

        # Project to Q, K, V
        q = self.q_proj(x)
        k = self.k_proj(x)
        v = self.v_proj(x)

        # Reshape for multi-head attention: (batch, seq, heads, d_head)
        q = q.view(batch_size, seq_len, self.num_heads, self.d_head)
        k = k.view(batch_size, seq_len, self.num_heads, self.d_head)
        v = v.view(batch_size, seq_len, self.num_heads, self.d_head)

        # Transpose to (batch, heads, seq, d_head) for PyTorch's SDPA
        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)

        # FlashAttention via scaled_dot_product_attention
        # Automatically uses FlashAttention when conditions are met
        attn_output = F.scaled_dot_product_attention(
            q,
            k,
            v,
            dropout_p=self.dropout if self.training else 0.0,
            is_causal=causal,
        )

        # Reshape back: (batch, seq, d_model)
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, self.d_model)

        # Final projection
        return self.o_proj(attn_output)


# Demonstrate the layer
model = FlashAttentionLayer(d_model=512, num_heads=8)