best available flash attention implementation for encoder like models for sm120?
Seven evidence rows as of September 1, 2026 — recommendations, one RTX 5090 forward benchmark, FlashAttention source-tree maturity facts, and one architecture caveat — cover attention backends for bidirectional, non-causal encoder workloads on NVIDIA SM120 (consumer/workstation Blackwell, e.g. RTX 5090). Throughput is in TFLOP/s, BF16.
SM120 production path; exact non-causal encoder attention and training through PyTorch autograd.
Force CUDNN_ATTENTION to test it, but permit automatic fallback if a mask or shape is unsupported.
Use recent CUDA-enabled PyTorch/cuDNN; call F.scaled_dot_product_attention with is_causal=False. docs.nvidia.com
Native SM120 build target; non-causal capability should be validated on the model shape.
Main repository setup now emits sm_120 code; benchmark against cuDNN before adopting.
CUDA 12.8+ and an SM120-targeted source build. github.com
Readiness matrix: each row is an implementation path, each column a capability an encoder stack needs on SM120. Filled squares are supported with cited evidence; hollow squares are unproven or not applicable on SM120. The bar on the right is the one measured RTX 5090 number against its stated roofline — forward pass only, one shape (B=1, H=8, Q=4096, KV=8192, BF16, CUDA 12.9). Hover any mark for the source detail.
# exact non-causal encoder attention, trains through autograd out = F.scaled_dot_product_attention(q, k, v, is_causal=False) # when benchmarking, pin the backend; in production keep automatic fallback with sdpa_kernel(SDPBackend.CUDNN_ATTENTION): out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, is_causal=False)
| Tier | Implementation | Result / support | Guidance | Source |
|---|---|---|---|---|
| Recommended default | PyTorch SDPA → cuDNN Attention | cudnn SDPA operation requires SM80 (Ampere) or newer architectures and cuda toolkit 12.x or newer. | Force CUDNN_ATTENTION to test it, but permit automatic fallback if a mask or shape is unsupported | docs.nvidia.com |
| Benchmark evidence | PyTorch SDPA → cuDNN Attention | 203.61 TFLOPS (97.19% of SOL) | 203.61 TFLOP/s, 97.19% of the stated BF16 roofline at B=1, H=8, Q=4096, KV=8192 | gau-nernst.github.io |
| Bleeding-edge alternative | Dao-AILab FlashAttention CuTe SM120 path | supported under compute capability 120 (sm_120) with CUDA 12.8 and above | Main repository setup now emits sm_120 code; benchmark against cuDNN before adopting | github.com |
| FlashAttention maturity evidence | Dao-AILab FlashAttention CuTe SM120 forward | Merged | Forward support is merged | github.com |
| FlashAttention maturity evidence | Dao-AILab FlashAttention CuTe SM120 backward | SM120 backward implementation committed | Training support exists in source, but validate correctness and speed for the exact workload | github.com |
| Do not confuse architectures | FlashAttention-4 | 1.3× speedup over cuDNN 9.13 and 2.7× over Triton on B200 GPUs with BF16, reaching up to 1613 TFLOPs/s | Do not infer RTX 5090 / SM120 compatibility from B200 results | proceedings.mlsys.org |
| Framework corroboration | cuDNN Attention backend | CUDNN_ATTN | vLLM-Omni reports 2× end-to-end versus generic SDPA on a mask-heavy DiT workload | docs.vllm.ai |
Recommendation and evidence table for SM120 (consumer Blackwell) encoder attention: 7 rows, one per recommendation, benchmark, maturity fact or architecture caveat, compiled as of 2026-09-01. The single throughput number measures BF16 forward attention TFLOP/s on an RTX 5090 at one shape and is not a training or multi-shape result. Requirement details are shortened in the table; full text is at each source.