[Hierarchical Compile] Handle autocast ctx manager

Summary: add a param to specify to the storage writer how to save tensors. Write now the only options are safetensors and torch.save.

Test Plan:
(lintrunner) [ankitageorge@devgpu003.cco3 /data/users/ankitageorge/fbsource/fbcode/caffe2 (1d57cb27b)]$ buck2 test 'fbcode//mode/opt' fbcode//caffe2/test/distributed/checkpoint:test_hf_storage
File changed: fbcode//caffe2/torch/distributed/checkpoint/filesystem.py
Buck UI: https://www.internalfb.com/buck2/e80cc963-e34a-4876-b6f4-7ce2794e48dd
Test UI: https://www.internalfb.com/intern/testinfra/testrun/3659174965882569
Network: Up: 32KiB  Down: 1.9KiB  (reSessionID-ef9fa764-a40a-451b-ab58-08eabe7a9422)
Executing actions. Remaining     0/4                                                                                             3.4s exec time total
Command: test.     Finished 2 local
Time elapsed: 19.6s
Tests finished: Pass 4. Fail 0. Fatal 0. Skip 0. Build failure 0

Reviewed By: saumishr

Differential Revision: D70271943

Pull Request resolved: https://github.com/pytorch/pytorch/pull/150025
Approved by: https://github.com/saumishr

.ci/aarch64_linux/aarch64_ci_build.sh

@@ -3,8 +3,11 @@ set -eux -o pipefail
 
 GPU_ARCH_VERSION=${GPU_ARCH_VERSION:-}
 
-# cuda arm build for Grace Hopper solely
-export TORCH_CUDA_ARCH_LIST="9.0"
+if [[ "$GPU_ARCH_VERSION" == *"12.6"* ]]; then
+    export TORCH_CUDA_ARCH_LIST="9.0"
+elif [[ "$GPU_ARCH_VERSION" == *"12.8"* ]]; then
+    export TORCH_CUDA_ARCH_LIST="9.0;10.0;12.0"
+fi
 
 SCRIPTPATH="$( cd -- "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P )"
 source $SCRIPTPATH/aarch64_ci_setup.sh
@@ -17,7 +20,7 @@ cd /
 # on the mounted pytorch repo
 git config --global --add safe.directory /pytorch
 pip install -r /pytorch/requirements.txt
-pip install auditwheel
+pip install auditwheel==6.2.0
 if [ "$DESIRED_CUDA" = "cpu" ]; then
     echo "BASE_CUDA_VERSION is not set. Building cpu wheel."
     #USE_PRIORITIZED_TEXT_FOR_LD for enable linker script optimization https://github.com/pytorch/pytorch/pull/121975/files

.ci/aarch64_linux/aarch64_ci_setup.sh

@@ -5,16 +5,14 @@ set -eux -o pipefail
 # By creating symlinks from desired /opt/python to /usr/local/bin/
 
 NUMPY_VERSION=2.0.2
-PYGIT2_VERSION=1.15.1
-if [[ "$DESIRED_PYTHON"  == "3.13" ]]; then
+if [[ "$DESIRED_PYTHON"  == "3.13" || "$DESIRED_PYTHON" == "3.13t" ]]; then
     NUMPY_VERSION=2.1.2
-    PYGIT2_VERSION=1.16.0
 fi
 
 SCRIPTPATH="$( cd "$(dirname "$0")" ; pwd -P )"
 source $SCRIPTPATH/../manywheel/set_desired_python.sh
 
-pip install -q numpy==${NUMPY_VERSION} pyyaml==6.0.2 scons==4.7.0 ninja==1.11.1 patchelf==0.17.2 pygit2==${PYGIT2_VERSION}
+pip install -q numpy==${NUMPY_VERSION} pyyaml==6.0.2 scons==4.7.0 ninja==1.11.1 patchelf==0.17.2
 
 for tool in python python3 pip pip3 ninja scons patchelf; do
     ln -sf ${DESIRED_PYTHON_BIN_DIR}/${tool} /usr/local/bin;

.ci/aarch64_linux/aarch64_wheel_ci_build.py

@@ -4,12 +4,9 @@
 import os
 import shutil
 from subprocess import check_call, check_output
-from typing import List
-
-from pygit2 import Repository
 
 
-def list_dir(path: str) -> List[str]:
+def list_dir(path: str) -> list[str]:
     """'
     Helper for getting paths for Python
     """
@@ -42,7 +39,7 @@ def build_ArmComputeLibrary() -> None:
             "clone",
             "https://github.com/ARM-software/ComputeLibrary.git",
             "-b",
-            "v24.09",
+            "v25.02",
             "--depth",
             "1",
             "--shallow-submodules",
@@ -58,7 +55,7 @@ def build_ArmComputeLibrary() -> None:
         shutil.copytree(f"{acl_checkout_dir}/{d}", f"{acl_install_dir}/{d}")
 
 
-def update_wheel(wheel_path) -> None:
+def update_wheel(wheel_path, desired_cuda) -> None:
     """
     Update the cuda wheel libraries
     """
@@ -80,7 +77,6 @@ def update_wheel(wheel_path) -> None:
         "/usr/local/cuda/lib64/libnvToolsExt.so.1",
         "/usr/local/cuda/lib64/libnvJitLink.so.12",
         "/usr/local/cuda/lib64/libnvrtc.so.12",
-        "/usr/local/cuda/lib64/libnvrtc-builtins.so.12.6",
         "/usr/local/cuda/lib64/libcudnn_adv.so.9",
         "/usr/local/cuda/lib64/libcudnn_cnn.so.9",
         "/usr/local/cuda/lib64/libcudnn_graph.so.9",
@@ -100,6 +96,18 @@ def update_wheel(wheel_path) -> None:
             "/usr/local/lib/libnvpl_lapack_core.so.0",
             "/usr/local/lib/libnvpl_blas_core.so.0",
         ]
+        if "126" in desired_cuda:
+            libs_to_copy += [
+                "/usr/local/cuda/lib64/libnvrtc-builtins.so.12.6",
+                "/usr/local/cuda/lib64/libcufile.so.0",
+                "/usr/local/cuda/lib64/libcufile_rdma.so.1",
+            ]
+        elif "128" in desired_cuda:
+            libs_to_copy += [
+                "/usr/local/cuda/lib64/libnvrtc-builtins.so.12.8",
+                "/usr/local/cuda/lib64/libcufile.so.0",
+                "/usr/local/cuda/lib64/libcufile_rdma.so.1",
+            ]
     else:
         libs_to_copy += [
             "/opt/OpenBLAS/lib/libopenblas.so.0",
@@ -171,22 +179,22 @@ if __name__ == "__main__":
     args = parse_arguments()
     enable_mkldnn = args.enable_mkldnn
     enable_cuda = args.enable_cuda
-    repo = Repository("/pytorch")
-    branch = repo.head.name
-    if branch == "HEAD":
-        branch = "master"
+    branch = check_output(
+        ["git", "rev-parse", "--abbrev-ref", "HEAD"], cwd="/pytorch"
+    ).decode()
 
     print("Building PyTorch wheel")
     build_vars = "MAX_JOBS=5 CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000 "
     os.system("cd /pytorch; python setup.py clean")
 
     override_package_version = os.getenv("OVERRIDE_PACKAGE_VERSION")
+    desired_cuda = os.getenv("DESIRED_CUDA")
     if override_package_version is not None:
         version = override_package_version
         build_vars += (
             f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={version} PYTORCH_BUILD_NUMBER=1 "
         )
-    elif branch in ["nightly", "master"]:
+    elif branch in ["nightly", "main"]:
         build_date = (
             check_output(["git", "log", "--pretty=format:%cs", "-1"], cwd="/pytorch")
             .decode()
@@ -196,12 +204,11 @@ if __name__ == "__main__":
             check_output(["cat", "version.txt"], cwd="/pytorch").decode().strip()[:-2]
         )
         if enable_cuda:
-            desired_cuda = os.getenv("DESIRED_CUDA")
             build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={version}.dev{build_date}+{desired_cuda} PYTORCH_BUILD_NUMBER=1 "
         else:
             build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={version}.dev{build_date} PYTORCH_BUILD_NUMBER=1 "
     elif branch.startswith(("v1.", "v2.")):
-        build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={branch[1:branch.find('-')]} PYTORCH_BUILD_NUMBER=1 "
+        build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={branch[1 : branch.find('-')]} PYTORCH_BUILD_NUMBER=1 "
 
     if enable_mkldnn:
         build_ArmComputeLibrary()
@@ -225,6 +232,6 @@ if __name__ == "__main__":
         print("Updating Cuda Dependency")
         filename = os.listdir("/pytorch/dist/")
         wheel_path = f"/pytorch/dist/{filename[0]}"
-        update_wheel(wheel_path)
+        update_wheel(wheel_path, desired_cuda)
     pytorch_wheel_name = complete_wheel("/pytorch/")
     print(f"Build Complete. Created {pytorch_wheel_name}..")

.ci/aarch64_linux/build_aarch64_wheel.py

@@ -12,22 +12,22 @@ import os
 import subprocess
 import sys
 import time
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Optional, Union
 
 import boto3
 
 
 # AMI images for us-east-1, change the following based on your ~/.aws/config
 os_amis = {
-    "ubuntu18_04": "ami-078eece1d8119409f",  # login_name: ubuntu
     "ubuntu20_04": "ami-052eac90edaa9d08f",  # login_name: ubuntu
     "ubuntu22_04": "ami-0c6c29c5125214c77",  # login_name: ubuntu
     "redhat8": "ami-0698b90665a2ddcf1",  # login_name: ec2-user
 }
-ubuntu18_04_ami = os_amis["ubuntu18_04"]
+
+ubuntu20_04_ami = os_amis["ubuntu20_04"]
 
 
-def compute_keyfile_path(key_name: Optional[str] = None) -> Tuple[str, str]:
+def compute_keyfile_path(key_name: Optional[str] = None) -> tuple[str, str]:
     if key_name is None:
         key_name = os.getenv("AWS_KEY_NAME")
         if key_name is None:
@@ -57,7 +57,7 @@ def ec2_instances_by_id(instance_id):
 
 
 def start_instance(
-    key_name, ami=ubuntu18_04_ami, instance_type="t4g.2xlarge", ebs_size: int = 50
+    key_name, ami=ubuntu20_04_ami, instance_type="t4g.2xlarge", ebs_size: int = 50
 ):
     inst = ec2.create_instances(
         ImageId=ami,
@@ -96,7 +96,7 @@ class RemoteHost:
         self.keyfile_path = keyfile_path
         self.login_name = login_name
 
-    def _gen_ssh_prefix(self) -> List[str]:
+    def _gen_ssh_prefix(self) -> list[str]:
         return [
             "ssh",
             "-o",
@@ -108,13 +108,13 @@ class RemoteHost:
         ]
 
     @staticmethod
-    def _split_cmd(args: Union[str, List[str]]) -> List[str]:
+    def _split_cmd(args: Union[str, list[str]]) -> list[str]:
         return args.split() if isinstance(args, str) else args
 
-    def run_ssh_cmd(self, args: Union[str, List[str]]) -> None:
+    def run_ssh_cmd(self, args: Union[str, list[str]]) -> None:
         subprocess.check_call(self._gen_ssh_prefix() + self._split_cmd(args))
 
-    def check_ssh_output(self, args: Union[str, List[str]]) -> str:
+    def check_ssh_output(self, args: Union[str, list[str]]) -> str:
         return subprocess.check_output(
             self._gen_ssh_prefix() + self._split_cmd(args)
         ).decode("utf-8")
@@ -157,7 +157,7 @@ class RemoteHost:
     def using_docker(self) -> bool:
         return self.container_id is not None
 
-    def run_cmd(self, args: Union[str, List[str]]) -> None:
+    def run_cmd(self, args: Union[str, list[str]]) -> None:
         if not self.using_docker():
             return self.run_ssh_cmd(args)
         assert self.container_id is not None
@@ -178,7 +178,7 @@ class RemoteHost:
         if rc != 0:
             raise subprocess.CalledProcessError(rc, docker_cmd)
 
-    def check_output(self, args: Union[str, List[str]]) -> str:
+    def check_output(self, args: Union[str, list[str]]) -> str:
         if not self.using_docker():
             return self.check_ssh_output(args)
         assert self.container_id is not None
@@ -230,7 +230,7 @@ class RemoteHost:
             )
         self.download_file(remote_file, local_file)
 
-    def list_dir(self, path: str) -> List[str]:
+    def list_dir(self, path: str) -> list[str]:
         return self.check_output(["ls", "-1", path]).split("\n")
 
 
@@ -327,7 +327,7 @@ def build_ArmComputeLibrary(host: RemoteHost, git_clone_flags: str = "") -> None
         ]
     )
     host.run_cmd(
-        f"git clone https://github.com/ARM-software/ComputeLibrary.git -b v24.09 {git_clone_flags}"
+        f"git clone https://github.com/ARM-software/ComputeLibrary.git -b v25.02 {git_clone_flags}"
     )
 
     host.run_cmd(f"cd ComputeLibrary && scons Werror=1 -j8 {acl_build_flags}")
@@ -358,7 +358,7 @@ def checkout_repo(
     branch: str = "main",
     url: str,
     git_clone_flags: str,
-    mapping: Dict[str, Tuple[str, str]],
+    mapping: dict[str, tuple[str, str]],
 ) -> Optional[str]:
     for prefix in mapping:
         if not branch.startswith(prefix):
@@ -657,18 +657,6 @@ def configure_system(
             "sudo apt-get install -y python3-dev python3-yaml python3-setuptools python3-wheel python3-pip"
         )
     host.run_cmd("pip3 install dataclasses typing-extensions")
-    # Install and switch to gcc-8 on Ubuntu-18.04
-    if not host.using_docker() and host.ami == ubuntu18_04_ami and compiler == "gcc-8":
-        host.run_cmd("sudo apt-get install -y g++-8 gfortran-8")
-        host.run_cmd(
-            "sudo update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-8 100"
-        )
-        host.run_cmd(
-            "sudo update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-8 100"
-        )
-        host.run_cmd(
-            "sudo update-alternatives --install /usr/bin/gfortran gfortran /usr/bin/gfortran-8 100"
-        )
     if not use_conda:
         print("Installing Cython + numpy from PyPy")
         host.run_cmd("sudo pip3 install Cython")
@@ -681,7 +669,7 @@ def build_domains(
     branch: str = "main",
     use_conda: bool = True,
     git_clone_flags: str = "",
-) -> Tuple[str, str, str, str]:
+) -> tuple[str, str, str, str]:
     vision_wheel_name = build_torchvision(
         host, branch=branch, use_conda=use_conda, git_clone_flags=git_clone_flags
     )
@@ -708,7 +696,7 @@ def start_build(
     pytorch_build_number: Optional[str] = None,
     shallow_clone: bool = True,
     enable_mkldnn: bool = False,
-) -> Tuple[str, str, str, str, str]:
+) -> tuple[str, str, str, str, str]:
     git_clone_flags = " --depth 1 --shallow-submodules" if shallow_clone else ""
     if host.using_docker() and not use_conda:
         print("Auto-selecting conda option for docker images")
@@ -759,7 +747,7 @@ def start_build(
         version = host.check_output("cat pytorch/version.txt").strip()[:-2]
         build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={version}.dev{build_date} PYTORCH_BUILD_NUMBER=1"
     if branch.startswith(("v1.", "v2.")):
-        build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={branch[1:branch.find('-')]} PYTORCH_BUILD_NUMBER=1"
+        build_vars += f"BUILD_TEST=0 PYTORCH_BUILD_VERSION={branch[1 : branch.find('-')]} PYTORCH_BUILD_NUMBER=1"
     if host.using_docker():
         build_vars += " CMAKE_SHARED_LINKER_FLAGS=-Wl,-z,max-page-size=0x10000"
     if enable_mkldnn:
@@ -932,9 +920,9 @@ def parse_arguments():
     parser.add_argument("--debug", action="store_true")
     parser.add_argument("--build-only", action="store_true")
     parser.add_argument("--test-only", type=str)
-    parser.add_argument(
-        "--os", type=str, choices=list(os_amis.keys()), default="ubuntu20_04"
-    )
+    group = parser.add_mutually_exclusive_group()
+    group.add_argument("--os", type=str, choices=list(os_amis.keys()))
+    group.add_argument("--ami", type=str)
     parser.add_argument(
         "--python-version",
         type=str,
@@ -964,7 +952,13 @@ def parse_arguments():
 
 if __name__ == "__main__":
     args = parse_arguments()
-    ami = os_amis[args.os]
+    ami = (
+        args.ami
+        if args.ami is not None
+        else os_amis[args.os]
+        if args.os is not None
+        else ubuntu20_04_ami
+    )
     keyfile_path, key_name = compute_keyfile_path(args.key_name)
 
     if args.list_instances:
@@ -1018,7 +1012,7 @@ if __name__ == "__main__":
         install_condaforge_python(host, args.python_version)
         sys.exit(0)
 
-    python_version = args.python_version if args.python_version is not None else "3.8"
+    python_version = args.python_version if args.python_version is not None else "3.9"
 
     if args.use_torch_from_pypi:
         configure_system(host, compiler=args.compiler, python_version=python_version)

.ci/docker/almalinux/Dockerfile

@@ -44,6 +44,8 @@ FROM base as cuda
 ARG CUDA_VERSION=12.4
 RUN rm -rf /usr/local/cuda-*
 ADD ./common/install_cuda.sh install_cuda.sh
+COPY ./common/install_nccl.sh install_nccl.sh
+COPY ./ci_commit_pins/nccl-cu* /ci_commit_pins/
 ENV CUDA_HOME=/usr/local/cuda-${CUDA_VERSION}
 # Preserve CUDA_VERSION for the builds
 ENV CUDA_VERSION=${CUDA_VERSION}

.ci/docker/build.sh

@@ -1,4 +1,8 @@
 #!/bin/bash
+# The purpose of this script is to:
+# 1. Extract the set of parameters to be used for a docker build based on the provided image name.
+# 2. Run docker build with the parameters found in step 1.
+# 3. Run the built image and print out the expected and actual versions of packages installed.
 
 set -ex
 
@@ -86,32 +90,21 @@ CMAKE_VERSION=3.18.5
 
 _UCX_COMMIT=7bb2722ff2187a0cad557ae4a6afa090569f83fb
 _UCC_COMMIT=20eae37090a4ce1b32bcce6144ccad0b49943e0b
+if [[ "$image" == *rocm* ]]; then
+  _UCX_COMMIT=cc312eaa4655c0cc5c2bcd796db938f90563bcf6
+  _UCC_COMMIT=0c0fc21559835044ab107199e334f7157d6a0d3d
+fi
 
 # It's annoying to rename jobs every time you want to rewrite a
 # configuration, so we hardcode everything here rather than do it
 # from scratch
 case "$image" in
-  pytorch-linux-focal-cuda12.4-cudnn9-py3-gcc9)
-    CUDA_VERSION=12.4.1
+  pytorch-linux-focal-cuda12.6-cudnn9-py3-gcc11)
+    CUDA_VERSION=12.6.3
     CUDNN_VERSION=9
     ANACONDA_PYTHON_VERSION=3.10
-    GCC_VERSION=9
+    GCC_VERSION=11
     PROTOBUF=yes
-    DB=yes
-    VISION=yes
-    KATEX=yes
-    UCX_COMMIT=${_UCX_COMMIT}
-    UCC_COMMIT=${_UCC_COMMIT}
-    CONDA_CMAKE=yes
-    TRITON=yes
-    ;;
-  pytorch-linux-focal-cuda12.1-cudnn9-py3-gcc9)
-    CUDA_VERSION=12.1.1
-    CUDNN_VERSION=9
-    ANACONDA_PYTHON_VERSION=3.10
-    GCC_VERSION=9
-    PROTOBUF=yes
-    DB=yes
     VISION=yes
     KATEX=yes
     UCX_COMMIT=${_UCX_COMMIT}
@@ -125,37 +118,6 @@ case "$image" in
     ANACONDA_PYTHON_VERSION=3.10
     GCC_VERSION=9
     PROTOBUF=yes
-    DB=yes
-    VISION=yes
-    KATEX=yes
-    UCX_COMMIT=${_UCX_COMMIT}
-    UCC_COMMIT=${_UCC_COMMIT}
-    CONDA_CMAKE=yes
-    TRITON=yes
-    INDUCTOR_BENCHMARKS=yes
-    ;;
-  pytorch-linux-focal-cuda12.1-cudnn9-py3-gcc9-inductor-benchmarks)
-    CUDA_VERSION=12.1.1
-    CUDNN_VERSION=9
-    ANACONDA_PYTHON_VERSION=3.10
-    GCC_VERSION=9
-    PROTOBUF=yes
-    DB=yes
-    VISION=yes
-    KATEX=yes
-    UCX_COMMIT=${_UCX_COMMIT}
-    UCC_COMMIT=${_UCC_COMMIT}
-    CONDA_CMAKE=yes
-    TRITON=yes
-    INDUCTOR_BENCHMARKS=yes
-    ;;
-  pytorch-linux-focal-cuda12.1-cudnn9-py3.12-gcc9-inductor-benchmarks)
-    CUDA_VERSION=12.1.1
-    CUDNN_VERSION=9
-    ANACONDA_PYTHON_VERSION=3.12
-    GCC_VERSION=9
-    PROTOBUF=yes
-    DB=yes
     VISION=yes
     KATEX=yes
     UCX_COMMIT=${_UCX_COMMIT}
@@ -170,7 +132,6 @@ case "$image" in
     ANACONDA_PYTHON_VERSION=3.12
     GCC_VERSION=9
     PROTOBUF=yes
-    DB=yes
     VISION=yes
     KATEX=yes
     UCX_COMMIT=${_UCX_COMMIT}
@@ -185,7 +146,61 @@ case "$image" in
     ANACONDA_PYTHON_VERSION=3.13
     GCC_VERSION=9
     PROTOBUF=yes
-    DB=yes
+    VISION=yes
+    KATEX=yes
+    UCX_COMMIT=${_UCX_COMMIT}
+    UCC_COMMIT=${_UCC_COMMIT}
+    CONDA_CMAKE=yes
+    TRITON=yes
+    INDUCTOR_BENCHMARKS=yes
+    ;;
+  pytorch-linux-focal-cuda12.6-cudnn9-py3-gcc9)
+    CUDA_VERSION=12.6.3
+    CUDNN_VERSION=9
+    ANACONDA_PYTHON_VERSION=3.10
+    GCC_VERSION=9
+    PROTOBUF=yes
+    VISION=yes
+    KATEX=yes
+    UCX_COMMIT=${_UCX_COMMIT}
+    UCC_COMMIT=${_UCC_COMMIT}
+    CONDA_CMAKE=yes
+    TRITON=yes
+    ;;
+  pytorch-linux-focal-cuda12.6-cudnn9-py3-gcc9-inductor-benchmarks)
+    CUDA_VERSION=12.6.3
+    CUDNN_VERSION=9
+    ANACONDA_PYTHON_VERSION=3.10
+    GCC_VERSION=9
+    PROTOBUF=yes
+    VISION=yes
+    KATEX=yes
+    UCX_COMMIT=${_UCX_COMMIT}
+    UCC_COMMIT=${_UCC_COMMIT}
+    CONDA_CMAKE=yes
+    TRITON=yes
+    INDUCTOR_BENCHMARKS=yes
+    ;;
+  pytorch-linux-focal-cuda12.6-cudnn9-py3.12-gcc9-inductor-benchmarks)
+    CUDA_VERSION=12.6.3
+    CUDNN_VERSION=9
+    ANACONDA_PYTHON_VERSION=3.12
+    GCC_VERSION=9
+    PROTOBUF=yes
+    VISION=yes
+    KATEX=yes
+    UCX_COMMIT=${_UCX_COMMIT}
+    UCC_COMMIT=${_UCC_COMMIT}
+    CONDA_CMAKE=yes
+    TRITON=yes
+    INDUCTOR_BENCHMARKS=yes
+    ;;
+  pytorch-linux-focal-cuda12.6-cudnn9-py3.13-gcc9-inductor-benchmarks)
+    CUDA_VERSION=12.6.3
+    CUDNN_VERSION=9
+    ANACONDA_PYTHON_VERSION=3.13
+    GCC_VERSION=9
+    PROTOBUF=yes
     VISION=yes
     KATEX=yes
     UCX_COMMIT=${_UCX_COMMIT}
@@ -200,21 +215,6 @@ case "$image" in
     ANACONDA_PYTHON_VERSION=3.10
     GCC_VERSION=9
     PROTOBUF=yes
-    DB=yes
-    VISION=yes
-    KATEX=yes
-    UCX_COMMIT=${_UCX_COMMIT}
-    UCC_COMMIT=${_UCC_COMMIT}
-    CONDA_CMAKE=yes
-    TRITON=yes
-    ;;
-  pytorch-linux-focal-cuda12.1-cudnn9-py3-gcc9)
-    CUDA_VERSION=12.1.1
-    CUDNN_VERSION=9
-    ANACONDA_PYTHON_VERSION=3.10
-    GCC_VERSION=9
-    PROTOBUF=yes
-    DB=yes
     VISION=yes
     KATEX=yes
     UCX_COMMIT=${_UCX_COMMIT}
@@ -226,7 +226,6 @@ case "$image" in
     ANACONDA_PYTHON_VERSION=3.9
     CLANG_VERSION=10
     PROTOBUF=yes
-    DB=yes
     VISION=yes
     CONDA_CMAKE=yes
     ONNX=yes
@@ -235,10 +234,7 @@ case "$image" in
     ANACONDA_PYTHON_VERSION=3.9
     CLANG_VERSION=10
     PROTOBUF=yes
-    DB=yes
     VISION=yes
-    VULKAN_SDK_VERSION=1.2.162.1
-    SWIFTSHADER=yes
     CONDA_CMAKE=yes
     TRITON=yes
     ;;
@@ -246,10 +242,7 @@ case "$image" in
     ANACONDA_PYTHON_VERSION=3.11
     CLANG_VERSION=10
     PROTOBUF=yes
-    DB=yes
     VISION=yes
-    VULKAN_SDK_VERSION=1.2.162.1
-    SWIFTSHADER=yes
     CONDA_CMAKE=yes
     TRITON=yes
     ;;
@@ -257,38 +250,42 @@ case "$image" in
     ANACONDA_PYTHON_VERSION=3.9
     GCC_VERSION=9
     PROTOBUF=yes
-    DB=yes
     VISION=yes
     CONDA_CMAKE=yes
     TRITON=yes
     ;;
   pytorch-linux-focal-rocm-n-1-py3)
     ANACONDA_PYTHON_VERSION=3.10
-    GCC_VERSION=9
+    GCC_VERSION=11
     PROTOBUF=yes
-    DB=yes
     VISION=yes
     ROCM_VERSION=6.2.4
     NINJA_VERSION=1.9.0
     CONDA_CMAKE=yes
     TRITON=yes
+    KATEX=yes
+    UCX_COMMIT=${_UCX_COMMIT}
+    UCC_COMMIT=${_UCC_COMMIT}
+    INDUCTOR_BENCHMARKS=yes
     ;;
   pytorch-linux-focal-rocm-n-py3)
     ANACONDA_PYTHON_VERSION=3.10
-    GCC_VERSION=9
+    GCC_VERSION=11
     PROTOBUF=yes
-    DB=yes
     VISION=yes
     ROCM_VERSION=6.3
     NINJA_VERSION=1.9.0
     CONDA_CMAKE=yes
     TRITON=yes
+    KATEX=yes
+    UCX_COMMIT=${_UCX_COMMIT}
+    UCC_COMMIT=${_UCC_COMMIT}
+    INDUCTOR_BENCHMARKS=yes
     ;;
   pytorch-linux-jammy-xpu-2024.0-py3)
     ANACONDA_PYTHON_VERSION=3.9
     GCC_VERSION=11
     PROTOBUF=yes
-    DB=yes
     VISION=yes
     XPU_VERSION=0.5
     NINJA_VERSION=1.9.0
@@ -299,7 +296,6 @@ case "$image" in
     ANACONDA_PYTHON_VERSION=3.9
     GCC_VERSION=11
     PROTOBUF=yes
-    DB=yes
     VISION=yes
     XPU_VERSION=2025.0
     NINJA_VERSION=1.9.0
@@ -310,7 +306,6 @@ case "$image" in
     ANACONDA_PYTHON_VERSION=3.9
     GCC_VERSION=11
     PROTOBUF=yes
-    DB=yes
     VISION=yes
     KATEX=yes
     CONDA_CMAKE=yes
@@ -324,7 +319,6 @@ case "$image" in
     CUDNN_VERSION=9
     CLANG_VERSION=12
     PROTOBUF=yes
-    DB=yes
     VISION=yes
     TRITON=yes
     ;;
@@ -332,7 +326,6 @@ case "$image" in
     ANACONDA_PYTHON_VERSION=3.9
     CLANG_VERSION=12
     PROTOBUF=yes
-    DB=yes
     VISION=yes
     CONDA_CMAKE=yes
     TRITON=yes
@@ -353,7 +346,6 @@ case "$image" in
     ANACONDA_PYTHON_VERSION=3.9
     GCC_VERSION=11
     PROTOBUF=yes
-    DB=yes
     VISION=yes
     KATEX=yes
     CONDA_CMAKE=yes
@@ -368,7 +360,7 @@ case "$image" in
     EXECUTORCH=yes
     ;;
   pytorch-linux-jammy-py3.12-halide)
-    CUDA_VERSION=12.4
+    CUDA_VERSION=12.6
     ANACONDA_PYTHON_VERSION=3.12
     GCC_VERSION=11
     CONDA_CMAKE=yes
@@ -376,7 +368,7 @@ case "$image" in
     TRITON=yes
     ;;
   pytorch-linux-jammy-py3.12-triton-cpu)
-    CUDA_VERSION=12.4
+    CUDA_VERSION=12.6
     ANACONDA_PYTHON_VERSION=3.12
     GCC_VERSION=11
     CONDA_CMAKE=yes
@@ -386,20 +378,19 @@ case "$image" in
     # TODO: Use 3.9 here because of this issue https://github.com/python/mypy/issues/13627.
     # We will need to update mypy version eventually, but that's for another day. The task
     # would be to upgrade mypy to 1.0.0 with Python 3.11
-    ANACONDA_PYTHON_VERSION=3.9
-    CONDA_CMAKE=yes
+    PYTHON_VERSION=3.9
+    PIP_CMAKE=yes
     ;;
   pytorch-linux-jammy-cuda11.8-cudnn9-py3.9-linter)
-    ANACONDA_PYTHON_VERSION=3.9
+    PYTHON_VERSION=3.9
     CUDA_VERSION=11.8
-    CONDA_CMAKE=yes
+    PIP_CMAKE=yes
     ;;
   pytorch-linux-jammy-aarch64-py3.10-gcc11)
     ANACONDA_PYTHON_VERSION=3.10
     GCC_VERSION=11
     ACL=yes
     PROTOBUF=yes
-    DB=yes
     VISION=yes
     CONDA_CMAKE=yes
     # snadampal: skipping llvm src build install because the current version
@@ -411,7 +402,6 @@ case "$image" in
     GCC_VERSION=11
     ACL=yes
     PROTOBUF=yes
-    DB=yes
     VISION=yes
     CONDA_CMAKE=yes
     # snadampal: skipping llvm src build install because the current version
@@ -422,7 +412,6 @@ case "$image" in
   *)
     # Catch-all for builds that are not hardcoded.
     PROTOBUF=yes
-    DB=yes
     VISION=yes
     echo "image '$image' did not match an existing build configuration"
     if [[ "$image" == *py* ]]; then
@@ -471,14 +460,21 @@ if [[ "$image" == *cuda*  && ${OS} == "ubuntu" ]]; then
   fi
 fi
 
