FlashAttention

Overview

FlashAttention is a GPU kernel that computes exact transformer attention faster and with far less memory than the standard method. It does not approximate anything. It gives the same answer as ordinary attention, up to small rounding differences, but it reorganizes the work so the large attention score matrix is never written out to the GPU's main memory. This cuts memory use from growing with the square of the sequence length to growing in proportion to it, and it speeds up training and inference, especially for long sequences.

The key insight is that modern GPUs are often limited by how fast they can move data, not by how fast they can do math. FlashAttention reduces data movement by processing attention in small tiles that fit in fast on-chip memory. This page explains the idea with everyday analogies, walks through a small numeric example, compares FlashAttention with other attention methods, and shows why it matters for models that read long tables with many rows and columns.

What it is

FlashAttention is an algorithm and a set of fused GPU kernels for the attention operation in transformers. Attention takes three matrices, called queries, keys and values, and produces an output where each position is a weighted mix of the values. The weights come from comparing queries with keys.

In the standard approach, three steps are done one after another, each as a separate kernel:

1. Multiply queries by the transposed keys to get a score matrix with one entry for every pair of positions.
2. Apply softmax to turn scores into weights.
3. Multiply the weights by the values.

The score matrix has a size equal to the sequence length squared. For a sequence of tens of thousands of tokens, that matrix is enormous. It is written to GPU memory after step 1, read back for step 2, written again, and read again for step 3. Those trips to memory are slow.

FlashAttention fuses the three steps into one kernel. It loads small blocks of queries, keys and values into fast on-chip memory, does all the math for those blocks, and writes only the final output. The big score matrix never exists in main memory.

What you get

• Exact results, not an approximation.
• Memory use that grows linearly with sequence length for the attention step.
• Speedups that are larger for longer sequences.
• Support for causal masks, padding masks, multi-query and grouped-query attention, dropout and other common options.
• Availability inside popular frameworks, including the attention function of PyTorch and many libraries built on Transformers.

What it is not

It is not a new model or a new type of attention. It does not change what a model learns. It is a faster way to compute the same function. It also does not remove the quadratic growth of computation: the number of operations still grows with the square of the sequence length. It removes the quadratic memory and reduces memory traffic.

Where it came from

Attention became the standard layer in language models after the 2017 paper on the transformer architecture. As models were trained on longer sequences, two problems grew: the quadratic memory for the score matrix and the time spent moving it.

Many earlier papers tried to approximate attention to make it cheaper, using sparse patterns, low-rank projections or kernel tricks. These cut the number of operations on paper, but in practice they often did not run faster on real GPUs, and they changed the results.

The original FlashAttention paper, published in 2022, took a different view. It argued that the real cost was memory traffic, and it called the approach IO-aware. By computing exact attention in tiles and carefully controlling reads and writes between the GPU's large main memory and its small fast memory, it achieved large speedups with no approximation. It was quickly adopted in many training codes.

A second version, published in 2023, reorganized the work across thread blocks and warps to improve how fully the GPU's compute units were used, and it reduced work that was not matrix multiplication. A third version, published in 2024, targeted newer GPU architectures, using asynchronous data movement and lower-precision number formats to push utilization higher. Later work and related kernels added support for decoding with long caches, variable-length batches and more attention variants. The exact feature list depends on the release, so check the project documentation for the version you use.

The method also inspired an ecosystem of compatible kernels and compilers, including those written with Triton, and its tiling idea is now part of the attention function in major frameworks.

How it works

To understand FlashAttention, you need a picture of GPU memory.

A GPU has two main kinds of memory. High-bandwidth memory is large, measured in tens of gigabytes, but relatively slow to reach. On-chip memory, often called shared memory or SRAM, is tiny, measured in kilobytes per compute unit, but very fast. Think of a cook in a kitchen. The pantry is the large memory: it holds everything but walking there takes time. The cutting board is on-chip memory: small, but everything on it is instantly at hand. A smart cook brings a small batch of ingredients to the board, finishes all the work on that batch, and only then returns to the pantry. A careless cook walks to the pantry for every single step.

