FlashAttention 2: faster attention starts with moving less data

Fewer operations did not mean faster training

Tri Dao had spent years working on approximate attention when large-scale model trainers challenged his assumptions: methods that performed fewer floating-point operations often failed to run faster. Memory traffic and parallelism mattered alongside arithmetic. In the workloads he describes in this July 2023 Latent Space conversation, reading and writing memory consumed most of the attention runtime.

Dense attention computes pairwise similarities between sequence elements. Sparse approaches reduce work by selecting only some pairs, effectively treating omitted entries as zero. Dao says ignoring those interactions tends to reduce model quality. [FlashAttention](https://arxiv.org/abs/2205.14135) preserves the attention computation without approximation, changing how it executes on hardware.

The crucial distinction is that FlashAttention’s memory usage is linear in sequence length; its computation remains quadratic. Dao reports roughly 2–4× wall-clock speedups and says these enabled training with 2–4× longer sequences at the same cost. He gives no hardware, workload or baseline details for those figures, so they describe reported benefits rather than a general conversion between speedup and affordable context length.

The mathematical rewrite that makes fusion possible

Two classical systems techniques provide the starting point. Kernel fusion combines operations so data can be loaded once, processed through several steps and then written out, avoiding intermediate writes and reloads. Tiling divides computation into blocks that fit in fast memory near the processor. Together, they let more work happen before data moves through slower memory.

Attention introduces a dependency that makes this difficult: softmax, the normalization operation, requires a sum across an attention-matrix row. Each block contains only part of that row. Straightforward tiling therefore cannot finish the normalization independently within each block.

Online softmax allows computation in smaller pieces, with rescaling that preserves the correct result. FlashAttention combines this existing mathematical idea with fusion: the rewrite makes blockwise attention possible, and the implementation uses those blocks to reduce memory reads and writes. The gain comes from making the mathematics and the execution strategy work together.

Fusion also reduces flexibility. A researcher who inserts an extra operation inside attention may no longer be able to use the existing fused routine. A compiler could recover flexibility by capturing a PyTorch computation graph and generating fused code. In July 2023, that worked for some computations, but attention’s softmax rewrite remained difficult; automatic support within a year or two was still a hope.

Small, fast memory changes the algorithm

Avoiding intermediate writes to GPU memory still requires somewhere to keep the data. In the simplified hierarchy used here, high-bandwidth memory (HBM) is the large memory pool—such as an A100’s 40 or 80 GB—while SRAM is much smaller, faster on-chip memory beside the compute units. CPU caches versus main memory offer a useful analogy for this capacity–speed tradeoff.

The illustrative figures are approximately 40 GB of HBM versus 20 MB of SRAM, with bandwidths around 1.5 TB/s and 19 TB/s respectively. These describe the example hierarchy, not universal GPU specifications. Tiling makes the small, fast tier useful by keeping each block’s working data small enough to fit.

This design problem predates GPU attention. Banks sorting transactions overnight had to coordinate large, slow disks with smaller, faster main memory. The technologies changed, but the question survived: how much useful work can happen before data must cross the expensive memory boundary again?

Limited area near compute units constrains on-chip memory capacity. The forecast in the conversation is that SRAM will grow relatively little while off-chip HBM becomes larger and faster. That makes input/output awareness—designing around data movement—a durable principle. Fusion and tiling could remain valuable across new devices and model architectures, even after today’s FlashAttention code is replaced.

Research improves when applications push back

FlashAttention’s popularity exceeded expectations, but fast, usable code was a priority during the paper’s development. Releasing that code gave users something to adopt and improve through feedback. Hazy Research, the Stanford group led by Chris Re, supported a similar loop by bringing together researchers focused on algorithms, systems and applications.

Medical imaging and seizure detection show why that mix matters. A collaborator committed to seizure detection keeps asking why a method fails on that task, pushing researchers to investigate rather than move to a more favorable application. Re also encouraged textbook study: understanding older ideas can reveal opportunities when the technical landscape changes.

The discussion presents academia and industry as complementary. Industry has advantages in the compute, infrastructure and engineering required to scale models. Academia can contribute through: - Evaluation: characterizing model capabilities and deciding what to measure. - Understanding: investigating how models produce their behavior. - Riskier research: pursuing uncertain directions beyond near-term company objectives.

The respondent describes plans to become a Princeton professor the following year while maintaining industry relationships to anticipate problems two or three years ahead. Evaluation benefits from the same connection: Chatbot Arena, which compares chatbot models, exemplifies measurement shaped by an emerging application. Useful evaluations deserve investment even when they draw less attention than new model releases.

FlashAttention 2 and the hardware lottery

[FlashAttention 2](https://crfm.stanford.edu/2023/07/17/flash2.html) began as a refactor to try NVIDIA’s [CUTLASS 3](https://github.com/NVIDIA/cutlass/discussions/787), released that January. The library supplied building blocks for efficient GPU matrix multiplication and memory loading. Several months of working with those primitives exposed further inefficiencies in the original FlashAttention.

The reported result was attention running roughly twice as fast as its predecessor, approaching the efficiency of highly optimized matrix multiplication. The comparison comes without workload shapes, precision, hardware or measurement details; it is an attention implementation result, not a specified end-to-end Transformer speedup.

The released code used NVIDIA libraries and ran on NVIDIA GPUs. The underlying memory-aware approach reached further: an AMD implementation of FlashAttention and GGML implementations for Mac and CPU execution are cited as examples. These were implementations of the broader idea, not identified as FlashAttention 2 ports.

Portability exposes the [hardware lottery](https://hardwarelottery.github.io/), extended here to software frameworks. Architectures that exploit existing machines gain adoption; adoption attracts engineering investment; hardware and frameworks then evolve to favor those architectures. Transformers benefited from this reinforcing loop, while recurrent alternatives faced a harder path to efficient execution. A promising model can be impractical to scale because it maps poorly to the available machines.

Compiler efforts such as Mojo offer a possible way to make the same code efficient across devices. New hardware makes different bets: fast on-chip memory could reduce transfers, with Tesla’s Dojo cited as an example. The respondent, explicitly speaking as a nonexpert in hardware, identifies higher manufacturing cost per gigabyte as a tradeoff.

Software support is another hurdle. Cerebras has its own stack and compiler, while broad PyTorch compatibility requires supporting many operations. Hardware development is estimated at three to five years, forcing chip designers to commit before they know which model architectures will dominate.

Making attention faster while testing alternatives

Researchers face a related timing problem: foundational work may take years while architectures change rapidly. The respondent recalls early PhD research on which matrices admit theoretically fast matrix–vector multiplication. That work developed mathematical maturity and research judgment useful for later problems. The advice is to mix foundational exploration with current methods, assess the results with an advisor and adjust, without prescribing a fixed allocation.

Transformer alternatives put that approach into practice. State-space methods are introduced as candidates for capturing long-range information with subquadratic scaling: computation grows more slowly than the square of sequence length. Recurrent neural networks (RNNs) revisit an older approach, with RWKV as a contemporary language-model example. A 14-billion-parameter RWKV model is described as apparently competitive with Transformers, without benchmark conditions. The larger question is whether attention is necessary for these capabilities; answering it requires serious effort to implement and scale competing architectures.

One potential advantage is long-context processing. At sequence lengths such as 16K, 32K or 100K tokens, FlashAttention still leaves quadratic attention computation. Better scaling could also benefit high-resolution images, audio and video. This was prospective: the speaker had not seen validation such as a released recurrent model with roughly 100K context.

Another opportunity is generation throughput, especially for offline document processing or servers batching many requests. A Transformer’s KV cache retains information from previous tokens and can consume enough memory to limit batch size. An RNN instead summarizes the past in a fixed-size state that does not grow with the history.

The proposed benefit follows a concrete chain: less memory per sequence permits larger batches, which could improve accelerator utilization and throughput. The speaker is optimistic but says this advantage had not, to their knowledge, been validated at scale. Making attention and its alternatives efficient enough to compare is therefore one research program: push both, scale both, and investigate which mechanisms produce useful capabilities.

Open weights, open data and room to experiment

The host separates openness into access to pretraining data, visibility into training, availability of weights and restrictions on use. In that framing, RedPajama exposes data and training work, Stable LM offers weights with closed data, and Llama 2 combines closed data with restricted weights. The host considers Llama 2 beneficial while rejecting its description as fully open source.

The response emphasizes practical consequences: broader commercial availability could encourage companies to fine-tune and serve their own Llama 2 derivatives, shifting some usage away from closed-model APIs and reducing concentrated decision-making power. Available weights also enable experiments such as the cited community work extending Llama 1 from a 2K context to roughly 8K.

Libraries for hosting and fine-tuning helped turn releases into usable services. Providers including Together reportedly announced Llama 2 access or hosting within hours of release. The value of available weights includes the software and services others can build around them.

Open datasets face weaker incentives. Legal complications impede publication, datasets attract less attention than models, and companies may regard expensive human-annotated instruction data as a competitive advantage. Together’s RedPajama and Cerebras’s cleaned, deduplicated derivative SlimPajama are examples of progress in open pretraining data. The host argues that reusable datasets such as the Pile and the Stack can have more impact than another underperforming model release.

A proposed alternative resembles Wikipedia: volunteers contribute instruction-data annotations and receive recognition that encourages participation. Collaborative model development might resemble Linux patches, though the implementation path remains unclear.

Asked about joining Together, the respondent points to its open models and datasets, trusted colleagues, and alignment with research on efficient training and inference. The plan was to lead a research team and make its work openly available. Openness also serves a practical purpose: releasing code that others can use and modify had already brought valuable feedback to the team’s research.

Capabilities still need explanations

Asked which capability arrived sooner than expected, the respondent chooses understanding jokes, attributing the surprise to larger models trained on large amounts of data. Reasoning remains a deeper open question: models exhibit behavior that looks like reasoning, but its mechanism is poorly understood. Explicit reasoning modules are proposed as a future architectural need, a prediction rather than an established requirement.

The closing advice returns to the combination behind FlashAttention: understand both the algorithm and the system it runs on. At large model scales, the mathematics and the execution machinery both matter. Their intersection offers a productive place to work.