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

use capability in the same way as on master

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

6
nanochat/flash_attention.py
View File

@ -25,11 +25,11 @@ def _load_flash_attention_3():
if not torch.cuda.is_available():
return None
try:
cap = torch.cuda.get_device_capability()
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
# varunneal kernel obtains better results for H100
hf_kernel = "varunneal/flash-attention-3" if cap == (9, 0) else "kernels-community/flash-attn3"
# 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"