Standard attention is the careless cook. FlashAttention is the smart one.

Step 1: split the work into tiles

Queries, keys and values are split into blocks along the sequence axis. A block of queries is paired with each block of keys and values, one pair at a time. The block sizes are chosen so that the pieces fit in on-chip memory.

Step 2: compute scores for one tile

For one block of queries and one block of keys, the kernel computes a small score matrix. This lives only in on-chip memory. It is small, so it never needs to go to main memory.

Step 3: the softmax problem

Softmax needs to see all scores in a row before it can normalize, because it divides each exponent by the sum of all exponents in that row. If you only see part of the row, how can you normalize?

The answer is the online softmax trick. As the kernel processes one block of keys after another, it keeps two running numbers for each query row: the largest score seen so far, and the running sum of exponents relative to that maximum. When a new block arrives with a larger maximum, the old sum and the old partial output are rescaled by a correction factor. At the end, the result is exactly what full softmax would give.

An everyday analogy: you are tallying votes from many ballot boxes, and you want each candidate's share. You can keep a running total and a running scale, and when a box arrives with a bigger count you adjust the earlier totals proportionally. After the last box, the shares are correct without ever laying out every ballot at once.

Step 4: accumulate the output

For each tile, the kernel multiplies the local weights by the matching block of values and adds the result into a running output for that block of queries. It uses the correction factors from the online softmax to keep the sum consistent. After the last key block, the output is divided by the final normalizer and written to main memory once.

Step 5: save a small summary for the backward pass

Training needs gradients, and the backward pass of standard attention uses the stored score matrix. FlashAttention does not store it. Instead it saves only the output and a small per-row statistic, such as the log of the softmax normalizer. During the backward pass, it recomputes the score tiles from queries and keys, which is cheap compared to reading a huge matrix from memory. This is a form of recomputation, in the same spirit as gradient checkpointing, but at the level of one kernel.

Step 6: handle masks and options

For causal attention, where a position can only look at earlier positions, tiles that lie entirely in the future are skipped, which saves about half the work. Padding masks, sliding windows and other patterns are handled by skipping or masking tiles. Dropout is applied inside the kernel with a saved random state so the backward pass can reproduce it.

Step 7: parallelize across the GPU

A GPU has many compute units. The kernel assigns different query blocks, batch items and heads to different units. For long sequences with a small batch, splitting the work along the sequence axis keeps the GPU full. Later versions improved this scheduling so that fewer cycles were wasted on bookkeeping and more went into matrix multiplication.

Step 8: use low precision wisely

Matrix multiplications run in 16-bit formats such as bfloat16 or float16, with accumulation in higher precision. The softmax statistics stay in float32. This is part of the standard recipe of mixed precision. Newer versions can use 8-bit floating point for some operations on supported hardware, with extra care to keep accuracy.

Step 9: decoding with a cache

When a language model generates text, it keeps a cache of past keys and values and processes one new token at a time. A single query attending to a long cache does not offer much parallelism. Special decoding kernels split the cache along the sequence axis, compute partial results, and combine them. Serving engines such as vLLM rely on related kernels.

Why this is faster

The number of math operations is about the same as in standard attention. The gain comes from fewer reads and writes to main memory. Because attention is memory-bound for typical head sizes, cutting memory traffic cuts time. And because the score matrix is never stored, memory use drops sharply, which lets you use larger batches or longer sequences.

Key concepts and terms

A worked example

This section has a small numeric walk-through of the online softmax, then code that shows how to call fast attention in practice.

Online softmax by hand

Suppose one query has four scores against four keys: 1, 3, 2 and 5. Standard softmax computes the exponent of each, sums them, and divides.

Now split the keys into two blocks: [1, 3] and [2, 5].

Block one:

• The maximum is 3.
• The exponents relative to the maximum are exp(1 - 3) = 0.135 and exp(3 - 3) = 1.0.
• The running sum is 1.135. The running maximum is 3.

Block two:

