Fix ptxas compilation on sm103 for triton kernels (#2539)

tdophung · web-flow · commit 97a09c29cffa · 2025-12-22T13:27:51.000-08:00
* add triton ptxas path  for gb300 to find where it is to avoid compilation errors

Signed-off-by: tdophung &lt;tdophung@nvidia.com&gt;

* add these flags in advance to preven future breaks when ops are extended to multi gpus

Signed-off-by: tdophung &lt;tdophung@nvidia.com&gt;

* add this also to L1

Signed-off-by: tdophung &lt;tdophung@nvidia.com&gt;

---------

Signed-off-by: tdophung &lt;tdophung@nvidia.com&gt;
diff --git a/qa/L0_jax_distributed_unittest/test.sh b/qa/L0_jax_distributed_unittest/test.sh
@@ -1,6 +1,7 @@
 # Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # See LICENSE for license information.
+export TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas
 
 function error_exit() {
     echo "Error: $1"
diff --git a/qa/L0_jax_unittest/test.sh b/qa/L0_jax_unittest/test.sh
@@ -1,6 +1,7 @@
 # Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # See LICENSE for license information.
+export TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas
 
 set -x
 
diff --git a/qa/L1_jax_distributed_unittest/test.sh b/qa/L1_jax_distributed_unittest/test.sh
@@ -1,6 +1,7 @@
 # Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # See LICENSE for license information.
+export TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas
 
 function test_fail() {
     RET=1
diff --git a/qa/L2_jax_distributed_unittest/test.sh b/qa/L2_jax_distributed_unittest/test.sh
@@ -1,6 +1,7 @@
 # Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # See LICENSE for license information.
+export TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas
 
 set -xe
 
diff --git a/qa/L2_jax_unittest/test.sh b/qa/L2_jax_unittest/test.sh
@@ -1,6 +1,7 @@
 # Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # See LICENSE for license information.
+export TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas
 
 set -x
 

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,7 @@`
`1`	`1`	`# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.`
`2`	`2`	`#`
`3`	`3`	`# See LICENSE for license information.`
	`4`	`+export TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas`
`4`	`5`
`5`	`6`	`function error_exit() {`
`6`	`7`	`echo "Error: $1"`