From 1e12db623cd3d0d4ebbe9ed14d1c7892324473cb Mon Sep 17 00:00:00 2001 From: Xu Zhao Date: Wed, 24 Sep 2025 13:36:38 -0400 Subject: [PATCH 1/2] install cu130 --- docker/tritonbench-nightly.dockerfile | 2 +- tools/cuda_utils.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/docker/tritonbench-nightly.dockerfile b/docker/tritonbench-nightly.dockerfile index 3a4156e3e..7ac1ea9b6 100644 --- a/docker/tritonbench-nightly.dockerfile +++ b/docker/tritonbench-nightly.dockerfile @@ -32,7 +32,7 @@ RUN cd /workspace/pytorch-ci; wget https://raw.githubusercontent.com/pytorch/pyt wget https://raw.githubusercontent.com/pytorch/pytorch/main/.ci/docker/common/install_cusparselt.sh && \ mkdir ci_commit_pins && cd ci_commit_pins && \ wget https://raw.githubusercontent.com/pytorch/pytorch/main/.ci/docker/ci_commit_pins/nccl-cu12.txt -RUN sudo bash -c "set -x;export OVERRIDE_GENCODE=\"${OVERRIDE_GENCODE}\" OVERRIDE_GENCODE_CUDNN=\"${OVERRIDE_GENCODE_CUDNN}\"; cd /workspace/pytorch-ci; bash install_cuda.sh 12.8" +RUN sudo bash -c "set -x;export OVERRIDE_GENCODE=\"${OVERRIDE_GENCODE}\" OVERRIDE_GENCODE_CUDNN=\"${OVERRIDE_GENCODE_CUDNN}\"; cd /workspace/pytorch-ci; bash install_cuda.sh 13.0" # Install miniconda RUN wget -q https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /workspace/Miniconda3-latest-Linux-x86_64.sh diff --git a/tools/cuda_utils.py b/tools/cuda_utils.py index 8fa97d65b..4169a481f 100644 --- a/tools/cuda_utils.py +++ b/tools/cuda_utils.py @@ -7,13 +7,17 @@ from .torch_utils import install_pytorch_nightly # defines the default CUDA version to compile against -DEFAULT_CUDA_VERSION = "12.8" +DEFAULT_CUDA_VERSION = "13.0" CUDA_VERSION_MAP = { "12.8": { "pytorch_url": "cu128", "jax": "jax[cuda12]", }, + "13.0": { + "pytorch_url": "cu130", + "jax": "jax[cuda13]", + }, } From 3587bd4dfd7eac388833b8cfd509e00f598cb649 Mon Sep 17 00:00:00 2001 From: Xu Zhao Date: Wed, 24 Sep 2025 15:46:32 -0400 Subject: [PATCH 2/2] fix nccl --- docker/tritonbench-nightly.dockerfile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docker/tritonbench-nightly.dockerfile b/docker/tritonbench-nightly.dockerfile index 7ac1ea9b6..1dd5c1e36 100644 --- a/docker/tritonbench-nightly.dockerfile +++ b/docker/tritonbench-nightly.dockerfile @@ -31,7 +31,7 @@ RUN cd /workspace/pytorch-ci; wget https://raw.githubusercontent.com/pytorch/pyt wget https://raw.githubusercontent.com/pytorch/pytorch/main/.ci/docker/common/install_nccl.sh && \ wget https://raw.githubusercontent.com/pytorch/pytorch/main/.ci/docker/common/install_cusparselt.sh && \ mkdir ci_commit_pins && cd ci_commit_pins && \ - wget https://raw.githubusercontent.com/pytorch/pytorch/main/.ci/docker/ci_commit_pins/nccl-cu12.txt + wget https://raw.githubusercontent.com/pytorch/pytorch/main/.ci/docker/ci_commit_pins/nccl-cu13.txt RUN sudo bash -c "set -x;export OVERRIDE_GENCODE=\"${OVERRIDE_GENCODE}\" OVERRIDE_GENCODE_CUDNN=\"${OVERRIDE_GENCODE_CUDNN}\"; cd /workspace/pytorch-ci; bash install_cuda.sh 13.0" # Install miniconda