• The block maximum is 5, which is larger than 3. The new running maximum is 5.
• The old sum must be rescaled by exp(3 - 5) = 0.135. The old sum becomes 1.135 times 0.135, which is about 0.154.
• The new exponents relative to 5 are exp(2 - 5) = 0.050 and exp(5 - 5) = 1.0, which sum to 1.050.
• The final sum is 0.154 + 1.050 = 1.204.

Check against the full computation. Relative to the maximum 5, the exponents of all four scores are exp(-4) = 0.018, exp(-2) = 0.135, exp(-3) = 0.050 and exp(0) = 1.0. Their sum is 1.203. The small difference is only rounding. The two methods agree, and the second never needed all four scores at once. The same rescaling is applied to the partly built output vector, which is why the final answer is exact.

Memory arithmetic

For one head with sequence length 16,384, the score matrix has 16,384 times 16,384 entries, about 268 million. At 2 bytes each, that is roughly 0.5 GB for a single head in a single sequence. With 32 heads and a batch of 8, that is on the order of 128 GB just for scores in one layer, which is impossible to store. FlashAttention keeps only blocks of the matrix in on-chip memory, so this large cost disappears, and the attention step needs memory in proportion to the sequence length times the head size.

Calling fast attention from PyTorch

Modern PyTorch exposes a scaled dot product attention function that can pick a FlashAttention-style kernel when the hardware and the data type allow it.

import torch
import torch.nn.functional as F

batch, heads, seq, dim = 4, 16, 8192, 64
q = torch.randn(batch, heads, seq, dim, device="cuda", dtype=torch.bfloat16)
k = torch.randn(batch, heads, seq, dim, device="cuda", dtype=torch.bfloat16)
v = torch.randn(batch, heads, seq, dim, device="cuda", dtype=torch.bfloat16)

# The framework chooses a fused kernel when possible.
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
print(out.shape)

Using it in a Transformers model

Many models in the Transformers library let you select the attention implementation when you load them.

from transformers import AutoModelForCausalLM
import torch

model = AutoModelForCausalLM.from_pretrained(
    "example-org/example-model",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
).to("cuda")

If the kernel is not installed or the GPU is not supported, loading raises an error or falls back to another implementation, depending on the setting. Check the message.

Checking correctness

To check that a fast kernel matches the reference, compare it against a plain implementation in float32 on a small input, and use a tolerance.

def reference_attention(q, k, v, causal=True):
    scale = q.shape[-1] ** -0.5
    scores = (q.float() @ k.float().transpose(-2, -1)) * scale
    if causal:
        n = scores.shape[-1]
        mask = torch.triu(torch.ones(n, n, device=q.device), 1).bool()
        scores = scores.masked_fill(mask, float("-inf"))
    return (scores.softmax(-1) @ v.float())

ref = reference_attention(q[:1, :2, :512], k[:1, :2, :512], v[:1, :2, :512])
fast = F.scaled_dot_product_attention(
    q[:1, :2, :512], k[:1, :2, :512], v[:1, :2, :512], is_causal=True
)
print((ref - fast.float()).abs().max())

The maximum error should be small, on the order of the rounding error of the lower precision format.

Measuring the gain

Time the forward and backward pass with and without the fused kernel at several sequence lengths. Expect small gains at short lengths and larger gains at long lengths. Also record peak memory. These two curves show the benefit more clearly than a single number.

Backward pass in a bit more detail

The backward pass needs gradients with respect to queries, keys and values. In standard attention these come from the stored weight matrix. In the tiled version, the kernel reloads blocks of queries, keys and values, recomputes the local scores, rebuilds the local weights using the saved row statistics, and then computes the local gradient pieces. Gradients for keys and values accumulate across query blocks, and gradients for queries accumulate across key blocks. The kernel orders the loops so that one of these accumulations stays in fast memory while the other is added to main memory with atomic or carefully ordered updates. The extra arithmetic from recomputation is small compared with the memory traffic saved, so the backward pass is still faster overall.

A table-shaped example of the cost