+no_cache_flag=""
+progress_flag=""
+# Do not use cache and progress=plain when in CI
+if [[ -n "${CI:-}" ]]; then
+  no_cache_flag="--no-cache"
+  progress_flag="--progress=plain"
+fi
+
 # Build image
 docker build \
-       --no-cache \
-       --progress=plain \
+       ${no_cache_flag} \
+       ${progress_flag} \
        --build-arg "BUILD_ENVIRONMENT=${image}" \
        --build-arg "PROTOBUF=${PROTOBUF:-}" \
        --build-arg "LLVMDEV=${LLVMDEV:-}" \
-       --build-arg "DB=${DB:-}" \
        --build-arg "VISION=${VISION:-}" \
        --build-arg "UBUNTU_VERSION=${UBUNTU_VERSION}" \
        --build-arg "CENTOS_VERSION=${CENTOS_VERSION}" \
@@ -486,13 +482,12 @@ docker build \
        --build-arg "GLIBC_VERSION=${GLIBC_VERSION}" \
        --build-arg "CLANG_VERSION=${CLANG_VERSION}" \
        --build-arg "ANACONDA_PYTHON_VERSION=${ANACONDA_PYTHON_VERSION}" \
+       --build-arg "PYTHON_VERSION=${PYTHON_VERSION}" \
        --build-arg "GCC_VERSION=${GCC_VERSION}" \
        --build-arg "CUDA_VERSION=${CUDA_VERSION}" \
        --build-arg "CUDNN_VERSION=${CUDNN_VERSION}" \
        --build-arg "TENSORRT_VERSION=${TENSORRT_VERSION}" \
        --build-arg "GRADLE_VERSION=${GRADLE_VERSION}" \
-       --build-arg "VULKAN_SDK_VERSION=${VULKAN_SDK_VERSION}" \
-       --build-arg "SWIFTSHADER=${SWIFTSHADER}" \
        --build-arg "CMAKE_VERSION=${CMAKE_VERSION:-}" \
        --build-arg "NINJA_VERSION=${NINJA_VERSION:-}" \
        --build-arg "KATEX=${KATEX:-}" \
@@ -502,6 +497,7 @@ docker build \
        --build-arg "UCX_COMMIT=${UCX_COMMIT}" \
        --build-arg "UCC_COMMIT=${UCC_COMMIT}" \
        --build-arg "CONDA_CMAKE=${CONDA_CMAKE}" \
+       --build-arg "PIP_CMAKE=${PIP_CMAKE}" \
        --build-arg "TRITON=${TRITON}" \
        --build-arg "TRITON_CPU=${TRITON_CPU}" \
        --build-arg "ONNX=${ONNX}" \
@@ -527,7 +523,7 @@ docker build \
 UBUNTU_VERSION=$(echo ${UBUNTU_VERSION} | sed 's/-rc$//')
 
 function drun() {
-  docker run --rm "$tmp_tag" $*
+  docker run --rm "$tmp_tag" "$@"
 }
 
 if [[ "$OS" == "ubuntu" ]]; then
@@ -575,3 +571,14 @@ if [ -n "$KATEX" ]; then
     exit 1
   fi
 fi
+
+HAS_TRITON=$(drun python -c "import triton" > /dev/null 2>&1 && echo "yes" || echo "no")
+if [[ -n "$TRITON" || -n "$TRITON_CPU" ]]; then
+  if [ "$HAS_TRITON" = "no" ]; then
+    echo "expecting triton to be installed, but it is not"
+    exit 1
+  fi
+elif [ "$HAS_TRITON" = "yes" ]; then
+  echo "expecting triton to not be installed, but it is"
+  exit 1
+fi

.ci/docker/centos-rocm/Dockerfile

@@ -55,13 +55,6 @@ RUN if [ -n "${PROTOBUF}" ]; then bash ./install_protobuf.sh; fi
 RUN rm install_protobuf.sh
 ENV INSTALLED_PROTOBUF ${PROTOBUF}
 
-# (optional) Install database packages like LMDB and LevelDB
-ARG DB
-COPY ./common/install_db.sh install_db.sh
-RUN if [ -n "${DB}" ]; then bash ./install_db.sh; fi
-RUN rm install_db.sh
-ENV INSTALLED_DB ${DB}
-
 # (optional) Install vision packages like OpenCV
 ARG VISION
 COPY ./common/install_vision.sh ./common/cache_vision_models.sh ./common/common_utils.sh ./
@@ -75,7 +68,7 @@ COPY ./common/install_rocm.sh install_rocm.sh
 RUN bash ./install_rocm.sh
 RUN rm install_rocm.sh
 COPY ./common/install_rocm_magma.sh install_rocm_magma.sh
-RUN bash ./install_rocm_magma.sh
+RUN bash ./install_rocm_magma.sh ${ROCM_VERSION}
 RUN rm install_rocm_magma.sh
 COPY ./common/install_amdsmi.sh install_amdsmi.sh
 RUN bash ./install_amdsmi.sh

.ci/docker/ci_commit_pins/executorch.txt

@@ -1 +1 @@
-a29b208a06ab378bb29ab1aa68932e412f8e09f1
+7e487c24e1c20c3f4606c2d8aca2778873b00b4c

.ci/docker/ci_commit_pins/nccl-cu11.txt

@@ -0,0 +1 @@
+v2.21.5-1

.ci/docker/ci_commit_pins/nccl-cu12.txt

@@ -0,0 +1 @@
+v2.26.2-1

.ci/docker/ci_commit_pins/timm.txt

@@ -1 +1 @@
-ac3470188b914c5d7a5058a7e28b9eb685a62427
+5d535d7a2d4b435b1b5c1177fd8f04a12b942b9a

.ci/docker/ci_commit_pins/triton-xpu.txt

@@ -1 +1 @@
-e98b6fcb8df5b44eb0d0addb6767c573d37ba024
+0bcc8265e677e5321606a3311bf71470f14456a8

.ci/docker/ci_commit_pins/triton.txt

@@ -1 +1 @@
-0d4682f073ded4d1a8260dd4208a43d735ae3a2b
+96316ce50fade7e209553aba4898cd9b82aab83b

.ci/docker/common/install_acl.sh

@@ -1,6 +1,6 @@
 set -euo pipefail
 
-readonly version=v24.04
+readonly version=v25.02
 readonly src_host=https://github.com/ARM-software
 readonly src_repo=ComputeLibrary
 

.ci/docker/common/install_base.sh

