flash_attn_func crashes with cudaErrorInvalidValue (SMEM overflow) at head_dim=256, causal, bf16 on sm_86/sm_89 (SM80 tile heuristic not head-dim-aware)

[arbi-dev]

Summary

flash_attn.cute.flash_attn_func crashes with a raw CUDA shared-memory overflow on SM80-family consumer GPUs (sm_86/sm_89) at head_dim=256, causal, bf16. Confirmed on an RTX 4090 (sm_89).

Environment

• flash-attn-4: built from main @ c75d019 (2026-07-29)
• nvidia-cutlass-dsl: 4.6.0.dev0
• torch: 2.13.0+cu13.0
• CUDA: 12.5 (driver), GPU reports CUDA 13.0 runtime
• GPU: NVIDIA GeForce RTX 4090 (sm_89)

Repro

import torch
from flash_attn.cute.interface import flash_attn_func

q = torch.randn(1, 256, 2, 256, dtype=torch.bfloat16, device="cuda")
k = torch.randn(1, 256, 2, 256, dtype=torch.bfloat16, device="cuda")
v = torch.randn(1, 256, 2, 256, dtype=torch.bfloat16, device="cuda")
flash_attn_func(q, k, v, causal=True)

error: cudaErrorInvalidValue (error code: 1)
...
Architecture: Ada (sm_89)
Compatible SM archs: sm_89, sm_86
...
Error: kernel '...FlashAttentionForwardSm80...' launch shared memory exceeds current GPU
arch sm_89 allowed. Allocated: 131072 bytes. Max: 101376 bytes.

head_dim=64 (or 128, or 192) with the same shapes/settings works fine — only large head dims trigger it.

Root cause

_flash_attn_fwd in flash_attn/cute/interface.py:539 hardcodes the same forward tile config for every head_dim on SM80:

elif arch // 10 == 8:
    fwd_cfg = FwdConfig(128, 64, True, True)  # SM80, should tune

Unlike SM90 (_tile_size_fwd_sm90, same file), SM80 has no head-dim-aware branching at all. With tile_m=128, tile_n=64, num_stages=1, the static SMEM usage at head_dim=head_dim_v=256 is:

2*(tile_m*head_dim) + 2*tile_n*(head_dim+head_dim_v)
= 2*(128*256) + 2*64*(256+256)
= 65536 + 65536 = 131072 bytes

sm_80 (A100) has 166912 B of static SMEM, so this fits there and the bug doesn't show up on A100. But sm_86/sm_89 (RTX 30xx/40xx) are capped at 101376 B (SMEM_CAPACITY_MAP in cutlass.utils), so 131072 > 101376 and the launch fails.

The kernel classes already have a can_implement() static method that computes exactly this SMEM check (flash_fwd.py:159-171), but it is never invoked at the _flash_attn_fwd/_flash_attn_bwd dispatch sites — both have a literal # TODO: check @can_implement comment (interface.py:982 and :1959). So instead of a clean, early RuntimeError, the failure surfaces as a raw CUDA launch error deep in compiled code.

There's also a second, compounding bug in can_implement() itself: it computes the SMEM budget via a hardcoded utils_basic.get_smem_capacity_in_bytes("sm_80") (166912 B) rather than the actual detected minor architecture (flash_fwd.py:168-169, flagged by the function's own # TODO: sm86 and sm89 comment). So even if wired up as-is, can_implement would incorrectly pass this exact D=256/SM80 config on sm_86/sm_89, since it'd check against the larger A100 budget.

PR #2609 (open, unmerged) wires up the can_implement() calls at the dispatch sites, which would at least turn this into a clean RuntimeError instead of an opaque CUDA error — but it doesn't fix the sm_80-hardcoded capacity constant above, so as currently written it would not actually catch this case on sm_86/sm_89 (left a comment on that PR with specifics). The ideal fix is giving SM80 a real head-dim-aware tile heuristic, like SM90 already has via _tile_size_fwd_sm90, so head_dim=256 picks a tile config that actually fits in 101376 B rather than merely failing cleanly.

Related

• PR Wire can_implement checks before JIT compile in interface.py #2609 — wires up can_implement() guards (fixes the "opaque error" part, not the capacity constant or the tile heuristic).