Imagine a table model that stacks 4,096 context rows and 64 columns, and applies attention along the row axis for each column. For each column, the sequence length is 4,096. The score matrix per head is 4,096 by 4,096, about 16.8 million entries. Across 64 columns and, say, 8 heads, that is more than 8.6 billion score entries in one layer for a single table. At 2 bytes each, this is over 17 GB, far too much to hold for every layer in training. A fused kernel avoids storing it. The model can then use more context rows or larger batches of tables. The numbers here are an illustration of the arithmetic, not a measurement of any specific model.

Choosing between kernels in a framework

Frameworks usually expose a single attention call and choose among several backends: a FlashAttention style kernel, a memory-efficient kernel, and a plain math path. The choice depends on data type, head size, device and whether you need features such as arbitrary masks or returning the attention weights. If you ask for the attention weights themselves, for example to draw a heat map, the fused path cannot supply them because it never builds them, and the framework will fall back to the plain path. Plan for this when you want interpretability, and extract weights only for a few small examples.

Long context beyond a single GPU

Even with linear memory, a very long sequence may exceed what one GPU can hold in activations or time. Context parallelism splits the sequence across GPUs, and ring-style schemes pass blocks of keys and values around the group while each device computes partial attention with a fused kernel and merges the partial results with the same online softmax rule. The idea is the same as the tile loop, applied across devices instead of within one. See Megatron-LM for how this fits with other forms of parallelism.

Strengths and limitations

Strengths

• Exact: no change to model outputs beyond rounding.
• Lower memory: the attention step no longer needs a quadratic score matrix.
• Faster: fewer trips to main memory, so speedups on both training and inference.
• Longer contexts: the memory saving lets you train and run with longer sequences or bigger batches.
• Broad support: causal masks, variable lengths, grouped-query attention and other common options.
• Easy to adopt: often a single flag or the default path in a framework.
• Composable with other methods: mixed precision, sharded training, tensor parallelism and many others.

Limitations

• Hardware dependence. The best kernels target specific GPU generations, and older or non-NVIDIA devices may use other kernels or fall back to slower paths.
• Compute still grows with the square of sequence length. Very long contexts remain expensive.
• Head size limits. Some kernels support only certain head dimensions.
• Unusual attention patterns may not be supported, such as arbitrary dense masks or custom score modifications, although newer tools make this easier.
• Building from source can be slow and sensitive to compiler and CUDA versions.
• Numerical results can differ slightly from a reference because the order of summation differs.
• Debugging is harder than with plain operations because the intermediate score matrix is not available to inspect.

How it compares with alternatives

The important point is that FlashAttention is not in competition with changes to the model. It is a faster way to compute standard attention. You can also combine it with sliding windows or grouped-query attention to get further savings.

Role in tabular foundation models

Tabular foundation models learn from raw tables. A table has rows and columns, cells of mixed types such as numbers, categories, dates and text, missing values, and sometimes columns without useful names. A model like TabPFN takes a set of labeled rows as context and predicts labels for new rows in a single forward pass, with no weight updates at prediction time. Attention is the engine of that process, so the cost of attention shapes what such models can do.

Attention over cells, rows and columns

A table can be seen as a grid of cells. A model needs to relate cells within a row, such as how age and income relate. It also needs to relate cells within a column, such as how one row's value compares with the distribution of others. Many designs apply attention along both axes. Attention along the column axis, across rows, is the one that implements in-context learning: a test row looks at training rows and borrows their label information. The number of rows in context can be thousands, so this attention step runs over long sequences. Fast exact attention makes it practical to use more rows, which usually gives better predictions.

Cost of long row contexts

If a table has many rows and many columns, a model that flattened all cells into one sequence would produce a very long input. Even with axis-wise attention, long context windows are common. Without a kernel like FlashAttention, the score matrices would dominate memory. With it, memory scales linearly, and the model can handle larger tables or larger batches of tables on the same hardware.

Exactness matters for numbers

Table predictions often involve fine numeric differences, for example a probability that sits close to a decision threshold. An approximate attention method might shift outputs in ways that are hard to audit. FlashAttention's exactness means that results match the reference implementation up to rounding, which makes it a safe drop-in for models whose behavior was validated with standard attention.