@@ -32,8 +32,12 @@ install_ubuntu() {
 
   # HACK: UCC testing relies on libnccl library from NVIDIA repo, and version 2.16 crashes
   # See https://github.com/pytorch/pytorch/pull/105260#issuecomment-1673399729
+  # TODO: Eliminate this hack, we should not relay on apt-get installation
+  # See https://github.com/pytorch/pytorch/issues/144768
   if [[ "$UBUNTU_VERSION" == "20.04"* && "$CUDA_VERSION" == "11.8"* ]]; then
     maybe_libnccl_dev="libnccl2=2.15.5-1+cuda11.8 libnccl-dev=2.15.5-1+cuda11.8 --allow-downgrades --allow-change-held-packages"
+  elif [[ "$UBUNTU_VERSION" == "20.04"* && "$CUDA_VERSION" == "12.4"* ]]; then
+    maybe_libnccl_dev="libnccl2=2.26.2-1+cuda12.4 libnccl-dev=2.26.2-1+cuda12.4 --allow-downgrades --allow-change-held-packages"
   else
     maybe_libnccl_dev=""
   fi

.ci/docker/common/install_cache.sh

@@ -9,7 +9,7 @@ install_ubuntu() {
   # Instead use lib and headers from OpenSSL1.1 installed in `install_openssl.sh``
   apt-get install -y cargo
   echo "Checking out sccache repo"
-  git clone https://github.com/mozilla/sccache -b v0.9.0
+  git clone https://github.com/mozilla/sccache -b v0.9.1
   cd sccache
   echo "Building sccache"
   cargo build --release

.ci/docker/common/install_clang.sh

@@ -4,16 +4,10 @@ set -ex
 
 if [ -n "$CLANG_VERSION" ]; then
 
-  if [[ $CLANG_VERSION == 9 && $UBUNTU_VERSION == 18.04 ]]; then
-    sudo apt-get update
-    # gpg-agent is not available by default on 18.04
-    sudo apt-get install  -y --no-install-recommends gpg-agent
-    wget --no-check-certificate -O - https://apt.llvm.org/llvm-snapshot.gpg.key | sudo apt-key add  -
-    apt-add-repository "deb http://apt.llvm.org/bionic/ llvm-toolchain-bionic-${CLANG_VERSION} main"
-  elif [[ $UBUNTU_VERSION == 22.04 ]]; then
+  if [[ $UBUNTU_VERSION == 22.04 ]]; then
     # work around ubuntu apt-get conflicts
     sudo apt-get -y -f install
-    wget --no-check-certificate -O - https://apt.llvm.org/llvm-snapshot.gpg.key | sudo apt-key add  -
+    wget --no-check-certificate -O - https://apt.llvm.org/llvm-snapshot.gpg.key | sudo apt-key add -
     if [[ $CLANG_VERSION == 18 ]]; then
       apt-add-repository "deb http://apt.llvm.org/jammy/ llvm-toolchain-jammy-18 main"
     fi
@@ -41,7 +35,7 @@ if [ -n "$CLANG_VERSION" ]; then
   # clang's packaging is a little messed up (the runtime libs aren't
   # added into the linker path), so give it a little help
   clang_lib=("/usr/lib/llvm-$CLANG_VERSION/lib/clang/"*"/lib/linux")
-  echo "$clang_lib" > /etc/ld.so.conf.d/clang.conf
+  echo "$clang_lib" >/etc/ld.so.conf.d/clang.conf
   ldconfig
 
   # Cleanup package manager

.ci/docker/common/install_conda.sh

@@ -62,11 +62,11 @@ if [ -n "$ANACONDA_PYTHON_VERSION" ]; then
 
   # libstdcxx from conda default channels are too old, we need GLIBCXX_3.4.30
   # which is provided in libstdcxx 12 and up.
-  conda_install libstdcxx-ng=12.3.0 -c conda-forge
+  conda_install libstdcxx-ng=12.3.0 --update-deps -c conda-forge
 
   # Install PyTorch conda deps, as per https://github.com/pytorch/pytorch README
   if [[ $(uname -m) == "aarch64" ]]; then
-    conda_install "openblas==0.3.28=*openmp*"
+    conda_install "openblas==0.3.29=*openmp*"
   else
     conda_install "mkl=2021.4.0 mkl-include=2021.4.0"
   fi

.ci/docker/common/install_cpython.sh

@@ -7,7 +7,7 @@ PYTHON_DOWNLOAD_GITHUB_BRANCH=https://github.com/python/cpython/archive/refs/hea
 GET_PIP_URL=https://bootstrap.pypa.io/get-pip.py
 
 # Python versions to be installed in /opt/$VERSION_NO
-CPYTHON_VERSIONS=${CPYTHON_VERSIONS:-"3.8.1 3.9.0 3.10.1 3.11.0 3.12.0 3.13.0 3.13.0t"}
+CPYTHON_VERSIONS=${CPYTHON_VERSIONS:-"3.9.0 3.10.1 3.11.0 3.12.0 3.13.0 3.13.0t"}
 
 function check_var {
     if [ -z "$1" ]; then

.ci/docker/common/install_cuda.sh

@@ -2,7 +2,6 @@
 
 set -ex
 
-NCCL_VERSION=v2.21.5-1
 CUDNN_VERSION=9.5.1.17
 
 function install_cusparselt_040 {
@@ -16,17 +15,6 @@ function install_cusparselt_040 {
     rm -rf tmp_cusparselt
 }
 
-function install_cusparselt_052 {
-    # cuSparseLt license: https://docs.nvidia.com/cuda/cusparselt/license.html
-    mkdir tmp_cusparselt && pushd tmp_cusparselt
-    wget -q https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-x86_64/libcusparse_lt-linux-x86_64-0.5.2.1-archive.tar.xz
-    tar xf libcusparse_lt-linux-x86_64-0.5.2.1-archive.tar.xz
-    cp -a libcusparse_lt-linux-x86_64-0.5.2.1-archive/include/* /usr/local/cuda/include/
-    cp -a libcusparse_lt-linux-x86_64-0.5.2.1-archive/lib/* /usr/local/cuda/lib64/
-    popd
-    rm -rf tmp_cusparselt
-}
-
 function install_cusparselt_062 {
     # cuSparseLt license: https://docs.nvidia.com/cuda/cusparselt/license.html
     mkdir tmp_cusparselt && pushd tmp_cusparselt
@@ -51,7 +39,7 @@ function install_cusparselt_063 {
 
 function install_118 {
     CUDNN_VERSION=9.1.0.70
-    echo "Installing CUDA 11.8 and cuDNN ${CUDNN_VERSION} and NCCL ${NCCL_VERSION} and cuSparseLt-0.4.0"
+    echo "Installing CUDA 11.8 and cuDNN ${CUDNN_VERSION} and NCCL and cuSparseLt-0.4.0"
     rm -rf /usr/local/cuda-11.8 /usr/local/cuda
     # install CUDA 11.8.0 in the same container
     wget -q https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run
@@ -69,56 +57,16 @@ function install_118 {
     cd ..
     rm -rf tmp_cudnn
 
-    # NCCL license: https://docs.nvidia.com/deeplearning/nccl/#licenses
-    # Follow build: https://github.com/NVIDIA/nccl/tree/master?tab=readme-ov-file#build
-    git clone -b $NCCL_VERSION --depth 1 https://github.com/NVIDIA/nccl.git
-    cd nccl && make -j src.build
-    cp -a build/include/* /usr/local/cuda/include/
-    cp -a build/lib/* /usr/local/cuda/lib64/
-    cd ..
-    rm -rf nccl
+    CUDA_VERSION=11.8 bash install_nccl.sh
 
     install_cusparselt_040
 
     ldconfig
 }
 
-function install_121 {
-    echo "Installing CUDA 12.1 and cuDNN ${CUDNN_VERSION} and NCCL ${NCCL_VERSION} and cuSparseLt-0.5.2"
-    rm -rf /usr/local/cuda-12.1 /usr/local/cuda
-    # install CUDA 12.1.0 in the same container
-    wget -q https://developer.download.nvidia.com/compute/cuda/12.1.1/local_installers/cuda_12.1.1_530.30.02_linux.run
-    chmod +x cuda_12.1.1_530.30.02_linux.run
-    ./cuda_12.1.1_530.30.02_linux.run --toolkit --silent
-    rm -f cuda_12.1.1_530.30.02_linux.run
-    rm -f /usr/local/cuda && ln -s /usr/local/cuda-12.1 /usr/local/cuda
-
-    # cuDNN license: https://developer.nvidia.com/cudnn/license_agreement
-    mkdir tmp_cudnn && cd tmp_cudnn
-    wget -q https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive.tar.xz -O cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive.tar.xz
-    tar xf cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive.tar.xz
-    cp -a cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive/include/* /usr/local/cuda/include/
-    cp -a cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive/lib/* /usr/local/cuda/lib64/
-    cd ..
-    rm -rf tmp_cudnn
-
-    # NCCL license: https://docs.nvidia.com/deeplearning/nccl/#licenses
-    # Follow build: https://github.com/NVIDIA/nccl/tree/master?tab=readme-ov-file#build
-    git clone -b $NCCL_VERSION --depth 1 https://github.com/NVIDIA/nccl.git
-    cd nccl && make -j src.build
-    cp -a build/include/* /usr/local/cuda/include/
-    cp -a build/lib/* /usr/local/cuda/lib64/
-    cd ..
-    rm -rf nccl
-
-    install_cusparselt_052
-
-    ldconfig
-}
-
 function install_124 {
   CUDNN_VERSION=9.1.0.70
-  echo "Installing CUDA 12.4.1 and cuDNN ${CUDNN_VERSION} and NCCL ${NCCL_VERSION} and cuSparseLt-0.6.2"
+  echo "Installing CUDA 12.4.1 and cuDNN ${CUDNN_VERSION} and NCCL and cuSparseLt-0.6.2"
   rm -rf /usr/local/cuda-12.4 /usr/local/cuda
   # install CUDA 12.4.1 in the same container
   wget -q https://developer.download.nvidia.com/compute/cuda/12.4.1/local_installers/cuda_12.4.1_550.54.15_linux.run
@@ -136,14 +84,7 @@ function install_124 {
   cd ..
   rm -rf tmp_cudnn
 
-  # NCCL license: https://docs.nvidia.com/deeplearning/nccl/#licenses
-  # Follow build: https://github.com/NVIDIA/nccl/tree/master?tab=readme-ov-file#build
-  git clone -b $NCCL_VERSION --depth 1 https://github.com/NVIDIA/nccl.git
-  cd nccl && make -j src.build
-  cp -a build/include/* /usr/local/cuda/include/
-  cp -a build/lib/* /usr/local/cuda/lib64/
-  cd ..
-  rm -rf nccl
+  CUDA_VERSION=12.4 bash install_nccl.sh
 
   install_cusparselt_062
 
@@ -151,7 +92,7 @@ function install_124 {
 }
 
 function install_126 {
-  echo "Installing CUDA 12.6.3 and cuDNN ${CUDNN_VERSION} and NCCL ${NCCL_VERSION} and cuSparseLt-0.6.3"
+  echo "Installing CUDA 12.6.3 and cuDNN ${CUDNN_VERSION} and NCCL and cuSparseLt-0.6.3"
   rm -rf /usr/local/cuda-12.6 /usr/local/cuda
   # install CUDA 12.6.3 in the same container
   wget -q https://developer.download.nvidia.com/compute/cuda/12.6.3/local_installers/cuda_12.6.3_560.35.05_linux.run
@@ -169,14 +110,7 @@ function install_126 {
   cd ..
   rm -rf tmp_cudnn
 
-  # NCCL license: https://docs.nvidia.com/deeplearning/nccl/#licenses
-  # Follow build: https://github.com/NVIDIA/nccl/tree/master?tab=readme-ov-file#build
-  git clone -b $NCCL_VERSION --depth 1 https://github.com/NVIDIA/nccl.git
-  cd nccl && make -j src.build
-  cp -a build/include/* /usr/local/cuda/include/
-  cp -a build/lib/* /usr/local/cuda/lib64/
-  cd ..
-  rm -rf nccl
+  CUDA_VERSION=12.6 bash install_nccl.sh
 
   install_cusparselt_063
 
@@ -214,37 +148,6 @@ function prune_118 {
     rm -rf $CUDA_BASE/libnvvp $CUDA_BASE/nsightee_plugins $CUDA_BASE/nsight-compute-2022.3.0 $CUDA_BASE/nsight-systems-2022.4.2/
 }
 
-function prune_121 {
-  echo "Pruning CUDA 12.1"
-  #####################################################################################
-  # CUDA 12.1 prune static libs
-  #####################################################################################
-    export NVPRUNE="/usr/local/cuda-12.1/bin/nvprune"
-    export CUDA_LIB_DIR="/usr/local/cuda-12.1/lib64"
-
-    export GENCODE="-gencode arch=compute_50,code=sm_50 -gencode arch=compute_60,code=sm_60 -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90"
-    export GENCODE_CUDNN="-gencode arch=compute_50,code=sm_50 -gencode arch=compute_60,code=sm_60 -gencode arch=compute_61,code=sm_61 -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90"
-
-    if [[ -n "$OVERRIDE_GENCODE" ]]; then
-        export GENCODE=$OVERRIDE_GENCODE
-    fi
-
-    # all CUDA libs except CuDNN and CuBLAS
-    ls $CUDA_LIB_DIR/ | grep "\.a" | grep -v "culibos" | grep -v "cudart" | grep -v "cudnn" | grep -v "cublas" | grep -v "metis"  \
-      | xargs -I {} bash -c \
-                "echo {} && $NVPRUNE $GENCODE $CUDA_LIB_DIR/{} -o $CUDA_LIB_DIR/{}"
-
-    # prune CuDNN and CuBLAS
-    $NVPRUNE $GENCODE_CUDNN $CUDA_LIB_DIR/libcublas_static.a -o $CUDA_LIB_DIR/libcublas_static.a
-    $NVPRUNE $GENCODE_CUDNN $CUDA_LIB_DIR/libcublasLt_static.a -o $CUDA_LIB_DIR/libcublasLt_static.a
-
-    #####################################################################################
-    # CUDA 12.1 prune visual tools
-    #####################################################################################
-    export CUDA_BASE="/usr/local/cuda-12.1/"
-    rm -rf $CUDA_BASE/libnvvp $CUDA_BASE/nsightee_plugins $CUDA_BASE/nsight-compute-2023.1.0 $CUDA_BASE/nsight-systems-2023.1.2/
-}
-
 function prune_124 {
   echo "Pruning CUDA 12.4"
   #####################################################################################
@@ -313,18 +216,45 @@ function prune_126 {
   rm -rf $CUDA_BASE/libnvvp $CUDA_BASE/nsightee_plugins $CUDA_BASE/nsight-compute-2024.3.2 $CUDA_BASE/nsight-systems-2024.5.1/
 }
 
+function install_128 {
+  CUDNN_VERSION=9.8.0.87
+  echo "Installing CUDA 12.8.0 and cuDNN ${CUDNN_VERSION} and NCCL and cuSparseLt-0.6.3"
+  rm -rf /usr/local/cuda-12.8 /usr/local/cuda
+  # install CUDA 12.8.0 in the same container
+  wget -q https://developer.download.nvidia.com/compute/cuda/12.8.0/local_installers/cuda_12.8.0_570.86.10_linux.run
+  chmod +x cuda_12.8.0_570.86.10_linux.run
+  ./cuda_12.8.0_570.86.10_linux.run --toolkit --silent
+  rm -f cuda_12.8.0_570.86.10_linux.run
+  rm -f /usr/local/cuda && ln -s /usr/local/cuda-12.8 /usr/local/cuda
+
+  # cuDNN license: https://developer.nvidia.com/cudnn/license_agreement
+  mkdir tmp_cudnn && cd tmp_cudnn
+  wget -q https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive.tar.xz -O cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive.tar.xz
+  tar xf cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive.tar.xz
+  cp -a cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive/include/* /usr/local/cuda/include/
+  cp -a cudnn-linux-x86_64-${CUDNN_VERSION}_cuda12-archive/lib/* /usr/local/cuda/lib64/
+  cd ..
+  rm -rf tmp_cudnn
+
+  CUDA_VERSION=12.8 bash install_nccl.sh
+
+  install_cusparselt_063
+
+  ldconfig
+}
+
 # idiomatic parameter and option handling in sh
 while test $# -gt 0
 do
     case "$1" in
     11.8) install_118; prune_118
         ;;
-    12.1) install_121; prune_121
-        ;;
     12.4) install_124; prune_124
         ;;
     12.6) install_126; prune_126
         ;;
+    12.8) install_128;
+        ;;
     *) echo "bad argument $1"; exit 1
         ;;
     esac

.ci/docker/common/install_cuda_aarch64.sh

@@ -3,19 +3,7 @@
 
 set -ex
 
-NCCL_VERSION=v2.21.5-1
-CUDNN_VERSION=9.5.1.17
-
-function install_cusparselt_062 {
-    # cuSparseLt license: https://docs.nvidia.com/cuda/cusparselt/license.html
-    mkdir tmp_cusparselt && pushd tmp_cusparselt
-    wget -q https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-sbsa/libcusparse_lt-linux-sbsa-0.6.2.3-archive.tar.xz
-    tar xf libcusparse_lt-linux-sbsa-0.6.2.3-archive.tar.xz
-    cp -a libcusparse_lt-linux-sbsa-0.6.2.3-archive/include/* /usr/local/cuda/include/
-    cp -a libcusparse_lt-linux-sbsa-0.6.2.3-archive/lib/* /usr/local/cuda/lib64/
-    popd
-    rm -rf tmp_cusparselt
-}
+CUDNN_VERSION=9.8.0.87
 
 function install_cusparselt_063 {
     # cuSparseLt license: https://docs.nvidia.com/cuda/cusparselt/license.html
@@ -28,16 +16,15 @@ function install_cusparselt_063 {
     rm -rf tmp_cusparselt
 }
 
-function install_124 {
-  CUDNN_VERSION=9.1.0.70
-  echo "Installing CUDA 12.4.1 and cuDNN ${CUDNN_VERSION} and NCCL ${NCCL_VERSION} and cuSparseLt-0.6.2"
-  rm -rf /usr/local/cuda-12.4 /usr/local/cuda
-  # install CUDA 12.4.1 in the same container
-  wget -q https://developer.download.nvidia.com/compute/cuda/12.4.1/local_installers/cuda_12.4.1_550.54.15_linux_sbsa.run
-  chmod +x cuda_12.4.1_550.54.15_linux_sbsa.run
-  ./cuda_12.4.1_550.54.15_linux_sbsa.run --toolkit --silent
-  rm -f cuda_12.4.1_550.54.15_linux_sbsa.run
-  rm -f /usr/local/cuda && ln -s /usr/local/cuda-12.4 /usr/local/cuda
+function install_128 {
+  echo "Installing CUDA 12.8.0 and cuDNN ${CUDNN_VERSION} and NCCL and cuSparseLt-0.6.3"
+  rm -rf /usr/local/cuda-12.8 /usr/local/cuda
+  # install CUDA 12.8.0 in the same container
+  wget -q https://developer.download.nvidia.com/compute/cuda/12.8.0/local_installers/cuda_12.8.0_570.86.10_linux_sbsa.run
+  chmod +x cuda_12.8.0_570.86.10_linux_sbsa.run
+  ./cuda_12.8.0_570.86.10_linux_sbsa.run --toolkit --silent
+  rm -f cuda_12.8.0_570.86.10_linux_sbsa.run
+  rm -f /usr/local/cuda && ln -s /usr/local/cuda-12.8 /usr/local/cuda
 
   # cuDNN license: https://developer.nvidia.com/cudnn/license_agreement
   mkdir tmp_cudnn && cd tmp_cudnn
@@ -48,125 +35,18 @@ function install_124 {
   cd ..
   rm -rf tmp_cudnn
 
-  # NCCL license: https://docs.nvidia.com/deeplearning/nccl/#licenses
-  # Follow build: https://github.com/NVIDIA/nccl/tree/master?tab=readme-ov-file#build
-  git clone -b ${NCCL_VERSION} --depth 1 https://github.com/NVIDIA/nccl.git
-  cd nccl && make -j src.build
-  cp -a build/include/* /usr/local/cuda/include/
-  cp -a build/lib/* /usr/local/cuda/lib64/
-  cd ..
-  rm -rf nccl
-
-  install_cusparselt_062
-
-  ldconfig
-}
-
-function prune_124 {
-  echo "Pruning CUDA 12.4"
-  #####################################################################################
-  # CUDA 12.4 prune static libs
-  #####################################################################################
-  export NVPRUNE="/usr/local/cuda-12.4/bin/nvprune"
-  export CUDA_LIB_DIR="/usr/local/cuda-12.4/lib64"
-
-  export GENCODE="-gencode arch=compute_50,code=sm_50 -gencode arch=compute_60,code=sm_60 -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90"
-  export GENCODE_CUDNN="-gencode arch=compute_50,code=sm_50 -gencode arch=compute_60,code=sm_60 -gencode arch=compute_61,code=sm_61 -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90"
-
-  if [[ -n "$OVERRIDE_GENCODE" ]]; then
-      export GENCODE=$OVERRIDE_GENCODE
-  fi
-
-  # all CUDA libs except CuDNN and CuBLAS
-  ls $CUDA_LIB_DIR/ | grep "\.a" | grep -v "culibos" | grep -v "cudart" | grep -v "cudnn" | grep -v "cublas" | grep -v "metis"  \
-      | xargs -I {} bash -c \
-                "echo {} && $NVPRUNE $GENCODE $CUDA_LIB_DIR/{} -o $CUDA_LIB_DIR/{}"
-
-  # prune CuDNN and CuBLAS
-  $NVPRUNE $GENCODE_CUDNN $CUDA_LIB_DIR/libcublas_static.a -o $CUDA_LIB_DIR/libcublas_static.a
-  $NVPRUNE $GENCODE_CUDNN $CUDA_LIB_DIR/libcublasLt_static.a -o $CUDA_LIB_DIR/libcublasLt_static.a
-
-  #####################################################################################
-  # CUDA 12.4 prune visual tools
-  #####################################################################################
-  export CUDA_BASE="/usr/local/cuda-12.4/"
-  rm -rf $CUDA_BASE/libnvvp $CUDA_BASE/nsightee_plugins $CUDA_BASE/nsight-compute-2024.1.0 $CUDA_BASE/nsight-systems-2023.4.4/
-}
-
-function install_126 {
-  echo "Installing CUDA 12.6.3 and cuDNN ${CUDNN_VERSION} and NCCL ${NCCL_VERSION} and cuSparseLt-0.6.3"
-  rm -rf /usr/local/cuda-12.6 /usr/local/cuda
-  # install CUDA 12.6.3 in the same container
-  wget -q https://developer.download.nvidia.com/compute/cuda/12.6.3/local_installers/cuda_12.6.3_560.35.05_linux_sbsa.run
-  chmod +x cuda_12.6.3_560.35.05_linux_sbsa.run
-  ./cuda_12.6.3_560.35.05_linux_sbsa.run --toolkit --silent
-  rm -f cuda_12.6.3_560.35.05_linux_sbsa.run
-  rm -f /usr/local/cuda && ln -s /usr/local/cuda-12.6 /usr/local/cuda
-
-  # cuDNN license: https://developer.nvidia.com/cudnn/license_agreement
-  mkdir tmp_cudnn && cd tmp_cudnn
-  wget -q https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-sbsa/cudnn-linux-sbsa-${CUDNN_VERSION}_cuda12-archive.tar.xz -O cudnn-linux-sbsa-${CUDNN_VERSION}_cuda12-archive.tar.xz
-  tar xf cudnn-linux-sbsa-${CUDNN_VERSION}_cuda12-archive.tar.xz
-  cp -a cudnn-linux-sbsa-${CUDNN_VERSION}_cuda12-archive/include/* /usr/local/cuda/include/
-  cp -a cudnn-linux-sbsa-${CUDNN_VERSION}_cuda12-archive/lib/* /usr/local/cuda/lib64/
-  cd ..
-  rm -rf tmp_cudnn
-
-  # NCCL license: https://docs.nvidia.com/deeplearning/nccl/#licenses
-  # Follow build: https://github.com/NVIDIA/nccl/tree/master?tab=readme-ov-file#build
-  git clone -b ${NCCL_VERSION} --depth 1 https://github.com/NVIDIA/nccl.git
-  cd nccl && make -j src.build
-  cp -a build/include/* /usr/local/cuda/include/
-  cp -a build/lib/* /usr/local/cuda/lib64/
-  cd ..
-  rm -rf nccl
+  CUDA_VERSION=12.8 bash install_nccl.sh
 
   install_cusparselt_063
 
   ldconfig
 }
 
-function prune_126 {
-  echo "Pruning CUDA 12.6"
-  #####################################################################################
-  # CUDA 12.6 prune static libs
-  #####################################################################################
-  export NVPRUNE="/usr/local/cuda-12.6/bin/nvprune"
-  export CUDA_LIB_DIR="/usr/local/cuda-12.6/lib64"
-
-  export GENCODE="-gencode arch=compute_50,code=sm_50 -gencode arch=compute_60,code=sm_60 -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90"
-  export GENCODE_CUDNN="-gencode arch=compute_50,code=sm_50 -gencode arch=compute_60,code=sm_60 -gencode arch=compute_61,code=sm_61 -gencode arch=compute_70,code=sm_70 -gencode arch=compute_75,code=sm_75 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_86,code=sm_86 -gencode arch=compute_90,code=sm_90"
-
-  if [[ -n "$OVERRIDE_GENCODE" ]]; then
-      export GENCODE=$OVERRIDE_GENCODE
-  fi
-  if [[ -n "$OVERRIDE_GENCODE_CUDNN" ]]; then
-      export GENCODE_CUDNN=$OVERRIDE_GENCODE_CUDNN
-  fi
-
-  # all CUDA libs except CuDNN and CuBLAS
-  ls $CUDA_LIB_DIR/ | grep "\.a" | grep -v "culibos" | grep -v "cudart" | grep -v "cudnn" | grep -v "cublas" | grep -v "metis"  \
-      | xargs -I {} bash -c \
-                "echo {} && $NVPRUNE $GENCODE $CUDA_LIB_DIR/{} -o $CUDA_LIB_DIR/{}"
-
-  # prune CuDNN and CuBLAS
-  $NVPRUNE $GENCODE_CUDNN $CUDA_LIB_DIR/libcublas_static.a -o $CUDA_LIB_DIR/libcublas_static.a
-  $NVPRUNE $GENCODE_CUDNN $CUDA_LIB_DIR/libcublasLt_static.a -o $CUDA_LIB_DIR/libcublasLt_static.a
-
-  #####################################################################################
-  # CUDA 12.6 prune visual tools
-  #####################################################################################
-  export CUDA_BASE="/usr/local/cuda-12.6/"
-  rm -rf $CUDA_BASE/libnvvp $CUDA_BASE/nsightee_plugins $CUDA_BASE/nsight-compute-2024.3.2 $CUDA_BASE/nsight-systems-2024.5.1/
-}
-
 # idiomatic parameter and option handling in sh
 while test $# -gt 0
 do
     case "$1" in
-    12.4) install_124; prune_124
-        ;;
-    12.6) install_126; prune_126
+    12.8) install_128;
         ;;
     *) echo "bad argument $1"; exit 1
         ;;

.ci/docker/common/install_cudnn.sh

@@ -4,7 +4,9 @@ if [[ -n "${CUDNN_VERSION}" ]]; then
     # cuDNN license: https://developer.nvidia.com/cudnn/license_agreement
     mkdir tmp_cudnn
     pushd tmp_cudnn
-    if [[ ${CUDA_VERSION:0:4} == "12.6" ]]; then
+    if [[ ${CUDA_VERSION:0:4} == "12.8" ]]; then
+        CUDNN_NAME="cudnn-linux-x86_64-9.8.0.87_cuda12-archive"
+    elif [[ ${CUDA_VERSION:0:4} == "12.6" ]]; then
         CUDNN_NAME="cudnn-linux-x86_64-9.5.1.17_cuda12-archive"
     elif [[ ${CUDA_VERSION:0:2} == "12" ]]; then
         CUDNN_NAME="cudnn-linux-x86_64-9.1.0.70_cuda12-archive"

.ci/docker/common/install_cusparselt.sh

@@ -5,7 +5,15 @@ set -ex
 # cuSPARSELt license: https://docs.nvidia.com/cuda/cusparselt/license.html
 mkdir tmp_cusparselt && cd tmp_cusparselt
 
-if [[ ${CUDA_VERSION:0:4} =~ ^12\.[2-6]$ ]]; then
+if [[ ${CUDA_VERSION:0:4} =~ ^12\.[5-8]$ ]]; then
+    arch_path='sbsa'
+    export TARGETARCH=${TARGETARCH:-$(uname -m)}
+    if [ ${TARGETARCH} = 'amd64' ] || [ "${TARGETARCH}" = 'x86_64' ]; then
+        arch_path='x86_64'
+    fi
+    CUSPARSELT_NAME="libcusparse_lt-linux-${arch_path}-0.6.3.2-archive"
+    curl --retry 3 -OLs https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-${arch_path}/${CUSPARSELT_NAME}.tar.xz
+elif [[ ${CUDA_VERSION:0:4} == "12.4" ]]; then
     arch_path='sbsa'
     export TARGETARCH=${TARGETARCH:-$(uname -m)}
     if [ ${TARGETARCH} = 'amd64' ] || [ "${TARGETARCH}" = 'x86_64' ]; then
@@ -13,17 +21,11 @@ if [[ ${CUDA_VERSION:0:4} =~ ^12\.[2-6]$ ]]; then
     fi
     CUSPARSELT_NAME="libcusparse_lt-linux-${arch_path}-0.6.2.3-archive"
     curl --retry 3 -OLs https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-${arch_path}/${CUSPARSELT_NAME}.tar.xz
-elif [[ ${CUDA_VERSION:0:4} == "12.1" ]]; then
-    arch_path='sbsa'
-    export TARGETARCH=${TARGETARCH:-$(uname -m)}
-    if [ ${TARGETARCH} = 'amd64' ] || [ "${TARGETARCH}" = 'x86_64' ]; then
-        arch_path='x86_64'
-    fi
-    CUSPARSELT_NAME="libcusparse_lt-linux-${arch_path}-0.5.2.1-archive"
-    curl --retry 3 -OLs https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-${arch_path}/${CUSPARSELT_NAME}.tar.xz
 elif [[ ${CUDA_VERSION:0:4} == "11.8" ]]; then
     CUSPARSELT_NAME="libcusparse_lt-linux-x86_64-0.4.0.7-archive"
     curl --retry 3 -OLs https://developer.download.nvidia.com/compute/cusparselt/redist/libcusparse_lt/linux-x86_64/${CUSPARSELT_NAME}.tar.xz
+else
+    echo "Not sure which libcusparselt version to install for this ${CUDA_VERSION}"
 fi
 
 tar xf ${CUSPARSELT_NAME}.tar.xz

.ci/docker/common/install_db.sh

@@ -1,38 +0,0 @@
-#!/bin/bash
-
-set -ex
-
-install_ubuntu() {
-  apt-get update
-
-  # Cleanup
-  apt-get autoclean && apt-get clean
-  rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
-}
-
-install_centos() {
-  # Need EPEL for many packages we depend on.
-  # See http://fedoraproject.org/wiki/EPEL
-  yum --enablerepo=extras install -y epel-release
-
-  # Cleanup
-  yum clean all
-  rm -rf /var/cache/yum
-  rm -rf /var/lib/yum/yumdb
-  rm -rf /var/lib/yum/history
-}
-
-# Install base packages depending on the base OS
-ID=$(grep -oP '(?<=^ID=).+' /etc/os-release | tr -d '"')
-case "$ID" in
-  ubuntu)
-    install_ubuntu
-    ;;
-  centos)
-    install_centos
-    ;;
-  *)
-    echo "Unable to determine OS..."
-    exit 1
-    ;;
-esac

.ci/docker/common/install_executorch.sh

@@ -37,7 +37,12 @@ install_conda_dependencies() {
 
 install_pip_dependencies() {
   pushd executorch
-  as_jenkins bash install_requirements.sh --pybind xnnpack
+  as_jenkins bash install_executorch.sh
+
+  # A workaround, ExecuTorch has moved to numpy 2.0 which is not compatible with the current
+  # numba and scipy version used in PyTorch CI
+  conda_run pip uninstall -y numba scipy
+
   popd
 }
 
@@ -45,10 +50,9 @@ setup_executorch() {
   pushd executorch
 
   export PYTHON_EXECUTABLE=python
-  export EXECUTORCH_BUILD_PYBIND=ON
-  export CMAKE_ARGS="-DEXECUTORCH_BUILD_XNNPACK=ON -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON"
+  export CMAKE_ARGS="-DEXECUTORCH_BUILD_PYBIND=ON -DEXECUTORCH_BUILD_XNNPACK=ON -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON"
 
-  as_jenkins .ci/scripts/setup-linux.sh cmake || true
+  as_jenkins .ci/scripts/setup-linux.sh --build-tool cmake || true
   popd
 }
 

.ci/docker/common/install_halide.sh

@@ -35,7 +35,9 @@ git clone https://github.com/halide/Halide.git
 pushd Halide
 git checkout ${COMMIT} && git submodule update --init --recursive
 pip_install -r requirements.txt
-cmake -G Ninja -DCMAKE_BUILD_TYPE=Release -S . -B build
+# NOTE: pybind has a requirement for cmake > 3.5 so set the minimum cmake version here with a flag
+#       Context: https://github.com/pytorch/pytorch/issues/150420
+cmake -G Ninja -DCMAKE_POLICY_VERSION_MINIMUM=3.5 -DCMAKE_BUILD_TYPE=Release -S . -B build
 cmake --build build
 test -e ${CONDA_PREFIX}/lib/python3 || ln -s python${ANACONDA_PYTHON_VERSION} ${CONDA_PREFIX}/lib/python3
 cmake --install build --prefix ${CONDA_PREFIX}

.ci/docker/common/install_linter.sh

@@ -2,8 +2,6 @@
 
 set -ex
 
-source "$(dirname "${BASH_SOURCE[0]}")/common_utils.sh"
-
 if [ -n "${UBUNTU_VERSION}" ]; then
   apt update
   apt-get install -y clang doxygen git graphviz nodejs npm libtinfo5
@@ -15,8 +13,8 @@ chown -R jenkins pytorch
 
 pushd pytorch
 # Install all linter dependencies
-pip_install -r requirements.txt
-conda_run lintrunner init
+pip install -r requirements.txt
+lintrunner init
 
 # Cache .lintbin directory as part of the Docker image
 cp -r .lintbin /tmp

.ci/docker/common/install_nccl.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+
+set -ex
+
+NCCL_VERSION=""
+if [[ ${CUDA_VERSION:0:2} == "11" ]]; then
+  NCCL_VERSION=$(cat ci_commit_pins/nccl-cu11.txt)
+elif [[ ${CUDA_VERSION:0:2} == "12" ]]; then
+  NCCL_VERSION=$(cat ci_commit_pins/nccl-cu12.txt)
+else
+  echo "Unexpected CUDA_VERSION ${CUDA_VERSION}"
+  exit 1
+fi
+
+if [[ -n "${NCCL_VERSION}" ]]; then
+  # NCCL license: https://docs.nvidia.com/deeplearning/nccl/#licenses
+  # Follow build: https://github.com/NVIDIA/nccl/tree/master?tab=readme-ov-file#build
+  git clone -b $NCCL_VERSION --depth 1 https://github.com/NVIDIA/nccl.git
+  pushd nccl
+  make -j src.build
+  cp -a build/include/* /usr/local/cuda/include/
+  cp -a build/lib/* /usr/local/cuda/lib64/
+  popd
+  rm -rf nccl
+  ldconfig
+fi

.ci/docker/common/install_ninja.sh

@@ -4,10 +4,15 @@ set -ex
 
 [ -n "$NINJA_VERSION" ]
 
-url="https://github.com/ninja-build/ninja/releases/download/v${NINJA_VERSION}/ninja-linux.zip"
+arch=$(uname -m)
+if [ "$arch" == "aarch64" ]; then
+    url="https://github.com/ninja-build/ninja/releases/download/v${NINJA_VERSION}/ninja-linux-aarch64.zip"
+else
+    url="https://github.com/ninja-build/ninja/releases/download/v${NINJA_VERSION}/ninja-linux.zip"
+fi
 
 pushd /tmp
 wget --no-verbose --output-document=ninja-linux.zip "$url"
 unzip ninja-linux.zip -d /usr/local/bin
 rm -f ninja-linux.zip
-popd
+popd

.ci/docker/common/install_onnx.sh

@@ -31,15 +31,15 @@ pip_install \
 pip_install coloredlogs packaging
 
 pip_install onnxruntime==1.18.1
-pip_install onnx==1.16.2
-pip_install onnxscript==0.1.0.dev20241124 --no-deps
+pip_install onnx==1.17.0
+pip_install onnxscript==0.2.2 --no-deps
 # required by onnxscript
 pip_install ml_dtypes
 
 # Cache the transformers model to be used later by ONNX tests. We need to run the transformers
 # package to download the model. By default, the model is cached at ~/.cache/huggingface/hub/
 IMPORT_SCRIPT_FILENAME="/tmp/onnx_import_script.py"
-as_jenkins echo 'import transformers; transformers.AutoModel.from_pretrained("sshleifer/tiny-gpt2"); transformers.AutoTokenizer.from_pretrained("sshleifer/tiny-gpt2"); transformers.AutoModelForSpeechSeq2Seq.from_pretrained("openai/whisper-large-v3");' > "${IMPORT_SCRIPT_FILENAME}"
+as_jenkins echo 'import transformers; transformers.GPTJForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gptj");' > "${IMPORT_SCRIPT_FILENAME}"
 
 # Need a PyTorch version for transformers to work
 pip_install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu

.ci/docker/common/install_openblas.sh

@@ -4,7 +4,7 @@
 set -ex
 
 cd /
-git clone https://github.com/OpenMathLib/OpenBLAS.git -b v0.3.28 --depth 1 --shallow-submodules
+git clone https://github.com/OpenMathLib/OpenBLAS.git -b v0.3.29 --depth 1 --shallow-submodules
 
 
 OPENBLAS_BUILD_FLAGS="

.ci/docker/common/install_python.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+set -ex
+
+apt-get update
+# Use deadsnakes in case we need an older python version
+sudo add-apt-repository ppa:deadsnakes/ppa
+apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev python3-pip python${PYTHON_VERSION}-venv
+
+# Use a venv because uv and some other package managers don't support --user install
+ln -s /usr/bin/python${PYTHON_VERSION} /usr/bin/python
+python -m venv /var/lib/jenkins/ci_env
+source /var/lib/jenkins/ci_env/bin/activate
+
+python -mpip install --upgrade pip
+python -mpip install -r /opt/requirements-ci.txt
+if [ -n "${PIP_CMAKE}" ]; then
+  python -mpip install cmake==3.31.6
+fi

.ci/docker/common/install_rocm.sh

@@ -8,10 +8,6 @@ ver() {
 
 install_ubuntu() {
     apt-get update
-    if [[ $UBUNTU_VERSION == 18.04 ]]; then
-      # gpg-agent is not available by default on 18.04
-      apt-get install -y --no-install-recommends gpg-agent
-    fi
     if [[ $UBUNTU_VERSION == 20.04 ]]; then
       # gpg-agent is not available by default on 20.04
       apt-get install -y --no-install-recommends gpg-agent

.ci/docker/common/install_rocm_drm.sh

@@ -25,7 +25,9 @@ python3 -m pip install meson ninja
 ###########################
 ### clone repo
 ###########################
-GIT_SSL_NO_VERIFY=true git clone https://gitlab.freedesktop.org/mesa/drm.git
+# TEMPORARY FIX: https://gitlab.freedesktop.org/mesa/drm.git is down until 2025/03/22
+# GIT_SSL_NO_VERIFY=true git clone https://gitlab.freedesktop.org/mesa/drm.git
+GIT_SSL_NO_VERIFY=true git clone git://anongit.freedesktop.org/mesa/drm
 pushd drm
 
 ###########################
@@ -115,7 +117,7 @@ index a5007ffc..13fa07fc 100644
  	if (!fp) {
 -		fprintf(stderr, "%s: %s\n", AMDGPU_ASIC_ID_TABLE,
 -			strerror(errno));
-+		fprintf(stderr, "amdgpu.ids: No such file or directory\n");
++		//fprintf(stderr, "amdgpu.ids: No such file or directory\n");
  		return;
  	}
 

.ci/docker/common/install_rocm_magma.sh

@@ -1,50 +1,28 @@
-#!/bin/bash
-# Script used in CI and CD pipeline
+#!/usr/bin/env bash
+# Script used only in CD pipeline
 
-set -ex
+set -eou pipefail
 
-# Magma build scripts need `python`
-ln -sf /usr/bin/python3 /usr/bin/python
+function do_install() {
+    rocm_version=$1
+    rocm_version_nodot=${1//./}
 
-ID=$(grep -oP '(?<=^ID=).+' /etc/os-release | tr -d '"')
-case "$ID" in
-  almalinux)
-    yum install -y gcc-gfortran
-    ;;
-  *)
-    echo "No preinstalls to build magma..."
-    ;;
-esac
+    # Version 2.7.2 + ROCm related updates
+    MAGMA_VERSION=a1625ff4d9bc362906bd01f805dbbe12612953f6
+    magma_archive="magma-rocm${rocm_version_nodot}-${MAGMA_VERSION}-1.tar.bz2"
 
-MKLROOT=${MKLROOT:-/opt/conda/envs/py_$ANACONDA_PYTHON_VERSION}
+    rocm_dir="/opt/rocm"
+    (
+        set -x
+        tmp_dir=$(mktemp -d)
+        pushd ${tmp_dir}
+        curl -OLs https://ossci-linux.s3.us-east-1.amazonaws.com/${magma_archive}
+        tar -xvf "${magma_archive}"
+        mkdir -p "${rocm_dir}/magma"
+        mv include "${rocm_dir}/magma/include"
+        mv lib "${rocm_dir}/magma/lib"
+        popd
+    )
+}
 
-# "install" hipMAGMA into /opt/rocm/magma by copying after build
-git clone https://bitbucket.org/icl/magma.git
-pushd magma
-
-# Version 2.7.2 + ROCm related updates
-git checkout a1625ff4d9bc362906bd01f805dbbe12612953f6
-
-cp make.inc-examples/make.inc.hip-gcc-mkl make.inc
-echo 'LIBDIR += -L$(MKLROOT)/lib' >> make.inc
-if [[ -f "${MKLROOT}/lib/libmkl_core.a" ]]; then
-    echo 'LIB = -Wl,--start-group -lmkl_gf_lp64 -lmkl_gnu_thread -lmkl_core -Wl,--end-group -lpthread -lstdc++ -lm -lgomp -lhipblas -lhipsparse' >> make.inc
-fi
-echo 'LIB += -Wl,--enable-new-dtags -Wl,--rpath,/opt/rocm/lib -Wl,--rpath,$(MKLROOT)/lib -Wl,--rpath,/opt/rocm/magma/lib -ldl' >> make.inc
-echo 'DEVCCFLAGS += --gpu-max-threads-per-block=256' >> make.inc
-export PATH="${PATH}:/opt/rocm/bin"
-if [[ -n "$PYTORCH_ROCM_ARCH" ]]; then
-  amdgpu_targets=`echo $PYTORCH_ROCM_ARCH | sed 's/;/ /g'`
-else
-  amdgpu_targets=`rocm_agent_enumerator | grep -v gfx000 | sort -u | xargs`
-fi
-for arch in $amdgpu_targets; do
-  echo "DEVCCFLAGS += --offload-arch=$arch" >> make.inc
-done
-# hipcc with openmp flag may cause isnan() on __device__ not to be found; depending on context, compiler may attempt to match with host definition
-sed -i 's/^FOPENMP/#FOPENMP/g' make.inc
-make -f make.gen.hipMAGMA -j $(nproc)
-LANG=C.UTF-8 make lib/libmagma.so -j $(nproc) MKLROOT="${MKLROOT}"
-make testing/testing_dgemm -j $(nproc) MKLROOT="${MKLROOT}"
-popd
-mv magma /opt/rocm
+do_install $1

.ci/docker/common/install_swiftshader.sh

@@ -1,24 +0,0 @@
-#!/bin/bash
-
-set -ex
-
-[ -n "${SWIFTSHADER}" ]
-
-retry () {
-    $*  || (sleep 1 && $*) || (sleep 2 && $*) || (sleep 4 && $*) || (sleep 8 && $*)
-}
-
-_https_amazon_aws=https://ossci-android.s3.amazonaws.com
-
-# SwiftShader
-_swiftshader_dir=/var/lib/jenkins/swiftshader
-_swiftshader_file_targz=swiftshader-abe07b943-prebuilt.tar.gz
-mkdir -p $_swiftshader_dir
-_tmp_swiftshader_targz="/tmp/${_swiftshader_file_targz}"
-
-curl --silent --show-error --location --fail --retry 3 \
-  --output "${_tmp_swiftshader_targz}" "$_https_amazon_aws/${_swiftshader_file_targz}"
-
-tar -C "${_swiftshader_dir}" -xzf "${_tmp_swiftshader_targz}"
-
-export VK_ICD_FILENAMES="${_swiftshader_dir}/build/Linux/vk_swiftshader_icd.json"

.ci/docker/common/install_triton.sh

@@ -2,6 +2,12 @@
 
 set -ex
 
+mkdir -p /opt/triton
+if [ -z "${TRITON}" ] && [ -z "${TRITON_CPU}" ]; then
+  echo "TRITON and TRITON_CPU are not set. Exiting..."
+  exit 0
+fi
+
 source "$(dirname "${BASH_SOURCE[0]}")/common_utils.sh"
 
 get_conda_version() {
@@ -52,6 +58,7 @@ cd triton
 as_jenkins git checkout ${TRITON_PINNED_COMMIT}
 as_jenkins git submodule update --init --recursive
 cd python
+pip_install pybind11==2.13.6
 
 # TODO: remove patch setup.py once we have a proper fix for https://github.com/triton-lang/triton/issues/4527
 as_jenkins sed -i -e 's/https:\/\/tritonlang.blob.core.windows.net\/llvm-builds/https:\/\/oaitriton.blob.core.windows.net\/public\/llvm-builds/g' setup.py
@@ -60,17 +67,22 @@ if [ -n "${UBUNTU_VERSION}" ] && [ -n "${GCC_VERSION}" ] && [[ "${GCC_VERSION}"
   # Triton needs at least gcc-9 to build
   apt-get install -y g++-9
 
-  CXX=g++-9 pip_install -e .
+  CXX=g++-9 conda_run python setup.py bdist_wheel
 elif [ -n "${UBUNTU_VERSION}" ] && [ -n "${CLANG_VERSION}" ]; then
   # Triton needs <filesystem> which surprisingly is not available with clang-9 toolchain
   add-apt-repository -y ppa:ubuntu-toolchain-r/test
   apt-get install -y g++-9
 
-  CXX=g++-9 pip_install -e .
+  CXX=g++-9 conda_run python setup.py bdist_wheel
 else
-  pip_install -e .
+  conda_run python setup.py bdist_wheel
 fi
 
+# Copy the wheel to /opt for multi stage docker builds
+cp dist/*.whl /opt/triton
+# Install the wheel for docker builds that don't use multi stage
+pip_install dist/*.whl
+
 if [ -n "${CONDA_CMAKE}" ]; then
   # TODO: This is to make sure that the same cmake and numpy version from install conda
   # script is used. Without this step, the newer cmake version (3.25.2) downloaded by

.ci/docker/common/install_ucc.sh

@@ -8,6 +8,12 @@ else
   with_cuda=no
 fi
 
+if [[ -d "/opt/rocm" ]]; then
+  with_rocm=/opt/rocm
+else
+  with_rocm=no
+fi
+
 function install_ucx() {
   set -ex
   git clone --recursive https://github.com/openucx/ucx.git
@@ -19,6 +25,7 @@ function install_ucx() {
   ./configure --prefix=$UCX_HOME      \
       --enable-mt                     \
       --with-cuda=$with_cuda          \
+      --with-rocm=$with_rocm          \
       --enable-profiling              \
       --enable-stats
   time make -j
@@ -36,12 +43,29 @@ function install_ucc() {
   git submodule update --init --recursive
 
   ./autogen.sh
+
   # We only run distributed tests on Tesla M60 and A10G
   NVCC_GENCODE="-gencode=arch=compute_52,code=sm_52 -gencode=arch=compute_86,code=compute_86"
+
+  if [[ -n "$ROCM_VERSION" ]]; then
+    if [[ -n "$PYTORCH_ROCM_ARCH" ]]; then
+      amdgpu_targets=`echo $PYTORCH_ROCM_ARCH | sed 's/;/ /g'`
+    else
+      amdgpu_targets=`rocm_agent_enumerator | grep -v gfx000 | sort -u | xargs`
+    fi
+    for arch in $amdgpu_targets; do
+      HIP_OFFLOAD="$HIP_OFFLOAD --offload-arch=$arch"
+    done
+  else
+    HIP_OFFLOAD="all-arch-no-native"
+  fi
+
   ./configure --prefix=$UCC_HOME          \
     --with-ucx=$UCX_HOME                  \
     --with-cuda=$with_cuda                \
-    --with-nvcc-gencode="${NVCC_GENCODE}"
+    --with-nvcc-gencode="${NVCC_GENCODE}" \
+    --with-rocm=$with_rocm                \
+    --with-rocm-arch="${HIP_OFFLOAD}"
   time make -j
   sudo make install
 

.ci/docker/common/install_vulkan_sdk.sh

@@ -1,24 +0,0 @@
-#!/bin/bash
-
-set -ex
-
-[ -n "${VULKAN_SDK_VERSION}" ]
-
-retry () {
-    $*  || (sleep 1 && $*) || (sleep 2 && $*) || (sleep 4 && $*) || (sleep 8 && $*)
-}
-
-_vulkansdk_dir=/var/lib/jenkins/vulkansdk
-_tmp_vulkansdk_targz=/tmp/vulkansdk.tar.gz
-
-curl \
-  --silent \
-  --show-error \
-  --location \
-  --fail \
-  --retry 3 \
-  --output "${_tmp_vulkansdk_targz}" "https://ossci-android.s3.amazonaws.com/vulkansdk-linux-x86_64-${VULKAN_SDK_VERSION}.tar.gz"
-
-mkdir -p "${_vulkansdk_dir}"
-tar -C "${_vulkansdk_dir}" -xzf "${_tmp_vulkansdk_targz}" --strip-components 1
-rm -rf "${_tmp_vulkansdk_targz}"

.ci/docker/libtorch/Dockerfile

@@ -49,6 +49,8 @@ RUN bash ./install_mkl.sh && rm install_mkl.sh
 FROM cpu as cuda
 ADD ./common/install_cuda.sh install_cuda.sh
 ADD ./common/install_magma.sh install_magma.sh
+COPY ./common/install_nccl.sh install_nccl.sh
+COPY ./ci_commit_pins/nccl-cu* /ci_commit_pins/
 ENV CUDA_HOME /usr/local/cuda
 
 FROM cuda as cuda11.8
@@ -56,11 +58,6 @@ RUN bash ./install_cuda.sh 11.8
 RUN bash ./install_magma.sh 11.8
 RUN ln -sf /usr/local/cuda-11.8 /usr/local/cuda
 
-FROM cuda as cuda12.1
-RUN bash ./install_cuda.sh 12.1
-RUN bash ./install_magma.sh 12.1
-RUN ln -sf /usr/local/cuda-12.1 /usr/local/cuda
-
 FROM cuda as cuda12.4
 RUN bash ./install_cuda.sh 12.4
 RUN bash ./install_magma.sh 12.4
@@ -71,7 +68,13 @@ RUN bash ./install_cuda.sh 12.6
 RUN bash ./install_magma.sh 12.6
 RUN ln -sf /usr/local/cuda-12.6 /usr/local/cuda
 
+FROM cuda as cuda12.8
+RUN bash ./install_cuda.sh 12.8
+RUN bash ./install_magma.sh 12.8
+RUN ln -sf /usr/local/cuda-12.8 /usr/local/cuda
+
 FROM cpu as rocm
+ARG ROCM_VERSION
 ARG PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH ${PYTORCH_ROCM_ARCH}
 ENV MKLROOT /opt/intel
@@ -90,7 +93,7 @@ RUN apt-get update -y && \
     apt-get clean
 
 RUN bash ./install_rocm_drm.sh && rm install_rocm_drm.sh
-RUN bash ./install_rocm_magma.sh && rm install_rocm_magma.sh
+RUN bash ./install_rocm_magma.sh ${ROCM_VERSION} && rm install_rocm_magma.sh
 
 FROM ${BASE_TARGET} as final
 COPY --from=openssl            /opt/openssl           /opt/openssl

.ci/docker/libtorch/build.sh

@@ -39,8 +39,8 @@ case ${GPU_ARCH_TYPE} in
         BASE_TARGET=rocm
         DOCKER_TAG=rocm${GPU_ARCH_VERSION}
         GPU_IMAGE=rocm/dev-ubuntu-20.04:${GPU_ARCH_VERSION}-complete
-        PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908;gfx90a;gfx1030;gfx1100;gfx1101;gfx942"
-        DOCKER_GPU_BUILD_ARG="--build-arg PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH}"
+        PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1200;gfx1201"
+        DOCKER_GPU_BUILD_ARG="--build-arg PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} --build-arg ROCM_VERSION=${GPU_ARCH_VERSION}"
         ;;
     *)
         echo "ERROR: Unrecognized GPU_ARCH_TYPE: ${GPU_ARCH_TYPE}"

.ci/docker/linter-cuda/Dockerfile

@@ -18,28 +18,30 @@ COPY ./common/install_user.sh install_user.sh
 RUN bash ./install_user.sh && rm install_user.sh
 
 # Install conda and other packages (e.g., numpy, pytest)
-ARG ANACONDA_PYTHON_VERSION
-ARG CONDA_CMAKE
-ENV ANACONDA_PYTHON_VERSION=$ANACONDA_PYTHON_VERSION
-ENV PATH /opt/conda/envs/py_$ANACONDA_PYTHON_VERSION/bin:/opt/conda/bin:$PATH
-COPY requirements-ci.txt /opt/conda/requirements-ci.txt
-COPY ./common/install_conda.sh install_conda.sh
-COPY ./common/common_utils.sh common_utils.sh
-COPY ./common/install_magma_conda.sh install_magma_conda.sh
-RUN bash ./install_conda.sh && rm install_conda.sh install_magma_conda.sh common_utils.sh /opt/conda/requirements-ci.txt
+ARG PYTHON_VERSION
+ARG PIP_CMAKE
+# Put venv into the env vars so users don't need to activate it
+ENV PATH /var/lib/jenkins/ci_env/bin:$PATH
+ENV VIRTUAL_ENV /var/lib/jenkins/ci_env
+COPY requirements-ci.txt /opt/requirements-ci.txt
+COPY ./common/install_python.sh install_python.sh
+RUN bash ./install_python.sh && rm install_python.sh /opt/requirements-ci.txt
 
 # Install cuda and cudnn
 ARG CUDA_VERSION
 COPY ./common/install_cuda.sh install_cuda.sh
-RUN bash ./install_cuda.sh ${CUDA_VERSION} && rm install_cuda.sh
+COPY ./common/install_nccl.sh install_nccl.sh
+COPY ./ci_commit_pins/nccl-cu* /ci_commit_pins/
+RUN bash ./install_cuda.sh ${CUDA_VERSION} && rm install_cuda.sh install_nccl.sh /ci_commit_pins/nccl-cu*
 ENV DESIRED_CUDA ${CUDA_VERSION}
 ENV PATH /usr/local/nvidia/bin:/usr/local/cuda/bin:$PATH
 
 # Note that Docker build forbids copying file outside the build context
 COPY ./common/install_linter.sh install_linter.sh
-COPY ./common/common_utils.sh common_utils.sh
 RUN bash ./install_linter.sh
-RUN rm install_linter.sh common_utils.sh
+RUN rm install_linter.sh
+
+RUN chown -R jenkins:jenkins /var/lib/jenkins/ci_env
 
 USER jenkins
 CMD ["bash"]

.ci/docker/linter/Dockerfile

@@ -15,20 +15,18 @@ COPY ./common/install_user.sh install_user.sh
 RUN bash ./install_user.sh && rm install_user.sh
 
 # Install conda and other packages (e.g., numpy, pytest)
-ARG ANACONDA_PYTHON_VERSION
-ARG CONDA_CMAKE
-ENV ANACONDA_PYTHON_VERSION=$ANACONDA_PYTHON_VERSION
-ENV PATH /opt/conda/envs/py_$ANACONDA_PYTHON_VERSION/bin:/opt/conda/bin:$PATH
-COPY requirements-ci.txt /opt/conda/requirements-ci.txt
-COPY ./common/install_conda.sh install_conda.sh
-COPY ./common/common_utils.sh common_utils.sh
-RUN bash ./install_conda.sh && rm install_conda.sh common_utils.sh /opt/conda/requirements-ci.txt
+ARG PYTHON_VERSION
+ARG PIP_CMAKE
+ENV PATH /var/lib/jenkins/ci_env/bin:$PATH
+ENV VIRTUAL_ENV /var/lib/jenkins/ci_env
+COPY requirements-ci.txt /opt/requirements-ci.txt
+COPY ./common/install_python.sh install_python.sh
+RUN bash ./install_python.sh && rm install_python.sh /opt/requirements-ci.txt
 
 # Note that Docker build forbids copying file outside the build context
 COPY ./common/install_linter.sh install_linter.sh
-COPY ./common/common_utils.sh common_utils.sh
 RUN bash ./install_linter.sh
-RUN rm install_linter.sh common_utils.sh
+RUN rm install_linter.sh
 
 USER jenkins
 CMD ["bash"]

.ci/docker/manywheel/Dockerfile

@@ -64,7 +64,9 @@ FROM base as cuda
 ARG BASE_CUDA_VERSION=10.2
 # Install CUDA
 ADD ./common/install_cuda.sh install_cuda.sh
-RUN bash ./install_cuda.sh ${BASE_CUDA_VERSION} && rm install_cuda.sh
+COPY ./common/install_nccl.sh install_nccl.sh
+COPY ./ci_commit_pins/nccl-cu* /ci_commit_pins/
+RUN bash ./install_cuda.sh ${BASE_CUDA_VERSION} && rm install_cuda.sh install_nccl.sh /ci_commit_pins/nccl-cu*
 
 FROM base as intel
 # MKL
@@ -195,6 +197,6 @@ RUN bash ./install_rocm_drm.sh && rm install_rocm_drm.sh
 # cmake3 is needed for the MIOpen build
 RUN ln -sf /usr/local/bin/cmake /usr/bin/cmake3
 ADD ./common/install_rocm_magma.sh install_rocm_magma.sh
-RUN bash ./install_rocm_magma.sh && rm install_rocm_magma.sh
+RUN bash ./install_rocm_magma.sh ${ROCM_VERSION} && rm install_rocm_magma.sh
 ADD ./common/install_miopen.sh install_miopen.sh
 RUN bash ./install_miopen.sh ${ROCM_VERSION} && rm install_miopen.sh

.ci/docker/manywheel/Dockerfile_2014

@@ -1,153 +0,0 @@
-# syntax = docker/dockerfile:experimental
-ARG ROCM_VERSION=3.7
-ARG BASE_CUDA_VERSION=10.2
-ARG GPU_IMAGE=nvidia/cuda:${BASE_CUDA_VERSION}-devel-centos7
-FROM quay.io/pypa/manylinux2014_x86_64 as base
-
-ENV LC_ALL en_US.UTF-8
-ENV LANG en_US.UTF-8
-ENV LANGUAGE en_US.UTF-8
-
-RUN sed -i s/mirror.centos.org/vault.centos.org/g /etc/yum.repos.d/*.repo
-RUN sed -i s/^#.*baseurl=http/baseurl=http/g /etc/yum.repos.d/*.repo
-RUN sed -i s/^mirrorlist=http/#mirrorlist=http/g /etc/yum.repos.d/*.repo
-RUN yum install -y wget curl perl util-linux xz bzip2 git patch which perl zlib-devel
-RUN yum install -y yum-utils centos-release-scl sudo
-RUN yum-config-manager --enable rhel-server-rhscl-7-rpms
-RUN yum install -y devtoolset-7-gcc devtoolset-7-gcc-c++ devtoolset-7-gcc-gfortran devtoolset-7-binutils
-ENV PATH=/opt/rh/devtoolset-7/root/usr/bin:$PATH
-ENV LD_LIBRARY_PATH=/opt/rh/devtoolset-7/root/usr/lib64:/opt/rh/devtoolset-7/root/usr/lib:$LD_LIBRARY_PATH
-
-# cmake
-RUN yum install -y cmake3 && \
-    ln -s /usr/bin/cmake3 /usr/bin/cmake
-FROM base as openssl
-# Install openssl (this must precede `build python` step)
-# (In order to have a proper SSL module, Python is compiled
-# against a recent openssl [see env vars above], which is linked
-# statically. We delete openssl afterwards.)
-ADD ./common/install_openssl.sh install_openssl.sh
-RUN bash ./install_openssl.sh && rm install_openssl.sh
-
-
-
-# remove unncessary python versions
-RUN rm -rf /opt/python/cp26-cp26m /opt/_internal/cpython-2.6.9-ucs2
-RUN rm -rf /opt/python/cp26-cp26mu /opt/_internal/cpython-2.6.9-ucs4
-RUN rm -rf /opt/python/cp33-cp33m /opt/_internal/cpython-3.3.6
-RUN rm -rf /opt/python/cp34-cp34m /opt/_internal/cpython-3.4.6
-
-FROM base as cuda
-ARG BASE_CUDA_VERSION=10.2
-# Install CUDA
-ADD ./common/install_cuda.sh install_cuda.sh
-RUN bash ./install_cuda.sh ${BASE_CUDA_VERSION} && rm install_cuda.sh
-
-FROM base as intel
-# MKL
-ADD ./common/install_mkl.sh install_mkl.sh
-RUN bash ./install_mkl.sh && rm install_mkl.sh
-
-FROM base as magma
-ARG BASE_CUDA_VERSION=10.2
-# Install magma
-ADD ./common/install_magma.sh install_magma.sh
-RUN bash ./install_magma.sh ${BASE_CUDA_VERSION} && rm install_magma.sh
-
-FROM base as jni
-# Install java jni header
-ADD ./common/install_jni.sh install_jni.sh
-ADD ./java/jni.h jni.h
-RUN bash ./install_jni.sh && rm install_jni.sh
-
-FROM base as libpng
-# Install libpng
-ADD ./common/install_libpng.sh install_libpng.sh
-RUN bash ./install_libpng.sh && rm install_libpng.sh
-
-FROM ${GPU_IMAGE} as common
-RUN sed -i s/mirror.centos.org/vault.centos.org/g /etc/yum.repos.d/*.repo
-RUN sed -i s/^#.*baseurl=http/baseurl=http/g /etc/yum.repos.d/*.repo
-RUN sed -i s/^mirrorlist=http/#mirrorlist=http/g /etc/yum.repos.d/*.repo
-ENV LC_ALL en_US.UTF-8
-ENV LANG en_US.UTF-8
-ENV LANGUAGE en_US.UTF-8
-RUN yum install -y \
-        aclocal \
-        autoconf \
-        automake \
-        bison \
-        bzip2 \
-        curl \
-        diffutils \
-        file \
-        git \
-        make \
-        patch \
-        perl \
-        unzip \
-        util-linux \
-        wget \
-        which \
-        xz \
-        yasm
-RUN yum install -y \
-    https://repo.ius.io/ius-release-el7.rpm \
-    https://ossci-linux.s3.amazonaws.com/epel-release-7-14.noarch.rpm
-
-RUN yum swap -y git git236-core
-# git236+ would refuse to run git commands in repos owned by other users
-# Which causes version check to fail, as pytorch repo is bind-mounted into the image
-# Override this behaviour by treating every folder as safe
-# For more details see https://github.com/pytorch/pytorch/issues/78659#issuecomment-1144107327
-RUN git config --global --add safe.directory "*"
-
-ENV SSL_CERT_FILE=/opt/_internal/certs.pem
-# Install LLVM version
-COPY --from=openssl            /opt/openssl                          /opt/openssl
-COPY --from=base               /opt/python                           /opt/python
-COPY --from=base               /opt/_internal                        /opt/_internal
-COPY --from=base               /usr/local/bin/auditwheel             /usr/local/bin/auditwheel
-COPY --from=intel              /opt/intel                            /opt/intel
-COPY --from=base               /usr/local/bin/patchelf               /usr/local/bin/patchelf
-COPY --from=libpng             /usr/local/bin/png*                   /usr/local/bin/
-COPY --from=libpng             /usr/local/bin/libpng*                /usr/local/bin/
-COPY --from=libpng             /usr/local/include/png*               /usr/local/include/
-COPY --from=libpng             /usr/local/include/libpng*            /usr/local/include/
-COPY --from=libpng             /usr/local/lib/libpng*                /usr/local/lib/
-COPY --from=libpng             /usr/local/lib/pkgconfig              /usr/local/lib/pkgconfig
-COPY --from=jni                /usr/local/include/jni.h              /usr/local/include/jni.h
-
-FROM common as cpu_final
-ARG BASE_CUDA_VERSION=10.2
-RUN yum install -y yum-utils centos-release-scl
-RUN yum-config-manager --enable rhel-server-rhscl-7-rpms
-RUN yum install -y devtoolset-7-gcc devtoolset-7-gcc-c++ devtoolset-7-gcc-gfortran devtoolset-7-binutils
-ENV PATH=/opt/rh/devtoolset-7/root/usr/bin:$PATH
-ENV LD_LIBRARY_PATH=/opt/rh/devtoolset-7/root/usr/lib64:/opt/rh/devtoolset-7/root/usr/lib:$LD_LIBRARY_PATH
-
-# cmake
-RUN yum install -y cmake3 && \
-    ln -s /usr/bin/cmake3 /usr/bin/cmake
-
-# ninja
-RUN yum install -y http://repo.okay.com.mx/centos/7/x86_64/release/okay-release-1-1.noarch.rpm
-RUN yum install -y ninja-build
-
-FROM cpu_final as cuda_final
-RUN rm -rf /usr/local/cuda-${BASE_CUDA_VERSION}
-COPY --from=cuda     /usr/local/cuda-${BASE_CUDA_VERSION}  /usr/local/cuda-${BASE_CUDA_VERSION}
-COPY --from=magma    /usr/local/cuda-${BASE_CUDA_VERSION}  /usr/local/cuda-${BASE_CUDA_VERSION}
-
-FROM common as rocm_final
-ARG ROCM_VERSION=3.7
-# Install ROCm
-ADD ./common/install_rocm.sh install_rocm.sh
-RUN bash ./install_rocm.sh ${ROCM_VERSION} && rm install_rocm.sh
-# cmake is already installed inside the rocm base image, but both 2 and 3 exist
-# cmake3 is needed for the later MIOpen custom build, so that step is last.
-RUN yum install -y cmake3 && \
-    rm -f /usr/bin/cmake && \
-    ln -s /usr/bin/cmake3 /usr/bin/cmake
-ADD ./common/install_miopen.sh install_miopen.sh
-RUN bash ./install_miopen.sh ${ROCM_VERSION} && rm install_miopen.sh

.ci/docker/manywheel/Dockerfile_2_28

@@ -36,7 +36,9 @@ FROM base as cuda
 ARG BASE_CUDA_VERSION=11.8
 # Install CUDA
 ADD ./common/install_cuda.sh install_cuda.sh
-RUN bash ./install_cuda.sh ${BASE_CUDA_VERSION} && rm install_cuda.sh
+COPY ./common/install_nccl.sh install_nccl.sh
+COPY ./ci_commit_pins/nccl-cu* /ci_commit_pins/
+RUN bash ./install_cuda.sh ${BASE_CUDA_VERSION} && rm install_cuda.sh install_nccl.sh ci_commit_pins/nccl-cu*
 
 FROM base as intel
 # MKL
@@ -158,7 +160,7 @@ ADD ./common/install_rocm_drm.sh install_rocm_drm.sh
 RUN bash ./install_rocm_drm.sh && rm install_rocm_drm.sh
 ENV MKLROOT /opt/intel
 ADD ./common/install_rocm_magma.sh install_rocm_magma.sh
-RUN bash ./install_rocm_magma.sh && rm install_rocm_magma.sh
+RUN bash ./install_rocm_magma.sh ${ROCM_VERSION} && rm install_rocm_magma.sh
 ADD ./common/install_miopen.sh install_miopen.sh
 RUN bash ./install_miopen.sh ${ROCM_VERSION} && rm install_miopen.sh
 

.ci/docker/manywheel/Dockerfile_2_28_aarch64

@@ -38,6 +38,12 @@ RUN yum install -y \
   sudo \
   gcc-toolset-${GCCTOOLSET_VERSION}-toolchain
 
+# (optional) Install non-default Ninja version
+ARG NINJA_VERSION
+COPY ./common/install_ninja.sh install_ninja.sh
+RUN if [ -n "${NINJA_VERSION}" ]; then bash ./install_ninja.sh; fi
+RUN rm install_ninja.sh
+
 # Ensure the expected devtoolset is used
 ENV PATH=/opt/rh/gcc-toolset-${GCCTOOLSET_VERSION}/root/usr/bin:$PATH
 ENV LD_LIBRARY_PATH=/opt/rh/gcc-toolset-${GCCTOOLSET_VERSION}/root/usr/lib64:/opt/rh/gcc-toolset-${GCCTOOLSET_VERSION}/root/usr/lib:$LD_LIBRARY_PATH

.ci/docker/manywheel/Dockerfile_cuda_aarch64

@@ -67,7 +67,9 @@ FROM base as cuda
 ARG BASE_CUDA_VERSION
 # Install CUDA
 ADD ./common/install_cuda_aarch64.sh install_cuda_aarch64.sh
-RUN bash ./install_cuda_aarch64.sh ${BASE_CUDA_VERSION} && rm install_cuda_aarch64.sh
+COPY ./common/install_nccl.sh install_nccl.sh
+COPY ./ci_commit_pins/nccl-cu* /ci_commit_pins/
+RUN bash ./install_cuda_aarch64.sh ${BASE_CUDA_VERSION} && rm install_cuda_aarch64.sh install_nccl.sh ci_commit_pins/nccl-cu*
 
 FROM base as magma
 ARG BASE_CUDA_VERSION

.ci/docker/manywheel/Dockerfile_s390x

@@ -42,6 +42,7 @@ RUN yum install -y \
   llvm-devel \
   libzstd-devel \
   python3.12-devel \
+  python3.12-test \
   python3.12-setuptools \
   python3.12-pip \
   python3-virtualenv \
@@ -101,24 +102,33 @@ CMD ["/bin/bash"]
 
 # install test dependencies:
 # - grpcio requires system openssl, bundled crypto fails to build
-# - ml_dtypes 0.4.0 requires some fixes provided in later commits to build
 RUN dnf install -y \
   protobuf-devel \
   protobuf-c-devel \
   protobuf-lite-devel \
-  wget \
-  patch
+  hdf5-devel \
+  python3-h5py \
+  git
 
-RUN env GRPC_PYTHON_BUILD_SYSTEM_OPENSSL=True pip3 install grpcio==1.65.4
-RUN cd ~ && \
-  git clone https://github.com/jax-ml/ml_dtypes && \
-  cd ml_dtypes && \
-  git checkout v0.4.0 && \
+RUN env GRPC_PYTHON_BUILD_SYSTEM_OPENSSL=True pip3 install grpcio
+
+# cmake-3.28.0 from pip for onnxruntime
+RUN python3 -mpip install cmake==3.28.0
+
+# build onnxruntime 1.21.0 from sources.
+# it is not possible to build it from sources using pip,
+# so just build it from upstream repository.
+# h5py is dependency of onnxruntime_training.
+# h5py==3.11.0 builds with hdf5-devel 1.10.5 from repository.
+# install newest flatbuffers version first:
+# for some reason old version is getting pulled in otherwise.
+# packaging package is required for onnxruntime wheel build.
+RUN pip3 install flatbuffers && \
+  pip3 install h5py==3.11.0 && \
+  pip3 install packaging && \
+  git clone https://github.com/microsoft/onnxruntime && \
+  cd onnxruntime && git checkout v1.21.0 && \
   git submodule update --init --recursive && \
-  wget https://github.com/jax-ml/ml_dtypes/commit/b969f76914d6b30676721bc92bf0f6021a0d1321.patch && \
-  wget https://github.com/jax-ml/ml_dtypes/commit/d4e6d035ecda073eab8bcf60f4eef572ee7087e6.patch && \
-  patch -p1 < b969f76914d6b30676721bc92bf0f6021a0d1321.patch && \
-  patch -p1 < d4e6d035ecda073eab8bcf60f4eef572ee7087e6.patch && \
-  python3 setup.py bdist_wheel && \
-  pip3 install dist/*.whl && \
-  rm -rf ml_dtypes
+  ./build.sh --config Release --parallel 0 --enable_pybind --build_wheel --enable_training --enable_training_apis --enable_training_ops --skip_tests --allow_running_as_root && \
+  pip3 install ./build/Linux/Release/dist/onnxruntime_training-*.whl && \
+  cd .. && /bin/rm -rf ./onnxruntime

.ci/docker/manywheel/build.sh

@@ -48,7 +48,7 @@ case ${GPU_ARCH_TYPE} in
         TARGET=final
         DOCKER_TAG=cpu-aarch64
         GPU_IMAGE=arm64v8/almalinux:8
-        DOCKER_GPU_BUILD_ARG=" --build-arg DEVTOOLSET_VERSION=11"
+        DOCKER_GPU_BUILD_ARG=" --build-arg DEVTOOLSET_VERSION=11 --build-arg NINJA_VERSION=1.12.1"
         MANY_LINUX_VERSION="2_28_aarch64"
         ;;
     cpu-cxx11-abi)
@@ -97,7 +97,7 @@ case ${GPU_ARCH_TYPE} in
             DEVTOOLSET_VERSION="11"
             GPU_IMAGE=rocm/dev-almalinux-8:${GPU_ARCH_VERSION}-complete
         fi
-        PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101"
+        PYTORCH_ROCM_ARCH="gfx900;gfx906;gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1200;gfx1201"
         DOCKER_GPU_BUILD_ARG="--build-arg ROCM_VERSION=${GPU_ARCH_VERSION} --build-arg PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} --build-arg DEVTOOLSET_VERSION=${DEVTOOLSET_VERSION}"
         ;;
     xpu)
@@ -121,7 +121,8 @@ fi
 (
     set -x
 
-    if [ "$(uname -m)" != "s390x" ]; then
+    # Only activate this if in CI
+    if [ "$(uname -m)" != "s390x" ] && [ -v CI ]; then
         # TODO: Remove LimitNOFILE=1048576 patch once https://github.com/pytorch/test-infra/issues/5712
         # is resolved. This patch is required in order to fix timing out of Docker build on Amazon Linux 2023.
         sudo sed -i s/LimitNOFILE=infinity/LimitNOFILE=1048576/ /usr/lib/systemd/system/docker.service
@@ -139,7 +140,7 @@ fi
         "${TOPDIR}/.ci/docker/"
 )
 
-GITHUB_REF=${GITHUB_REF:-$(git symbolic-ref -q HEAD || git describe --tags --exact-match)}
+GITHUB_REF=${GITHUB_REF:-"dev")}
 GIT_BRANCH_NAME=${GITHUB_REF##*/}
 GIT_COMMIT_SHA=${GITHUB_SHA:-$(git rev-parse HEAD)}
 DOCKER_IMAGE_BRANCH_TAG=${DOCKER_IMAGE}-${GIT_BRANCH_NAME}

.ci/docker/manywheel/build_scripts/build_utils.sh

@@ -3,7 +3,7 @@
 # Script used only in CD pipeline
 
 OPENSSL_DOWNLOAD_URL=https://www.openssl.org/source/old/1.1.1/
-CURL_DOWNLOAD_URL=https://curl.askapache.com/download
+CURL_DOWNLOAD_URL=https://curl.se/download
 
 AUTOCONF_DOWNLOAD_URL=https://ftp.gnu.org/gnu/autoconf
 

.ci/docker/requirements-ci.txt

@@ -41,11 +41,14 @@ fbscribelogger==0.1.7
 #Pinned versions: 0.1.6
 #test that import:
 
-flatbuffers==2.0
+flatbuffers==2.0 ; platform_machine != "s390x"
 #Description: cross platform serialization library
 #Pinned versions: 2.0
 #test that import:
 
+flatbuffers ; platform_machine == "s390x"
+#Description: cross platform serialization library; Newer version is required on s390x for new python version
+
 hypothesis==5.35.1
 # Pin hypothesis to avoid flakiness: https://github.com/pytorch/pytorch/issues/31136
 #Description: advanced library for generating parametrized tests
@@ -90,10 +93,10 @@ librosa>=0.6.2 ; python_version < "3.11"
 #Pinned versions:
 #test that import:
 
-mypy==1.13.0
+mypy==1.14.0
 # Pin MyPy version because new errors are likely to appear with each release
 #Description: linter
-#Pinned versions: 1.10.0
+#Pinned versions: 1.14.0
 #test that import: test_typing.py, test_type_hints.py
 
 networkx==2.8.8
@@ -102,10 +105,10 @@ networkx==2.8.8
 #Pinned versions: 2.8.8
 #test that import: functorch
 
-#ninja
-#Description: build system.  Note that it install from
-#here breaks things so it is commented out
-#Pinned versions: 1.10.0.post1
+ninja==1.11.1.3
+#Description: build system. Used in some tests. Used in build to generate build
+#time tracing information
+#Pinned versions: 1.11.1.3
 #test that import: run_test.py, test_cpp_extensions_aot.py,test_determination.py
 
 numba==0.49.0 ; python_version < "3.9"
@@ -294,7 +297,7 @@ ghstack==0.8.0
 #Pinned versions: 0.8.0
 #test that import:
 
-jinja2==3.1.5
+jinja2==3.1.6
 #Description: jinja2 template engine
 #Pinned versions: 3.1.4
 #test that import:
@@ -329,7 +332,7 @@ lxml==5.3.0
 
 PyGithub==2.3.0
 
-sympy==1.13.1 ; python_version >= "3.9"
+sympy==1.13.3
 #Description: Required by coremltools, also pinned in .github/requirements/pip-requirements-macOS.txt
 #Pinned versions:
 #test that import:
@@ -339,7 +342,7 @@ onnx==1.17.0
 #Pinned versions:
 #test that import:
 
-onnxscript==0.1.0.dev20240817
+onnxscript==0.2.2
 #Description: Required by mypy and test_public_bindings.py when checking torch.onnx._internal
 #Pinned versions:
 #test that import:
@@ -353,7 +356,7 @@ parameterized==0.8.1
 #Pinned versions: 1.24.0
 #test that import: test_sac_estimator.py
 
-pwlf==2.2.1 ; python_version >= "3.8"
+pwlf==2.2.1
 #Description: required for testing torch/distributed/_tools/sac_estimator.py
 #Pinned versions: 2.2.1
 #test that import: test_sac_estimator.py
@@ -362,12 +365,17 @@ pwlf==2.2.1 ; python_version >= "3.8"
 # To build PyTorch itself
 astunparse
 PyYAML
+pyzstd
 setuptools
 
-ninja==1.11.1 ; platform_machine == "aarch64"
 scons==4.5.2 ; platform_machine == "aarch64"
 
-pulp==2.9.0 ; python_version >= "3.8"
+pulp==2.9.0
 #Description: required for testing ilp formulaiton under torch/distributed/_tools
 #Pinned versions: 2.9.0
 #test that import: test_sac_ilp.py
+
+dataclasses_json==0.6.7
+#Description: required for data pipeline and scripts under tools/stats
+#Pinned versions: 0.6.7
+#test that import:

.ci/docker/triton_version.txt

@@ -1 +1 @@
-3.2.0
+3.3.0

.ci/docker/ubuntu-cuda/Dockerfile

@@ -2,7 +2,7 @@ ARG UBUNTU_VERSION
 ARG CUDA_VERSION
 ARG IMAGE_NAME
 
-FROM ${IMAGE_NAME}
+FROM ${IMAGE_NAME} as base
 
 ARG UBUNTU_VERSION
 ARG CUDA_VERSION
@@ -50,13 +50,6 @@ RUN if [ -n "${PROTOBUF}" ]; then bash ./install_protobuf.sh; fi
 RUN rm install_protobuf.sh
 ENV INSTALLED_PROTOBUF ${PROTOBUF}
 
-# (optional) Install database packages like LMDB and LevelDB
-ARG DB
-COPY ./common/install_db.sh install_db.sh
-RUN if [ -n "${DB}" ]; then bash ./install_db.sh; fi
-RUN rm install_db.sh
-ENV INSTALLED_DB ${DB}
-
 # (optional) Install vision packages like OpenCV
 ARG VISION
 COPY ./common/install_vision.sh ./common/cache_vision_models.sh ./common/common_utils.sh ./
@@ -97,14 +90,20 @@ RUN if [ -n "${CMAKE_VERSION}" ]; then bash ./install_cmake.sh; fi
 RUN rm install_cmake.sh
 
 ARG TRITON
+
+FROM base as triton-builder
 # Install triton, this needs to be done before sccache because the latter will
 # try to reach out to S3, which docker build runners don't have access
 COPY ./common/install_triton.sh install_triton.sh
 COPY ./common/common_utils.sh common_utils.sh
 COPY ci_commit_pins/triton.txt triton.txt
 COPY triton_version.txt triton_version.txt
-RUN if [ -n "${TRITON}" ]; then bash ./install_triton.sh; fi
-RUN rm install_triton.sh common_utils.sh triton.txt triton_version.txt
+RUN bash ./install_triton.sh
+
+FROM base as final
+COPY --from=triton-builder /opt/triton /opt/triton
+RUN if [ -n "${TRITON}" ]; then pip install /opt/triton/*.whl; chown -R jenkins:jenkins /opt/conda; fi
+RUN rm -rf /opt/triton
 
 ARG HALIDE
 # Build and install halide
@@ -159,6 +158,16 @@ COPY ./common/install_cusparselt.sh install_cusparselt.sh
 RUN bash install_cusparselt.sh
 RUN rm install_cusparselt.sh
 
+# Install NCCL
+ARG CUDA_VERSION
+COPY ./common/install_nccl.sh install_nccl.sh
+COPY ./ci_commit_pins/nccl-cu* /ci_commit_pins/
+RUN bash install_nccl.sh
+RUN rm install_nccl.sh /ci_commit_pins/nccl-cu*
+ENV USE_SYSTEM_NCCL=1
+ENV NCCL_INCLUDE_DIR="/usr/local/cuda/include/"
+ENV NCCL_LIB_DIR="/usr/local/cuda/lib64/"
+
 # Install CUDSS
 ARG CUDA_VERSION
 COPY ./common/install_cudss.sh install_cudss.sh

.ci/docker/ubuntu-rocm/Dockerfile

@@ -14,21 +14,20 @@ ENV PYTORCH_ROCM_ARCH ${PYTORCH_ROCM_ARCH}
 COPY ./common/install_base.sh install_base.sh
 RUN bash ./install_base.sh && rm install_base.sh
 
-# Install clang
-ARG LLVMDEV
-ARG CLANG_VERSION
-COPY ./common/install_clang.sh install_clang.sh
-RUN bash ./install_clang.sh && rm install_clang.sh
-
 # Install user
 COPY ./common/install_user.sh install_user.sh
 RUN bash ./install_user.sh && rm install_user.sh
 
+# Install katex
+ARG KATEX
+COPY ./common/install_docs_reqs.sh install_docs_reqs.sh
+RUN bash ./install_docs_reqs.sh && rm install_docs_reqs.sh
+
 # Install conda and other packages (e.g., numpy, pytest)
 ARG ANACONDA_PYTHON_VERSION
-ARG CONDA_CMAKE
 ENV ANACONDA_PYTHON_VERSION=$ANACONDA_PYTHON_VERSION
 ENV PATH /opt/conda/envs/py_$ANACONDA_PYTHON_VERSION/bin:/opt/conda/bin:$PATH
+ARG CONDA_CMAKE
 COPY requirements-ci.txt /opt/conda/requirements-ci.txt
 COPY ./common/install_conda.sh install_conda.sh
 COPY ./common/common_utils.sh common_utils.sh
@@ -39,6 +38,11 @@ ARG GCC_VERSION
 COPY ./common/install_gcc.sh install_gcc.sh
 RUN bash ./install_gcc.sh && rm install_gcc.sh
 
+# Install clang
+ARG CLANG_VERSION
+COPY ./common/install_clang.sh install_clang.sh
+RUN bash ./install_clang.sh && rm install_clang.sh
+
 # (optional) Install protobuf for ONNX
 ARG PROTOBUF
 COPY ./common/install_protobuf.sh install_protobuf.sh
@@ -46,13 +50,6 @@ RUN if [ -n "${PROTOBUF}" ]; then bash ./install_protobuf.sh; fi
 RUN rm install_protobuf.sh
 ENV INSTALLED_PROTOBUF ${PROTOBUF}
 
-# (optional) Install database packages like LMDB and LevelDB
-ARG DB
-COPY ./common/install_db.sh install_db.sh
-RUN if [ -n "${DB}" ]; then bash ./install_db.sh; fi
-RUN rm install_db.sh
-ENV INSTALLED_DB ${DB}
-
 # (optional) Install vision packages like OpenCV
 ARG VISION
 COPY ./common/install_vision.sh ./common/cache_vision_models.sh ./common/common_utils.sh ./
@@ -66,7 +63,7 @@ COPY ./common/install_rocm.sh install_rocm.sh
 RUN bash ./install_rocm.sh
 RUN rm install_rocm.sh
 COPY ./common/install_rocm_magma.sh install_rocm_magma.sh
-RUN bash ./install_rocm_magma.sh
+RUN bash ./install_rocm_magma.sh ${ROCM_VERSION}
 RUN rm install_rocm_magma.sh
 ADD ./common/install_miopen.sh install_miopen.sh
 RUN bash ./install_miopen.sh ${ROCM_VERSION} && rm install_miopen.sh
@@ -85,6 +82,32 @@ COPY ./common/install_amdsmi.sh install_amdsmi.sh
 RUN bash ./install_amdsmi.sh
 RUN rm install_amdsmi.sh
 
+# (optional) Install UCC
+ARG UCX_COMMIT
+ARG UCC_COMMIT
+ENV UCX_COMMIT $UCX_COMMIT
+ENV UCC_COMMIT $UCC_COMMIT
+ENV UCX_HOME /usr
+ENV UCC_HOME /usr
+ADD ./common/install_ucc.sh install_ucc.sh
+RUN if [ -n "${UCX_COMMIT}" ] && [ -n "${UCC_COMMIT}" ]; then bash ./install_ucc.sh; fi
+RUN rm install_ucc.sh
+
+COPY ./common/install_openssl.sh install_openssl.sh
+ENV OPENSSL_ROOT_DIR /opt/openssl
+RUN bash ./install_openssl.sh
+ENV OPENSSL_DIR /opt/openssl
+
+ARG INDUCTOR_BENCHMARKS
+ARG ANACONDA_PYTHON_VERSION
+ENV ANACONDA_PYTHON_VERSION=$ANACONDA_PYTHON_VERSION
+COPY ./common/install_inductor_benchmark_deps.sh install_inductor_benchmark_deps.sh
+COPY ./common/common_utils.sh common_utils.sh
+COPY ci_commit_pins/huggingface.txt huggingface.txt
+COPY ci_commit_pins/timm.txt timm.txt
+RUN if [ -n "${INDUCTOR_BENCHMARKS}" ]; then bash ./install_inductor_benchmark_deps.sh; fi
+RUN rm install_inductor_benchmark_deps.sh common_utils.sh timm.txt huggingface.txt
+
 # (optional) Install non-default CMake version
 ARG CMAKE_VERSION
 COPY ./common/install_cmake.sh install_cmake.sh
@@ -107,17 +130,17 @@ COPY triton_version.txt triton_version.txt
 RUN if [ -n "${TRITON}" ]; then bash ./install_triton.sh; fi
 RUN rm install_triton.sh common_utils.sh triton.txt triton_version.txt
 
-# This is needed by sccache
-COPY ./common/install_openssl.sh install_openssl.sh
-ENV OPENSSL_ROOT_DIR /opt/openssl
-RUN bash ./install_openssl.sh
-ENV OPENSSL_DIR /opt/openssl
 
 # Install ccache/sccache (do this last, so we get priority in PATH)
 COPY ./common/install_cache.sh install_cache.sh
 ENV PATH /opt/cache/bin:$PATH
 RUN bash ./install_cache.sh && rm install_cache.sh
 
+# Install Open MPI for ROCm
+COPY ./common/install_openmpi.sh install_openmpi.sh
+RUN if [ -n "${CUDA_VERSION}" ]; then bash install_openmpi.sh; fi
+RUN rm install_openmpi.sh
+
 # Include BUILD_ENVIRONMENT environment variable in image
 ARG BUILD_ENVIRONMENT
 ENV BUILD_ENVIRONMENT ${BUILD_ENVIRONMENT}

.ci/docker/ubuntu-xpu/Dockerfile

@@ -77,13 +77,6 @@ COPY triton_version.txt triton_version.txt
 RUN if [ -n "${TRITON}" ]; then bash ./install_triton.sh; fi
 RUN rm install_triton.sh common_utils.sh triton-xpu.txt triton_version.txt
 
-# (optional) Install database packages like LMDB and LevelDB
-ARG DB
-COPY ./common/install_db.sh install_db.sh
-RUN if [ -n "${DB}" ]; then bash ./install_db.sh; fi
-RUN rm install_db.sh
-ENV INSTALLED_DB ${DB}
-
 # (optional) Install vision packages like OpenCV
 ARG VISION
 COPY ./common/install_vision.sh ./common/cache_vision_models.sh ./common/common_utils.sh ./

.ci/docker/ubuntu/Dockerfile

@@ -1,6 +1,6 @@
 ARG UBUNTU_VERSION
 
-FROM ubuntu:${UBUNTU_VERSION}
+FROM ubuntu:${UBUNTU_VERSION} as base
 
 ARG UBUNTU_VERSION
 
@@ -52,9 +52,16 @@ RUN  bash ./install_lcov.sh && rm install_lcov.sh
 # Install cuda and cudnn
 ARG CUDA_VERSION
 COPY ./common/install_cuda.sh install_cuda.sh
-RUN bash ./install_cuda.sh ${CUDA_VERSION} && rm install_cuda.sh
+COPY ./common/install_nccl.sh install_nccl.sh
+COPY ./ci_commit_pins/nccl-cu* /ci_commit_pins/
+RUN bash ./install_cuda.sh ${CUDA_VERSION} && rm install_cuda.sh install_nccl.sh /ci_commit_pins/nccl-cu*
 ENV DESIRED_CUDA ${CUDA_VERSION}
 ENV PATH /usr/local/nvidia/bin:/usr/local/cuda/bin:$PATH
+# No effect if cuda not installed
+ENV USE_SYSTEM_NCCL=1
+ENV NCCL_INCLUDE_DIR="/usr/local/cuda/include/"
+ENV NCCL_LIB_DIR="/usr/local/cuda/lib64/"
+
 
 # (optional) Install UCC
 ARG UCX_COMMIT
@@ -74,13 +81,6 @@ RUN if [ -n "${PROTOBUF}" ]; then bash ./install_protobuf.sh; fi
 RUN rm install_protobuf.sh
 ENV INSTALLED_PROTOBUF ${PROTOBUF}
 
-# (optional) Install database packages like LMDB and LevelDB
-ARG DB
-COPY ./common/install_db.sh install_db.sh
-RUN if [ -n "${DB}" ]; then bash ./install_db.sh; fi
-RUN rm install_db.sh
-ENV INSTALLED_DB ${DB}
-
 # (optional) Install vision packages like OpenCV
 ARG VISION
 COPY ./common/install_vision.sh ./common/cache_vision_models.sh ./common/common_utils.sh ./
@@ -88,18 +88,6 @@ RUN if [ -n "${VISION}" ]; then bash ./install_vision.sh; fi
 RUN rm install_vision.sh cache_vision_models.sh common_utils.sh
 ENV INSTALLED_VISION ${VISION}
 
-# (optional) Install Vulkan SDK
-ARG VULKAN_SDK_VERSION
-COPY ./common/install_vulkan_sdk.sh install_vulkan_sdk.sh
-RUN if [ -n "${VULKAN_SDK_VERSION}" ]; then bash ./install_vulkan_sdk.sh; fi
-RUN rm install_vulkan_sdk.sh
-
-# (optional) Install swiftshader
-ARG SWIFTSHADER
-COPY ./common/install_swiftshader.sh install_swiftshader.sh
-RUN if [ -n "${SWIFTSHADER}" ]; then bash ./install_swiftshader.sh; fi
-RUN rm install_swiftshader.sh
-
 # (optional) Install non-default CMake version
 ARG CMAKE_VERSION
 COPY ./common/install_cmake.sh install_cmake.sh
@@ -127,20 +115,21 @@ RUN if [ -n "${INDUCTOR_BENCHMARKS}" ]; then bash ./install_inductor_benchmark_d
 RUN rm install_inductor_benchmark_deps.sh common_utils.sh timm.txt huggingface.txt
 
 ARG TRITON
-# Install triton, this needs to be done before sccache because the latter will
-# try to reach out to S3, which docker build runners don't have access
+ARG TRITON_CPU
+
+# Create a separate stage for building Triton and Triton-CPU.  install_triton
+# will check for the presence of env vars
+FROM base as triton-builder
 COPY ./common/install_triton.sh install_triton.sh
 COPY ./common/common_utils.sh common_utils.sh
 COPY ci_commit_pins/triton.txt triton.txt
-RUN if [ -n "${TRITON}" ]; then bash ./install_triton.sh; fi
-RUN rm install_triton.sh common_utils.sh triton.txt
-
-ARG TRITON_CPU
-COPY ./common/install_triton.sh install_triton.sh
-COPY ./common/common_utils.sh common_utils.sh
 COPY ci_commit_pins/triton-cpu.txt triton-cpu.txt
-RUN if [ -n "${TRITON_CPU}" ]; then bash ./install_triton.sh; fi
-RUN rm install_triton.sh common_utils.sh triton-cpu.txt
+RUN bash ./install_triton.sh
+
+FROM base as final
+COPY --from=triton-builder /opt/triton /opt/triton
+RUN if [ -n "${TRITON}" ] || [ -n "${TRITON_CPU}" ]; then pip install /opt/triton/*.whl; chown -R jenkins:jenkins /opt/conda; fi
+RUN rm -rf /opt/triton
 
 ARG EXECUTORCH
 # Build and install executorch

.ci/magma-rocm/.gitignore

@@ -0,0 +1,2 @@
+output/
+magma-rocm*/

.ci/magma-rocm/Makefile

@@ -0,0 +1,35 @@
+SHELL=/usr/bin/env bash
+
+DOCKER_CMD ?= docker
+DESIRED_ROCM ?= 6.3
+DESIRED_ROCM_SHORT = $(subst .,,$(DESIRED_ROCM))
+PACKAGE_NAME = magma-rocm
+# inherit this from underlying docker image, do not pass this env var to docker
+#PYTORCH_ROCM_ARCH ?= gfx900;gfx906;gfx908;gfx90a;gfx942;gfx1030;gfx1100;gfx1101;gfx1102;gfx1200;gfx1201
+
+DOCKER_RUN = set -eou pipefail; ${DOCKER_CMD} run --rm -i \
+	-v $(shell git rev-parse --show-toplevel)/.ci:/builder \
+	-w /builder \
+	-e PACKAGE_NAME=${PACKAGE_NAME}${DESIRED_ROCM_SHORT} \
+	-e DESIRED_ROCM=${DESIRED_ROCM} \
+	"pytorch/manylinux2_28-builder:rocm${DESIRED_ROCM}-main" \
+	magma-rocm/build_magma.sh
+
+.PHONY: all
+all: magma-rocm63
+all: magma-rocm624
+
+.PHONY:
+clean:
+	$(RM) -r magma-*
+	$(RM) -r output
+
+.PHONY: magma-rocm63
+magma-rocm63: DESIRED_ROCM := 6.3
+magma-rocm63:
+	$(DOCKER_RUN)
+
+.PHONY: magma-rocm624
+magma-rocm624: DESIRED_ROCM := 6.2.4
+magma-rocm624:
+	$(DOCKER_RUN)

.ci/magma-rocm/README.md

@@ -0,0 +1,48 @@
+# Magma ROCm
+
+This folder contains the scripts and configurations to build libmagma.so, linked for various versions of ROCm.
+
+## Building
+
+Look in the `Makefile` for available targets to build. To build any target, for example `magma-rocm63`, run
+
+```
+# Using `docker`
+make magma-rocm63
+
+# Using `podman`
+DOCKER_CMD=podman make magma-rocm63
+```
+
+This spawns a `pytorch/manylinux-rocm<version>` docker image, which has the required `devtoolset` and ROCm versions installed.
+Within the docker image, it runs `build_magma.sh` with the correct environment variables set, which package the necessary files
+into a tarball, with the following structure:
+
+```
+.
+├── include       # header files
+├── lib           # libmagma.so
+├── info
+│   ├── licenses  # license file
+│   └── recipe    # build script
+```
+
+More specifically, `build_magma.sh` copies over the relevant files from the `package_files` directory depending on the ROCm version.
+Outputted binaries should be in the `output` folder.
+
+
+## Pushing
+
+Packages can be uploaded to an S3 bucket using:
+
+```
+aws s3 cp output/*/magma-cuda*.bz2 <bucket-with-path>
+```
+
+If you do not have upload permissions, please ping @seemethere or @soumith to gain access
+
+## New versions
+
+New ROCm versions can be added by creating a new make target with the next desired version. For ROCm version N.n, the target should be named `magma-rocmNn`.
+
+Make sure to edit the appropriate environment variables (e.g., DESIRED_ROCM) in the `Makefile` accordingly. Remember also to check `build_magma.sh` to ensure the logic for copying over the files remains correct.

.ci/magma-rocm/build_magma.sh

@@ -0,0 +1,42 @@
+#!/usr/bin/env bash
+
+set -eou pipefail
+
+# Environment variables
+# The script expects DESIRED_CUDA and PACKAGE_NAME to be set
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+
+# Version 2.7.2 + ROCm related updates
+MAGMA_VERSION=a1625ff4d9bc362906bd01f805dbbe12612953f6
+
+# Folders for the build
+PACKAGE_FILES=${ROOT_DIR}/magma-rocm/package_files # metadata
+PACKAGE_DIR=${ROOT_DIR}/magma-rocm/${PACKAGE_NAME} # build workspace
+PACKAGE_OUTPUT=${ROOT_DIR}/magma-rocm/output # where tarballs are stored
+PACKAGE_BUILD=${PACKAGE_DIR} # where the content of the tarball is prepared
+PACKAGE_RECIPE=${PACKAGE_BUILD}/info/recipe
+PACKAGE_LICENSE=${PACKAGE_BUILD}/info/licenses
+mkdir -p ${PACKAGE_DIR} ${PACKAGE_OUTPUT}/linux-64 ${PACKAGE_BUILD} ${PACKAGE_RECIPE} ${PACKAGE_LICENSE}
+
+# Fetch magma sources and verify checksum
+pushd ${PACKAGE_DIR}
+git clone https://bitbucket.org/icl/magma.git
+pushd magma
+git checkout ${MAGMA_VERSION}
+popd
+popd
+
+# build
+pushd ${PACKAGE_DIR}/magma
+# The build.sh script expects to be executed from the sources root folder
+INSTALL_DIR=${PACKAGE_BUILD} ${PACKAGE_FILES}/build.sh
+popd
+
+# Package recipe, license and tarball
+# Folder and package name are backward compatible for the build workflow
+cp ${PACKAGE_FILES}/build.sh ${PACKAGE_RECIPE}/build.sh
+cp ${PACKAGE_DIR}/magma/COPYRIGHT ${PACKAGE_LICENSE}/COPYRIGHT
+pushd ${PACKAGE_BUILD}
+tar cjf ${PACKAGE_OUTPUT}/linux-64/${PACKAGE_NAME}-${MAGMA_VERSION}-1.tar.bz2 include lib info
+echo Built in ${PACKAGE_OUTPUT}/linux-64/${PACKAGE_NAME}-${MAGMA_VERSION}-1.tar.bz2
+popd

.ci/magma-rocm/package_files/build.sh

@@ -0,0 +1,38 @@
+# Magma build scripts need `python`
+ln -sf /usr/bin/python3 /usr/bin/python
+
+ID=$(grep -oP '(?<=^ID=).+' /etc/os-release | tr -d '"')
+case "$ID" in
+  almalinux)
+    yum install -y gcc-gfortran
+    ;;
+  *)
+    echo "No preinstalls to build magma..."
+    ;;
+esac
+
+MKLROOT=${MKLROOT:-/opt/conda/envs/py_$ANACONDA_PYTHON_VERSION}
+
+cp make.inc-examples/make.inc.hip-gcc-mkl make.inc
+echo 'LIBDIR += -L$(MKLROOT)/lib' >> make.inc
+if [[ -f "${MKLROOT}/lib/libmkl_core.a" ]]; then
+    echo 'LIB = -Wl,--start-group -lmkl_gf_lp64 -lmkl_gnu_thread -lmkl_core -Wl,--end-group -lpthread -lstdc++ -lm -lgomp -lhipblas -lhipsparse' >> make.inc
+fi
+echo 'LIB += -Wl,--enable-new-dtags -Wl,--rpath,/opt/rocm/lib -Wl,--rpath,$(MKLROOT)/lib -Wl,--rpath,/opt/rocm/magma/lib -ldl' >> make.inc
+echo 'DEVCCFLAGS += --gpu-max-threads-per-block=256' >> make.inc
+export PATH="${PATH}:/opt/rocm/bin"
+if [[ -n "$PYTORCH_ROCM_ARCH" ]]; then
+  amdgpu_targets=`echo $PYTORCH_ROCM_ARCH | sed 's/;/ /g'`
+else
+  amdgpu_targets=`rocm_agent_enumerator | grep -v gfx000 | sort -u | xargs`
+fi
+for arch in $amdgpu_targets; do
+  echo "DEVCCFLAGS += --offload-arch=$arch" >> make.inc
+done
+# hipcc with openmp flag may cause isnan() on __device__ not to be found; depending on context, compiler may attempt to match with host definition
+sed -i 's/^FOPENMP/#FOPENMP/g' make.inc
+make -f make.gen.hipMAGMA -j $(nproc)
+LANG=C.UTF-8 make lib/libmagma.so -j $(nproc) MKLROOT="${MKLROOT}"
+make testing/testing_dgemm -j $(nproc) MKLROOT="${MKLROOT}"
+cp -R lib ${INSTALL_DIR}
+cp -R include ${INSTALL_DIR}

.ci/magma/Makefile

@@ -12,13 +12,13 @@ DOCKER_RUN = set -eou pipefail; ${DOCKER_CMD} run --rm -i \
 	-e PACKAGE_NAME=${PACKAGE_NAME}${DESIRED_CUDA_SHORT} \
 	-e DESIRED_CUDA=${DESIRED_CUDA} \
 	-e CUDA_ARCH_LIST="${CUDA_ARCH_LIST}" \
-	"pytorch/manylinux-builder:cuda${DESIRED_CUDA}-main" \
+	"pytorch/manylinux2_28-builder:cuda${DESIRED_CUDA}-main" \
 	magma/build_magma.sh
 
 .PHONY: all
+all: magma-cuda128
 all: magma-cuda126
 all: magma-cuda124
-all: magma-cuda121
 all: magma-cuda118
 
 .PHONY:
@@ -26,6 +26,12 @@ clean:
 	$(RM) -r magma-*
 	$(RM) -r output
 
+.PHONY: magma-cuda128
+magma-cuda128: DESIRED_CUDA := 12.8
+magma-cuda128: CUDA_ARCH_LIST += -gencode arch=compute_100,code=sm_100 -gencode arch=compute_120,code=sm_120
+magma-cuda128:
+	$(DOCKER_RUN)
+
 .PHONY: magma-cuda126
 magma-cuda126: DESIRED_CUDA := 12.6
 magma-cuda126:
@@ -36,11 +42,6 @@ magma-cuda124: DESIRED_CUDA := 12.4
 magma-cuda124:
 	$(DOCKER_RUN)
 
-.PHONY: magma-cuda121
-magma-cuda121: DESIRED_CUDA := 12.1
-magma-cuda121:
-	$(DOCKER_RUN)
-
 .PHONY: magma-cuda118
 magma-cuda118: DESIRED_CUDA := 11.8
 magma-cuda118: CUDA_ARCH_LIST += -gencode arch=compute_37,code=sm_37

.ci/manywheel/build_common.sh

@@ -111,12 +111,6 @@ case ${DESIRED_PYTHON} in
     ;;
 esac
 
-if [[ "$DESIRED_DEVTOOLSET" == *"cxx11-abi"* ]]; then
-    export _GLIBCXX_USE_CXX11_ABI=1
-else
-    export _GLIBCXX_USE_CXX11_ABI=0
-fi
-
 if [[ "$DESIRED_CUDA" == *"rocm"* ]]; then
     echo "Calling build_amd.py at $(date)"
     python tools/amd_build/build_amd.py
@@ -209,12 +203,6 @@ if [[ -n "$BUILD_PYTHONLESS" ]]; then
 
     mkdir -p /tmp/$LIBTORCH_HOUSE_DIR
 
-    if [[ "$DESIRED_DEVTOOLSET" == *"cxx11-abi"* ]]; then
-        LIBTORCH_ABI="cxx11-abi-"
-    else
-        LIBTORCH_ABI=
-    fi
-
     zip -rq /tmp/$LIBTORCH_HOUSE_DIR/libtorch-$LIBTORCH_ABI$LIBTORCH_VARIANT-$PYTORCH_BUILD_VERSION.zip libtorch
     cp /tmp/$LIBTORCH_HOUSE_DIR/libtorch-$LIBTORCH_ABI$LIBTORCH_VARIANT-$PYTORCH_BUILD_VERSION.zip \
        /tmp/$LIBTORCH_HOUSE_DIR/libtorch-$LIBTORCH_ABI$LIBTORCH_VARIANT-latest.zip

.ci/manywheel/build_cuda.sh

@@ -14,6 +14,7 @@ export USE_CUDA_STATIC_LINK=1
 export INSTALL_TEST=0 # dont install test binaries into site-packages
 export USE_CUPTI_SO=0
 export USE_CUSPARSELT=${USE_CUSPARSELT:-1} # Enable if not disabled by libtorch build
+export USE_CUFILE=${USE_CUFILE:-1}
 
 # Keep an array of cmake variables to add to
 if [[ -z "$CMAKE_ARGS" ]]; then
@@ -52,8 +53,12 @@ cuda_version_nodot=$(echo $CUDA_VERSION | tr -d '.')
 
 TORCH_CUDA_ARCH_LIST="5.0;6.0;7.0;7.5;8.0;8.6"
 case ${CUDA_VERSION} in
+    12.8)
+        TORCH_CUDA_ARCH_LIST="7.5;8.0;8.6;9.0;10.0;12.0+PTX" #removing sm_50-sm_70 as these architectures are deprecated in CUDA 12.8 and will be removed in future releases
+        EXTRA_CAFFE2_CMAKE_FLAGS+=("-DATEN_NO_TEST=ON")
+        ;;
     12.6)
