FA3 loader reports success on Blackwell (sm_120) then dies at first kernel launch

FA3 loader reports success on Blackwell (sm_120) then dies at first kernel launch

Summary

On RTX 50-series (Blackwell, sm_120), _load_flash_attention_3() successfully imports an FA3 kernel that has no sm_120 image. USE_FA3 is set to True, the run prints the success message, and then crashes at the first attention call. The SDPA fallback exists and works, but is never reached.

Environment

GPU     : NVIDIA GeForce RTX 5090 (32 GB), compute capability (12, 0)
Driver  : 580.95.05, CUDA 13.0
torch   : 2.9.1+cu128
Python  : 3.10, uv-managed venv

Reproduction

torchrun --standalone --nproc_per_node=1 -m scripts.base_train -- \
    --depth=24 --target-param-data-ratio=8 --device-batch-size=2 --fp8

Output (trimmed):

GPU: NVIDIA GeForce RTX 5090 | Peak FLOPS (BF16): 2.10e+14
COMPUTE_DTYPE: torch.bfloat16 (auto-detected: CUDA SM 120 (bf16 supported))
✓ Using Flash Attention 3: efficient, new and awesome.
...
Total batch size 1,048,576 => gradient accumulation steps: 256
CUDA error (/build/source/flash-attn/flash_fwd_launch_template.h:192): no kernel image is available for execution on the device

Note the path in the error — /build/source/flash-attn/... — is from the kernel's build machine, not the local checkout. Model construction, FP8 conversion and dataloader setup all succeed; the failure is at the first flash_attn_func call, ~12 seconds in.

Cause

nanochat/flash_attention.py, _load_flash_attention_3():

        major, _ = torch.cuda.get_device_capability()
        # FA3 kernels are currently compiled for Hopper (sm90), Ada (sm89) and Ampere (sm80/sm86)
        # Blackwell (sm100) needs SDPA fallback until FA3 is recompiled or FA4 is released
        ...
        if major == 9:
            hf_kernel = "varunneal/flash-attention-3"
            return get_kernel(hf_kernel).flash_attn_interface
        else:
            hf_kernel = "kernels-community/flash-attn3"
            if has_kernel(hf_kernel):
                return get_kernel(hf_kernel).flash_attn_interface
            else:
                return None

The comment anticipates exactly this case, but the code only special-cases major == 9. Blackwell (major == 12) falls into the else, and has_kernel("kernels-community/flash-attn3") returns True — it appears to report repository-level availability rather than per-architecture. The kernel loads, HAS_FA3 becomes True, and the failure is deferred to launch time.

Suggested fix

One line, immediately after the capability check:

        major, _ = torch.cuda.get_device_capability()
        if major >= 10:
            return None   # Blackwell sm100/sm120: no FA3 kernels available

This returns before the kernels import, so it also avoids downloading a kernel that cannot run. Verified on the affected hardware:

python -c "from nanochat.flash_attention import USE_FA3, HAS_FA3; print(USE_FA3, HAS_FA3)"
# False False

After the patch the run proceeds normally on the SDPA path and completes a full d24 pretrain (details in the linked discussion post, if opened).

>= 10 rather than == 12 covers sm_100 (datacenter Blackwell) as well, matching the existing comment. If has_kernel() is later made architecture-aware upstream, this guard becomes redundant and can be dropped.

Workaround without patching

The file already has a test hook:

_override_impl = None   # set to 'sdpa'

Setting it to 'sdpa' produces the same result via _resolve_use_fa3(). Uninstalling the kernels package also works (the ImportError is caught by the existing except Exception: return None), which is presumably why older environments without kernels installed never hit this.

Note on impact

With --window-pattern L the SDPA path stays on its fused branch (at Tq == Tk with a full-context window, _sdpa_attention returns via F.scaled_dot_product_attention(..., is_causal=True) with no explicit mask), so the performance cost on Blackwell is modest. With the SSSL default, SDPA must build an explicit mask and the existing "GPU utilization will be terrible" warning in base_train.py applies.

---

Disclosure per repo AI policy: this report was drafted with LLM assistance. The bug, environment, logs and the tested fix are from my own run.