tacit/nanochat
1
0
Fork 0
mirror of https://github.com/karpathy/nanochat.git synced 2026-04-01 21:25:21 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

small fix

Browse Source
This commit is contained in:
svlandeg 2026-03-11 16:49:09 +01:00
parent 5985cd867b
commit ff833f8137
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
nanochat/flash_attention.py
View File

@ -29,10 +29,10 @@ def _load_flash_attention_3():
# 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
# varunneal kernel obtains better results for H100/Hopper
hf_kernel = "varunneal/flash-attention-3" if major == 9 else "kernels-community/flash-attn3"
from kernels import get_kernel
import os
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
from kernels import get_kernel
hf_kernel = "varunneal/flash-attention-3" if major == 9 else "kernels-community/flash-attn3"
return get_kernel(hf_kernel).flash_attn_interface
except Exception:
return None