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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user