Skip to content

Flash Attention: Complete Technical Guide

Overview

Flash Attention is a fast and memory-efficient attention algorithm that reduces the time and memory complexity of transformer attention through I/O-aware computation. It achieves 2-3x speedup and uses significantly less memory than standard attention.

  • Paper: "Flash Attention: Fast and Memory-Efficient Exact Attention with IO-Awareness"
  • Authors: Dao et al., Stanford/MIT-IBM (2022)
  • Key Innovation: Minimize data movement (I/O), not just FLOPs
  • Impact: 2-3x faster attention, reduced memory footprint
  • Adoption: Standard in vLLM, HuggingFace, PyTorch
  • Versions: Flash-1, Flash-2 (2x faster), Flash-3 (in progress)

The Problem: Standard Attention Bottleneck

Standard Attention Computation

Input:
Q: Query matrix (N × d_k)
K: Key matrix (N × d_k)
V: Value matrix (N × d_v)
N: Sequence length
d_k: Key dimension
d_v: Value dimension

Standard Algorithm:
1. S = Q @ K^T                    # (N × N) attention scores
2. S_scaled = S / sqrt(d_k)      # Scale
3. P = softmax(S_scaled, dim=-1)  # (N × N) attention weights
4. O = P @ V                      # (N × d_v) output

Memory usage:
  - Q, K, V: O(N × d)
  - S: O(N²) ← HUGE! This is the problem!
  - P: O(N²)
  - O: O(N × d)

For N=4096 sequence, d=64:
  - S matrix: 4096² × 4 bytes = 67MB
  - P matrix: 4096² × 4 bytes = 67MB
  - Total for just attention: 134MB!

For N=16384 (longer sequences):
  - S matrix: 16384² × 4 bytes = 1GB!
  - This alone can exceed GPU memory!

Where Time is Spent

Attention computation breakdown:

For N=4096, d=64:

Operation            FLOPs      Memory(GB)   Time      Bottleneck
─────────────────────────────────────────────────────────────────
Q @ K^T             2N²d        0.03         Compute  Memory
Softmax             N²          0.13         Memory   I/O
P @ V               2N²d        0.03         Compute  Memory
Total               4N²d        0.20         Variable I/O-bound!

Key insight:
  - FLOPs are reasonable (matrix multiply)
  - But moving data between GPU memory and cache is slow
  - Softmax requires reading ENTIRE (N × N) matrix
  - This is I/O-bound, not compute-bound!

GPU memory hierarchy:
Registers (32 KB, fast)
  - Can keep small blocks
L2 Cache (40 MB, medium)
  - Can fit maybe 512×512 values
HBM (40 GB, slow)
  - Where S and P matrices live
  - Very slow to access!

Problem: S matrix (N² elements) doesn't fit in cache!
→ Constant back-and-forth to HBM
→ Memory becomes bottleneck

Roofline Analysis

GPU A100 Characteristics:
  - Peak compute: 312 TFLOPS
  - Memory bandwidth: 2 TB/s
  - Compute-to-bandwidth ratio: 156 FLOP/byte

For standard attention:
Bytes moved: O(N²)
Computation: O(N² × d)
Compute intensity: (N² × d) / N² = d = 64

But A100 can sustain ~312 TB/s * (156 FLOP/byte) / d = 
    312 TB/s * 156 / 64 = 760 TFLOPS (if we had good I/O)

In reality:
  - Need to write/read S, P from HBM
  - Massive I/O overhead
  - Actual throughput: ~50-100 TFLOPS
  - Terrible! We're wasting 60-80% of GPU potential!

Flash Attention:
  - Blocks computation to fit in cache
  - Reduces I/O from O(N²) to O(N)
  - Actual throughput: ~200 TFLOPS
  - 2-4x improvement!

Flash Attention: The Solution

Core Idea: Tiling and Recomputation

Key insight from operating systems:
"Memory hierarchies are ubiquitous - design around them!"

