{"record":{"id":"4a736d4e2496d138","repo":"labmlai/annotated_deep_learning_paper_implementations","slug":"head-size-self-d-head-too-large-for-flash-atten","errorCode":null,"errorMessage":"Head size ${self.d_head} too large for Flash Attention","messagePattern":"Head size (.+?) too large for Flash Attention","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"labml_nn/diffusion/stable_diffusion/model/unet_attention.py","lineNumber":219,"sourceCode":"        # Get batch size and number of elements along sequence axis (`width * height`)\n        batch_size, seq_len, _ = q.shape\n\n        # Stack `q`, `k`, `v` vectors for flash attention, to get a single tensor of\n        # shape `[batch_size, seq_len, 3, n_heads * d_head]`\n        qkv = torch.stack((q, k, v), dim=2)\n        # Split the heads\n        qkv = qkv.view(batch_size, seq_len, 3, self.n_heads, self.d_head)\n\n        # Flash attention works for head sizes `32`, `64` and `128`, so we have to pad the heads to\n        # fit this size.\n        if self.d_head <= 32:\n            pad = 32 - self.d_head\n        elif self.d_head <= 64:\n            pad = 64 - self.d_head\n        elif self.d_head <= 128:\n            pad = 128 - self.d_head\n        else:\n            raise ValueError(f'Head size ${self.d_head} too large for Flash Attention')\n\n        # Pad the heads\n        if pad:\n            qkv = torch.cat((qkv, qkv.new_zeros(batch_size, seq_len, 3, self.n_heads, pad)), dim=-1)\n\n        # Compute attention\n        # $$\\underset{seq}{softmax}\\Bigg(\\frac{Q K^\\top}{\\sqrt{d_{key}}}\\Bigg)V$$\n        # This gives a tensor of shape `[batch_size, seq_len, n_heads, d_padded]`\n        out, _ = self.flash(qkv)\n        # Truncate the extra head size\n        out = out[:, :, :, :self.d_head]\n        # Reshape to `[batch_size, seq_len, n_heads * d_head]`\n        out = out.reshape(batch_size, seq_len, self.n_heads * self.d_head)\n\n        # Map to `[batch_size, height * width, d_model]` with a linear layer\n        return self.to_out(out)\n\n    def normal_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):","sourceCodeStart":201,"sourceCodeEnd":237,"githubUrl":"https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/33ab02281c2b928e6b32792909cc79cbdcfe1d6a/labml_nn/diffusion/stable_diffusion/model/unet_attention.py#L201-L237","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Increase n_heads or decrease d_model so that d_model // n_heads <= 128","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","Sanity-check the config before building the model: assert d_model % n_heads == 0 and d_model // n_heads <= 128"],"exampleFix":"# before: d_head = 1024 // 4 = 256 -> raises\nattn = FlashAttention(1024, 4)\n\n# after: d_head = 1024 // 16 = 64 -> OK\nattn = FlashAttention(1024, 16)","handlingStrategy":"validation","validationCode":"d_model, n_heads = 1024, 16\nd_head = d_model // n_heads\nassert d_model % n_heads == 0, 'd_model must be divisible by n_heads'\nassert d_head <= 128, f'd_head={d_head} exceeds flash-attention limit of 128'","typeGuard":"def flash_attention_head_size_ok(d_model: int, n_heads: int) -> bool:\n    return d_model % n_heads == 0 and d_model // n_heads <= 128","tryCatchPattern":"try:\n    out = attn.forward(q, k, v)\nexcept ValueError as e:\n    if 'too large for Flash Attention' in str(e):\n        raise SystemExit('Reduce d_head: increase n_heads or disable flash attention')\n    raise","preventionTips":["Always compute d_head = d_model // n_heads when editing attention configs","Keep d_head a power of two <= 128 for flash-attention compatibility","Unit-test model construction with the exact production config before launching training"],"tags":["pytorch","flash-attention","stable-diffusion","config"],"backgroundTag":"attention-head-dim-limit","analyzedSha":"33ab02281c2b928e6b32792909cc79cbdcfe1d6a","analyzedAt":"2026-08-25T10:30:27.743Z","schemaVersion":2},"datasetVersion":"2026-08-25T11:17:15.655Z"}