[Feat] Enable PDL automatically on Hopper architecture (#5981)

This commit is contained in:
Huapeng Zhou
2025-06-01 12:30:17 -07:00
committed by GitHub
parent c6a0cacc35
commit 2f7420bc84
2 changed files with 28 additions and 9 deletions

View File

@@ -39,3 +39,10 @@ def _to_tensor_scalar_tuple(x):
return (x, 0)
else:
return (None, x)
def is_hopper_arch() -> bool:
# Hopper arch's compute capability == 9.0
device = torch.cuda.current_device()
major, minor = torch.cuda.get_device_capability(device)
return major == 9