Skip to content

Commit

Permalink
add sm_100 and sm_120 to cuda 12.8 domain builds (#6273)
Browse files Browse the repository at this point in the history
current torchvision/audio builds missing sm_100 and sm_120

adding to cu128 builds

`cuobjdump -ptx
/usr/local/lib/python3.12/dist-packages/torchvision/_C.so | grep "arch =
sm_" | sort --unique
arch = sm_50
arch = sm_60
arch = sm_70
arch = sm_75
arch = sm_80
arch = sm_86
arch = sm_90`

cc @atalman @ptrblck @nWEIdia
  • Loading branch information
tinglvv authored Feb 10, 2025
1 parent dbde372 commit e8c908c
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion tools/pkg-helpers/pytorch_pkg_helpers/cuda.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,11 @@


def get_cuda_arch_list(sanitized_version: str) -> str:
base_arch_list = "5.0+PTX;6.0;7.0;7.5;8.0;8.6;9.0"
if float(sanitized_version) >= 12.0:
return "5.0+PTX;6.0;7.0;7.5;8.0;8.6;9.0"
if sanitized_version == "12.8":
return base_arch_list + ";10.0;12.0"
return base_arch_list
if float(sanitized_version) > 11.3:
return "3.5;5.0+PTX;6.0;7.0;7.5;8.0;8.6"
# mainly for cuda 10.2
Expand Down

0 comments on commit e8c908c

Please sign in to comment.