labmlai/annotated_deep_learning_paper_implementations · error · ValueError

Head size ${self.d_head} too large for Flash Attention

Error message

Head size ${self.d_head} too large for Flash Attention

What it means

The Stable Diffusion UNet attention module wraps a Flash Attention kernel that only supports per-head dimensions up to 128. Before calling the kernel, the code pads the head dimension up to the next supported bucket (32, 64, or 128). If d_head exceeds 128 there is no bucket to pad to, so the wrapper raises ValueError rather than silently computing wrong results.

Source

Thrown at labml_nn/diffusion/stable_diffusion/model/unet_attention.py:219

        # Get batch size and number of elements along sequence axis (`width * height`)
        batch_size, seq_len, _ = q.shape

        # Stack `q`, `k`, `v` vectors for flash attention, to get a single tensor of
        # shape `[batch_size, seq_len, 3, n_heads * d_head]`
        qkv = torch.stack((q, k, v), dim=2)
        # Split the heads
        qkv = qkv.view(batch_size, seq_len, 3, self.n_heads, self.d_head)

        # Flash attention works for head sizes `32`, `64` and `128`, so we have to pad the heads to
        # fit this size.
        if self.d_head <= 32:
            pad = 32 - self.d_head
        elif self.d_head <= 64:
            pad = 64 - self.d_head
        elif self.d_head <= 128:
            pad = 128 - self.d_head
        else:
            raise ValueError(f'Head size ${self.d_head} too large for Flash Attention')

        # Pad the heads
        if pad:
            qkv = torch.cat((qkv, qkv.new_zeros(batch_size, seq_len, 3, self.n_heads, pad)), dim=-1)

        # Compute attention
        # $$\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_{key}}}\Bigg)V$$
        # This gives a tensor of shape `[batch_size, seq_len, n_heads, d_padded]`
        out, _ = self.flash(qkv)
        # Truncate the extra head size
        out = out[:, :, :, :self.d_head]
        # Reshape to `[batch_size, seq_len, n_heads * d_head]`
        out = out.reshape(batch_size, seq_len, self.n_heads * self.d_head)

        # Map to `[batch_size, height * width, d_model]` with a linear layer
        return self.to_out(out)

    def normal_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):

View on GitHub (pinned to 33ab02281c)

Solutions

1. Increase n_heads or decrease d_model so that d_model // n_heads <= 128
2. If you need large heads, bypass the flash path: monkeypatch/replace flash_attention with torch.nn.functional.scaled_dot_product_attention or the plain attention implementation
3. Sanity-check the config before building the model: assert d_model % n_heads == 0 and d_model // n_heads <= 128

Example fix

# before: d_head = 1024 // 4 = 256 -> raises
attn = FlashAttention(1024, 4)

# after: d_head = 1024 // 16 = 64 -> OK
attn = FlashAttention(1024, 16)

d_model, n_heads = 1024, 16
d_head = d_model // n_heads
assert d_model % n_heads == 0, 'd_model must be divisible by n_heads'
assert d_head <= 128, f'd_head={d_head} exceeds flash-attention limit of 128'

def flash_attention_head_size_ok(d_model: int, n_heads: int) -> bool:
    return d_model % n_heads == 0 and d_model // n_heads <= 128

try:
    out = attn.forward(q, k, v)
except ValueError as e:
    if 'too large for Flash Attention' in str(e):
        raise SystemExit('Reduce d_head: increase n_heads or disable flash attention')
    raise

When it happens

Trigger: Constructing/running UNetAttentionModule or CrossAttention where d_model // n_heads > 128; calling forward() or compute_flash_attention() on such a module; any custom config that sets n_heads low relative to d_model (e.g. d_model=1024 with n_heads=4 gives d_head=256).

Common situations: Porting a UNet or attention config from another repo that used PyTorch SDPA (no head-dim limit); experimenting with fewer heads for memory; typo in n_heads; upgrading labml-nn where flash attention became the default path.

Related errors

• Head size {d_k} too large for flash attention
• Unknown variant {configs.glu_variant}
• Please install `bitsandbytes` with `pip install bitsandbytes
• Invalid learning rate: {lr}
• Invalid epsilon value: {eps}

AI-assisted analysis of labmlai/annotated_deep_learning_paper_implementations@33ab02281c (2026-08-25). Data as JSON: /api/errors/4a736d4e2496d138. Report an issue: GitHub.