From d25ee7f8df53f53e170d95592208f31102b29c71 Mon Sep 17 00:00:00 2001 From: Anatoly Myachev Date: Tue, 3 Dec 2024 23:41:14 +0100 Subject: [PATCH] Search for 'ptxas' only for cuda backend in 'supports_tma' function Signed-off-by: Anatoly Myachev --- python/triton/_internal_testing.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/triton/_internal_testing.py b/python/triton/_internal_testing.py index 5ba0be1e34f9..377eed877a51 100644 --- a/python/triton/_internal_testing.py +++ b/python/triton/_internal_testing.py @@ -142,11 +142,13 @@ def to_numpy(x): def supports_tma(byval_only=False): + if not is_cuda(): + return False _, cuda_version = _path_to_binary("ptxas") min_cuda_version = (12, 0) if byval_only else (12, 3) cuda_version_tuple = tuple(map(int, cuda_version.split("."))) assert len(cuda_version_tuple) == 2, cuda_version_tuple - return is_cuda() and torch.cuda.get_device_capability()[0] >= 9 and cuda_version_tuple >= min_cuda_version + return torch.cuda.get_device_capability()[0] >= 9 and cuda_version_tuple >= min_cuda_version def tma_skip_msg(byval_only=False):