nvcc segfaults building flash-attn 2.x for SM 120 (RTX 50-series)

[bmdhodl]

Summary

Building flash-attn from source with FLASH_ATTN_CUDA_ARCHS=120 crashes nvcc during compilation of backward kernels. This blocks all RTX 50-series (SM 12.0) users from using flash-attn 2.x.

Environment

• GPU: RTX 5070 Ti (SM 12.0)
• CUDA: 12.8 and 12.9 (both fail)
• PyTorch: 2.9.1+cu128
• WSL2 Ubuntu 22.04

Reproduction

FLASH_ATTN_CUDA_ARCHS=120 pip install flash-attn --no-build-isolation

nvcc segfaults on backward pass kernels (e.g. flash_bwd_hdim256_bf16_sm80.cu, flash_bwd_hdim128_bf16_causal_sm80.cu). Tested with both CUDA 12.8 and 12.9.

Prebuilt wheels

The v2.8.3 release only has wheels up to torch 2.8. The torch 2.8 wheel has ABI incompatibility with torch 2.9 (undefined symbol: _ZNK3c106SymInt6sym_neERKS0_). A torch 2.9 + cu128 wheel for sm_120 would unblock this.

Context

• SM 120 uses MMAv2 tensor cores (same as Ampere), so flash-attn 2.x MMAv2 kernels should be compatible at the hardware level
• PyTorch SDPA achieves ~160 TFLOPS on the 5070 Ti using cuDNN flash attention, so the performance gap may be small
• The related Triton sm_120 codegen fix was just merged: [NVIDIA] Fix PTX codegen segfaults on consumer Blackwell (sm_120) triton-lang/triton#9734
• See also FA4 SM120 support? #2307 for FA4 sm_120 status