[FA4, CuTe, SM100] `_flash_attn_fwd` fails at launch for `head_dim < 32` on B200
On B200 (SM100), the FA4 cute forward kernel raises `RuntimeError: CUDA Error: cudaErrorInvalidValue` at kernel launch when `head_dim < 32`, even though `_validate_head_dims` advertises `8 <= head_dim <= 128` as supported on SM100.
## Script to reproduce the error
Pinned versions:
```bash
uv venv --python=3.12
uv pip install torch==2.12.0 "nvidia-cutlass-dsl[cu13]==4.4.2" "flash-attn-4[cu13]==4.0.0b14"
uv run python fa4_head_dim_repro.py
```
Python script:
```python
import math
import sys
import torch
from flash_attn.cute.interface import _flash_attn_fwd
def fa4_call(head_dim: int):
batch, seqlen, num_heads = 1, 2048, 2
dtype = torch.bfloat16
device = "cuda"
q = torch.randn(batch, seqlen, num_heads, head_dim, dtype=dtype, device=device)
k = torch.randn(batch, seqlen, num_heads, head_dim, dtype=dtype, device=device)
v = torch.randn(batch, seqlen, num_heads, head_dim, dtype=dtype, device=device)
torch.cuda.synchronize()
try:
out, lse = _flash_attn_fwd(
q,
k,
v,
softmax_scale=1.0 / math.sqrt(head_dim),
causal=True,
return_lse=True,
)
torch.cuda.synchronize()
except Exception as exc:
return f"head_dim={head_dim:>3}: FAIL {type(exc).__name__}: {exc}"
return (
f"head_dim={head_dim:>3}: OK "
f"out.shape={tuple(out.shape)} lse.shape={tuple(lse.shape)}"
)
def main():
cap = torch.cuda.get_device_capability()
print(f"python: {sys.version.split()[0]}")
print(f"torch: {torch.__version__}")
print(
f"device: {torch.cuda.get_device_name(0)} "
f"sm_{cap[0]}{cap[1]} (capability={cap})"
)
print("Calling flash_attn.cute._flash_attn_fwd with varying head_dim:")
print("(batch=1, seqlen=2048, num_heads=2, bf16, causal=True)")
results = [fa4_call(hd) for hd in (8, 16, 32, 64, 96, 128)]
print("\n".join(results))
if __name__ == "__main__":
main()
```
### Output
```
python: 3.12.3
torch: 2.12.0+cu130
device: NVIDIA B200 sm_100 (capability=(10, 0))
Calling flash_attn.cute._flash_attn_fwd with varying head_dim:
(batch=1, seqlen=2048, num_heads=2, bf16, causal=True)
head_dim= 8: FAIL RuntimeError: CUDA Error: cudaErrorInvalidValue
head_dim= 16: FAIL RuntimeError: CUDA Error: cudaErrorInvalidValue
head_dim= 32: OK out.shape=(1, 2048, 2, 32) lse.shape=(1, 2, 2048)
head_dim= 64: OK out.shape=(1, 2048, 2, 64) lse.shape=(1, 2, 2048)
head_dim= 96: OK out.shape=(1, 2048, 2, 96) lse.shape=(1, 2, 2048)
head_dim=128: OK out.shape=(1, 2048, 2, 128) lse.shape=(1, 2, 2048)
```
1 条评论