Skip to content

Commit df930ef

Browse files
authored
[CUDA][PTXAS] Make TRITON_PTXAS_BLACKWELL_PATH work to overwrite ptxas-blackwell binary (#8945)
Mainly for THOR as by default ptxas-blackwell (cuda 12.9) would be used and it would go "I don't understand sm110a". Context: There is an environment variable TRITON_PTXAS_BLACKWELL_PATH that was designed for users to overwrite this ptxas-blackwell, via the lines noted in #8945 (review). But: `export TRITON_PTXAS-BLACKWELL_PATH=/usr/local/cuda/bin/ptxas` is NOT a valid bash script, but `export TRITON_PTXAS_BLACKWELL_PATH` is. So this PR maps `ptxas-blackwell` binary to look for `TRITON_PTXAS_BLACKWELL_PATH`. IF set, the ptxas-blackwell would use the provided "path" (nit: it is call a "path" but we should supply the binary, like "TRITON_PTXAS_BLACKWELL_PATH_AND_BINARY".
1 parent 3a731bd commit df930ef

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

python/triton/knobs.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -196,7 +196,8 @@ def __init__(self, binary: str) -> None:
196196
binary += sysconfig.get_config_var("EXE")
197197
self.binary = binary
198198
self.default_path = os.path.join(os.path.dirname(__file__), "backends", "nvidia", "bin", binary)
199-
super().__init__(f"TRITON_{binary.upper()}_PATH")
199+
# Convert ptxas-blackwell to PTXAS_BLACKWELL, not PTXAS-BLACKWELL
200+
super().__init__(f"TRITON_{binary.upper().replace('-', '_')}_PATH")
200201

201202
def get(self) -> NvidiaTool:
202203
return self.transform(getenv(self.key))

0 commit comments

Comments
 (0)