PyTorch 2.13, released July 8, brought the fused FlexAttention kernel to Apple Silicon's Metal (MPS) backend. Before this, custom attention masks on a Mac fell back to dense scaled-dot-product attention — full quadratic work no matter how much your mask threw away. Now the sparse blocks are skipped, and PyTorch reports up to a ~12x speedup over SDPA on sparse patterns.

This is the practical walkthrough: the three masks you'll actually use, as copy-paste code that runs on device="mps". If you want the strategic why, read what PyTorch 2.13 changes for founders running models on a Mac. This piece is just the code.

The two hooks#

FlexAttention gives you two places to inject logic, and they do different jobs:

You can pass either, both, or neither.

Setup#

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

# Wrap once at module load so the mask + attention fuse into one kernel.
flex_attention = torch.compile(flex_attention)

device = "mps"
B, H, S, D = 1, 16, 8192, 64  # batch, heads, seq len, head dim
q = torch.randn(B, H, S, D, device=device, dtype=torch.float16)
k = torch.randn(B, H, S, D, device=device, dtype=torch.float16)
v = torch.randn(B, H, S, D, device=device, dtype=torch.float16)

Note the shape: (batch, heads, sequence, head_dim) — identical to scaled_dot_product_attention, so swapping SDPA out needs no reshape.

Mask 1: causal#

The baseline. A query attends only to keys at or before its position.

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

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

B=None, H=None broadcasts the same mask across the batch and all heads — build it once, reuse it every forward pass. A dense causal mask is only mildly sparse, so this is where FlexAttention roughly matches SDPA. The wins come next.

Mask 2: sliding-window causal#

A token attends only to the last W positions. Over a long context this is very sparse — most of the grid is dark — which is exactly the pattern the MPS kernel now skips.

W = 1024  # window size

def sliding_window_causal(b, h, q_idx, kv_idx):
    causal_mask = q_idx >= kv_idx
    window_mask = q_idx - kv_idx < W
    return causal_mask & window_mask

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

A 1k-token window over an 8k context leaves roughly seven-eighths of the blocks fully masked. Those blocks are never computed — that's the ~12x territory.

Mask 3: document-packed (no cross-contamination)#

When you pack several short documents into one sequence to fill the context window, tokens must not attend across document boundaries. Give each token a document id, then AND it into the causal rule.

# document_id[i] = which document token i belongs to
document_id = torch.zeros(S, dtype=torch.int32, device=device)
document_id[2048:] = 1
document_id[5000:] = 2  # three packed documents

def document_causal(b, h, q_idx, kv_idx):
    causal_mask = q_idx >= kv_idx
    same_doc = document_id[q_idx] == document_id[kv_idx]
    return causal_mask & same_doc

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

Because the mask is block-diagonal, cross-document blocks are skipped wholesale — packing four documents is close to a quarter of the attention cost of treating them as one long sequence.

Bonus: a score_mod for soft-capping#

Some models cap logits with tanh to keep attention scores bounded. That's a score_mod, not a mask — it changes the value, not whether a position is live:

SOFTCAP = 30.0

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

out = flex_attention(q, k, v, block_mask=block_mask, score_mod=soft_cap)

Pass block_mask and score_mod together — FlexAttention applies the mask for sparsity and the score_mod on what survives.

The two mistakes that drop you to the slow path#

  1. You forgot torch.compile. Eager flex_attention is correct but unfused; the sparse-pattern speedup is a compiled-kernel property. Wrap it once at load.
  2. Your mask isn't actually sparse. FlexAttention can only skip blocks that are fully masked. A dense causal mask leaves most blocks partially live, so there's little to skip and you'll see numbers close to SDPA. The ~12x is for genuinely sparse shapes — small windows, block-diagonal document packing — not for "causal" alone.

Build the block mask once, reuse it across the run, keep the compile warm, and the Mac on your desk does work you used to rent a GPU for.