feat: use fa3 mla by default on hopper (#5210)

Co-authored-by: yundai424 <yundai424@gmail.com>
Co-authored-by: hebiao064 <hebiaobuaa@gmail.com>
This commit is contained in:
Yineng Zhang
2025-04-12 01:09:25 -07:00
committed by GitHub
parent 115ae2e728
commit 57de7c6b5f
3 changed files with 42 additions and 11 deletions

View File

@@ -1828,3 +1828,12 @@ def fast_topk(values, topk, dim):
else:
# Use topk for efficiency with larger k values
return torch.topk(values, topk, dim=dim)
def is_hopper_with_cuda_12_3():
if not is_cuda():
return False
is_hopper = torch.cuda.get_device_capability()[0] == 9
cuda_version = torch.version.cuda.split(".")
is_cuda_compatible = int(cuda_version[0]) == 12 and int(cuda_version[1]) >= 3
return is_hopper and is_cuda_compatible