labmlai/annotated_deep_learning_paper_implementations · error · ValueError

Head size {d_k} too large for flash attention

Error message

Head size {d_k} too large for flash attention

What it means

GPT-NeoX's compute_flash_attention pads the key/query head dimension d_k to the next supported size (32/64/128) before invoking the fused flash attention kernel. d_k above 128 exceeds the kernel's maximum supported head size, so the code raises ValueError instead of calling the kernel with an unsupported shape.

Source

Thrown at labml_nn/neox/model.py:308

        # Reshape from `[batch_size, seq_len, n_heads, d_k] to `[batch_size, seq_len, n_hidden]`
        output = output.reshape(*x.shape)

        # Final linear layer
        return self.output(output)

    def compute_flash_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
        # Stack them into shape `[batch_size, seq_len, 3, n_heads, d_k]`
        qkv = torch.stack((q, k, v), dim=2)
        d_k = qkv.shape[-1]
        if d_k <= 32:
            pad = 32 - d_k
        elif d_k <= 64:
            pad = 64 - d_k
        elif d_k <= 128:
            pad = 128 - d_k
        else:
            raise ValueError(f'Head size {d_k} too large for flash attention')

        if pad > 0:
            qkv = torch.cat((qkv, qkv.new_zeros(*qkv.shape[:-1], pad)), dim=-1)

        output, _ = self.flash_attention(qkv, causal=True)
        # The output is of shape `[batch_size, seq_len, n_heads, d_k + padding]`
        output = output[:, :, :, :d_k]

        return output

    def compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
        # Disable auto-casting to fp16 for attention computation
        with autocast(enabled=False):
            if q.dtype == torch.float16:
                # Convert to fp32 if the current dtype is fp16
                attn = torch.einsum('bihk,bjhk->bijh', q.float(), k.float())
            else:
                # Do not cast for bfloat

View on GitHub (pinned to 33ab02281c)

Solutions

1. Change n_heads so that d_model // n_heads <= 128
2. Fall back to standard attention: replace compute_flash_attention with a plain softmax(QK^T/sqrt(d_k))V implementation for those layers
3. Validate the config at load time: assert d_model % n_heads == 0 and d_model // n_heads <= 128

Example fix

# before: d_model=8192, n_heads=32 -> d_k=256 -> raises
model = NeoX(d_model=8192, n_heads=32)

# after: d_k=64 -> OK
model = NeoX(d_model=8192, n_heads=128)

d_model, n_heads = 8192, 128
assert d_model % n_heads == 0 and d_model // n_heads <= 128, \n    f'd_k={d_model // n_heads} exceeds flash-attention limit of 128'

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

try:
    logits = model(idx, target)
except ValueError as e:
    if 'too large for flash attention' in str(e):
        raise SystemExit('Fix NeoX config: d_model // n_heads must be <= 128')
    raise

When it happens

Trigger: Running NeoX model forward() with d_model // n_heads > 128 (e.g. d_model=12288 with 96 heads gives d_head=128 which is fine, but d_model=12288 with 48 heads gives 256 which raises); directly calling compute_flash_attention with an oversized qkv tensor.

Common situations: Scaling down n_heads when adapting NeoX configs; using configs copied from models that relied on PyTorch's built-in SDPA which supports larger head dims; mismatched d_model/n_heads after editing a YAML config.

Related errors

• Head size ${self.d_head} 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/988d779d5929f8e5. Report an issue: GitHub.