FlexAttention in PyTorch: Custom and Sparse Attention Without Writing Kernels

Every serious LLM or vision project eventually needs an attention variant that FlashAttention does not ship: a sliding window with a few global sink tokens, document-boundary masking so packed sequences do not attend across examples, logit soft-capping, ALiBi or a custom positional bias, prefix-LM masking, or one of the 2025-era sparse patterns that only attend to a selected subset of blocks. Historically you had three bad options: fall back to a dense F.scaled_dot_product_attention with an explicit [B, H, Q, K] boolean mask and watch memory explode quadratically, hand-write a Triton kernel and own it forever, or wait for someone upstream to add your variant.

torch.nn.attention.flex_attention is the fourth option. You write the mask and the score modification as small Python functions, PyTorch compiles them into a fused FlashAttention-style kernel, and the sparsity you describe is actually skipped rather than computed and masked out. This tutorial shows how it works, the patterns that matter in production, the performance numbers to expect, and the sharp edges we hit on client engagements.

The two hooks: score_mod and mask_mod

FlexAttention gives you two places to inject logic, and the difference between them is the whole point.

A score_mod transforms the attention logit after the query-key dot product:

def score_mod(score, b, h, q_idx, kv_idx):
    return score  # modify and return

A mask_mod returns a boolean: may this query position attend to this key position?

def causal(b, h, q_idx, kv_idx):
    return q_idx >= kv_idx

Both are traced per-element but never actually run per-element; torch.compile lowers them into the body of a fused kernel. The critical distinction: a mask_mod can be analyzed ahead of time into a BlockMask, a coarse map (default 128x128 blocks) of which blocks are entirely empty, which are partially masked, and which are full. Empty blocks are never scheduled. That is where the speedup comes from. A score_mod that returns -inf produces the same numerics but the kernel still does the work, so always express structure as a mask, not as a penalty.

import torch
from torch.nn.attention.flex_attention import flex_attention, create_block_mask

flex_attention = torch.compile(flex_attention, dynamic=False)
create_block_mask = torch.compile(create_block_mask)

B, H, S, D = 4, 16, 8192, 64
q, k, v = (torch.randn(B, H, S, D, device="cuda", dtype=torch.bfloat16) for _ in range(3))

block_mask = create_block_mask(causal, B=None, H=None, Q_LEN=S, KV_LEN=S, device="cuda")
out = flex_attention(q, k, v, block_mask=block_mask)

B=None / H=None tell the builder the mask does not depend on batch or head, so one small mask is broadcast instead of materialized per batch element. Build the BlockMask once per shape, outside the training step, and cache it. Rebuilding it every forward pass is the single most common reason a FlexAttention migration ends up slower than what it replaced.

Pattern 1: document masking for packed training

If you pack multiple documents into a fixed-length training sequence to avoid padding waste, tokens from document 3 must not attend to document 2. Done with a dense mask this costs S^2 memory; done with a mask_mod it costs nothing and the off-diagonal blocks are skipped entirely.

def make_doc_mask(doc_id):            # doc_id: [S] int tensor, cumulative doc index
    def mask_mod(b, h, q_idx, kv_idx):
        return (q_idx >= kv_idx) & (doc_id[q_idx] == doc_id[kv_idx])
    return mask_mod

block_mask = create_block_mask(make_doc_mask(doc_id), B=None, H=None,
                               Q_LEN=S, KV_LEN=S, device="cuda")

With many short documents per packed sequence, the block sparsity is extreme — we routinely see 85-95% of blocks eliminated, and wall-clock attention time drops roughly in proportion. Packing plus document masking is usually a bigger end-to-end training win than any kernel micro-optimization, because it removes padding waste and quadratic cost at the same time.

Pattern 2: sliding window with attention sinks

Long-context models increasingly interleave full-attention layers with windowed layers, plus a handful of "sink" tokens at the start that every position may attend to (this stabilizes streaming generation).

WINDOW = 1024
SINKS = 4