Standard approach wastes cache:
  - Load K, V to process all Q
  - Compute full S matrix (doesn't fit in cache!)
  - Compute full P matrix (doesn't fit in cache!)
  - Lots of I/O

Flash Attention approach:
Divide computation into blocks ("tiles") that fit in cache!

1. Process Q in tiles (e.g., 64 rows at a time)
2. For each Q tile, iterate through K, V tiles
3. Accumulate results incrementally
4. Keep intermediate results in fast cache (SRAM)
5. Write only final results to slow HBM

Benefits:
  - Fit intermediate S, P in fast cache
  - Reduce I/O to HBM
  - Trade: More compute for less I/O (good trade!)

Visual Algorithm

Standard Attention (wasteful I/O):
- ┌─────────────────────────────────┐
    - HBM (GPU Memory - SLOW)         │
    - ┌─────────────────────────────┐ │
        - Q (all N tokens)            │ │ Read once
        - K (all N tokens)            │ │ Read many times
        - V (all N tokens)            │ │
        - S[N × N] (all scores)       │ │ ← Write, Read, Write (wasteful!)
        - P[N × N] (all weights)      │ │ ← Write, Read (wasteful!)
        - O[N × d] (output)           │ │ Write once
      - ┘ │
    - ↑        ↑      ↓        │
  - ┘
           I/O-bottleneck!

Flash Attention (efficient I/O):
- ┌─────────────────────────────────┐
    - HBM (GPU Memory)                │
    - ┌──────────┬──────────┐         │
          - Q-block1 │ Q-block2 │ ...     │ Read once
          - K (all)  │          │         │
          - V (all)  │          │         │
          - O-block1 │ O-block2 │ ...     │ Write once
      - ┴──────────┘         │
    - ↓             ↓              │
  - ┘
  SRAM (GPU Cache - FAST)
- ┌─────────────────────────┐
    - Q-block, K-block, V-blk │ Compute here!
    - Intermediate S, P       │ (fits in cache)
    - Incremental accum       │
  - ┘

I/O Complexity Analysis

Standard Attention I/O

Work analysis (FLOPs):
S = Q @ K^T:     O(N² × d) FLOPs
softmax(S):      O(N²) FLOPs
P = P @ V:       O(N² × d) FLOPs
Total:           O(N² × d) FLOPs

Memory I/O analysis:
Q ∈ HBM → registers: O(N × d) I/O
K ∈ HBM → registers: O(N × d) I/O (for each of N blocks)
V ∈ HBM → registers: O(N × d) I/O (for each of N blocks)
S ∈ HBM: O(N²) I/O (write and read!)
P ∈ HBM: O(N²) I/O (write and read!)
O ∈ HBM: O(N × d) I/O (write once)

Total I/O: O(N²) ← HUGE!

I/O to Compute Ratio:
Bytes/FLOP = O(N²) / O(N² × d) = 1/d ≈ 1/64

This means for every 64 FLOPs, we move 1 byte.
GPU memory bandwidth: 2TB/s
Effective throughput: 2TB/s × 64 ops/byte = 128 TFLOPS
But peak is 312 TFLOPS → 40% utilization (bad!)

Flash Attention I/O

Work analysis (FLOPs):
Still O(N² × d) FLOPs (exact same computation!)

Memory I/O analysis (with M = cache size in elements):
Q ∈ HBM → SRAM: O(N × d) I/O
K ∈ HBM → SRAM: O(N × d) I/O
V ∈ HBM → SRAM: O(N × d) I/O
O ∈ HBM: O(N × d) I/O (write once)

No need to write/read S, P (keep in SRAM)!

Total I/O: O(N × d) ← HUGE IMPROVEMENT!

I/O to Compute Ratio:
Bytes/FLOP = O(N × d) / O(N² × d) = 1/N ≈ 1/4096 (much better!)

Effective throughput: 2TB/s × 4096 ops/byte = 8.2 PETA-FLOPS!
But peak is 312 TFLOPS → this analysis is simplified
In practice: ~2-3x speedup (I/O reduced from N² to N terms)

Comparison Table

Complexity Metric        Standard    Flash-Attn   Improvement
──────────────────────────────────────────────────────────────
FLOPs                    O(N²d)      O(N²d)       Same
Total I/O                O(N²)       O(Nd)        O(N) factor!
Memory peak              O(N²)       O(M)         Huge!
For N=4096:
  Standard I/O           67 MB       0.25 MB      267x less!
  Standard memory peak   67 MB       varies
  Flash peak (M=4MB)     4 MB        constant!

M = SRAM size (typically 2-4 MB on A100)

Flash Attention Algorithm in Detail

High-Level Algorithm

Algorithm: FlashAttention

Input:
  Q: (N, d) query matrix
  K: (N, d) key matrix
  V: (N, d) value matrix

Config:
  B_r: block size for rows (e.g., 64)
  B_c: block size for columns (e.g., 64)

def flash_attention(Q, K, V, B_r, B_c):
    N = Q.shape[0]
    O = zeros((N, d))  # Output
    l = zeros(N)       # Softmax denominator
    m = -inf * ones(N) # Softmax max (for numerical stability)

    # Main loop: process Q in blocks
    for i in range(0, N, B_r):
        i_end = min(i + B_r, N)
        Q_block = Q[i:i_end]  # Load Q block to SRAM

        # Sub-loop: compute attention with all K, V blocks
        for j in range(0, N, B_c):
            j_end = min(j + B_c, N)
            K_block = K[j:j_end]  # Load K block to SRAM
            V_block = V[j:j_end]  # Load V block to SRAM

            # Compute attention for this block pair (stays in SRAM)
            S_block = Q_block @ K_block.T  # (B_r, B_c)

            # Numerical stability: track maximum
            m_new = max(m[i:i_end], max_row(S_block))

            # Softmax trick: compute in log-space
            P_block = exp(S_block - m_new)
            l_new = exp(m[i:i_end] - m_new) * l[i:i_end] + sum_row(P_block)

            # Accumulate output
            O[i:i_end] += (P_block @ V_block)

            # Update tracking variables
            m[i:i_end] = m_new
            l[i:i_end] = l_new

    # Normalize output
    O = O / l  # Divide by softmax denominator

    return O

Key Technique: Softmax with Reductions

Challenge: Softmax is not associative (can't compute in blocks)

Standard softmax:
P[i] = exp(S[i]) / sum(exp(S))

Problem: Don't know the denominator until we've seen all S!

Flash Attention solution: Online softmax (numerically stable)

For a single block:
1. Find max: m = max(S)
2. Compute: exp_S = exp(S - m)  ← Subtract max for stability
3. Sum: l = sum(exp_S)
4. Result: P = exp_S / l

For multiple blocks (online update):
Block 1: m_1, l_1
Block 2: m_2, l_2

Updated max: m_new = max(m_1, m_2)
Updated sum: l_new = exp(m_1 - m_new) * l_1 + exp(m_2 - m_new) * l_2

This is EXACT softmax, just computed incrementally!

Accumulated output:
O_new = (exp(m_1 - m_new) * O_1 + exp(m_2 - m_new) * (P_2 @ V))

Memory Layout

SRAM allocation (Typical A100: 4 MB = 4M bytes):

For 64-token blocks with 64-dim:
  - Q_block: 64 × 64 × 2 bytes = 8 KB
  - K_block: 64 × 64 × 2 bytes = 8 KB
  - V_block: 64 × 64 × 2 bytes = 8 KB
  - S_block: 64 × 64 × 4 bytes = 16 KB (larger, float32)
  - P_block: 64 × 64 × 2 bytes = 8 KB
  - O_block: 64 × 64 × 2 bytes = 8 KB
  - Scalars (m, l): ~1 KB
  - Total: ~57 KB (well within 4 MB SRAM!)

This is why tiling works!
All intermediate values fit comfortably in SRAM.

Flash Attention vs Standard: Step-by-Step Example

Concrete Example: N=256, d=64

Standard Attention:

Step 1: Q @ K^T
  - Load Q from HBM: 256 × 64 × 2 = 32 KB (fast, once)
  - Load K from HBM: 256 × 64 × 2 = 32 KB
  - Compute: S = Q @ K^T (256 × 256 × 64 = 4M FLOPs)
  - Write S to HBM: 256 × 256 × 4 = 256 KB (slow!)
  - Data movement: 32 + 32 + 256 = 320 KB

Step 2: Softmax(S)
  - Load S from HBM: 256 KB (slow)
  - Compute softmax: 256 × 256 = 65K FLOPs
  - Write P to HBM: 256 KB (slow)
  - Data movement: 256 + 256 = 512 KB

Step 3: P @ V
  - Load P from HBM: 256 KB (slow)
  - Load V from HBM: 32 KB
  - Compute: O = P @ V (256 × 256 × 64 = 4M FLOPs)
  - Write O to HBM: 32 KB
  - Data movement: 256 + 32 + 32 = 320 KB

Total HBM I/O: 32 + 32 + 256 + 256 + 256 + 32 + 32 = 896 KB

Flash Attention (block_size = 64):

Outer loop: i = 0, 64, 128, 192 (4 blocks)
  Inner loop for each i: j = 0, 64, 128, 192 (4 blocks)
    Load Q_i (64 × 64 × 2 = 8 KB) - once per i
    Load K_j, V_j (16 KB) - 4 times per i
    Compute S, P in SRAM (no I/O)
    Accumulate O (no write until end)

Total HBM I/O:
  - Load Q: 4 × 8 KB = 32 KB (same as standard)
  - Load K: 4 × 4 × 8 KB = 128 KB (vs 32 KB standard - wait, more?)
  - Load V: 4 × 4 × 8 KB = 128 KB
  - Write O: 4 × 8 KB = 32 KB (same as standard)
  - Total: 320 KB (less than standard's 896 KB!)

Savings: (896 - 320) / 896 = 64% I/O reduction!

But actual speedup is 2-3x because:
1. S, P no longer ping-pong to/from HBM (huge savings!)
2. Better cache locality
3. Overlapping computation and I/O

Performance Improvements

Benchmark Results

Measurement: Attention latency (lower = better)

Model: BERT-base (sequence length 512)
Hardware: A100 GPU

Method                  Time (ms)  Speedup  Memory (GB)
─────────────────────────────────────────────────────
Standard Attention      4.2        1x       0.35
PyTorch Optimized       3.8        1.1x     0.30
Flash Attention v1      2.0        2.1x     0.05
Flash Attention v2      1.5        2.8x     0.05

Measurement: Full transformer layer

Model: GPT-2 (1.5B params)
Sequence length: 1024
Hardware: A100

Layer Component         Standard   Flash-v2  Speedup
─────────────────────────────────────────────────────
Attention              12.5 ms    4.3 ms    2.9x
Feedforward            15.2 ms    15.1 ms   1.0x
Total layer            27.7 ms    19.4 ms   1.4x

Memory Usage Reduction

Sequence length: 4096
Hidden dimension: 768
Batch size: 1

                    Standard    Flash-Attn  Reduction
─────────────────────────────────────────────────────
Attention scores    67 MB       0 MB        100%
Attention weights   67 MB       0 MB        100%
Activations saved   320 MB      200 MB      37%
KV cache            24 MB       24 MB       0%

Total peak memory   478 MB      224 MB      53%

This means:
- Can fit 2x larger batch on same GPU
- Or 2x longer sequence

Flash Attention Versions

Flash Attention v1 (Original, 2022)

Features:
  - IO-aware attention
  - Exact computation (same output as standard)
  - 2x speedup, 60% memory reduction
  - For forward pass

Limitations:
  - Backward pass not optimized
  - Limited to specific block sizes
  - Some numerical precision issues

Code:
```python
import torch
from flash_attn import flash_attn_func

# Flash Attention v1
output = flash_attn_func(q, k, v)
### Flash Attention v2 (2023, Improved)
Major improvements over v1: - Backward pass also optimized - 2.8x speedup (vs 2x in v1) - Supports variable sequence lengths - Better numerical stability - Supports head_dim up to 256

Typical speedup: For sequence_length = 1024, head_dim = 64: - v1: 2.0x faster - v2: 2.8x faster

Code:

from flash_attn import flash_attn_func

# Flash Attention v2 (automatic)
output = flash_attn_func(q, k, v)

New features: - flash_attn_kvpacked_func: KV packed format - flash_attn_qkvpacked_func: All packed - Support for variable sequence lengths in batch

### Flash Attention v3 (2024, In Progress)
Planned improvements: - Support for FP8 quantization - Grouped Query Attention (GQA) optimization - Multi-GPU coordination - Speculative decoding integration

Expected: - 3-4x speedup - Better for quantized models - Seamless integration with vLLM

---

## Implementation & Integration

### PyTorch Integration

```python
# Option 1: Using Flash Attention directly
import torch
from flash_attn import flash_attn_func

Q = torch.randn(batch, seq_len, num_heads, head_dim, device='cuda')
K = torch.randn(batch, seq_len, num_heads, head_dim, device='cuda')
V = torch.randn(batch, seq_len, num_heads, head_dim, device='cuda')

# Flash attention (with causal mask for autoregressive)
output = flash_attn_func(
    Q, K, V,
    causal=True,  # For training/generation
    dropout_p=0.1,
    softmax_scale=None  # Auto: 1/sqrt(head_dim)
)

# Option 2: Using scaled_dot_product_attention (PyTorch 2.0+)
output = torch.nn.functional.scaled_dot_product_attention(
    Q, K, V,
    is_causal=True,
    scale=1.0 / (head_dim ** 0.5)
)
# PyTorch automatically uses Flash Attention if available!

Transformer Integration

class FlashAttention(nn.Module):
    def __init__(self, dim, num_heads=8, dropout=0.0):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.dropout = dropout

        self.qkv = nn.Linear(dim, 3 * dim)
        self.out = nn.Linear(dim, dim)

    def forward(self, x, causal=False):
        B, N, C = x.shape

        # Project to Q, K, V
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
        q, k, v = qkv.permute(2, 0, 3, 1, 4)  # (3, B, H, N, D)

        # Flash Attention
        from flash_attn import flash_attn_func
        out = flash_attn_func(
            q, k, v,
            dropout_p=self.dropout,
            causal=causal
        )

        # Merge heads
        out = out.reshape(B, N, C)
        out = self.out(out)

        return out

HuggingFace Integration

# Most HuggingFace models automatically use Flash Attention!

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b",
    attn_implementation="flash_attention_2",  # Enable Flash Attention
    torch_dtype="float16"
)

# Or enable globally
import torch
torch.backends.cuda.enable_flash_sdp(True)

vLLM Integration

# vLLM automatically uses Flash Attention!

from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Llama-2-7b-hf",
    # Flash Attention automatically enabled
)

outputs = llm.generate(prompts, sampling_params)

Numerical Considerations

Numerical Stability

Challenge: Softmax can overflow/underflow

exp(1000) → ∞ (overflow)
exp(-1000) → 0 (underflow)

Standard softmax (numerically unstable):
P[i] = exp(S[i]) / sum(exp(S))

Flash Attention (numerically stable):
1. m = max(S)  # Find maximum
2. P[i] = exp(S[i] - m) / sum(exp(S[j] - m))  # Subtract max

Example:
S = [1000, 1001, 999.9]

Naive: exp(1000) = ∞ (overflow!)

Stable:
m = 1001
P = [exp(-1), exp(0), exp(-1.1)] / sum(...)
  = [0.368, 1.0, 0.333] / 1.70
  = [0.216, 0.588, 0.196] (correct!)

Gradient Stability

Flash Attention v1: Gradient computation could be unstable

Flash Attention v2: Improved gradient stability
  - Better numerical behavior in backward pass
  - More accurate gradients
  - Safer for training

For most use cases: v2 is now standard

Limitations and Trade-offs

What Flash Attention Does NOT Do

❌ Doesn't reduce algorithmic complexity
  - Still O(N²d) FLOPs
  - Just reduces I/O overhead

❌ Doesn't change model accuracy
  - Exact computation (bit-identical to standard)
  - No approximation introduced

❌ Doesn't help with very long sequences indefinitely
  - Still hits memory limits eventually
  - But pushes them much further

❌ Doesn't work on CPU
  - Requires GPU with good memory bandwidth
  - No SRAM on CPU to exploit

Hardware Requirements

Flash Attention works best on:
✓ A100 / H100 (excellent)
✓ RTX 4090 (very good)
✓ RTX 3090 (good)
✓ Any modern GPU with good bandwidth

Works less efficiently on:
⚠ Older GPUs (< 2020)
⚠ Mobile GPUs
⚠ Some edge cases

Requires:
  - CUDA (NVIDIA GPUs)
  - Recent GPU compute capability (SM 75+)
  - Properly compiled kernels

Practical Comparison: Real Examples

Example 1: Training a 3B Model

Setup:
  - Model: LLaMA 3B
  - Sequence length: 2048
  - Batch size: 8
  - Hardware: RTX 4090
  - Task: Fine-tuning on custom data

Without Flash Attention:
  - Memory peak: 20 GB
  - Training step: 2.5 seconds
  - Can't fit (RTX 4090 has 24GB)
  - Need gradient checkpointing

With Flash Attention:
  - Memory peak: 16 GB
  - Training step: 1.8 seconds
  - Fits comfortably!
  - Can increase batch size

Benefit: 40% memory reduction + 28% faster training!

Example 2: Inference Long Document QA

Setup:
  - Model: LLaMA 7B
  - Document: 8192 tokens
  - Question: 100 tokens
  - Hardware: A100 40GB

Without Flash Attention:
  - Memory needed: 14GB model + 40GB attention = 54GB ❌
  - Can't fit on single A100!
  - Need multi-GPU or truncate

With Flash Attention:
  - Memory needed: 14GB model + 2GB attention = 16GB ✓
  - Fits with room to spare!
  - Inference: 8 seconds for 8192-token context
  - Can process full document!

Benefit: Makes long-context inference feasible!

Example 3: Real-Time Streaming

Setup:
  - Model: Mistral 7B
  - User: Chat application
  - Sequence so far: 512 tokens
  - Hardware: RTX 3090 (24GB)

Without Flash Attention:
  - Attention memory: (512)² × 2 bytes = 512 KB (small)
  - But latency per token: 150 ms (slow for streaming)
  - User perceives lag

With Flash Attention:
  - Same memory: 512 KB
  - Latency per token: 50 ms ✓ (much better!)
  - Feels responsive!
  - Better user experience

Benefit: 3x faster for interactive applications!

Best Practices

✅ Do's

  1. Always use Flash Attention v2 if available (it's usually better)
  2. Enable it in HuggingFace models (attn_implementation="flash_attention_2")
  3. For training: Use with gradient checkpointing for even more memory savings
  4. For inference: Pair with KV cache quantization for further speedup
  5. Test on target hardware (sometimes different GPUs have different speedups)
  6. Monitor accuracy during training (usually identical but good to verify)
  7. Use for long contexts (Flash Attention's strength)
  8. Combine with other optimizations (quantization, batching, etc.)

❌ Don'ts

  1. ❌ Assume exact bit-identical outputs (might have minor numerical differences)
  2. ❌ Use for very short sequences < 256 tokens (overhead might not be worth it)
  3. ❌ Expect speedup on CPU (use GPU)
  4. ❌ Mix different precision types without care
  5. ❌ Forget to install correct version (CUDA compatibility matters)
  6. ❌ Assume works on all hardware (test first)
  7. ❌ Disable if model is working (don't fix what isn't broken)
  8. ❌ Expect magic accuracy improvements (it's for speed, not accuracy)

Installation and Setup

Install Flash Attention v2

# For CUDA 11.8+
pip install flash-attn

# For specific CUDA version
pip install flash-attn --no-build-isolation

# From source (if prebuilt doesn't work)
git clone https://github.com/Dao-AILab/flash-attention
cd flash-attention
python setup.py install

# Verify installation
python -c "from flash_attn import flash_attn_func; print('✓ Flash Attention installed')"

Enable in Your Project

# Method 1: Direct import
from flash_attn import flash_attn_func

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

# Method 2: PyTorch 2.0+ automatic
import torch
output = torch.nn.functional.scaled_dot_product_attention(q, k, v)
# Automatically uses Flash Attention if available!

# Method 3: HuggingFace
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    attn_implementation="flash_attention_2"
)

# Method 4: vLLM (automatic)
from vllm import LLM
llm = LLM(model_name)  # Automatically uses Flash Attention

Key Takeaways

🔑 Flash Attention reduces I/O complexity from O(N²) to O(N)
2.8x faster attention computation (Flash Attention v2)
💾 50-60% memory reduction for attention
🎯 Enables long-context processing (8K+ tokens)
📈 Exact computation (no approximation, no accuracy loss)
🚀 Becoming industry standard (PyTorch, HF, vLLM)


Performance Summary

Speedup Factors (Flash Attention v2 vs Standard):

Sequence Length      Speedup    Memory Reduction
─────────────────────────────────────────────────
256                  1.5x       30%
512                  2.0x       45%
1024                 2.5x       50%
2048                 2.8x       55%
4096+                2.8x       60%

Longer sequences see more benefit!

Further Reading

  • Flash Attention v1 Paper: "Flash Attention: Fast and Memory-Efficient Exact Attention with IO-Awareness"
  • Flash Attention v2 Paper: "Flash-2: Faster Causal Attention by Reducing IO with Minimal Overhead"
  • GitHub: github.com/Dao-AILab/flash-attention
  • Implementation: Optimized CUDA kernels