Missing values and masks

Missing cells are normal in tables. A model usually marks them with a special value or a mask. Attention kernels must support masks correctly, so that a missing cell does not contribute as if it had a real value. Padding masks for tables with different numbers of rows or columns are also required, since batches contain tables of different shapes. Variable-length kernels avoid wasting compute on padding by packing tables of different sizes into one batch with offsets.

No column names and permutation invariance

When column names are missing, the model must rely on the values alone. Good table models are built to be insensitive to the order of rows and sometimes columns. Attention has this property by nature, because it treats its input as a set unless position information is added. Fast attention does not change this property. It just makes the computation cheaper so that designs which avoid fixed positions remain affordable.

Pre-training on synthetic and real tables

Pre-training on large numbers of generated tables requires many forward and backward passes over attention-heavy models. Cutting attention time and memory directly cuts training cost. It also allows larger table sizes in the synthetic data distribution, which helps the model generalize to bigger real tables.

Imputation and anomaly detection

Imputation asks the model to fill hidden cells. Anomaly detection asks how unlikely a cell is given the rest of the row and the other rows. Both tasks rely on the same attention layers. With faster attention, you can score or fill many cells per second, which matters when a dataset has millions of cells. At serving time, lower memory use means more tables can be processed in parallel.

Continual learning

When a model keeps learning from new table families, the cost of each training pass matters, since passes are repeated. Efficient attention lowers the cost of replaying old data along with new data, which helps limit forgetting.

Semantic and sector classification

If a model reads column names and sample values as text to infer meaning or industry sector, the text portion of the input is a normal language-model sequence, and FlashAttention applies directly. Long headers, descriptions and sample values can be included without running out of memory.

Scaling across GPUs

At larger scales, attention kernels are combined with tensor parallelism, Megatron-LM style sequence and context parallelism, or sharded training with FSDP. The kernel handles the work inside one GPU, and the parallel library handles the split across GPUs. For architecture and research context, see /model/architecture, /model/research and /model/benchmarks.

Practical guidance and best practices

Use it by default for transformer training and serving

If your hardware supports it, enable fast attention. It is rarely worse than the standard approach, and it often brings large gains. Start with the framework's built-in attention function, which selects a suitable kernel automatically.

Use 16-bit data types

Fused attention kernels are designed for bfloat16 and float16 inputs. If you pass float32, the framework may fall back to a slower path. Use mixed precision with bfloat16 where the hardware supports it.

Check head size and shapes

Confirm that your head dimension is supported. Many kernels work best with head sizes that are multiples of eight. If you see a fallback warning, check shapes, strides and data types.

Prefer variable-length interfaces for ragged batches

If your batch contains sequences of different lengths, use packed or variable-length interfaces rather than padding to the longest. This is especially useful for table data, where sizes vary a lot.

Verify numerics once

When you switch attention implementations, compare outputs and losses against the old one on a small run. Differences should be tiny. If they are not, check masks, scaling and dropout.

Match versions

Match the kernel package, the framework, the CUDA toolkit and the driver. Mismatches are a common source of build and import errors. Using prebuilt wheels or container images avoids long compile times.

Combine with other memory savers

For very long sequences, combine fast attention with gradient checkpointing, sharded optimizers such as DeepSpeed ZeRO, and sequence or context parallelism. Attention is only one part of memory use. Feed-forward activations and optimizer state also matter.

Profile before and after

Measure step time and peak memory with a profiler. Do not assume the gain. For short sequences, other costs may dominate, and the improvement may be small.

Consider alternatives for special masks

If your model uses an unusual attention pattern, check whether your framework offers a flexible attention API that compiles custom score modifications into a fused kernel. This keeps the memory benefit without giving up the pattern.

Use the decoding kernels for serving

For text generation or any model with a cache, use decoding-specific kernels or a serving engine that includes them, such as vLLM, and consider paged cache layouts to cut waste.

Keep a fallback path

