FlashAttention

As explained in Transformers in Action

FlashAttention is an exact attention-computation technique that tiles queries, keys, and values to reduce memory traffic and linearize memory usage.

FlashAttention is an exact attention-computation technique for scaling transformers to longer inputs. It changes the implementation rather than the attention result, reducing memory traffic without approximation error. section 9.4.1

Standard attention materializes the full n × n matrix of QKᵀ, applies softmax, and multiplies by V. FlashAttention instead streams tiles of queries and keys through fast on-chip SRAM, computes partial products, performs numerically stable incremental normalization, and accumulates outputs in HBM. section 9.4.1 section 9.4.1

The implementation improves attention efficiency by:

• Avoiding materialization of the full attention matrix and reducing memory complexity from O(n²) to O(n). section 9.4.1
• Reducing redundant reads and writes between high-bandwidth memory and on-chip memory. section 9.4.1
• Delivering reported 2× to 4× speedups over naïve kernels while retaining exactness, unlike sparse or linear attention methods. section 9.4.1

FlashAttention-3 deployment

A Transformers model can enable FlashAttention-3 through the attention implementation setting `kernels-community/vllm-flash-attn3`. The variant targets NVIDIA Hopper hardware and uses features such as Warpgroup Matrix Multiply-Accumulate and the Tensor Memory Accelerator. section 9.4.1

Figure 9.8 Blockwise computation in FlashAttention. Rather than materializing the full n × n attention matrix, the algorithm streams tiles of Q and K into SRAM, applies softmax locally, rescales for global normalization, and accumulates outputs. This avoids expensive HBM traffic.

Use the blockwise-computation figure to trace how Q and K tiles move through SRAM, receive local softmax treatment, undergo global rescaling, and accumulate outputs without full-matrix materialization. section 9.4.1

Listing 9.3 Enabling FlashAttention-3

model_id = "<your model id>"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype="auto",
    attn_implementation="kernels-community/vllm-flash-attn3", #1
)

Use the configuration listing to inspect the model-loading change that enables the FlashAttention-3 implementation. section 9.4.1

How Transformers in Action covers it

How deep it goes

Explains tiled computation, memory movement, normalization, complexity, speedups, and FlashAttention-3 hardware features.

Examples

Demonstrates blockwise Q-K streaming, incremental output accumulation, and a Transformers configuration enabling FlashAttention-3.

What it assumes

Assumes familiarity with transformer attention, Q, K, V, softmax, and GPU memory while contrasting standard attention from first principles.

What you can do after

Operate/configure FlashAttention-3 in a Transformers model and evaluate its memory and throughput benefits.

How it connects

Connects attention with softmax, GPU memory hierarchy, sparse attention, linear attention, and inference optimization.

liveBookmanning.com

Read Transformers in Action on liveBook Get the book at manning.com

Related concepts

• kernel fusionuses · Kernel fusion is an optimization that combines multiple computational operations into fewer executable units to reduce launch overhead and intermediate data movement.
• tilinguses · Tiling is a technique that partitions data or computation into smaller blocks to improve memory locality, data reuse, and processing efficiency.
• memory boundaddresses · A memory-bound workload is one whose performance is limited primarily by the rate of accessing or transferring data rather than by computational throughput.

This explainer was generated by Manning's knowledge graph from the text of Transformers in Action. The section links point to the parts of the book it draws on.