def sliding_sink(b, h, q_idx, kv_idx):
    causal = q_idx >= kv_idx
    in_window = (q_idx - kv_idx) < WINDOW
    is_sink = kv_idx < SINKS
    return causal & (in_window | is_sink)

At 32k context with a 1k window this is ~97% block-sparse. Expressing the same thing as a score_mod returning -inf would be numerically identical and roughly 20x slower.

Pattern 3: score_mod for biases and soft-capping

Use score_mod only for things that genuinely change every surviving logit:

def alibi(score, b, h, q_idx, kv_idx):
    return score - alibi_slopes[h] * (q_idx - kv_idx)

def soft_cap(score, b, h, q_idx, kv_idx):
    return 30.0 * torch.tanh(score / 30.0)

Both compose with any block_mask, and both are differentiable — FlexAttention generates a fused backward pass too, so a learned relative-position bias trains without a hand-written gradient kernel. Compose multiple mask functions with and_masks / or_masks from torch.nn.attention.flex_attention rather than writing one giant predicate; the resulting code is far easier to unit-test.

Decoding and paged KV caches

During generation the query length is 1 and the KV length grows, so a mask built for training shapes is wrong. Use BlockMask.from_kv_blocks or rebuild against the current KV length, and keep the cache length bucketed (for example to multiples of 1024) so you are not recompiling on every step. FlexAttention also integrates with paged KV caches: you convert the logical mask into physical block indices via the page table, which is how it is used inside modern serving stacks to support arbitrary mask shapes without a bespoke kernel per variant.

For inference there is a dedicated flex_attention decoding path optimized for short queries; benchmark it against the vendor FlashAttention/FlashInfer kernel for your exact variant. Where a hand-tuned kernel exists for your pattern (plain causal, plain sliding window), it will usually still win by 5-15%. FlexAttention's value is covering the long tail of patterns where no hand-tuned kernel exists at all.

What to expect, and the sharp edges

From our engagements, realistic expectations on H100-class hardware:

• Within ~10% of FlashAttention-2/3 for patterns those kernels support natively.
• 2-10x faster than a dense masked SDPA fallback at 8k-32k context, with memory that scales with unmasked blocks rather than S^2.
• Backward pass typically a bit further from the hand-tuned ideal than forward; budget for that in training throughput estimates.

Sharp edges worth knowing before you commit:

1. It requires torch.compile. Eager FlexAttention exists but is a reference implementation and is slow. If some part of your model refuses to compile, fix that first.
2. Recompilation storms. Varying sequence lengths trigger recompiles. Bucket your lengths, set dynamic=False, and watch TORCH_LOGS="recompiles" during the first epoch.
3. Head dimension and dtype constraints. Stick to head dims that are multiples of 64 and to bf16/fp16; exotic shapes fall off the fast path quietly.
4. Block size is a tunable. The default 128x128 is good for long sequences; for shorter contexts or very fine-grained masks, smaller BLOCK_SIZE captures more sparsity at the cost of more blocks to schedule.
5. Verify numerics first. Write a test that compares your FlexAttention output against a dense reference with an explicit mask, at small S, in fp32, within 1e-2. Mask bugs are silent: the model still trains, just worse. We have seen a sign error in a document mask cost a client a full pretraining run before anyone noticed.
6. Check the mask visually. block_mask.to_string() prints an ASCII sparsity map and reports the sparsity percentage. If it does not look like what you drew on the whiteboard, stop.

Where this fits

FlexAttention changes the economics of attention research in production. Trying a new sparse pattern used to be a two-week Triton project with a maintenance tail; it is now a ten-line function and an afternoon of benchmarking. That matters most for teams doing long-context work, packed pretraining or domain-specific architectures — exactly the situations where the published kernels do not fit and the dense fallback is unaffordable.

The discipline that makes it pay off is boring: express structure as masks, cache BlockMasks, bucket your shapes, and always validate against a dense reference before you trust a training run to it.

If your team is fighting attention memory at long context, packing documents without masking them properly, or sitting on a hand-written Triton kernel nobody wants to maintain, our PyTorch consultants can audit the attention path and migrate it. Get in touch or see our PyTorch Performance Optimization service.