Keep the ability to switch to the standard implementation. It helps with debugging, since you can inspect the score matrix, and with running on hardware that lacks fused kernel support.

A simple benchmark script

Use this pattern to compare attention implementations on your own hardware. Warm up first, synchronize before timing, and report both time and peak memory.

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

def bench(fn, steps=20, warmup=5):
    for _ in range(warmup):
        fn()
    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    start = time.time()
    for _ in range(steps):
        fn()
    torch.cuda.synchronize()
    ms = (time.time() - start) / steps * 1000
    mem = torch.cuda.max_memory_allocated() / 1e9
    return ms, mem

for seq in (1024, 4096, 16384):
    q = torch.randn(2, 16, seq, 64, device="cuda", dtype=torch.bfloat16)
    k, v = torch.randn_like(q), torch.randn_like(q)
    ms, mem = bench(lambda: F.scaled_dot_product_attention(q, k, v, is_causal=True))
    print(seq, round(ms, 2), "ms", round(mem, 2), "GB")

Run the same loop with the standard implementation at short lengths. At long lengths the standard version may run out of memory, which is itself a useful result.

How tile size is chosen

The kernel picks block sizes so that a block of queries, a block of keys and a block of values fit in on-chip memory together with the running statistics. Larger tiles mean fewer trips to main memory and better reuse, but they need more on-chip space, and too large a tile reduces how many blocks run at once. This is a balancing act, and kernel authors tune it for each GPU generation and head size. As a user you rarely touch it, but it explains why some head sizes and data types are faster than others.

Variants you may meet

• Multi-query and grouped-query attention reduce the size of the key and value cache, which helps decoding. Fast kernels handle them without copying the keys and values.
• Sliding window attention limits each position to a window of recent tokens. Kernels skip tiles outside the window.
• Alibi style biases and relative position biases add a term to the scores. Many kernels support common forms.
• Rotary position embeddings are applied to queries and keys before attention, so they are handled outside the kernel, or fused with it in some libraries.
• Soft capping of scores, used in some models, requires explicit support in the kernel.
• Block-sparse masks skip entire tiles, which can give extra speed when the pattern is regular.

Troubleshooting guide

Reading a profile of an attention layer

In a profile of a standard attention layer, you see separate kernels for the matrix multiply, the softmax, the dropout and the second matrix multiply, with long gaps and large memory copies. In a profile with a fused kernel, those collapse into one block, and the memory traffic falls. If you still see many small kernels around attention, look at the projections, the position embedding and the layer norm, which may be unfused and are good targets for other fused kernels, for instance those written with Triton.

Where it fits in the training stack

At the bottom is the hardware and CUDA. Above that are kernels, including fused attention. Above those are the framework layers in PyTorch, then the model code in Transformers, then the distributed layer such as FSDP or DeepSpeed ZeRO. Each layer solves a different problem. The attention kernel reduces the cost on a single GPU. The distributed layer splits the job across many. Both are needed for large models.

Common mistakes

• Expecting a speedup at very short sequence lengths. The benefit grows with length.
• Passing float32 inputs and not noticing that a slow fallback is used.
• Believing it makes attention cost linear in compute. Only memory becomes linear.
• Installing a kernel build that does not match the CUDA version, then losing time on cryptic errors.
• Using an unsupported head size and silently getting a slower path.
• Forgetting masks for padded or missing positions, so padding leaks into the result.
• Comparing outputs with a tolerance that is too tight for 16-bit math.
• Treating it as an approximation and worrying about accuracy, when it is exact up to rounding.
• Turning off the fused kernel to debug and then forgetting to turn it back on in production.
• Assuming that all attention variants are supported, such as arbitrary dense masks or custom biases.
• Ignoring the rest of the memory budget. Long sequences also grow feed-forward activations.
• Using the kernel for tiny models where other costs dominate and complexity is not worth it.
• Asking for attention weights from a fused kernel and being surprised by a slower fallback.
• Changing sequence length or head size late in a project without re-checking kernel support.
• Judging the gain on a tiny test and not on the real sequence lengths used in production.