-        TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST};9.0+PTX"
+        TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST};9.0"
         EXTRA_CAFFE2_CMAKE_FLAGS+=("-DATEN_NO_TEST=ON")
         ;;
     12.4)
@@ -114,7 +119,16 @@ if [[ $USE_CUSPARSELT == "1" && $CUDA_VERSION == "11.8" ]]; then
         )
 fi
 
-if [[ $CUDA_VERSION == "12.4" || $CUDA_VERSION == "12.6" ]]; then
+
+# Turn USE_CUFILE off for CUDA 11.8, 12.4 since nvidia-cufile-cu11 and 1.9.0.20 are
+# not available in PYPI
+if [[ $CUDA_VERSION == "11.8" || $CUDA_VERSION == "12.4" ]]; then
+    export USE_CUFILE=0
+fi
+
+
+# CUDA_VERSION 12.4, 12.6, 12.8
+if [[ $CUDA_VERSION == 12* ]]; then
     export USE_STATIC_CUDNN=0
     # Try parallelizing nvcc as well
     export TORCH_NVCC_FLAGS="-Xfatbin -compress-all --threads 2"
@@ -155,6 +169,16 @@ if [[ $CUDA_VERSION == "12.4" || $CUDA_VERSION == "12.6" ]]; then
             "libnvrtc.so.12"
             "libnvrtc-builtins.so"
         )
+        if [[ $USE_CUFILE == 1 ]]; then
+            DEPS_LIST+=(
+                "/usr/local/cuda/lib64/libcufile.so.0"
+                "/usr/local/cuda/lib64/libcufile_rdma.so.1"
+            )
+            DEPS_SONAME+=(
+                "libcufile.so.0"
+                "libcufile_rdma.so.1"
+            )
+        fi
     else
         echo "Using nvidia libs from pypi."
         CUDA_RPATHS=(
@@ -171,6 +195,11 @@ if [[ $CUDA_VERSION == "12.4" || $CUDA_VERSION == "12.6" ]]; then
             '$ORIGIN/../../nvidia/nccl/lib'
             '$ORIGIN/../../nvidia/nvtx/lib'
         )
+        if [[ $USE_CUFILE == 1 ]]; then
+            CUDA_RPATHS+=(
+                '$ORIGIN/../../nvidia/cufile/lib'
+            )
+        fi
         CUDA_RPATHS=$(IFS=: ; echo "${CUDA_RPATHS[*]}")
         export C_SO_RPATH=$CUDA_RPATHS':$ORIGIN:$ORIGIN/lib'
         export LIB_SO_RPATH=$CUDA_RPATHS':$ORIGIN'

.ci/manywheel/build_libtorch.sh

@@ -95,12 +95,6 @@ python setup.py clean
 retry pip install -qr requirements.txt
 retry pip install -q numpy==2.0.1
 
-if [[ "$DESIRED_DEVTOOLSET" == *"cxx11-abi"* ]]; then
-    export _GLIBCXX_USE_CXX11_ABI=1
-else
-    export _GLIBCXX_USE_CXX11_ABI=0
-fi
-
 if [[ "$DESIRED_CUDA" == *"rocm"* ]]; then
     echo "Calling build_amd.py at $(date)"
     python tools/amd_build/build_amd.py
@@ -169,12 +163,6 @@ fi
 
 )
 
-if [[ "$DESIRED_DEVTOOLSET" == *"cxx11-abi"* ]]; then
-    LIBTORCH_ABI="cxx11-abi-"
-else
-    LIBTORCH_ABI=
-fi
-
 (
     set -x
 

.ci/pytorch/build.sh

@@ -35,7 +35,7 @@ if [[ "$BUILD_ENVIRONMENT" == *cuda* ]]; then
 fi
 
 if [[ "$BUILD_ENVIRONMENT" == *cuda11* ]]; then
-  if [[ "$BUILD_ENVIRONMENT" != *cuda11.3* && "$BUILD_ENVIRONMENT" != *clang* ]]; then
+  if [[ "$BUILD_ENVIRONMENT" != *clang* ]]; then
     # TODO: there is a linking issue when building with UCC using clang,
     # disable it for now and to be fix later.
     # TODO: disable UCC temporarily to enable CUDA 12.1 in CI
@@ -173,6 +173,7 @@ if [[ "$BUILD_ENVIRONMENT" == *xpu* ]]; then
   source /opt/intel/oneapi/compiler/latest/env/vars.sh
   # XPU kineto feature dependencies are not fully ready, disable kineto build as temp WA
   export USE_KINETO=0
+  export TORCH_XPU_ARCH_LIST=pvc
 fi
 
 # sccache will fail for CUDA builds if all cores are used for compiling
@@ -191,7 +192,7 @@ fi
 
 # We only build FlashAttention files for CUDA 8.0+, and they require large amounts of
 # memory to build and will OOM
-if [[ "$BUILD_ENVIRONMENT" == *cuda* ]] && [[ 1 -eq $(echo "${TORCH_CUDA_ARCH_LIST} >= 8.0" | bc) ]]; then
+if [[ "$BUILD_ENVIRONMENT" == *cuda* ]] && [[ 1 -eq $(echo "${TORCH_CUDA_ARCH_LIST} >= 8.0" | bc) ]] && [ -z "$MAX_JOBS_OVERRIDE" ]; then
   echo "WARNING: FlashAttention files require large amounts of memory to build and will OOM"
   echo "Setting MAX_JOBS=(nproc-2)/3 to reduce memory usage"
   export MAX_JOBS="$(( $(nproc --ignore=2) / 3 ))"
@@ -276,10 +277,8 @@ else
     # or building non-XLA tests.
     if [[ "$BUILD_ENVIRONMENT" != *rocm*  &&
           "$BUILD_ENVIRONMENT" != *xla* ]]; then
-      if [[ "$BUILD_ENVIRONMENT" != *py3.8* ]]; then
-        # Install numpy-2.0.2 for builds which are backward compatible with 1.X
-        python -mpip install numpy==2.0.2
-      fi
+      # Install numpy-2.0.2 for builds which are backward compatible with 1.X
+      python -mpip install numpy==2.0.2
 
       WERROR=1 python setup.py clean
 
@@ -377,8 +376,10 @@ else
     # This is an attempt to mitigate flaky libtorch build OOM error. By default, the build parallelization
     # is set to be the number of CPU minus 2. So, let's try a more conservative value here. A 4xlarge has
     # 16 CPUs
-    MAX_JOBS=$(nproc --ignore=4)
-    export MAX_JOBS
+    if [ -z "$MAX_JOBS_OVERRIDE" ]; then
+      MAX_JOBS=$(nproc --ignore=4)
+      export MAX_JOBS
+    fi
 
     # NB: Install outside of source directory (at the same level as the root
     # pytorch folder) so that it doesn't get cleaned away prior to docker push.

.ci/pytorch/check_binary.sh

@@ -59,78 +59,16 @@ else
   export install_root="$(dirname $(which python))/../lib/python${py_dot}/site-packages/torch/"
 fi
 
-###############################################################################
-# Setup XPU ENV
-###############################################################################
-if [[ "$DESIRED_CUDA" == 'xpu' ]]; then
-  set +u
-  # Refer https://www.intel.com/content/www/us/en/developer/articles/tool/pytorch-prerequisites-for-intel-gpus.html
-  source /opt/intel/oneapi/compiler/latest/env/vars.sh
-  source /opt/intel/oneapi/pti/latest/env/vars.sh
-fi
-
 ###############################################################################
 # Check GCC ABI
 ###############################################################################
 
-# NOTE [ Building libtorch with old vs. new gcc ABI ]
-#
-# Packages built with one version of ABI could not be linked against by client
-# C++ libraries that were compiled using the other version of ABI. Since both
-# gcc ABIs are still common in the wild, we need to support both ABIs. Currently:
-#
-# - All the nightlies built on CentOS 7 + devtoolset7 use the old gcc ABI.
-# - All the nightlies built on Ubuntu 16.04 + gcc 5.4 use the new gcc ABI.
+# NOTE: As of https://github.com/pytorch/pytorch/issues/126551 we only produce
+#       wheels with cxx11-abi
 
 echo "Checking that the gcc ABI is what we expect"
 if [[ "$(uname)" != 'Darwin' ]]; then
-  function is_expected() {
-    if [[ "$DESIRED_DEVTOOLSET" == *"cxx11-abi"* || "$DESIRED_CUDA" == *"rocm"* ]]; then
-      if [[ "$1" -gt 0 || "$1" == "ON " ]]; then
-        echo 1
-      fi
-    else
-      if [[ -z "$1" || "$1" == 0 || "$1" == "OFF" ]]; then
-        echo 1
-      fi
-    fi
-  }
-
-  # First we check that the env var in TorchConfig.cmake is correct
-
-  # We search for D_GLIBCXX_USE_CXX11_ABI=1 in torch/TorchConfig.cmake
-  torch_config="${install_root}/share/cmake/Torch/TorchConfig.cmake"
-  if [[ ! -f "$torch_config" ]]; then
-    echo "No TorchConfig.cmake found!"
-    ls -lah "$install_root/share/cmake/Torch"
-    exit 1
-  fi
-  echo "Checking the TorchConfig.cmake"
-  cat "$torch_config"
-
-  # The sed call below is
-  #   don't print lines by default (only print the line we want)
-  # -n
-  #   execute the following expression
-  # e
-  #   replace lines that match with the first capture group and print
-  # s/.*D_GLIBCXX_USE_CXX11_ABI=\(.\)".*/\1/p
-  #   any characters, D_GLIBCXX_USE_CXX11_ABI=, exactly one any character, a
-  #   quote, any characters
-  #   Note the exactly one single character after the '='. In the case that the
-  #     variable is not set the '=' will be followed by a '"' immediately and the
-  #     line will fail the match and nothing will be printed; this is what we
-  #     want.  Otherwise it will capture the 0 or 1 after the '='.
-  # /.*D_GLIBCXX_USE_CXX11_ABI=\(.\)".*/
-  #   replace the matched line with the capture group and print
-  # /\1/p
-  actual_gcc_abi="$(sed -ne 's/.*D_GLIBCXX_USE_CXX11_ABI=\(.\)".*/\1/p' < "$torch_config")"
-  if [[ "$(is_expected "$actual_gcc_abi")" != 1 ]]; then
-    echo "gcc ABI $actual_gcc_abi not as expected."
-    exit 1
-  fi
-
-  # We also check that there are [not] cxx11 symbols in libtorch
+  # We also check that there are cxx11 symbols in libtorch
   #
   echo "Checking that symbols in libtorch.so have the right gcc abi"
   python3 "$(dirname ${BASH_SOURCE[0]})/smoke_test/check_binary_symbols.py"
@@ -208,35 +146,11 @@ setup_link_flags () {
 
 TEST_CODE_DIR="$(dirname $(realpath ${BASH_SOURCE[0]}))/test_example_code"
 build_and_run_example_cpp () {
-  if [[ "$DESIRED_DEVTOOLSET" == *"cxx11-abi"* ]]; then
-    GLIBCXX_USE_CXX11_ABI=1
-  else
-    GLIBCXX_USE_CXX11_ABI=0
-  fi
   setup_link_flags
-  g++ ${TEST_CODE_DIR}/$1.cpp -I${install_root}/include -I${install_root}/include/torch/csrc/api/include -D_GLIBCXX_USE_CXX11_ABI=$GLIBCXX_USE_CXX11_ABI -std=gnu++17 -L${install_root}/lib ${REF_LIB} ${ADDITIONAL_LINKER_FLAGS} -ltorch $TORCH_CPU_LINK_FLAGS $TORCH_CUDA_LINK_FLAGS $C10_LINK_FLAGS -o $1
+  g++ ${TEST_CODE_DIR}/$1.cpp -I${install_root}/include -I${install_root}/include/torch/csrc/api/include -std=gnu++17 -L${install_root}/lib ${REF_LIB} ${ADDITIONAL_LINKER_FLAGS} -ltorch $TORCH_CPU_LINK_FLAGS $TORCH_CUDA_LINK_FLAGS $C10_LINK_FLAGS -o $1
   ./$1
 }
 
-build_example_cpp_with_incorrect_abi () {
-  if [[ "$DESIRED_DEVTOOLSET" == *"cxx11-abi"* ]]; then
-    GLIBCXX_USE_CXX11_ABI=0
-  else
-    GLIBCXX_USE_CXX11_ABI=1
-  fi
-  set +e
-  setup_link_flags
-  g++ ${TEST_CODE_DIR}/$1.cpp -I${install_root}/include -I${install_root}/include/torch/csrc/api/include -D_GLIBCXX_USE_CXX11_ABI=$GLIBCXX_USE_CXX11_ABI -std=gnu++17 -L${install_root}/lib ${REF_LIB} ${ADDITIONAL_LINKER_FLAGS} -ltorch $TORCH_CPU_LINK_FLAGS $TORCH_CUDA_LINK_FLAGS $C10_LINK_FLAGS -o $1
-  ERRCODE=$?
-  set -e
-  if [ "$ERRCODE" -eq "0" ]; then
-    echo "Building example with incorrect ABI didn't throw error. Aborting."
-    exit 1
-  else
-    echo "Building example with incorrect ABI throws expected error. Proceeding."
-  fi
-}
-
 ###############################################################################
 # Check simple Python/C++ calls
 ###############################################################################
@@ -246,11 +160,6 @@ if [[ "$PACKAGE_TYPE" == 'libtorch' ]]; then
     export LD_LIBRARY_PATH=/usr/local/cuda/lib64
   fi
   build_and_run_example_cpp simple-torch-test
-  # `_GLIBCXX_USE_CXX11_ABI` is always ignored by gcc in devtoolset7, so we test
-  # the expected failure case for Ubuntu 16.04 + gcc 5.4 only.
-  if [[ "$DESIRED_DEVTOOLSET" == *"cxx11-abi"* ]]; then
-    build_example_cpp_with_incorrect_abi simple-torch-test
-  fi
 else
   pushd /tmp
   python -c 'import torch'
@@ -385,10 +294,19 @@ except RuntimeError as e:
 fi
 
 ###############################################################################
-# Check for C++ ABI compatibility between gcc7 and gcc9 compiled binaries
+# Check for C++ ABI compatibility to GCC-11
 ###############################################################################
-if [[ "$(uname)" == 'Linux' && ("$PACKAGE_TYPE" == 'conda' || "$PACKAGE_TYPE" == 'manywheel')]]; then
+if [[ "$(uname)" == 'Linux' &&  "$PACKAGE_TYPE" == 'manywheel' ]]; then
   pushd /tmp
-  python -c "import torch; exit(0 if torch.compiled_with_cxx11_abi() else (0 if torch._C._PYBIND11_BUILD_ABI == '_cxxabi1011' else 1))"
+  # Per https://gcc.gnu.org/onlinedocs/gcc/C_002b_002b-Dialect-Options.html gcc-11 is ABI16
+  # Though manylinux_2.28 should have been build with gcc-14, per
+  # https://github.com/pypa/manylinux?tab=readme-ov-file#manylinux_2_28-almalinux-8-based
+  # On s390x gcc 14 is used because it contains fix for interaction
+  # between precompiled headers and vectorization builtins.
+  # This fix is not available in earlier gcc versions.
+  # gcc-14 uses ABI19.
+  if [[ "$(uname -m)" != "s390x" ]]; then
+    python -c "import torch; exit(0 if torch._C._PYBIND11_BUILD_ABI == '_cxxabi1016' else 1)"
+  fi
   popd
 fi

.ci/pytorch/common_utils.sh

@@ -169,24 +169,34 @@ function install_torchrec_and_fbgemm() {
   torchrec_commit=$(get_pinned_commit torchrec)
   local fbgemm_commit
   fbgemm_commit=$(get_pinned_commit fbgemm)
+  if [[ "$BUILD_ENVIRONMENT" == *rocm* ]] ; then
+    fbgemm_commit=$(get_pinned_commit fbgemm_rocm)
+  fi
   pip_uninstall torchrec-nightly
   pip_uninstall fbgemm-gpu-nightly
   pip_install setuptools-git-versioning scikit-build pyre-extensions
 
-  # TODO (huydhn): I still have no clue on why sccache doesn't work with only fbgemm_gpu here, but it
-  # seems to be an sccache-related issue
-  if [[ "$IS_A100_RUNNER" == "1" ]]; then
-    unset CMAKE_CUDA_COMPILER_LAUNCHER
-    sudo mv /opt/cache/bin /opt/cache/bin-backup
-  fi
+  if [[ "$BUILD_ENVIRONMENT" == *rocm* ]] ; then
+    # install torchrec first because it installs fbgemm nightly on top of rocm fbgemm
+    pip_install --no-use-pep517 --user "git+https://github.com/pytorch/torchrec.git@${torchrec_commit}"
+    pip_uninstall fbgemm-gpu-nightly
 
-  # See https://github.com/pytorch/pytorch/issues/106971
-  CUDA_PATH=/usr/local/cuda-12.1 pip_install --no-use-pep517 --user "git+https://github.com/pytorch/FBGEMM.git@${fbgemm_commit}#egg=fbgemm-gpu&subdirectory=fbgemm_gpu"
-  pip_install --no-use-pep517 --user "git+https://github.com/pytorch/torchrec.git@${torchrec_commit}"
-
-  if [[ "$IS_A100_RUNNER" == "1" ]]; then
-    export CMAKE_CUDA_COMPILER_LAUNCHER=/opt/cache/bin/sccache
-    sudo mv /opt/cache/bin-backup /opt/cache/bin
+    pip_install tabulate  # needed for newer fbgemm
+    pip_install patchelf  # needed for rocm fbgemm
+    git clone --recursive https://github.com/pytorch/fbgemm
+    pushd fbgemm/fbgemm_gpu
+    git checkout "${fbgemm_commit}"
+    python setup.py install \
+      --package_variant=rocm \
+      -DHIP_ROOT_DIR="${ROCM_PATH}" \
+      -DCMAKE_C_FLAGS="-DTORCH_USE_HIP_DSA" \
+      -DCMAKE_CXX_FLAGS="-DTORCH_USE_HIP_DSA"
+    popd
+    rm -rf fbgemm
+  else
+    # See https://github.com/pytorch/pytorch/issues/106971
+    CUDA_PATH=/usr/local/cuda-12.1 pip_install --no-use-pep517 --user "git+https://github.com/pytorch/FBGEMM.git@${fbgemm_commit}#egg=fbgemm-gpu&subdirectory=fbgemm_gpu"
+    pip_install --no-use-pep517 --user "git+https://github.com/pytorch/torchrec.git@${torchrec_commit}"
   fi
 }
 
@@ -216,6 +226,11 @@ function checkout_install_torchbench() {
     # to install and test other models
     python install.py --continue_on_fail
   fi
+
+  # TODO (huydhn): transformers-4.44.2 added by https://github.com/pytorch/benchmark/pull/2488
+  # is regressing speedup metric. This needs to be investigated further
+  pip install transformers==4.38.1
+
   echo "Print all dependencies after TorchBench is installed"
   python -mpip freeze
   popd

.ci/pytorch/macos-build.sh

@@ -33,55 +33,11 @@ if which sccache > /dev/null; then
   export PATH="${tmp_dir}:$PATH"
 fi
 
-cross_compile_arm64() {
-  # Cross compilation for arm64
-  # Explicitly set USE_DISTRIBUTED=0 to align with the default build config on mac. This also serves as the sole CI config that tests
-  # that building with USE_DISTRIBUTED=0 works at all. See https://github.com/pytorch/pytorch/issues/86448
-  USE_DISTRIBUTED=0 CMAKE_OSX_ARCHITECTURES=arm64 MACOSX_DEPLOYMENT_TARGET=11.0 USE_MKLDNN=OFF USE_QNNPACK=OFF WERROR=1 BUILD_TEST=OFF USE_PYTORCH_METAL=1 python setup.py bdist_wheel
-}
-
-compile_arm64() {
-  # Compilation for arm64
-  # TODO: Compile with OpenMP support (but this causes CI regressions as cross-compilation were done with OpenMP disabled)
-  USE_DISTRIBUTED=0 USE_OPENMP=1 MACOSX_DEPLOYMENT_TARGET=11.0 WERROR=1 BUILD_TEST=OFF USE_PYTORCH_METAL=1 python setup.py bdist_wheel
-}
-
-compile_x86_64() {
-  USE_DISTRIBUTED=0 WERROR=1 python setup.py bdist_wheel --plat-name=macosx_10_9_x86_64
-}
-
-build_lite_interpreter() {
-    echo "Testing libtorch (lite interpreter)."
-
-    CPP_BUILD="$(pwd)/../cpp_build"
-    # Ensure the removal of the tmp directory
-    trap 'rm -rfv ${CPP_BUILD}' EXIT
-    rm -rf "${CPP_BUILD}"
-    mkdir -p "${CPP_BUILD}/caffe2"
-
-    # It looks libtorch need to be built in "${CPP_BUILD}/caffe2 folder.
-    BUILD_LIBTORCH_PY=$PWD/tools/build_libtorch.py
-    pushd "${CPP_BUILD}/caffe2" || exit
-    VERBOSE=1 DEBUG=1 python "${BUILD_LIBTORCH_PY}"
-    popd || exit
-
-    "${CPP_BUILD}/caffe2/build/bin/test_lite_interpreter_runtime"
-}
-
 print_cmake_info
 
-if [[ ${BUILD_ENVIRONMENT} = *arm64* ]]; then
-  if [[ $(uname -m) == "arm64" ]]; then
-    compile_arm64
-  else
-    cross_compile_arm64
-  fi
-elif [[ ${BUILD_ENVIRONMENT} = *lite-interpreter* ]]; then
-  export BUILD_LITE_INTERPRETER=1
-  build_lite_interpreter
-else
-  compile_x86_64
-fi
+# Explicitly set USE_DISTRIBUTED=0 to align with the default build config on mac. This also serves as the sole CI config that tests
+# that building with USE_DISTRIBUTED=0 works at all. See https://github.com/pytorch/pytorch/issues/86448
+USE_DISTRIBUTED=0 USE_OPENMP=1 MACOSX_DEPLOYMENT_TARGET=11.0 WERROR=1 BUILD_TEST=OFF USE_PYTORCH_METAL=1 python setup.py bdist_wheel
 
 if which sccache > /dev/null; then
   print_sccache_stats

.ci/pytorch/macos-test.sh

@@ -18,6 +18,9 @@ if [[ ! $(python -c "import torch; print(int(torch.backends.openmp.is_available(
 fi
 popd
 
+# enable debug asserts in serialization
+export TORCH_SERIALIZATION_DEBUG=1
+
 setup_test_python() {
   # The CircleCI worker hostname doesn't resolve to an address.
   # This environment variable makes ProcessGroupGloo default to

.ci/pytorch/run_tests.sh

@@ -40,7 +40,7 @@ retry () {
 if [[ "$#" != 3 ]]; then
   if [[ -z "${DESIRED_PYTHON:-}" || -z "${DESIRED_CUDA:-}" || -z "${PACKAGE_TYPE:-}" ]]; then
     echo "USAGE: run_tests.sh  PACKAGE_TYPE  DESIRED_PYTHON  DESIRED_CUDA"
-    echo "The env variable PACKAGE_TYPE must be set to 'conda' or 'manywheel' or 'libtorch'"
+    echo "The env variable PACKAGE_TYPE must be set to 'manywheel' or 'libtorch'"
     echo "The env variable DESIRED_PYTHON must be set like '2.7mu' or '3.6m' etc"
     echo "The env variable DESIRED_CUDA must be set like 'cpu' or 'cu80' etc"
     exit 1

.ci/pytorch/smoke_test/check_binary_symbols.py

@@ -6,7 +6,7 @@ import itertools
 import os
 import re
 from pathlib import Path
-from typing import Any, List, Tuple
+from typing import Any
 
 
 # We also check that there are [not] cxx11 symbols in libtorch
@@ -46,17 +46,17 @@ LIBTORCH_PRE_CXX11_PATTERNS = _apply_libtorch_symbols(PRE_CXX11_SYMBOLS)
 
 
 @functools.lru_cache(100)
-def get_symbols(lib: str) -> List[Tuple[str, str, str]]:
+def get_symbols(lib: str) -> list[tuple[str, str, str]]:
     from subprocess import check_output
 
     lines = check_output(f'nm "{lib}"|c++filt', shell=True)
     return [x.split(" ", 2) for x in lines.decode("latin1").split("\n")[:-1]]
 
 
-def grep_symbols(lib: str, patterns: List[Any]) -> List[str]:
+def grep_symbols(lib: str, patterns: list[Any]) -> list[str]:
     def _grep_symbols(
-        symbols: List[Tuple[str, str, str]], patterns: List[Any]
-    ) -> List[str]:
+        symbols: list[tuple[str, str, str]], patterns: list[Any]
+    ) -> list[str]:
         rc = []
         for _s_addr, _s_type, s_name in symbols:
             for pattern in patterns:
@@ -80,7 +80,7 @@ def grep_symbols(lib: str, patterns: List[Any]) -> List[str]:
         return functools.reduce(list.__add__, (x.result() for x in tasks), [])
 
 
-def check_lib_symbols_for_abi_correctness(lib: str, pre_cxx11_abi: bool = True) -> None:
+def check_lib_symbols_for_abi_correctness(lib: str) -> None:
     print(f"lib: {lib}")
     cxx11_symbols = grep_symbols(lib, LIBTORCH_CXX11_PATTERNS)
     pre_cxx11_symbols = grep_symbols(lib, LIBTORCH_PRE_CXX11_PATTERNS)
@@ -88,28 +88,12 @@ def check_lib_symbols_for_abi_correctness(lib: str, pre_cxx11_abi: bool = True)
     num_pre_cxx11_symbols = len(pre_cxx11_symbols)
     print(f"num_cxx11_symbols: {num_cxx11_symbols}")
     print(f"num_pre_cxx11_symbols: {num_pre_cxx11_symbols}")
-    if pre_cxx11_abi:
-        if num_cxx11_symbols > 0:
-            raise RuntimeError(
-                f"Found cxx11 symbols, but there shouldn't be any, see: {cxx11_symbols[:100]}"
-            )
-        if num_pre_cxx11_symbols < 1000:
-            raise RuntimeError("Didn't find enough pre-cxx11 symbols.")
-        # Check for no recursive iterators, regression test for https://github.com/pytorch/pytorch/issues/133437
-        rec_iter_symbols = grep_symbols(
-            lib, [re.compile("std::filesystem::recursive_directory_iterator.*")]
+    if num_pre_cxx11_symbols > 0:
+        raise RuntimeError(
+            f"Found pre-cxx11 symbols, but there shouldn't be any, see: {pre_cxx11_symbols[:100]}"
         )
-        if len(rec_iter_symbols) > 0:
-            raise RuntimeError(
-                f"recursive_directory_iterator in used pre-CXX11 binaries, see; {rec_iter_symbols}"
-            )
-    else:
-        if num_pre_cxx11_symbols > 0:
-            raise RuntimeError(
-                f"Found pre-cxx11 symbols, but there shouldn't be any, see: {pre_cxx11_symbols[:100]}"
-            )
-        if num_cxx11_symbols < 100:
-            raise RuntimeError("Didn't find enought cxx11 symbols")
+    if num_cxx11_symbols < 100:
+        raise RuntimeError("Didn't find enought cxx11 symbols")
 
 
 def main() -> None:
@@ -121,9 +105,8 @@ def main() -> None:
         else:
             install_root = Path(distutils.sysconfig.get_python_lib()) / "torch"
 
-    libtorch_cpu_path = install_root / "lib" / "libtorch_cpu.so"
-    pre_cxx11_abi = "cxx11-abi" not in os.getenv("DESIRED_DEVTOOLSET", "")
-    check_lib_symbols_for_abi_correctness(libtorch_cpu_path, pre_cxx11_abi)
+    libtorch_cpu_path = str(install_root / "lib" / "libtorch_cpu.so")
+    check_lib_symbols_for_abi_correctness(libtorch_cpu_path)
 
 
 if __name__ == "__main__":

.ci/pytorch/smoke_test/max_autotune.py

@@ -46,7 +46,9 @@ def train(args, model, device, train_loader, optimizer, epoch):
         optimizer.step()
         if batch_idx % args.log_interval == 0:
             print(
-                f"Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} ({100. * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}"  # noqa: B950
+                f"Train Epoch: {epoch} "
+                f"[{batch_idx * len(data)}/{len(train_loader.dataset)} "
+                f"({100.0 * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}"
             )
             if args.dry_run:
                 break
@@ -71,7 +73,9 @@ def test(model, device, test_loader):
     test_loss /= len(test_loader.dataset)
 
     print(
-        f"\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({100. * correct / len(test_loader.dataset):.0f}%)\n"  # noqa: B950
+        f"\nTest set: Average loss: {test_loss:.4f}, "
+        f"Accuracy: {correct}/{len(test_loader.dataset)} "
+        f"({100.0 * correct / len(test_loader.dataset):.0f}%)\n"
     )
 
 

.ci/pytorch/smoke_test/smoke_test.py

@@ -6,6 +6,8 @@ import re
 import subprocess
 import sys
 from pathlib import Path
+from tempfile import NamedTemporaryFile
+from typing import Optional
 
 import torch
 import torch._dynamo
@@ -75,10 +77,13 @@ def read_release_matrix():
 
 
 def test_numpy():
-    import numpy as np
+    try:
+        import numpy as np
 
-    x = np.arange(5)
-    torch.tensor(x)
+        x = np.arange(5)
+        torch.tensor(x)
+    except ImportError:
+        print("Numpy check skipped. Numpy is not installed.")
 
 
 def check_version(package: str) -> None:
@@ -161,8 +166,71 @@ def test_cuda_runtime_errors_captured() -> None:
         raise RuntimeError("Expected CUDA RuntimeError but have not received!")
 
 
+def test_cuda_gds_errors_captured() -> None:
+    major_version = int(torch.version.cuda.split(".")[0])
+    minor_version = int(torch.version.cuda.split(".")[1])
+
+    if target_os == "windows":
+        print(f"{target_os} is not supported for GDS smoke test")
+        return
+
+    if major_version < 12 or (major_version == 12 and minor_version < 6):
+        print("CUDA version is not supported for GDS smoke test")
+        return
+
+    cuda_exception_missed = True
+    try:
+        print("Testing test_cuda_gds_errors_captured")
+        with NamedTemporaryFile() as f:
+            torch.cuda.gds.GdsFile(f.name, os.O_CREAT | os.O_RDWR)
+    except RuntimeError as e:
+        expected_error = "cuFileHandleRegister failed"
+        if re.search(expected_error, f"{e}"):
+            print(f"Caught CUDA exception with success: {e}")
+            cuda_exception_missed = False
+        else:
+            raise e
+    if cuda_exception_missed:
+        raise RuntimeError(
+            "Expected cuFileHandleRegister failed RuntimeError but have not received!"
+        )
+
+
+def find_pypi_package_version(package: str) -> Optional[str]:
+    from importlib import metadata
+
+    dists = metadata.distributions()
+    for dist in dists:
+        if dist.metadata["Name"].startswith(package):
+            return dist.version
+    return None
+
+
+def cudnn_to_version_str(cudnn_version: int) -> str:
+    patch = int(cudnn_version % 10)
+    minor = int((cudnn_version / 100) % 100)
+    major = int((cudnn_version / 10000) % 10000)
+    return f"{major}.{minor}.{patch}"
+
+
+def compare_pypi_to_torch_versions(
+    package: str, pypi_version: str, torch_version: str
+) -> None:
+    if pypi_version is None:
+        raise RuntimeError(f"Can't find {package} in PyPI for Torch: {torch_version}")
+    if pypi_version.startswith(torch_version):
+        print(f"Found matching {package}. Torch: {torch_version} PyPI {pypi_version}")
+    else:
+        raise RuntimeError(
+            f"Wrong {package} version. Torch: {torch_version} PyPI: {pypi_version}"
+        )
+
+
 def smoke_test_cuda(
-    package: str, runtime_error_check: str, torch_compile_check: str
+    package: str,
+    runtime_error_check: str,
+    torch_compile_check: str,
+    pypi_pkg_check: str,
 ) -> None:
     if not torch.cuda.is_available() and is_cuda_system:
         raise RuntimeError(f"Expected CUDA {gpu_arch_ver}. However CUDA is not loaded.")
@@ -192,20 +260,30 @@ def smoke_test_cuda(
             raise RuntimeError(
                 f"Wrong CUDA version. Loaded: {torch.version.cuda} Expected: {gpu_arch_ver}"
             )
-        print(f"torch cuda: {torch.version.cuda}")
-        # todo add cudnn version validation
-        print(f"torch cudnn: {torch.backends.cudnn.version()}")
-        print(f"cuDNN enabled? {torch.backends.cudnn.enabled}")
 
+        print(f"torch cuda: {torch.version.cuda}")
         torch.cuda.init()
         print("CUDA initialized successfully")
         print(f"Number of CUDA devices: {torch.cuda.device_count()}")
         for i in range(torch.cuda.device_count()):
             print(f"Device {i}: {torch.cuda.get_device_name(i)}")
 
-        # nccl is availbale only on Linux
+        print(f"cuDNN enabled? {torch.backends.cudnn.enabled}")
+        torch_cudnn_version = cudnn_to_version_str(torch.backends.cudnn.version())
+        print(f"Torch cuDNN version: {torch_cudnn_version}")
+
         if sys.platform in ["linux", "linux2"]:
-            print(f"torch nccl version: {torch.cuda.nccl.version()}")
+            torch_nccl_version = ".".join(str(v) for v in torch.cuda.nccl.version())
+            print(f"Torch nccl; version: {torch_nccl_version}")
+
+        # Pypi dependencies are installed on linux ony and nccl is availbale only on Linux.
+        if pypi_pkg_check == "enabled" and sys.platform in ["linux", "linux2"]:
+            compare_pypi_to_torch_versions(
+                "cudnn", find_pypi_package_version("nvidia-cudnn"), torch_cudnn_version
+            )
+            compare_pypi_to_torch_versions(
+                "nccl", find_pypi_package_version("nvidia-nccl"), torch_nccl_version
+            )
 
         if runtime_error_check == "enabled":
             test_cuda_runtime_errors_captured()
@@ -364,6 +442,13 @@ def parse_args():
         choices=["enabled", "disabled"],
         default="enabled",
     )
+    parser.add_argument(
+        "--pypi-pkg-check",
+        help="Check pypi package versions cudnn and nccl",
+        type=str,
+        choices=["enabled", "disabled"],
+        default="enabled",
+    )
     return parser.parse_args()
 
 
@@ -379,14 +464,19 @@ def main() -> None:
     smoke_test_conv2d()
     test_linalg()
     test_numpy()
+
     if is_cuda_system:
         test_linalg("cuda")
+        test_cuda_gds_errors_captured()
 
     if options.package == "all":
         smoke_test_modules()
 
     smoke_test_cuda(
-        options.package, options.runtime_error_check, options.torch_compile_check
+        options.package,
+        options.runtime_error_check,
+        options.torch_compile_check,
+        options.pypi_pkg_check,
     )
 
 

.ci/pytorch/test.sh

@@ -46,6 +46,9 @@ BUILD_BIN_DIR="$BUILD_DIR"/bin
 SHARD_NUMBER="${SHARD_NUMBER:=1}"
 NUM_TEST_SHARDS="${NUM_TEST_SHARDS:=1}"
 
+# enable debug asserts in serialization
+export TORCH_SERIALIZATION_DEBUG=1
+
 export VALGRIND=ON
 # export TORCH_INDUCTOR_INSTALL_GXX=ON
 if [[ "$BUILD_ENVIRONMENT" == *clang9* || "$BUILD_ENVIRONMENT" == *xpu* ]]; then
@@ -174,6 +177,9 @@ if [[ "$BUILD_ENVIRONMENT" == *rocm* ]]; then
   # Print GPU info
   rocminfo
   rocminfo | grep -E 'Name:.*\sgfx|Marketing'
+
+  # for benchmarks/dynamo/check_accuracy.py, we need to put results in a rocm specific directory to avoid clashes with cuda
+  MAYBE_ROCM="rocm/"
 fi
 
 if [[ "$BUILD_ENVIRONMENT" == *xpu* ]]; then
@@ -308,6 +314,13 @@ test_python() {
   assert_git_not_dirty
 }
 
+test_lazy_tensor_meta_reference_disabled() {
+  export TORCH_DISABLE_FUNCTIONALIZATION_META_REFERENCE=1
+  echo "Testing lazy tensor operations without meta reference"
+  time python test/run_test.py --include lazy/test_ts_opinfo.py --verbose
+  export -n TORCH_DISABLE_FUNCTIONALIZATION_META_REFERENCE
+}
+
 
 test_dynamo_wrapped_shard() {
   if [[ -z "$NUM_TEST_SHARDS" ]]; then
@@ -411,7 +424,10 @@ test_inductor_cpp_wrapper_shard() {
 
   # Run certain inductor unit tests with cpp wrapper. In the end state, we
   # should be able to run all the inductor unit tests with cpp_wrapper.
-  python test/run_test.py --include inductor/test_torchinductor --verbose
+  python test/run_test.py \
+    --include inductor/test_torchinductor inductor/test_max_autotune inductor/test_cpu_repro \
+    --verbose
+  python test/run_test.py --inductor --include test_torch -k 'take' --verbose
 
   # Run inductor benchmark tests with cpp wrapper.
   # Skip benchmark tests if it's in rerun-disabled-mode.
@@ -424,7 +440,7 @@ test_inductor_cpp_wrapper_shard() {
     --output "$TEST_REPORTS_DIR/inductor_cpp_wrapper_training.csv"
     python benchmarks/dynamo/check_accuracy.py \
       --actual "$TEST_REPORTS_DIR/inductor_cpp_wrapper_training.csv" \
-      --expected "benchmarks/dynamo/ci_expected_accuracy/inductor_timm_training.csv"
+      --expected "benchmarks/dynamo/ci_expected_accuracy/${MAYBE_ROCM}inductor_timm_training.csv"
 
     python benchmarks/dynamo/torchbench.py --device cuda --accuracy \
       --bfloat16 --inference --inductor --only hf_T5 --output "$TEST_REPORTS_DIR/inductor_cpp_wrapper_inference.csv"
@@ -434,7 +450,7 @@ test_inductor_cpp_wrapper_shard() {
       --bfloat16 --inference --inductor --only moco --output "$TEST_REPORTS_DIR/inductor_cpp_wrapper_inference.csv"
     python benchmarks/dynamo/check_accuracy.py \
       --actual "$TEST_REPORTS_DIR/inductor_cpp_wrapper_inference.csv" \
-      --expected "benchmarks/dynamo/ci_expected_accuracy/inductor_torchbench_inference.csv"
+      --expected "benchmarks/dynamo/ci_expected_accuracy/${MAYBE_ROCM}inductor_torchbench_inference.csv"
   fi
 }
 
@@ -467,6 +483,8 @@ elif [[ "${TEST_CONFIG}" == *aot_eager* ]]; then
   DYNAMO_BENCHMARK_FLAGS+=(--backend aot_eager)
 elif [[ "${TEST_CONFIG}" == *aot_inductor* ]]; then
   DYNAMO_BENCHMARK_FLAGS+=(--export-aot-inductor)
+elif [[ "${TEST_CONFIG}" == *max_autotune_inductor* ]]; then
+  DYNAMO_BENCHMARK_FLAGS+=(--inductor --inductor-compile-mode max-autotune)
 elif [[ "${TEST_CONFIG}" == *inductor* && "${TEST_CONFIG}" != *perf* ]]; then
   DYNAMO_BENCHMARK_FLAGS+=(--inductor)
 fi
@@ -481,6 +499,59 @@ else
   DYNAMO_BENCHMARK_FLAGS+=(--device cuda)
 fi
 
+test_cachebench() {
+  TEST_REPORTS_DIR=$(pwd)/test/test-reports
+  mkdir -p "$TEST_REPORTS_DIR"
+
+  local BENCHMARK
+  if [[ "${SHARD_NUMBER}" == 1 ]]; then
+    local BENCHMARK=torchbench
+  elif [[ "${SHARD_NUMBER}" == 2 ]]; then
+    local BENCHMARK=huggingface
+  else
+    echo "invalid SHARD_NUMBER: ${SHARD_NUMBER}"
+    exit 1
+  fi
+
+  local mode_options=("training" "inference")
+
+  for mode in "${mode_options[@]}"; do
+    $TASKSET python "benchmarks/dynamo/cachebench.py" \
+        --mode "$mode" \
+        --device cuda \
+        --benchmark "$BENCHMARK" \
+        --repeat 3 \
+        --output "$TEST_REPORTS_DIR/cachebench_${BENCHMARK}_${mode}.json"
+
+    $TASKSET python "benchmarks/dynamo/cachebench.py" \
+        --mode "$mode" \
+        --dynamic \
+        --device cuda \
+        --benchmark "$BENCHMARK" \
+        --repeat 3 \
+        --output "$TEST_REPORTS_DIR/cachebench_${BENCHMARK}_${mode}_dynamic.json"
+  done
+}
+
+test_verify_cachebench() {
+  TMP_TEST_REPORTS_DIR=$(mktemp -d)
+  TEST_OUTPUT="$TMP_TEST_REPORTS_DIR/test.json"
+
+  $TASKSET python "benchmarks/dynamo/cachebench.py" \
+      --mode training \
+      --device cpu \
+      --model nanogpt \
+      --benchmark torchbench \
+      --output "$TEST_OUTPUT"
+
+  # -s checks file exists and is non empty
+  if [[ ! -s "$TEST_OUTPUT" ]]; then
+    echo "Cachebench failed to produce an output."
+    echo "Run 'python benchmarks/dynamo/cachebench.py' to make sure it works"
+    exit 1
+  fi
+}
+
 test_perf_for_dashboard() {
   TEST_REPORTS_DIR=$(pwd)/test/test-reports
   mkdir -p "$TEST_REPORTS_DIR"
@@ -509,6 +580,10 @@ test_perf_for_dashboard() {
     test_inductor_set_cpu_affinity
   elif [[ "${TEST_CONFIG}" == *cuda_a10g* ]]; then
     device=cuda_a10g
+  elif [[ "${TEST_CONFIG}" == *h100* ]]; then
+    device=cuda_h100
+  elif [[ "${TEST_CONFIG}" == *rocm* ]]; then
+    device=rocm
   fi
 
   for mode in "${modes[@]}"; do
@@ -625,16 +700,16 @@ test_single_dynamo_benchmark() {
       TEST_CONFIG=${TEST_CONFIG//_avx512/}
     fi
     python "benchmarks/dynamo/$suite.py" \
-      --ci --accuracy --timing --explain \
+      --ci --accuracy --timing --explain --print-compilation-time \
       "${DYNAMO_BENCHMARK_FLAGS[@]}" \
       "$@" "${partition_flags[@]}" \
       --output "$TEST_REPORTS_DIR/${name}_${suite}.csv"
     python benchmarks/dynamo/check_accuracy.py \
       --actual "$TEST_REPORTS_DIR/${name}_$suite.csv" \
-      --expected "benchmarks/dynamo/ci_expected_accuracy/${TEST_CONFIG}_${name}.csv"
+      --expected "benchmarks/dynamo/ci_expected_accuracy/${MAYBE_ROCM}${TEST_CONFIG}_${name}.csv"
     python benchmarks/dynamo/check_graph_breaks.py \
       --actual "$TEST_REPORTS_DIR/${name}_$suite.csv" \
-      --expected "benchmarks/dynamo/ci_expected_accuracy/${TEST_CONFIG}_${name}.csv"
+      --expected "benchmarks/dynamo/ci_expected_accuracy/${MAYBE_ROCM}${TEST_CONFIG}_${name}.csv"
   fi
 }
 
@@ -657,7 +732,7 @@ test_inductor_halide() {
 }
 
 test_inductor_triton_cpu() {
-  python test/run_test.py --include inductor/test_triton_cpu_backend.py --verbose
+  python test/run_test.py --include inductor/test_triton_cpu_backend.py inductor/test_torchinductor_strided_blocks.py --verbose
   assert_git_not_dirty
 }
 
@@ -687,6 +762,8 @@ test_dynamo_benchmark() {
       fi
     elif [[ "${TEST_CONFIG}" == *aot_inductor* ]]; then
       test_single_dynamo_benchmark "inference" "$suite" "$shard_id" --inference --bfloat16 "$@"
+    elif [[ "${TEST_CONFIG}" == *max_autotune_inductor* ]]; then
+      test_single_dynamo_benchmark "inference" "$suite" "$shard_id" --inference --bfloat16 "$@"
     else
       test_single_dynamo_benchmark "inference" "$suite" "$shard_id" --inference --bfloat16 "$@"
       test_single_dynamo_benchmark "training" "$suite" "$shard_id" --training --amp "$@"
@@ -721,7 +798,7 @@ test_inductor_torchbench_smoketest_perf() {
       --only $test --output "$TEST_REPORTS_DIR/inductor_warm_start_smoketest_$test.csv"
     python benchmarks/dynamo/check_accuracy.py \
       --actual "$TEST_REPORTS_DIR/inductor_warm_start_smoketest_$test.csv" \
-      --expected "benchmarks/dynamo/ci_expected_accuracy/inductor_huggingface_training.csv"
+      --expected "benchmarks/dynamo/ci_expected_accuracy/${MAYBE_ROCM}inductor_huggingface_training.csv"
   done
 }
 
@@ -1096,8 +1173,9 @@ build_xla() {
   apply_patches
   SITE_PACKAGES="$(python -c 'from distutils.sysconfig import get_python_lib; print(get_python_lib())')"
   # These functions are defined in .circleci/common.sh in pytorch/xla repo
-  retry install_deps_pytorch_xla $XLA_DIR $USE_CACHE
+  retry install_pre_deps_pytorch_xla $XLA_DIR $USE_CACHE
   CMAKE_PREFIX_PATH="${SITE_PACKAGES}/torch:${CMAKE_PREFIX_PATH}" XLA_SANDBOX_BUILD=1 build_torch_xla $XLA_DIR
+  retry install_post_deps_pytorch_xla
   assert_git_not_dirty
 }
 
@@ -1397,14 +1475,13 @@ test_executorch() {
   pushd /executorch
 
   export PYTHON_EXECUTABLE=python
-  export EXECUTORCH_BUILD_PYBIND=ON
-  export CMAKE_ARGS="-DEXECUTORCH_BUILD_XNNPACK=ON -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON"
+  export CMAKE_ARGS="-DEXECUTORCH_BUILD_PYBIND=ON -DEXECUTORCH_BUILD_XNNPACK=ON -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON"
 
   # For llama3
   bash examples/models/llama3_2_vision/install_requirements.sh
   # NB: We need to rebuild ExecuTorch runner here because it depends on PyTorch
   # from the PR
-  bash .ci/scripts/setup-linux.sh cmake
+  bash .ci/scripts/setup-linux.sh --build-tool cmake
 
   echo "Run ExecuTorch unit tests"
   pytest -v -n auto
@@ -1428,7 +1505,7 @@ test_executorch() {
 test_linux_aarch64() {
   python test/run_test.py --include test_modules test_mkldnn test_mkldnn_fusion test_openmp test_torch test_dynamic_shapes \
         test_transformers test_multiprocessing test_numpy_interop test_autograd test_binary_ufuncs test_complex test_spectral_ops \
-        test_foreach test_reductions test_unary_ufuncs \
+        test_foreach test_reductions test_unary_ufuncs test_tensor_creation_ops test_ops \
         --shard "$SHARD_NUMBER" "$NUM_TEST_SHARDS" --verbose
 
   # Dynamo tests
@@ -1450,6 +1527,27 @@ test_linux_aarch64() {
        --shard "$SHARD_NUMBER" "$NUM_TEST_SHARDS" --verbose
 }
 
+test_operator_benchmark() {
+  TEST_REPORTS_DIR=$(pwd)/test/test-reports
+  mkdir -p "$TEST_REPORTS_DIR"
+  TEST_DIR=$(pwd)
+
+  test_inductor_set_cpu_affinity
+
+  cd benchmarks/operator_benchmark/pt_extension
+  python setup.py install
+
+  cd "${TEST_DIR}"/benchmarks/operator_benchmark
+  $TASKSET python -m benchmark_all_test --device "$1" --tag-filter "$2" \
+      --output-dir "${TEST_REPORTS_DIR}/operator_benchmark_eager_float32_cpu.csv"
+
+  pip_install pandas
+  python check_perf_csv.py \
+      --actual "${TEST_REPORTS_DIR}/operator_benchmark_eager_float32_cpu.csv" \
+      --expected "expected_ci_operator_benchmark_eager_float32_cpu.csv"
+}
+
+
 if ! [[ "${BUILD_ENVIRONMENT}" == *libtorch* || "${BUILD_ENVIRONMENT}" == *-bazel-* ]]; then
   (cd test && python -c "import torch; print(torch.__config__.show())")
   (cd test && python -c "import torch; print(torch.__config__.parallel_info())")
@@ -1480,6 +1578,19 @@ elif [[ "$TEST_CONFIG" == distributed ]]; then
   if [[ "${SHARD_NUMBER}" == 1 ]]; then
     test_rpc
   fi
+elif [[ "${TEST_CONFIG}" == *operator_benchmark* ]]; then
+  TEST_MODE="short"
+
+  if [[ "${TEST_CONFIG}" == *cpu* ]]; then
+    if [[ "${TEST_CONFIG}" == *long* ]]; then
+      TEST_MODE="long"
+    elif [[ "${TEST_CONFIG}" == *all* ]]; then
+      TEST_MODE="all"
+    fi
+
+    test_operator_benchmark cpu ${TEST_MODE}
+
+  fi
 elif [[ "${TEST_CONFIG}" == *inductor_distributed* ]]; then
   test_inductor_distributed
 elif [[ "${TEST_CONFIG}" == *inductor-halide* ]]; then
@@ -1496,6 +1607,16 @@ elif [[ "${TEST_CONFIG}" == *timm* ]]; then
   install_torchvision
   id=$((SHARD_NUMBER-1))
   test_dynamo_benchmark timm_models "$id"
+elif [[ "${TEST_CONFIG}" == cachebench ]]; then
+  install_torchaudio cuda
+  install_torchvision
+  checkout_install_torchbench nanogpt BERT_pytorch resnet50 hf_T5 llama moco
+  PYTHONPATH=$(pwd)/torchbench test_cachebench
+elif [[ "${TEST_CONFIG}" == verify_cachebench ]]; then
+  install_torchaudio cpu
+  install_torchvision
+  checkout_install_torchbench nanogpt
+  PYTHONPATH=$(pwd)/torchbench test_verify_cachebench
 elif [[ "${TEST_CONFIG}" == *torchbench* ]]; then
   if [[ "${TEST_CONFIG}" == *cpu* ]]; then
     install_torchaudio cpu
@@ -1532,6 +1653,7 @@ elif [[ "${TEST_CONFIG}" == *inductor_cpp_wrapper* ]]; then
   install_torchvision
   checkout_install_torchbench hf_T5 llama moco
   PYTHONPATH=$(pwd)/torchbench test_inductor_cpp_wrapper_shard "$SHARD_NUMBER"
+  test_inductor_aoti
 elif [[ "${TEST_CONFIG}" == *inductor* ]]; then
   install_torchvision
   test_inductor_shard "${SHARD_NUMBER}"
@@ -1551,6 +1673,7 @@ elif [[ "${BUILD_ENVIRONMENT}" == *rocm* && -n "$TESTS_TO_INCLUDE" ]]; then
   test_python_shard "$SHARD_NUMBER"
   test_aten
 elif [[ "${SHARD_NUMBER}" == 1 && $NUM_TEST_SHARDS -gt 1 ]]; then
+  test_lazy_tensor_meta_reference_disabled
   test_without_numpy
   install_torchvision
   test_python_shard 1

.ci/pytorch/test_example_code/cnn_smoke_win_arm64.py

@@ -0,0 +1,41 @@
+r"""
+It's used to check basic rnn features with cpu-only.
+For example, it would throw exception if some components are missing
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+
+
+class SimpleCNN(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = nn.Conv2d(1, 1, 3)
+        self.pool = nn.MaxPool2d(2, 2)
+
+    def forward(self, inputs):
+        output = self.pool(F.relu(self.conv(inputs)))
+        output = output.view(1)
+        return output
+
+
+try:
+    # Mock one infer
+    net = SimpleCNN()
+    net_inputs = torch.rand((1, 1, 5, 5))
+    outputs = net(net_inputs)
+    print(outputs)
+
+    criterion = nn.MSELoss()
+    optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.1)
+
+    # Mock one step training
+    label = torch.full((1,), 1.0, dtype=torch.float)
+    loss = criterion(outputs, label)
+    loss.backward()
+    optimizer.step()
+
+except Exception as e:
+    print(f"An error occurred: {e}")

.ci/pytorch/test_example_code/rnn_smoke_win_arm64.py

@@ -0,0 +1,13 @@
+r"""
+It's used to check basic rnn features with cpu-only.
+For example, it would throw exception if missing some components are missing
+"""
+
+import torch
+import torch.nn as nn
+
+
+rnn = nn.RNN(10, 20, 2)
+inputs = torch.randn(5, 3, 10)
+h0 = torch.randn(2, 3, 20)
+output, hn = rnn(inputs, h0)

.ci/pytorch/win-test.sh

@@ -18,6 +18,9 @@ export PYTORCH_FINAL_PACKAGE_DIR="${PYTORCH_FINAL_PACKAGE_DIR:-/c/w/build-result
 PYTORCH_FINAL_PACKAGE_DIR_WIN=$(cygpath -w "${PYTORCH_FINAL_PACKAGE_DIR}")
 export PYTORCH_FINAL_PACKAGE_DIR_WIN
 
+# enable debug asserts in serialization
+export TORCH_SERIALIZATION_DEBUG=1
+
 mkdir -p "$TMP_DIR"/build/torch
 
 export SCRIPT_HELPERS_DIR=$SCRIPT_PARENT_DIR/win-test-helpers

.ci/pytorch/windows/arm64/bootstrap_apl.bat

@@ -0,0 +1,31 @@
+@echo off
+
+echo Dependency ARM Performance Libraries (APL) installation started.
+
+:: Pre-check for downloads and dependencies folders
+if not exist "%DOWNLOADS_DIR%" mkdir %DOWNLOADS_DIR%
+if not exist "%DEPENDENCIES_DIR%" mkdir %DEPENDENCIES_DIR%
+
+:: Set download URL for the ARM Performance Libraries (APL)
+set DOWNLOAD_URL="https://developer.arm.com/-/cdn-downloads/permalink/Arm-Performance-Libraries/Version_24.10/arm-performance-libraries_24.10_Windows.msi"
+set INSTALLER_FILE=%DOWNLOADS_DIR%\arm-performance-libraries.msi
+
+:: Download installer
+echo Downloading ARM Performance Libraries (APL)...
+curl -L -o "%INSTALLER_FILE%" %DOWNLOAD_URL%
+
+:: Install ARM Performance Libraries (APL)
+echo Installing ARM Performance Libraries (APL)...
+msiexec /i "%INSTALLER_FILE%" /qn /norestart ACCEPT_EULA=1 INSTALLFOLDER="%DEPENDENCIES_DIR%"
+
+:: Check if installation was successful
+if %errorlevel% neq 0 (
+    echo "Failed to install ARM Performance Libraries (APL) components. (exitcode = %errorlevel%)"
+    exit /b 1
+)
+
+:: Add to environment
+echo ARMPL_DIR=%DEPENDENCIES_DIR%\armpl_24.10\>> %GITHUB_ENV%
+echo %DEPENDENCIES_DIR%\armpl_24.10\bin\>> %GITHUB_PATH%
+
+echo Dependency ARM Performance Libraries (APL) installation finished.

.ci/pytorch/windows/arm64/bootstrap_buildtools.bat

@@ -0,0 +1,41 @@
+@echo off
+
+echo Dependency MSVC Build Tools with C++ with ARM64/ARM64EC components installation started.
+
+:: Pre-check for downloads and dependencies folders
+if not exist "%DOWNLOADS_DIR%" mkdir "%DOWNLOADS_DIR%"
+if not exist "%DEPENDENCIES_DIR%" mkdir "%DEPENDENCIES_DIR%"
+
+:: Set download URL for the Visual Studio Installer
+set DOWNLOAD_URL=https://aka.ms/vs/17/release/vs_BuildTools.exe
+set INSTALLER_FILE=%DOWNLOADS_DIR%\vs_BuildTools.exe
+
+:: Download installer
+echo Downloading Visual Studio Build Tools with C++ installer...
+curl -L -o "%INSTALLER_FILE%" %DOWNLOAD_URL%
+
+:: Install the Visual Studio Build Tools with C++ components
+echo Installing Visual Studio Build Tools with C++ components...
+echo Installing MSVC %MSVC_VERSION%
+"%INSTALLER_FILE%" --norestart --quiet --wait --installPath "%DEPENDENCIES_DIR%\VSBuildTools" ^
+    --add Microsoft.VisualStudio.Workload.VCTools ^
+    --add Microsoft.VisualStudio.Component.Windows10SDK ^
+    --add Microsoft.VisualStudio.Component.Windows11SDK.22621 ^
+    --add Microsoft.VisualStudio.Component.VC.ASAN ^
+    --add Microsoft.VisualStudio.Component.VC.CMake.Project ^
+    --add Microsoft.VisualStudio.Component.VC.CoreBuildTools ^
+    --add Microsoft.VisualStudio.Component.VC.CoreIde ^
+    --add Microsoft.VisualStudio.Component.VC.Redist.14.Latest ^
+    --add Microsoft.VisualStudio.Component.VC.Tools.ARM64EC ^
+    --add Microsoft.VisualStudio.Component.VC.Tools.ARM64 ^
+    --add Microsoft.VisualStudio.Component.VC.Tools.x86.x64
+
+echo exitcode = %errorlevel%
+
+:: Check if installation was successful
+if %errorlevel% neq 0 (
+    echo Failed to install Visual Studio Build Tools with C++ components.
+    exit /b 1
+)
+
+echo Dependency Visual Studio Build Tools with C++ installation finished.

.ci/pytorch/windows/arm64/bootstrap_git.bat

@@ -0,0 +1,37 @@
+:: we need to install newer version of Git manually as "-submodules" function is not supported in the default version of runner.
+
+@echo off
+
+echo Dependency Git installation started.
+
+:: Pre-check for downloads and dependencies folders
+if not exist "%DOWNLOADS_DIR%" mkdir %DOWNLOADS_DIR%
+if not exist "%DEPENDENCIES_DIR%" mkdir %DEPENDENCIES_DIR%
+
+:: Set download URL for the Git
+set DOWNLOAD_URL="https://github.com/git-for-windows/git/releases/download/v2.46.0.windows.1/Git-2.46.0-64-bit.exe"
+set INSTALLER_FILE=%DOWNLOADS_DIR%\Git-2.46.0-64-bit.exe
+
+:: Download installer
+echo Downloading Git...
+curl -L -o "%INSTALLER_FILE%" %DOWNLOAD_URL%
+
+:: Install Git
+echo Installing Git...
+"%INSTALLER_FILE%" /VERYSILENT /DIR="%DEPENDENCIES_DIR%\git"
+
+dir %DEPENDENCIES_DIR%\git
+
+:: Check if installation was successful
+if %errorlevel% neq 0 (
+    echo "Failed to install Git. (exitcode = %errorlevel%)"
+    exit /b 1
+)
+
+:: Enable long paths
+call "%DEPENDENCIES_DIR%\git\cmd\git.exe" config --system core.longpaths true
+
+:: Add to PATH
+echo %DEPENDENCIES_DIR%\git\cmd\;%DEPENDENCIES_DIR%\git\bin\>> %GITHUB_PATH%
+
+echo Dependency Git installation finished.

.ci/pytorch/windows/arm64/bootstrap_libuv.bat

@@ -0,0 +1,33 @@
+@echo off
+
+echo Dependency libuv installation started.
+
+:: Pre-check for downloads and dependencies folders
+if not exist "%DOWNLOADS_DIR%" mkdir %DOWNLOADS_DIR%
+if not exist "%DEPENDENCIES_DIR%" mkdir %DEPENDENCIES_DIR%
+
+:: activate visual studio
+call "%DEPENDENCIES_DIR%\VSBuildTools\VC\Auxiliary\Build\vcvarsall.bat" arm64
+where cl.exe
+
+cd %DEPENDENCIES_DIR%
+git clone https://github.com/libuv/libuv.git -b v1.39.0
+
+echo Configuring libuv...
+mkdir libuv\build
+cd libuv\build
+cmake .. -DBUILD_TESTING=OFF
+
+echo Building libuv...
+cmake --build . --config Release
+
+echo Installing libuv...
+cmake --install . --prefix ../install
+
+:: Check if installation was successful
+if %errorlevel% neq 0 (
+    echo "Failed to install libuv. (exitcode = %errorlevel%)"
+    exit /b 1
+)
+
+echo Dependency libuv installation finished.

.ci/pytorch/windows/arm64/bootstrap_openblas.bat

@@ -0,0 +1,46 @@
+@echo off
+
+echo Dependency OpenBLAS installation started.
+
+:: Pre-check for downloads and dependencies folders
+if not exist "%DOWNLOADS_DIR%" mkdir %DOWNLOADS_DIR%
+if not exist "%DEPENDENCIES_DIR%" mkdir %DEPENDENCIES_DIR%
+
+:: activate visual studio
+call "%DEPENDENCIES_DIR%\VSBuildTools\VC\Auxiliary\Build\vcvarsall.bat" arm64
+where cl.exe
+
+:: Clone OpenBLAS
+cd %DEPENDENCIES_DIR%
+git clone https://github.com/OpenMathLib/OpenBLAS.git -b v0.3.29
+
+echo Configuring OpenBLAS...
+mkdir OpenBLAS\build
+cd OpenBLAS\build
+cmake .. -G Ninja ^
+  -DBUILD_TESTING=0 ^
+  -DBUILD_BENCHMARKS=0 ^
+  -DC_LAPACK=1 ^
+  -DNOFORTRAN=1 ^
+  -DDYNAMIC_ARCH=0 ^
+  -DARCH=arm64 ^
+  -DBINARY=64 ^
+  -DTARGET=GENERIC ^
+  -DUSE_OPENMP=1 ^
+  -DCMAKE_SYSTEM_PROCESSOR=ARM64 ^
+  -DCMAKE_SYSTEM_NAME=Windows ^
+  -DCMAKE_BUILD_TYPE=Release
+
+echo Building OpenBLAS...
+cmake --build . --config Release
+
+echo Installing OpenBLAS...
+cmake --install . --prefix ../install
+
+:: Check if installation was successful
+if %errorlevel% neq 0 (
+    echo "Failed to install OpenBLAS. (exitcode = %errorlevel%)"
+    exit /b 1
+)
+
+echo Dependency OpenBLAS installation finished.

.ci/pytorch/windows/arm64/bootstrap_python.bat

@@ -0,0 +1,44 @@
+@echo off
+
+echo Dependency Python installation started.
+
+:: Pre-check for downloads and dependencies folders
+if not exist "%DOWNLOADS_DIR%" mkdir %DOWNLOADS_DIR%
+if not exist "%DEPENDENCIES_DIR%" mkdir %DEPENDENCIES_DIR%
+
+if "%DESIRED_PYTHON%" == "3.13" (
+    echo Python version is set to 3.13
+    set DOWNLOAD_URL=https://www.python.org/ftp/python/3.13.2/python-3.13.2-arm64.exe
+) else if "%DESIRED_PYTHON%" == "3.12" (
+    echo Python version is set to 3.12
+    set DOWNLOAD_URL=https://www.python.org/ftp/python/3.12.7/python-3.12.7-arm64.exe
+) else if "%DESIRED_PYTHON%" == "3.11" (
+    echo Python version is set to 3.11
+    set DOWNLOAD_URL=https://www.python.org/ftp/python/3.11.9/python-3.11.9-arm64.exe
+) else (
+    echo DESIRED_PYTHON not defined, Python version is set to 3.12
+    set DOWNLOAD_URL=https://www.python.org/ftp/python/3.12.7/python-3.12.7-arm64.exe
+)
+
+set INSTALLER_FILE=%DOWNLOADS_DIR%\python-installer.exe
+
+:: Download installer
+echo Downloading Python...
+curl -L -o "%INSTALLER_FILE%" "%DOWNLOAD_URL%"
+
+:: Install Python
+echo Installing Python...
+"%INSTALLER_FILE%" /quiet Include_debug=1 TargetDir="%DEPENDENCIES_DIR%\Python"
+
+:: Check if installation was successful
+if %errorlevel% neq 0 (
+    echo "Failed to install Python. (exitcode = %errorlevel%)"
+    exit /b 1
+)
+
+:: Add to PATH
+echo %DEPENDENCIES_DIR%\Python\>> %GITHUB_PATH%
+echo %DEPENDENCIES_DIR%\Python\scripts\>> %GITHUB_PATH%
+echo %DEPENDENCIES_DIR%\Python\libs\>> %GITHUB_PATH%
+
+echo Dependency Python installation finished.

.ci/pytorch/windows/arm64/bootstrap_rust.bat

@@ -0,0 +1,33 @@
+@echo off
+
+echo Dependency Rust installation started.
+
+:: Pre-check for downloads and dependencies folders
+if not exist "%DOWNLOADS_DIR%" mkdir %DOWNLOADS_DIR%
+if not exist "%DEPENDENCIES_DIR%" mkdir %DEPENDENCIES_DIR%
+
+set DOWNLOAD_URL="https://static.rust-lang.org/rustup/dist/x86_64-pc-windows-msvc/rustup-init.exe"
+set INSTALLER_FILE=%DOWNLOADS_DIR%\rustup-init.exe
+set RUSTUP_HOME=%DEPENDENCIES_DIR%\rust
+set CARGO_HOME=%DEPENDENCIES_DIR%\cargo
+
+:: Download installer
+echo Downloading Rust...
+curl -L -o "%INSTALLER_FILE%" %DOWNLOAD_URL%
+
+:: Install APL
+echo Installing Rust...
+"%INSTALLER_FILE%" -q -y --default-host aarch64-pc-windows-msvc --default-toolchain stable --profile default
+
+:: Check if installation was successful
+if %errorlevel% neq 0 (
+    echo "Failed to install Rust. (exitcode = %errorlevel%)"
+    exit /b 1
+)
+
+:: Add to PATH
+echo %DEPENDENCIES_DIR%\cargo\bin\>> %GITHUB_PATH%
+echo RUSTUP_HOME=%DEPENDENCIES_DIR%\rust>> %GITHUB_ENV%
+echo CARGO_HOME=%DEPENDENCIES_DIR%\cargo>> %GITHUB_ENV%
+
+echo Dependency Rust installation finished.

.ci/pytorch/windows/arm64/bootstrap_sccache.bat

@@ -0,0 +1,33 @@
+@echo off
+
+echo Dependency sccache installation started.
+
+:: Pre-check for downloads and dependencies folders
+if not exist "%DOWNLOADS_DIR%" mkdir %DOWNLOADS_DIR%
+if not exist "%DEPENDENCIES_DIR%" mkdir %DEPENDENCIES_DIR%
+
+:: Set download URL for the sccache
+set DOWNLOAD_URL="https://github.com/mozilla/sccache/releases/download/v0.8.1/sccache-v0.8.1-x86_64-pc-windows-msvc.zip"
+set INSTALLER_FILE=%DOWNLOADS_DIR%\sccache.zip
+
+:: Download installer
+echo Downloading sccache.zip...
+curl -L -o "%INSTALLER_FILE%" %DOWNLOAD_URL%
+
+:: Install sccache
+echo Extracting sccache.zip...
+tar -xf "%INSTALLER_FILE%" -C %DEPENDENCIES_DIR%
+cd %DEPENDENCIES_DIR%
+ren sccache-v0.8.1-x86_64-pc-windows-msvc sccache
+cd ..
+
+:: Check if installation was successful
+if %errorlevel% neq 0 (
+    echo "Failed to install sccache. (exitcode = %errorlevel%)"
+    exit /b 1
+)
+
+:: Add to PATH
+echo %DEPENDENCIES_DIR%\sccache\>> %GITHUB_PATH%
+
+echo Dependency sccache installation finished.

.ci/pytorch/windows/arm64/bootstrap_tests.bat

@@ -0,0 +1,22 @@
+:: change to source directory
+cd %PYTORCH_ROOT%
+
+:: activate visual studio
+call "%DEPENDENCIES_DIR%\VSBuildTools\VC\Auxiliary\Build\vcvarsall.bat" arm64
+where cl.exe
+
+:: create virtual environment
+python -m venv .venv
+echo * > .venv\.gitignore
+call .\.venv\Scripts\activate
+where python
+
+:: install dependencies
+python -m pip install --upgrade pip
+pip install -r requirements.txt
+pip install pytest numpy protobuf expecttest hypothesis
+
+:: find file name for pytorch wheel
+for /f "delims=" %%f in ('dir /b "%PYTORCH_FINAL_PACKAGE_DIR%" ^| findstr "torch-"') do set "TORCH_WHEEL_FILENAME=%PYTORCH_FINAL_PACKAGE_DIR%\%%f"
+
+pip install %TORCH_WHEEL_FILENAME%

.ci/pytorch/windows/arm64/build_libtorch.bat

@@ -0,0 +1,101 @@
+@echo on
+
+:: environment variables
+set CMAKE_BUILD_TYPE=%BUILD_TYPE%
+set CMAKE_C_COMPILER_LAUNCHER=sccache
+set CMAKE_CXX_COMPILER_LAUNCHER=sccache
+set libuv_ROOT=%DEPENDENCIES_DIR%\libuv\install
+set MSSdk=1
+if defined PYTORCH_BUILD_VERSION (
+  set PYTORCH_BUILD_VERSION=%PYTORCH_BUILD_VERSION%
+  set PYTORCH_BUILD_NUMBER=1
+)
+
+:: Set BLAS type
+if %ENABLE_APL% == 1 (
+    set BLAS=APL
+    set USE_LAPACK=1
+) else if %ENABLE_OPENBLAS% == 1 (
+    set BLAS=OpenBLAS
+    set OpenBLAS_HOME=%DEPENDENCIES_DIR%\OpenBLAS\install
+)
+
+:: activate visual studio
+call "%DEPENDENCIES_DIR%\VSBuildTools\VC\Auxiliary\Build\vcvarsall.bat" arm64
+where cl.exe
+
+:: change to source directory
+cd %PYTORCH_ROOT%
+
+:: copy libuv.dll
+copy %libuv_ROOT%\lib\Release\uv.dll torch\lib\uv.dll
+
+:: create virtual environment
+python -m venv .venv
+echo * > .venv\.gitignore
+call .\.venv\Scripts\activate
+where python
+
+:: python install dependencies
+python -m pip install --upgrade pip
+pip install -r requirements.txt
+:: DISTUTILS_USE_SDK should be set after psutil dependency
+set DISTUTILS_USE_SDK=1
+
+:: start sccache server and reset sccache stats
+sccache --start-server
+sccache --zero-stats
+sccache --show-stats
+ 
+:: Prepare the environment
+mkdir libtorch
+mkdir libtorch\bin
+mkdir libtorch\cmake
+mkdir libtorch\include
+mkdir libtorch\lib
+mkdir libtorch\share
+mkdir libtorch\test
+
+:: Call LibTorch build script
+python ./tools/build_libtorch.py
+
+:: Check if there is an error
+IF ERRORLEVEL 1 exit /b 1
+IF NOT ERRORLEVEL 0 exit /b 1
+ 
+:: Move the files to the correct location
+move /Y torch\bin\*.* libtorch\bin\
+move /Y torch\cmake\*.* libtorch\cmake\
+robocopy /move /e torch\include\ libtorch\include\
+move /Y torch\lib\*.* libtorch\lib\
+robocopy /move /e torch\share\ libtorch\share\
+move /Y torch\test\*.* libtorch\test\
+move /Y libtorch\bin\*.dll libtorch\lib\
+
+:: Set version
+echo %PYTORCH_BUILD_VERSION% > libtorch\build-version
+git rev-parse HEAD > libtorch\build-hash
+
+:: Set LIBTORCH_PREFIX
+IF "%DEBUG%" == "" (
+    set LIBTORCH_PREFIX=libtorch-win-arm64-shared-with-deps
+) ELSE (
+    set LIBTORCH_PREFIX=libtorch-win-arm64-shared-with-deps-debug
+)
+
+:: Create output
+C:\Windows\System32\tar.exe -cvaf %LIBTORCH_PREFIX%-%PYTORCH_BUILD_VERSION%.zip -C libtorch *
+
+:: Copy output to target directory
+if not exist ..\output mkdir ..\output
+copy /Y "%LIBTORCH_PREFIX%-%PYTORCH_BUILD_VERSION%.zip" "%PYTORCH_FINAL_PACKAGE_DIR%\"
+copy /Y "%LIBTORCH_PREFIX%-%PYTORCH_BUILD_VERSION%.zip" "%PYTORCH_FINAL_PACKAGE_DIR%\%LIBTORCH_PREFIX%-latest.zip"
+
+:: Cleanup raw data to save space
+rmdir /s /q libtorch
+
+:: Check if installation was successful
+if %errorlevel% neq 0 (
+    echo "Failed on build_libtorch. (exitcode = %errorlevel%)"
+    exit /b 1
+)

.ci/pytorch/windows/arm64/build_pytorch.bat

@@ -0,0 +1,60 @@
+@echo on
+
+:: environment variables
+set CMAKE_BUILD_TYPE=%BUILD_TYPE%
+set CMAKE_C_COMPILER_LAUNCHER=sccache
+set CMAKE_CXX_COMPILER_LAUNCHER=sccache
+set libuv_ROOT=%DEPENDENCIES_DIR%\libuv\install
+set MSSdk=1
+if defined PYTORCH_BUILD_VERSION (
+  set PYTORCH_BUILD_VERSION=%PYTORCH_BUILD_VERSION%
+  set PYTORCH_BUILD_NUMBER=1
+)
+
+:: Set BLAS type
+if %ENABLE_APL% == 1 (
+    set BLAS=APL
+    set USE_LAPACK=1
+) else if %ENABLE_OPENBLAS% == 1 (
+    set BLAS=OpenBLAS
+    set OpenBLAS_HOME=%DEPENDENCIES_DIR%\OpenBLAS\install
+)
+
+:: activate visual studio
+call "%DEPENDENCIES_DIR%\VSBuildTools\VC\Auxiliary\Build\vcvarsall.bat" arm64
+where cl.exe
+
+:: change to source directory
+cd %PYTORCH_ROOT%
+
+:: copy libuv.dll
+copy %libuv_ROOT%\lib\Release\uv.dll torch\lib\uv.dll
+
+:: create virtual environment
+python -m venv .venv
+echo * > .venv\.gitignore
+call .\.venv\Scripts\activate
+where python
+
+:: python install dependencies
+python -m pip install --upgrade pip
+pip install -r requirements.txt
+:: DISTUTILS_USE_SDK should be set after psutil dependency
+set DISTUTILS_USE_SDK=1
+
+:: start sccache server and reset sccache stats
+sccache --start-server
+sccache --zero-stats
+sccache --show-stats
+
+:: Call PyTorch build script
+python setup.py bdist_wheel -d "%PYTORCH_FINAL_PACKAGE_DIR%"
+
+:: show sccache stats
+sccache --show-stats
+
+:: Check if installation was successful
+if %errorlevel% neq 0 (
+    echo "Failed on build_pytorch. (exitcode = %errorlevel%)"
+    exit /b 1
+)

.ci/pytorch/windows/arm64/smoke_test.bat

@@ -0,0 +1,49 @@
+@echo off
+setlocal
+
+if "%PACKAGE_TYPE%" == "wheel" goto wheel
+if "%PACKAGE_TYPE%" == "libtorch" goto libtorch
+
+echo "unknown package type"
+exit /b 1
+
+:wheel
+call %PYTORCH_ROOT%\.ci\pytorch\windows\arm64\bootstrap_tests.bat
+
+echo Running python rnn_smoke.py...
+python %PYTORCH_ROOT%\.ci\pytorch\test_example_code\rnn_smoke_win_arm64.py
+if errorlevel 1 exit /b 1
+
+echo Checking that basic CNN works...
+python %PYTORCH_ROOT%\.ci\pytorch\test_example_code\cnn_smoke_win_arm64.py
+if errorlevel 1 exit /b 1
+
+goto end
+
+:libtorch
+echo "install and test libtorch"
+
+if not exist tmp mkdir tmp
+
+for /F "delims=" %%i in ('where /R "%PYTORCH_FINAL_PACKAGE_DIR:/=\%" *-latest.zip') do C:\Windows\System32\tar.exe -xf "%%i" -C tmp
+if ERRORLEVEL 1 exit /b 1
+
+pushd tmp
+
+set VC_VERSION_LOWER=14
+set VC_VERSION_UPPER=36
+
+call "%DEPENDENCIES_DIR%\VSBuildTools\VC\Auxiliary\Build\vcvarsall.bat" arm64
+
+set install_root=%CD%
+set INCLUDE=%INCLUDE%;%install_root%\include;%install_root%\include\torch\csrc\api\include
+set LIB=%LIB%;%install_root%\lib
+set PATH=%PATH%;%install_root%\lib
+
+cl %PYTORCH_ROOT%\.ci\pytorch\test_example_code\simple-torch-test.cpp c10.lib torch_cpu.lib /EHsc /std:c++17
+if ERRORLEVEL 1 exit /b 1
+
+.\simple-torch-test.exe
+if ERRORLEVEL 1 exit /b 1
+
+:end

.ci/pytorch/windows/condaenv.bat

@@ -9,12 +9,13 @@ FOR %%v IN (%DESIRED_PYTHON%) DO (
     set PYTHON_VERSION_STR=%%v
     set PYTHON_VERSION_STR=!PYTHON_VERSION_STR:.=!
     conda remove -n py!PYTHON_VERSION_STR! --all -y || rmdir %CONDA_HOME%\envs\py!PYTHON_VERSION_STR! /s
-    if "%%v" == "3.8" call conda create -n py!PYTHON_VERSION_STR! -y -q numpy=1.11 pyyaml boto3 cmake ninja typing_extensions setuptools=72.1.0 python=%%v
-    if "%%v" == "3.9" call conda create -n py!PYTHON_VERSION_STR! -y -q numpy=2.0.1 pyyaml boto3 cmake ninja typing_extensions setuptools=72.1.0 python=%%v
-    if "%%v" == "3.10" call conda create -n py!PYTHON_VERSION_STR! -y -q -c=conda-forge numpy=2.0.1 pyyaml boto3 cmake ninja typing_extensions setuptools=72.1.0 python=%%v
-    if "%%v" == "3.11" call conda create -n py!PYTHON_VERSION_STR! -y -q -c=conda-forge numpy=2.0.1 pyyaml boto3 cmake ninja typing_extensions setuptools=72.1.0 python=%%v
-    if "%%v" == "3.12" call conda create -n py!PYTHON_VERSION_STR! -y -q -c=conda-forge numpy=2.0.1 pyyaml boto3 cmake ninja typing_extensions setuptools=72.1.0 python=%%v
-    if "%%v" == "3.13" call conda create -n py!PYTHON_VERSION_STR! -y -q -c=conda-forge numpy=2.1.2 pyyaml boto3 cmake ninja typing_extensions setuptools=72.1.0 python=%%v
+    if "%%v" == "3.9" call conda create -n py!PYTHON_VERSION_STR! -y numpy=2.0.1 boto3 cmake ninja typing_extensions setuptools=72.1.0 python=%%v
+    if "%%v" == "3.10" call conda create -n py!PYTHON_VERSION_STR! -y -c=conda-forge numpy=2.0.1  boto3 cmake ninja typing_extensions setuptools=72.1.0 python=%%v
+    if "%%v" == "3.11" call conda create -n py!PYTHON_VERSION_STR! -y -c=conda-forge numpy=2.0.1  boto3 cmake ninja typing_extensions setuptools=72.1.0 python=%%v
+    if "%%v" == "3.12" call conda create -n py!PYTHON_VERSION_STR! -y -c=conda-forge numpy=2.0.1  boto3 cmake ninja typing_extensions setuptools=72.1.0 python=%%v
+    if "%%v" == "3.13" call conda create -n py!PYTHON_VERSION_STR! -y -c=conda-forge numpy=2.1.2  boto3 cmake ninja typing_extensions setuptools=72.1.0 python=%%v
+    if "%%v" == "3.13t" call conda create -n py!PYTHON_VERSION_STR! -y -c=conda-forge numpy=2.1.2 boto3 cmake ninja typing_extensions setuptools=72.1.0 python-freethreading python=3.13
+    call conda run -n py!PYTHON_VERSION_STR! pip install pyyaml
     call conda run -n py!PYTHON_VERSION_STR! pip install mkl-include
     call conda run -n py!PYTHON_VERSION_STR! pip install mkl-static
 )

.ci/pytorch/windows/cuda128.bat

@@ -0,0 +1,59 @@
+@echo off
+
+set MODULE_NAME=pytorch
+
+IF NOT EXIST "setup.py" IF NOT EXIST "%MODULE_NAME%" (
+    call internal\clone.bat
+    cd %~dp0
+) ELSE (
+    call internal\clean.bat
+)
+IF ERRORLEVEL 1 goto :eof
+
+call internal\check_deps.bat
+IF ERRORLEVEL 1 goto :eof
+
+REM Check for optional components
+
+set USE_CUDA=
+set CMAKE_GENERATOR=Visual Studio 15 2017 Win64
+
+IF "%NVTOOLSEXT_PATH%"=="" (
+    IF EXIST "C:\Program Files\NVIDIA Corporation\NvToolsExt\lib\x64\nvToolsExt64_1.lib"  (
+        set NVTOOLSEXT_PATH=C:\Program Files\NVIDIA Corporation\NvToolsExt
+    ) ELSE (
+        echo NVTX ^(Visual Studio Extension ^for CUDA^) ^not installed, failing
+        exit /b 1
+    )
+)
+
+IF "%CUDA_PATH_V128%"=="" (
+    IF EXIST "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8\bin\nvcc.exe" (
+        set "CUDA_PATH_V128=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.8"
+    ) ELSE (
+        echo CUDA 12.8 not found, failing
+        exit /b 1
+    )
+)
+
+IF "%BUILD_VISION%" == "" (
+    set TORCH_CUDA_ARCH_LIST=5.0;6.0;6.1;7.0;7.5;8.0;8.6;9.0;10.0;12.0
+    set TORCH_NVCC_FLAGS=-Xfatbin -compress-all
+) ELSE (
+    set NVCC_FLAGS=-D__CUDA_NO_HALF_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_50,code=sm_50 -gencode=arch=compute_60,code=sm_60 -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_80,code=compute_80 -gencode=arch=compute_86,code=compute_86 -gencode=arch=compute_90,code=compute_90 -gencode=arch=compute_100,code=compute_100 -gencode=arch=compute_120,code=compute_120
+)
+
+set "CUDA_PATH=%CUDA_PATH_V128%"
+set "PATH=%CUDA_PATH_V128%\bin;%PATH%"
+
+:optcheck
+
+call internal\check_opts.bat
+IF ERRORLEVEL 1 goto :eof
+
+if exist "%NIGHTLIES_PYTORCH_ROOT%" cd %NIGHTLIES_PYTORCH_ROOT%\..
+call  %~dp0\internal\copy.bat
+IF ERRORLEVEL 1 goto :eof
+
+call  %~dp0\internal\setup.bat
+IF ERRORLEVEL 1 goto :eof

.ci/pytorch/windows/internal/cuda_install.bat

@@ -9,7 +9,8 @@ if "%CUDA_VERSION%" == "xpu" (
     exit /b 0
 )
 
-set SRC_DIR=%NIGHTLIES_PYTORCH_ROOT%
+set SRC_DIR=%~dp0\..
+
 if not exist "%SRC_DIR%\temp_build" mkdir "%SRC_DIR%\temp_build"
 
 set /a CUDA_VER=%CUDA_VERSION%
@@ -23,9 +24,9 @@ set CUDNN_LIB_FOLDER="lib\x64"
 if exist "C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v%CUDA_VERSION_STR%\bin\nvcc.exe" goto set_cuda_env_vars
 
 if %CUDA_VER% EQU 118 goto cuda118
-if %CUDA_VER% EQU 121 goto cuda121
 if %CUDA_VER% EQU 124 goto cuda124
 if %CUDA_VER% EQU 126 goto cuda126
+if %CUDA_VER% EQU 128 goto cuda128
 
 echo CUDA %CUDA_VERSION_STR% is not supported
 exit /b 1
@@ -111,6 +112,33 @@ xcopy /Y "%SRC_DIR%\temp_build\zlib\dll_x64\*.dll" "C:\Windows\System32"
 
 goto cuda_common
 
+:cuda128
+
+set CUDA_INSTALL_EXE=cuda_12.8.0_571.96_windows.exe
+if not exist "%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%" (
+    curl -k -L "https://ossci-windows.s3.amazonaws.com/%CUDA_INSTALL_EXE%" --output "%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%"
+    if errorlevel 1 exit /b 1
+    set "CUDA_SETUP_FILE=%SRC_DIR%\temp_build\%CUDA_INSTALL_EXE%"
+    set "ARGS=cuda_profiler_api_12.8 thrust_12.8 nvcc_12.8 cuobjdump_12.8 nvprune_12.8 nvprof_12.8 cupti_12.8 cublas_12.8 cublas_dev_12.8 cudart_12.8 cufft_12.8 cufft_dev_12.8 curand_12.8 curand_dev_12.8 cusolver_12.8 cusolver_dev_12.8 cusparse_12.8 cusparse_dev_12.8 npp_12.8 npp_dev_12.8 nvrtc_12.8 nvrtc_dev_12.8 nvml_dev_12.8 nvjitlink_12.8 nvtx_12.8"
+)
+
+set CUDNN_FOLDER=cudnn-windows-x86_64-9.7.0.66_cuda12-archive
+set CUDNN_LIB_FOLDER="lib"
+set "CUDNN_INSTALL_ZIP=%CUDNN_FOLDER%.zip"
+if not exist "%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%" (
+    curl -k -L "http://s3.amazonaws.com/ossci-windows/%CUDNN_INSTALL_ZIP%" --output "%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%"
+    if errorlevel 1 exit /b 1
+    set "CUDNN_SETUP_FILE=%SRC_DIR%\temp_build\%CUDNN_INSTALL_ZIP%"
+)
+
+@REM cuDNN 8.3+ required zlib to be installed on the path
+echo Installing ZLIB dlls
+curl -k -L "http://s3.amazonaws.com/ossci-windows/zlib123dllx64.zip" --output "%SRC_DIR%\temp_build\zlib123dllx64.zip"
+7z x "%SRC_DIR%\temp_build\zlib123dllx64.zip" -o"%SRC_DIR%\temp_build\zlib"
+xcopy /Y "%SRC_DIR%\temp_build\zlib\dll_x64\*.dll" "C:\Windows\System32"
+
+goto cuda_common
+
 :cuda_common
 :: NOTE: We only install CUDA if we don't have it installed already.
 :: With GHA runners these should be pre-installed as part of our AMI process

.ci/pytorch/windows/internal/smoke_test.bat

@@ -27,7 +27,6 @@ for /F "delims=" %%i in ('wmic path win32_VideoController get name') do (
 endlocal & set NVIDIA_GPU_EXISTS=%NVIDIA_GPU_EXISTS%
 
 if "%PACKAGE_TYPE%" == "wheel" goto wheel
-if "%PACKAGE_TYPE%" == "conda" goto conda
 if "%PACKAGE_TYPE%" == "libtorch" goto libtorch
 
 echo "unknown package type"
@@ -37,16 +36,23 @@ exit /b 1
 echo "install wheel package"
 
 set PYTHON_INSTALLER_URL=
+if "%DESIRED_PYTHON%" == "3.13t" set "PYTHON_INSTALLER_URL=https://www.python.org/ftp/python/3.13.0/python-3.13.0-amd64.exe"
 if "%DESIRED_PYTHON%" == "3.13" set "PYTHON_INSTALLER_URL=https://www.python.org/ftp/python/3.13.0/python-3.13.0-amd64.exe"
 if "%DESIRED_PYTHON%" == "3.12" set "PYTHON_INSTALLER_URL=https://www.python.org/ftp/python/3.12.0/python-3.12.0-amd64.exe"
 if "%DESIRED_PYTHON%" == "3.11" set "PYTHON_INSTALLER_URL=https://www.python.org/ftp/python/3.11.0/python-3.11.0-amd64.exe"
 if "%DESIRED_PYTHON%" == "3.10" set "PYTHON_INSTALLER_URL=https://www.python.org/ftp/python/3.10.0/python-3.10.0-amd64.exe"
 if "%DESIRED_PYTHON%" == "3.9" set "PYTHON_INSTALLER_URL=https://www.python.org/ftp/python/3.9.0/python-3.9.0-amd64.exe"
-if "%DESIRED_PYTHON%" == "3.8" set "PYTHON_INSTALLER_URL=https://www.python.org/ftp/python/3.8.2/python-3.8.2-amd64.exe"
 if "%PYTHON_INSTALLER_URL%" == "" (
     echo Python %DESIRED_PYTHON% not supported yet
 )
 
+set ADDITIONAL_OPTIONS=""
+set PYTHON_EXEC="python"
+if "%DESIRED_PYTHON%" == "3.13t" (
+    set ADDITIONAL_OPTIONS="Include_freethreaded=1"
+    set PYTHON_EXEC="python3.13t"
+)
+
 del python-amd64.exe
 curl --retry 3 -kL "%PYTHON_INSTALLER_URL%" --output python-amd64.exe
 if errorlevel 1 exit /b 1
@@ -55,85 +61,39 @@ if errorlevel 1 exit /b 1
 :: the installed Python to PATH system-wide. Even calling set PATH=%ORIG_PATH% later on won't make
 :: a change. As the builder directory will be removed after the smoke test, all subsequent non-binary
 :: jobs will fail to find any Python executable there
-start /wait "" python-amd64.exe /quiet InstallAllUsers=1 PrependPath=0 Include_test=0 TargetDir=%CD%\Python
+start /wait "" python-amd64.exe /quiet InstallAllUsers=1 PrependPath=0 Include_test=0 %ADDITIONAL_OPTIONS% TargetDir=%CD%\Python
 if errorlevel 1 exit /b 1
 
 set "PATH=%CD%\Python%PYTHON_VERSION%\Scripts;%CD%\Python;%PATH%"
-
-if "%DESIRED_PYTHON%" == "3.13" pip install -q --pre numpy==2.1.0 protobuf
-if "%DESIRED_PYTHON%" == "3.12" pip install -q --pre numpy==2.0.2 protobuf
-if "%DESIRED_PYTHON%" == "3.11" pip install -q --pre numpy==2.0.2 protobuf
-if "%DESIRED_PYTHON%" == "3.10" pip install -q --pre numpy==2.0.2 protobuf
-if "%DESIRED_PYTHON%" == "3.9" pip install -q --pre numpy==2.0.2 protobuf
-if "%DESIRED_PYTHON%" == "3.8" pip install -q numpy protobuf
+if "%DESIRED_PYTHON%" == "3.13t" %PYTHON_EXEC% -m pip install --pre numpy==2.2.1 protobuf
+if "%DESIRED_PYTHON%" == "3.13" %PYTHON_EXEC% -m pip install --pre numpy==2.1.2 protobuf
+if "%DESIRED_PYTHON%" == "3.12" %PYTHON_EXEC% -m pip install --pre numpy==2.0.2 protobuf
+if "%DESIRED_PYTHON%" == "3.11" %PYTHON_EXEC% -m pip install --pre numpy==2.0.2 protobuf
+if "%DESIRED_PYTHON%" == "3.10" %PYTHON_EXEC% -m pip install --pre numpy==2.0.2 protobuf
+if "%DESIRED_PYTHON%" == "3.9" %PYTHON_EXEC% -m pip install --pre numpy==2.0.2 protobuf networkx
 
 if errorlevel 1 exit /b 1
 
-for /F "delims=" %%i in ('where /R "%PYTORCH_FINAL_PACKAGE_DIR:/=\%" *.whl') do pip install "%%i"
+if "%PYTORCH_BUILD_VERSION:dev=%" NEQ "%PYTORCH_BUILD_VERSION%" (
+    set "CHANNEL=nightly"
+) else (
+    set "CHANNEL=test"
+)
+
+set "EXTRA_INDEX= "
+if "%CUDA_VERSION%" == "xpu" set "EXTRA_INDEX=--index-url https://download.pytorch.org/whl/%CHANNEL%/xpu"
+
+for /F "delims=" %%i in ('where /R "%PYTORCH_FINAL_PACKAGE_DIR:/=\%" *.whl') do %PYTHON_EXEC% -m pip install "%%i" %EXTRA_INDEX%
 if errorlevel 1 exit /b 1
 
 goto smoke_test
 
-:conda
-echo "install conda package"
-
-:: Install Miniconda3
-set "CONDA_HOME=%CD%\conda"
-set "tmp_conda=%CONDA_HOME%"
-set "miniconda_exe=%CD%\miniconda.exe"
-set "CONDA_EXTRA_ARGS=cpuonly -c pytorch-nightly"
-if "%CUDA_VERSION%" == "118" (
-    set "CONDA_EXTRA_ARGS=pytorch-cuda=11.8 -c nvidia -c pytorch-nightly"
-)
-if "%CUDA_VERSION%" == "121" (
-    set "CONDA_EXTRA_ARGS=pytorch-cuda=12.1 -c nvidia -c pytorch-nightly"
-)
-if "%CUDA_VERSION%" == "124" (
-    set "CONDA_EXTRA_ARGS=pytorch-cuda=12.4 -c nvidia -c pytorch-nightly"
-)
-if "%CUDA_VERSION%" == "126" (
-    set "CONDA_EXTRA_ARGS=pytorch-cuda=12.6 -c nvidia -c pytorch-nightly"
-)
-
-rmdir /s /q conda
-del miniconda.exe
-curl -k https://repo.anaconda.com/miniconda/Miniconda3-latest-Windows-x86_64.exe -o "%miniconda_exe%"
-start /wait "" "%miniconda_exe%" /S /InstallationType=JustMe /RegisterPython=0 /AddToPath=0 /D=%tmp_conda%
-if ERRORLEVEL 1 exit /b 1
-
-set "PATH=%CONDA_HOME%;%CONDA_HOME%\scripts;%CONDA_HOME%\Library\bin;%PATH%"
-
-conda create -qyn testenv python=%DESIRED_PYTHON%
-if errorlevel 1 exit /b 1
-call conda install -yq conda-build
-if errorlevel 1 exit /b 1
-call %CONDA_HOME%\condabin\activate.bat testenv
-if errorlevel 1 exit /b 1
-set "NO_ARCH_PATH=%PYTORCH_FINAL_PACKAGE_DIR:/=\%\noarch"
-mkdir %NO_ARCH_PATH%
-for /F "delims=" %%i in ('where /R "%PYTORCH_FINAL_PACKAGE_DIR:/=\%" *') do xcopy "%%i" %NO_ARCH_PATH% /Y
-if ERRORLEVEL 1 exit /b 1
-call conda index %PYTORCH_FINAL_PACKAGE_DIR%
-if errorlevel 1 exit /b 1
-call conda install -yq -c "file:///%PYTORCH_FINAL_PACKAGE_DIR%" pytorch==%PYTORCH_BUILD_VERSION% -c pytorch -c numba/label/dev -c nvidia
-if ERRORLEVEL 1 exit /b 1
-call conda install -yq numpy
-if ERRORLEVEL 1 exit /b 1
-
-set /a CUDA_VER=%CUDA_VERSION%
-set CUDA_VER_MAJOR=%CUDA_VERSION:~0,-1%
-set CUDA_VER_MINOR=%CUDA_VERSION:~-1,1%
-set CUDA_VERSION_STR=%CUDA_VER_MAJOR%.%CUDA_VER_MINOR%
-
-:: Install package we just build
-
-
 :smoke_test
-python -c "import torch"
+%PYTHON_EXEC% -c "import torch"
 if ERRORLEVEL 1 exit /b 1
 
 echo Checking that MKL is available
-python -c "import torch; exit(0 if torch.backends.mkl.is_available() else 1)"
+%PYTHON_EXEC% -c "import torch; exit(0 if torch.backends.mkl.is_available() else 1)"
 if ERRORLEVEL 1 exit /b 1
 
 if "%NVIDIA_GPU_EXISTS%" == "0" (
@@ -142,24 +102,24 @@ if "%NVIDIA_GPU_EXISTS%" == "0" (
 )
 
 echo Checking that CUDA archs are setup correctly
-python -c "import torch; torch.randn([3,5]).cuda()"
+%PYTHON_EXEC% -c "import torch; torch.randn([3,5]).cuda()"
 if ERRORLEVEL 1 exit /b 1
 
 echo Checking that magma is available
-python -c "import torch; torch.rand(1).cuda(); exit(0 if torch.cuda.has_magma else 1)"
+%PYTHON_EXEC% -c "import torch; torch.rand(1).cuda(); exit(0 if torch.cuda.has_magma else 1)"
 if ERRORLEVEL 1 exit /b 1
 
 echo Checking that CuDNN is available
-python -c "import torch; exit(0 if torch.backends.cudnn.is_available() else 1)"
+%PYTHON_EXEC% -c "import torch; exit(0 if torch.backends.cudnn.is_available() else 1)"
 if ERRORLEVEL 1 exit /b 1
 
 echo Checking that basic RNN works
-python %PYTORCH_ROOT%\.ci\pytorch\test_example_code\rnn_smoke.py
+%PYTHON_EXEC% %PYTORCH_ROOT%\.ci\pytorch\test_example_code\rnn_smoke.py
 
 if ERRORLEVEL 1 exit /b 1
 
 echo Checking that basic CNN works
-python %PYTORCH_ROOT%\.ci\pytorch\test_example_code\cnn_smoke.py
+%PYTHON_EXEC% %PYTORCH_ROOT%\.ci\pytorch\test_example_code\cnn_smoke.py
 if ERRORLEVEL 1 exit /b 1
 
 goto end
@@ -167,7 +127,6 @@ goto end
 :libtorch
 echo "install and test libtorch"
 
-if "%VC_YEAR%" == "2019" powershell internal\vs2019_install.ps1
 if "%VC_YEAR%" == "2022" powershell internal\vs2022_install.ps1
 
 if ERRORLEVEL 1 exit /b 1
@@ -179,10 +138,6 @@ pushd tmp\libtorch
 
 set VC_VERSION_LOWER=17
 set VC_VERSION_UPPER=18
-IF "%VC_YEAR%" == "2019" (
-    set VC_VERSION_LOWER=16
-    set VC_VERSION_UPPER=17
-)
 
 for /f "usebackq tokens=*" %%i in (`"%ProgramFiles(x86)%\Microsoft Visual Studio\Installer\vswhere.exe" -legacy -products * -version [%VC_VERSION_LOWER%^,%VC_VERSION_UPPER%^) -property installationPath`) do (
     if exist "%%i" if exist "%%i\VC\Auxiliary\Build\vcvarsall.bat" (

.ci/pytorch/windows/internal/static_lib_test.bat

@@ -70,7 +70,6 @@ echo "install and test libtorch"
 pip install cmake
 echo "installing cmake"
 
-if "%VC_YEAR%" == "2019" powershell internal\vs2019_install.ps1
 if "%VC_YEAR%" == "2022" powershell internal\vs2022_install.ps1
 
 if ERRORLEVEL 1 exit /b 1
@@ -83,10 +82,6 @@ pushd tmp\libtorch
 
 set VC_VERSION_LOWER=17
 set VC_VERSION_UPPER=18
-IF "%VC_YEAR%" == "2019" (
-    set VC_VERSION_LOWER=16
-    set VC_VERSION_UPPER=17
-)
 
 for /f "usebackq tokens=*" %%i in (`"%ProgramFiles(x86)%\Microsoft Visual Studio\Installer\vswhere.exe" -legacy -products * -version [%VC_VERSION_LOWER%^,%VC_VERSION_UPPER%^) -property installationPath`) do (
     if exist "%%i" if exist "%%i\VC\Auxiliary\Build\vcvarsall.bat" (

.ci/pytorch/windows/internal/vc_install_helper.bat

@@ -1,12 +1,8 @@
-if "%VC_YEAR%" == "2019" powershell windows/internal/vs2019_install.ps1
 if "%VC_YEAR%" == "2022" powershell windows/internal/vs2022_install.ps1
 
 set VC_VERSION_LOWER=17
 set VC_VERSION_UPPER=18
-if "%VC_YEAR%" == "2019" (
-    set VC_VERSION_LOWER=16
-    set VC_VERSION_UPPER=17
-)
+
 
 for /f "usebackq tokens=*" %%i in (`"%ProgramFiles(x86)%\Microsoft Visual Studio\Installer\vswhere.exe"  -products Microsoft.VisualStudio.Product.BuildTools -version [%VC_VERSION_LOWER%^,%VC_VERSION_UPPER%^) -property installationPath`) do (
     if exist "%%i" if exist "%%i\VC\Auxiliary\Build\vcvarsall.bat" (

.ci/pytorch/windows/internal/vs2019_install.ps1

@@ -1,48 +0,0 @@
-# https://developercommunity.visualstudio.com/t/install-specific-version-of-vs-component/1142479
-# https://docs.microsoft.com/en-us/visualstudio/releases/2019/history#release-dates-and-build-numbers
-
-# 16.8.6 BuildTools
-$VS_DOWNLOAD_LINK = "https://ossci-windows.s3.us-east-1.amazonaws.com/vs16.8.6_BuildTools.exe"
-$COLLECT_DOWNLOAD_LINK = "https://aka.ms/vscollect.exe"
-$VS_INSTALL_ARGS = @("--nocache","--quiet","--wait", "--add Microsoft.VisualStudio.Workload.VCTools",
-                                                     "--add Microsoft.Component.MSBuild",
-                                                     "--add Microsoft.VisualStudio.Component.Roslyn.Compiler",
-                                                     "--add Microsoft.VisualStudio.Component.TextTemplating",
-                                                     "--add Microsoft.VisualStudio.Component.VC.CoreIde",
-                                                     "--add Microsoft.VisualStudio.Component.VC.Redist.14.Latest",
-                                                     "--add Microsoft.VisualStudio.ComponentGroup.NativeDesktop.Core",
-                                                     "--add Microsoft.VisualStudio.Component.VC.Tools.x86.x64",
-                                                     "--add Microsoft.VisualStudio.ComponentGroup.NativeDesktop.Win81")
-
-curl.exe --retry 3 -kL $VS_DOWNLOAD_LINK --output vs_installer.exe
-if ($LASTEXITCODE -ne 0) {
-    echo "Download of the VS 2019 Version 16.8.5 installer failed"
-    exit 1
-}
-
-if (Test-Path "${env:ProgramFiles(x86)}\Microsoft Visual Studio\Installer\vswhere.exe") {
-    $existingPath = & "${env:ProgramFiles(x86)}\Microsoft Visual Studio\Installer\vswhere.exe" -products "Microsoft.VisualStudio.Product.BuildTools" -version "[16, 17)" -property installationPath
-    if ($existingPath -ne $null) {
-        if (!${env:CIRCLECI}) {
-            echo "Found correctly versioned existing BuildTools installation in $existingPath"
-            exit 0
-        }
-        echo "Found existing BuildTools installation in $existingPath, keeping it"
-    }
-}
-
-$process = Start-Process "${PWD}\vs_installer.exe" -ArgumentList $VS_INSTALL_ARGS -NoNewWindow -Wait -PassThru
-Remove-Item -Path vs_installer.exe -Force
-$exitCode = $process.ExitCode
-if (($exitCode -ne 0) -and ($exitCode -ne 3010)) {
-    echo "VS 2019 installer exited with code $exitCode, which should be one of [0, 3010]."
-    curl.exe --retry 3 -kL $COLLECT_DOWNLOAD_LINK --output Collect.exe
-    if ($LASTEXITCODE -ne 0) {
-        echo "Download of the VS Collect tool failed."
-        exit 1
-    }
-    Start-Process "${PWD}\Collect.exe" -NoNewWindow -Wait -PassThru
-    New-Item -Path "C:\w\build-results" -ItemType "directory" -Force
-    Copy-Item -Path "C:\Users\${env:USERNAME}\AppData\Local\Temp\vslogs.zip" -Destination "C:\w\build-results\"
-    exit 1
-}

.ci/pytorch/windows/internal/xpu_install.bat

@@ -47,9 +47,9 @@ set XPU_EXTRA_INSTALLED=0
 set XPU_EXTRA_UNINSTALL=0
 
 if not [%XPU_VERSION%]==[] if [%XPU_VERSION%]==[2025.0] (
-    set XPU_BUNDLE_URL=https://registrationcenter-download.intel.com/akdlm/IRC_NAS/efc86abd-cb77-452e-a03f-a741895b8ece/intel-deep-learning-essentials-2025.0.0.336_offline.exe
+    set XPU_BUNDLE_URL=https://registrationcenter-download.intel.com/akdlm/IRC_NAS/9d6d6c17-ca2d-4735-9331-99447e4a1280/intel-deep-learning-essentials-2025.0.1.28_offline.exe
     set XPU_BUNDLE_PRODUCT_NAME=intel.oneapi.win.deep-learning-essentials.product
-    set XPU_BUNDLE_VERSION=2025.0.0+335
+    set XPU_BUNDLE_VERSION=2025.0.1+20
     set XPU_BUNDLE_INSTALLED=0
     set XPU_BUNDLE_UNINSTALL=0
     set XPU_EXTRA_URL=NULL
@@ -104,14 +104,6 @@ goto xpu_install_end
 
 :xpu_bundle_install
 
-:: Install Level Zero SDK
-set XPU_EXTRA_LZ_URL=https://github.com/oneapi-src/level-zero/releases/download/v1.14.0/level-zero-sdk_1.14.0.zip
-curl -k -L %XPU_EXTRA_LZ_URL% --output "%SRC_DIR%\temp_build\level_zero_sdk.zip"
-echo "Installing level zero SDK..."
-7z x "%SRC_DIR%\temp_build\level_zero_sdk.zip" -o"%SRC_DIR%\temp_build\level_zero"
-set "INCLUDE=%SRC_DIR%\temp_build\level_zero\include;%INCLUDE%"
-
-:: Install Bundle
 curl -o xpu_bundle.exe --retry 3 --retry-all-errors -k %XPU_BUNDLE_URL%
 echo "XPU Bundle installing..."
 start /wait "Intel Pytorch Bundle Installer" "xpu_bundle.exe" --action=install --eula=accept --silent --log-dir install_bundle
@@ -128,3 +120,14 @@ if errorlevel 1 exit /b 1
 del xpu_extra.exe
 
 :xpu_install_end
+
+if not "%XPU_ENABLE_KINETO%"=="1" goto install_end
+:: Install Level Zero SDK
+set XPU_EXTRA_LZ_URL=https://github.com/oneapi-src/level-zero/releases/download/v1.14.0/level-zero-sdk_1.14.0.zip
+curl -k -L %XPU_EXTRA_LZ_URL% --output "%SRC_DIR%\temp_build\level_zero_sdk.zip"
+echo "Installing level zero SDK..."
+7z x "%SRC_DIR%\temp_build\level_zero_sdk.zip" -o"%SRC_DIR%\temp_build\level_zero"
+set "INCLUDE=%SRC_DIR%\temp_build\level_zero\include;%INCLUDE%"
+del "%SRC_DIR%\temp_build\level_zero_sdk.zip"
+
+:install_end

.ci/pytorch/windows/xpu.bat

@@ -28,11 +28,6 @@ call "%XPU_BUNDLE_ROOT%\compiler\latest\env\vars.bat"
 call "%XPU_BUNDLE_ROOT%\ocloc\latest\env\vars.bat"
 IF ERRORLEVEL 1 goto :eof
 
-:: Workaround for https://github.com/pytorch/pytorch/issues/134989
-set CMAKE_SHARED_LINKER_FLAGS=/FORCE:MULTIPLE
-set CMAKE_MODULE_LINKER_FLAGS=/FORCE:MULTIPLE
-set CMAKE_EXE_LINKER_FLAGS=/FORCE:MULTIPLE
-
 if exist "%NIGHTLIES_PYTORCH_ROOT%" cd %NIGHTLIES_PYTORCH_ROOT%\..
 call %~dp0\internal\copy_cpu.bat
 IF ERRORLEVEL 1 goto :eof