[fix] fix enable_pdl for blackwell (#9011)

This commit is contained in:
JieXin Liang
2025-08-19 01:16:08 +08:00
committed by GitHub
parent c480a3f6ea
commit 6cdcbcc674
2 changed files with 7 additions and 7 deletions

View File

@@ -43,8 +43,8 @@ def _to_tensor_scalar_tuple(x):
@functools.lru_cache(maxsize=1)
def is_hopper_arch() -> bool:
def is_arch_support_pdl() -> bool:
# Hopper arch's compute capability == 9.0
device = torch.cuda.current_device()
major, minor = torch.cuda.get_device_capability(device)
return major == 9
return major >= 9