Frequently asked questions

What is FlashAttention in simple terms?

FlashAttention is a faster and lighter way to compute the attention step of a transformer. It breaks the work into small tiles that fit in the GPU's fast memory, so the huge table of attention scores never has to be written to slow memory. The final answer is the same as with the normal method.

Is FlashAttention an approximation?

No. It computes exact attention. Results can differ from a reference by tiny rounding amounts, because the order of additions changes, but the method does not drop or approximate any interactions between positions.

Why is it faster if it does the same math?

Because on modern GPUs the slow part is often moving data between main memory and the compute units, not the arithmetic itself. By keeping intermediate results in fast on-chip memory and avoiding reads and writes of the big score matrix, FlashAttention spends far less time waiting for data.

Does it make attention linear in sequence length?

It makes the memory use of the attention step linear in sequence length, but the number of arithmetic operations still grows with the square of the length. Very long contexts remain costly in compute, even though they now fit in memory.

Which GPUs does it support?

The best-optimized kernels target recent NVIDIA GPU generations, and the project documentation lists the supported architectures for each release. Other hardware may be served by related kernels in frameworks, or by compilers that generate fused attention, so check the support matrix for your device.

Do I need to change my model to use it?

Usually not. Many frameworks route their standard attention function through a fused kernel automatically, and many model libraries let you select the implementation with a single option. The model weights and architecture stay the same.

Can I use it for training and inference?

Yes. It has forward and backward kernels for training, and it also has kernels tuned for inference, including decoding with a long cache of past keys and values. Gains are often largest for training with long sequences and for processing long prompts.

How does it help models that read tables?

Table models often need many rows in context, which makes the attention step long and memory hungry. FlashAttention keeps memory use linear and speeds up the step, so a model can look at more rows or process more tables at once. Because it is exact, it does not change the predictions a validated model was designed to make.

Does it work with masks and variable-length batches?

Yes. It supports causal masks, padding masks and several other patterns, and it has interfaces for packed batches of different lengths. This matters for tables, where each example can have a different number of rows and columns.

What is the difference between FlashAttention and memory-efficient attention?

Both avoid storing the full score matrix by processing it in chunks. FlashAttention is a specific family of kernels designed around the memory hierarchy of the GPU, with tuned versions for newer hardware. Other memory-efficient kernels use similar ideas with different implementations and wider hardware support. Frameworks often choose between them automatically.

Glossary

• Attention: a layer that mixes value vectors using weights computed from queries and keys.
• Causal attention: attention where each position sees only itself and earlier positions.
• Fused kernel: a single GPU program that performs several operations in one pass.
• Grouped-query attention: a variant where several query heads share one set of keys and values.
• High-bandwidth memory: the large main memory of a GPU.
• IO-aware: designed to reduce data movement between memory levels.
• Online softmax: a way to compute softmax in blocks using running statistics.
• Recomputation: recalculating values in the backward pass to avoid storing them.
• Score matrix: the matrix of query-key comparison values.
• SRAM: small, very fast on-chip memory.
• Softmax: a function that converts scores into weights that sum to one.
• Tile: a small block of a matrix processed at one time.
• Variable-length batch: a batch that packs sequences of different lengths without padding.

Summary

FlashAttention computes exact transformer attention with tiling, an online softmax and recomputation in the backward pass, so the large score matrix never goes to GPU main memory. The result is lower memory use, faster training and inference, and the ability to handle long sequences, with no change in what the model computes. Its limits are hardware support, head size constraints and the fact that compute still grows with the square of the length.

For tabular foundation models, it makes in-context learning over many rows affordable, supports masks for missing cells and variable table shapes, keeps predictions faithful to validated behavior, and lowers the cost of pre-training on synthetic and real tables, imputation, anomaly scoring and continual learning. Use it together with mixed precision, tensor parallelism, Megatron-LM and Triton based kernels where needed, and read the model-side context at /model/architecture, /model/research and /model/benchmarks.

Related

It belongs to Training tooling in my stack, under LLMs, retrieval, and evaluation. Others in the same group: