Update on "remove allocation of new unbacked symbols during mod eval"

When executing code like torch._check(numel % newsize == 0, ...), we previously allocated a new unbacked symbol due to #113165. However, this allocation is no longer necessary and can cause issues due to inconsistent behavior when tracing torch._check multiple times.

In particular, the allocation can lead to a memoization disaster where the previously allocated symbol is returned instead of a new one, causing unexpected behavior.

This PR removes the unnecessary allocation, ensuring consistent behavior and avoiding potential issues. The change is validated by the following code, which now compiles without issues:
```
import torch

def fn(x):
    i0 = x.nonzero().size(0)
    y = torch.zeros((i0, 192))
    return y.view([12, -1, 192])
with torch._dynamo.config.patch({"capture_dynamic_output_shape_ops": True}):
    torch.compile(fn, fullgraph=True)(torch.ones((12,)))
```

By removing this unnecessary allocation, we simplify the code and avoid potential issues."

[ghstack-poisoned]

.ci/docker/almalinux/Dockerfile

@@ -7,13 +7,13 @@ ENV LC_ALL en_US.UTF-8
 ENV LANG en_US.UTF-8
 ENV LANGUAGE en_US.UTF-8
 
-ARG DEVTOOLSET_VERSION=13
+ARG DEVTOOLSET_VERSION=11
 
 RUN yum -y update
 RUN yum -y install epel-release
 # install glibc-langpack-en make sure en_US.UTF-8 locale is available
 RUN yum -y install glibc-langpack-en
-RUN yum install -y sudo wget curl perl util-linux xz bzip2 git patch which perl zlib-devel openssl-devel yum-utils autoconf automake make gcc-toolset-${DEVTOOLSET_VERSION}-gcc gcc-toolset-${DEVTOOLSET_VERSION}-gcc-c++ gcc-toolset-${DEVTOOLSET_VERSION}-gcc-gfortran gcc-toolset-${DEVTOOLSET_VERSION}-gdb
+RUN yum install -y sudo wget curl perl util-linux xz bzip2 git patch which perl zlib-devel openssl-devel yum-utils autoconf automake make gcc-toolset-${DEVTOOLSET_VERSION}-toolchain
 # Just add everything as a safe.directory for git since these will be used in multiple places with git
 RUN git config --global --add safe.directory '*'
 ENV PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/bin:$PATH
@@ -41,7 +41,6 @@ RUN bash ./install_conda.sh && rm install_conda.sh
 # Install CUDA
 FROM base as cuda
 ARG CUDA_VERSION=12.6
-ARG DEVTOOLSET_VERSION=13
 RUN rm -rf /usr/local/cuda-*
 ADD ./common/install_cuda.sh install_cuda.sh
 COPY ./common/install_nccl.sh install_nccl.sh
@@ -51,8 +50,7 @@ ENV CUDA_HOME=/usr/local/cuda-${CUDA_VERSION}
 # Preserve CUDA_VERSION for the builds
 ENV CUDA_VERSION=${CUDA_VERSION}
 # Make things in our path by default
-ENV PATH=/usr/local/cuda-${CUDA_VERSION}/bin:/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/bin:$PATH
-
+ENV PATH=/usr/local/cuda-${CUDA_VERSION}/bin:$PATH
 
 FROM cuda as cuda12.6
 RUN bash ./install_cuda.sh 12.6
@@ -70,22 +68,8 @@ FROM cuda as cuda13.0
 RUN bash ./install_cuda.sh 13.0
 ENV DESIRED_CUDA=13.0
 
-FROM ${ROCM_IMAGE} as rocm_base
-ARG DEVTOOLSET_VERSION=13
-ENV LC_ALL en_US.UTF-8
-ENV LANG en_US.UTF-8
-ENV LANGUAGE en_US.UTF-8
-# Install devtoolset on ROCm base image
-RUN yum -y update && \
-    yum -y install epel-release && \
-    yum -y install glibc-langpack-en && \
-    yum install -y sudo wget curl perl util-linux xz bzip2 git patch which perl zlib-devel openssl-devel yum-utils autoconf automake make gcc-toolset-${DEVTOOLSET_VERSION}-gcc gcc-toolset-${DEVTOOLSET_VERSION}-gcc-c++ gcc-toolset-${DEVTOOLSET_VERSION}-gcc-gfortran gcc-toolset-${DEVTOOLSET_VERSION}-gdb
-RUN git config --global --add safe.directory '*'
-ENV PATH=/opt/rh/gcc-toolset-${DEVTOOLSET_VERSION}/root/usr/bin:$PATH
-
-FROM rocm_base as rocm
+FROM ${ROCM_IMAGE} as rocm
 ARG PYTORCH_ROCM_ARCH
-ARG DEVTOOLSET_VERSION=13
 ENV PYTORCH_ROCM_ARCH ${PYTORCH_ROCM_ARCH}
 ADD ./common/install_mkl.sh install_mkl.sh
 RUN bash ./install_mkl.sh && rm install_mkl.sh
@@ -104,7 +88,6 @@ COPY --from=cuda13.0  /usr/local/cuda-13.0 /usr/local/cuda-13.0
 
 # Final step
 FROM ${BASE_TARGET} as final
-ARG DEVTOOLSET_VERSION=13
 COPY --from=openssl            /opt/openssl           /opt/openssl
 COPY --from=patchelf           /patchelf              /usr/local/bin/patchelf
 COPY --from=conda              /opt/conda             /opt/conda

.ci/docker/almalinux/build.sh

@@ -63,7 +63,7 @@ docker build \
   --target final \
   --progress plain \
   --build-arg "BASE_TARGET=${BASE_TARGET}" \
-  --build-arg "DEVTOOLSET_VERSION=13" \
+  --build-arg "DEVTOOLSET_VERSION=11" \
   ${EXTRA_BUILD_ARGS} \
   -t ${tmp_tag} \
   $@ \

.ci/docker/build.sh

@@ -261,9 +261,9 @@ case "$tag" in
     PYTHON_VERSION=3.10
     CUDA_VERSION=12.8.1
     ;;
-  pytorch-linux-jammy-aarch64-py3.10-gcc13)
+  pytorch-linux-jammy-aarch64-py3.10-gcc11)
     ANACONDA_PYTHON_VERSION=3.10
-    GCC_VERSION=13
+    GCC_VERSION=11
     ACL=yes
     VISION=yes
     OPENBLAS=yes
@@ -271,19 +271,9 @@ case "$tag" in
     # from pytorch/llvm:9.0.1 is x86 specific
     SKIP_LLVM_SRC_BUILD_INSTALL=yes
     ;;
-  pytorch-linux-jammy-aarch64-py3.10-clang21)
+  pytorch-linux-jammy-aarch64-py3.10-gcc11-inductor-benchmarks)
     ANACONDA_PYTHON_VERSION=3.10
-    CLANG_VERSION=21
-    ACL=yes
-    VISION=yes
-    OPENBLAS=yes
-    # snadampal: skipping llvm src build install because the current version
-    # from pytorch/llvm:9.0.1 is x86 specific
-    SKIP_LLVM_SRC_BUILD_INSTALL=yes
-    ;;
-  pytorch-linux-jammy-aarch64-py3.10-gcc13-inductor-benchmarks)
-    ANACONDA_PYTHON_VERSION=3.10
-    GCC_VERSION=13
+    GCC_VERSION=11
     ACL=yes
     VISION=yes
     OPENBLAS=yes

.ci/docker/common/install_clang.sh

@@ -8,8 +8,8 @@ if [ -n "$CLANG_VERSION" ]; 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 -
-    if [[ $CLANG_VERSION -ge 18 ]]; then
-      apt-add-repository "deb http://apt.llvm.org/jammy/ llvm-toolchain-jammy-${CLANG_VERSION} main"
+    if [[ $CLANG_VERSION == 18 ]]; then
+      apt-add-repository "deb http://apt.llvm.org/jammy/ llvm-toolchain-jammy-18 main"
     fi
   fi
 

.ci/docker/common/install_cuda.sh

@@ -129,7 +129,7 @@ function install_129 {
 }
 
 function install_128 {
-  CUDNN_VERSION=9.8.0.87
+  CUDNN_VERSION=9.10.2.21
   echo "Installing CUDA 12.8.1 and cuDNN ${CUDNN_VERSION} and NVSHMEM and NCCL and cuSparseLt-0.7.1"
   # install CUDA 12.8.1 in the same container
   install_cuda 12.8.1 cuda_12.8.1_570.124.06_linux

.ci/docker/common/install_gcc.sh

@@ -7,11 +7,11 @@ if [ -n "$GCC_VERSION" ]; then
   # Need the official toolchain repo to get alternate packages
   add-apt-repository ppa:ubuntu-toolchain-r/test
   apt-get update
-  apt-get install -y g++-$GCC_VERSION gfortran-$GCC_VERSION
+  apt-get install -y g++-$GCC_VERSION
   update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-"$GCC_VERSION" 50
   update-alternatives --install /usr/bin/g++ g++ /usr/bin/g++-"$GCC_VERSION" 50
   update-alternatives --install /usr/bin/gcov gcov /usr/bin/gcov-"$GCC_VERSION" 50
-  update-alternatives --install /usr/bin/gfortran gfortran /usr/bin/gfortran-"$GCC_VERSION" 50
+
 
   # Cleanup package manager
   apt-get autoclean && apt-get clean

.ci/docker/common/install_libgomp.sh

@@ -1,56 +0,0 @@
-#!/bin/bash
-# Script used only in CD pipeline
-
-set -ex
-
-# install dependencies
-dnf -y install gmp-devel libmpc-devel texinfo flex bison
-
-cd /usr/local/src
-# fetch source for gcc 13
-git clone --depth 1 --single-branch -b releases/gcc-13.3.0 https://github.com/gcc-mirror/gcc.git gcc-13.3.0
-
-mkdir -p gcc-13.3.0/build-gomp
-cd gcc-13.3.0/build-gomp
-
-# configure gcc build
-# I got these flags by:
-# 1. downloading the source rpm for gcc-11 on AlmaLinux 8 container
-#    dnf install -y dnf-plugins-core rpmdevtools
-#   dnf download --source libgomp
-# 2. extracting the gcc.spec from the source.
-#    rpmdev-extract gcc-xx.src.rpm
-# 3. extracting optflags and ld_flags from gcc.spec:
-#    rpm --eval '%{optflags}'
-#    rpm --eval '%{build_ldflags}'
-#
-# I had to remove the following flags because they didn't compile for this version of libgomp:
-#   -Werror=format-security
-#   -specs=/usr/lib/rpm/redhat/redhat-hardened-cc1
-#   -specs=/usr/lib/rpm/redhat/redhat-annobin-cc1
-#
-# I added -march=armv8-a -mtune=generic to make them explicit. I don't think they're strictly needed.
-
-OPT_FLAGS='-O2 -march=armv8-a -mtune=generic'\
-' -fexceptions -g -grecord-gcc-switches -pipe -Wall'\
-' -Wp,-D_FORTIFY_SOURCE=2 -Wp,-D_GLIBCXX_ASSERTIONS'\
-' -fstack-protector-strong -fasynchronous-unwind-tables'\
-' -fstack-clash-protection'
-
-LDFLAGS='-Wl,-z,relro -Wl,--as-needed -Wl,-z,now'
-
-CFLAGS="$OPT_FLAGS" \
-CXXFLAGS="$OPT_FLAGS" \
-LDFLAGS="$LDFLAGS" \
-../configure \
-  --prefix=/usr \
-  --libdir=/usr/lib64 \
-  --enable-languages=c,c++ \
-  --disable-multilib \
-  --disable-bootstrap \
-  --enable-libgomp
-
-# only build libgomp
-make -j$(nproc) all-target-libgomp
-
-make install-target-libgomp

.ci/docker/common/install_openblas.sh

@@ -10,7 +10,6 @@ git clone https://github.com/OpenMathLib/OpenBLAS.git -b "${OPENBLAS_VERSION}" -
 
 OPENBLAS_CHECKOUT_DIR="OpenBLAS"
 OPENBLAS_BUILD_FLAGS="
-CC=gcc
 NUM_THREADS=128
 USE_OPENMP=1
 NO_SHARED=0

.ci/docker/manywheel/Dockerfile_2_28_aarch64

@@ -50,10 +50,6 @@ RUN rm install_ninja.sh
 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
 
-# Build a newer version of libgomp than that supported in in Almalinux 8.
-COPY ./common/install_libgomp.sh install_libgomp.sh
-RUN bash ./install_libgomp.sh && rm install_libgomp.sh
-
 # 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

.ci/docker/requirements-docs.txt

@@ -1,11 +1,15 @@
-sphinx==7.2.6
+sphinx==5.3.0
 #Description: This is used to generate PyTorch docs
-#Pinned versions: 7.2.6
+#Pinned versions: 5.3.0
 
-pytorch_sphinx_theme2==0.2.0
-#Description: This is needed to generate PyTorch docs
-#Pinned versions: 0.2.0
+standard-imghdr==3.13.0; python_version >= "3.13"
+#Description: This is needed by Sphinx, so it needs to be added here.
+# The reasons are as follows:
+# 1) This module has been removed from the Python standard library since Python 3.13(https://peps.python.org/pep-0594/#imghdr);
+# 2) The current version of Sphinx (5.3.0) is not compatible with Python 3.13.
+# Once Sphinx is upgraded to a version compatible with Python 3.13 or later, we can remove this dependency.
 
+-e git+https://github.com/pytorch/pytorch_sphinx_theme.git@71e55749be14ceb56e7f8211a9fb649866b87ad4#egg=pytorch_sphinx_theme2
 # TODO: sphinxcontrib.katex 0.9.0 adds a local KaTeX server to speed up pre-rendering
 # but it doesn't seem to work and hangs around idly. The initial thought that it is probably
 # something related to Docker setup. We can investigate this later.
@@ -32,17 +36,17 @@ tensorboard==2.18.0 ; python_version >= "3.13"
 #Description: This is used to generate PyTorch docs
 #Pinned versions: 2.13.0
 
-breathe==4.36.0
+breathe==4.34.0
 #Description: This is used to generate PyTorch C++ docs
-#Pinned versions: 4.36.0
+#Pinned versions: 4.34.0
 
-exhale==0.3.7
+exhale==0.2.3
 #Description: This is used to generate PyTorch C++ docs
-#Pinned versions: 0.3.7
+#Pinned versions: 0.2.3
 
-docutils==0.20
+docutils==0.16
 #Description: This is used to generate PyTorch C++ docs
-#Pinned versions: 0.20
+#Pinned versions: 0.16
 
 bs4==0.0.1
 #Description: This is used to generate PyTorch C++ docs
@@ -52,13 +56,13 @@ IPython==8.12.0
 #Description: This is used to generate PyTorch functorch docs
 #Pinned versions: 8.12.0
 
-myst-nb==1.3.0
+myst-nb==0.17.2
 #Description: This is used to generate PyTorch functorch and torch.compile docs.
-#Pinned versions: 1.3.0
+#Pinned versions: 0.17.2
 
 # The following are required to build torch.distributed.elastic.rendezvous.etcd* docs
 python-etcd==0.4.5
 sphinx-copybutton==0.5.0
-sphinx-design==0.6.1
+sphinx-design==0.4.0
 sphinxcontrib-mermaid==1.0.0
-myst-parser==4.0.1
+myst-parser==0.18.1

.ci/magma-rocm/build_magma.sh

@@ -6,8 +6,8 @@ set -eou pipefail
 # The script expects DESIRED_CUDA and PACKAGE_NAME to be set
 ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
 
-# https://github.com/icl-utk-edu/magma/pull/65
-MAGMA_VERSION=d6e4117bc88e73f06d26c6c2e14f064e8fc3d1ec
+# post merge of https://github.com/icl-utk-edu/magma/pull/65
+MAGMA_VERSION=c0792ae825fb36872784892ea643dd6f3456bc5f
 
 # Folders for the build
 PACKAGE_FILES=${ROOT_DIR}/magma-rocm/package_files # metadata
@@ -20,7 +20,7 @@ mkdir -p ${PACKAGE_DIR} ${PACKAGE_OUTPUT}/linux-64 ${PACKAGE_BUILD} ${PACKAGE_RE
 
 # Fetch magma sources and verify checksum
 pushd ${PACKAGE_DIR}
-git clone https://github.com/jeffdaily/magma
+git clone https://github.com/icl-utk-edu/magma
 pushd magma
 git checkout ${MAGMA_VERSION}
 popd

.ci/pytorch/python_doc_push_script.sh

@@ -89,41 +89,23 @@ if [ "$is_main_doc" = true ]; then
 
   make coverage
   # Now we have the coverage report, we need to make sure it is empty.
-  # Sphinx 7.2.6+ format: python.txt contains a statistics table with a TOTAL row
-  # showing the undocumented count in the third column.
-  # Example: | TOTAL | 99.83% | 2 |
+  # Count the number of lines in the file and turn that number into a variable
+  # $lines. The `cut -f1 ...` is to only parse the number, not the filename
+  # Skip the report header by subtracting 2: the header will be output even if
+  # there are no undocumented items.
   #
   # Also: see docs/source/conf.py for "coverage_ignore*" items, which should
   # be documented then removed from there.
-
-  # Extract undocumented count from TOTAL row in Sphinx 7.2.6 statistics table
-  # The table format is: | Module | Coverage | Undocumented |
-  # Extract the third column (undocumented count) from the TOTAL row
-  undocumented=$(grep "| TOTAL" build/coverage/python.txt | awk -F'|' '{print $4}' | tr -d ' ')
-
-  if [ -z "$undocumented" ] || ! [[ "$undocumented" =~ ^[0-9]+$ ]]; then
+  lines=$(wc -l build/coverage/python.txt 2>/dev/null |cut -f1 -d' ')
+  undocumented=$((lines - 2))
+  if [ $undocumented -lt 0 ]; then
     echo coverage output not found
     exit 1
-  elif [ "$undocumented" -gt 0 ]; then
-    set +x  # Disable command echoing for cleaner output
-    echo ""
-    echo "====================="
-    echo "UNDOCUMENTED OBJECTS:"
-    echo "====================="
-    echo ""
-    # Find the line number of the TOTAL row and print only what comes after it
-    total_line=$(grep -n "| TOTAL" build/coverage/python.txt | cut -d: -f1)
-    if [ -n "$total_line" ]; then
-      # Print only the detailed list (skip the statistics table)
-      tail -n +$((total_line + 2)) build/coverage/python.txt
-    else
-      # Fallback to showing entire file if TOTAL line not found
-      cat build/coverage/python.txt
-    fi
-    echo ""
+  elif [ $undocumented -gt 0 ]; then
+    echo undocumented objects found:
+    cat build/coverage/python.txt
     echo "Make sure you've updated relevant .rsts in docs/source!"
-    echo "You can reproduce locally by running 'cd docs && make coverage && tail -n +\$((grep -n \"| TOTAL\" build/coverage/python.txt | cut -d: -f1) + 2)) build/coverage/python.txt'"
-    set -x  # Re-enable command echoing
+    echo "You can reproduce locally by running 'cd docs && make coverage && cat build/coverage/python.txt'"
     exit 1
   fi
 else

.ci/pytorch/smoke_test/smoke_test.py

@@ -272,6 +272,18 @@ def smoke_test_cuda(
         torch_cudnn_version = cudnn_to_version_str(torch.backends.cudnn.version())
         print(f"Torch cuDNN version: {torch_cudnn_version}")
 
+        torch_cudnn_compile_version = torch._C._cudnn.getCompileVersion()
+        print(f"Torch cuDNN compile-time version: {torch_cudnn_compile_version}")
+        torch_cudnn_runtime_version = tuple(
+            [int(x) for x in torch_cudnn_version.split(".")]
+        )
+        if torch_cudnn_runtime_version != torch_cudnn_compile_version:
+            raise RuntimeError(
+                "cuDNN runtime version doesn't match comple version. "
+                f"Loaded: {torch_cudnn_runtime_version} "
+                f"Expected: {torch_cudnn_compile_version}"
+            )
+
         if sys.platform in ["linux", "linux2"]:
             torch_nccl_version = ".".join(str(v) for v in torch.cuda.nccl.version())
             print(f"Torch nccl; version: {torch_nccl_version}")

.ci/pytorch/test.sh

@@ -337,7 +337,7 @@ test_python() {
 
 test_python_smoke() {
   # Smoke tests for H100/B200
-  time python test/run_test.py --include test_matmul_cuda test_scaled_matmul_cuda inductor/test_fp8 inductor/test_max_autotune inductor/test_cutedsl_grouped_mm $PYTHON_TEST_EXTRA_OPTION --upload-artifacts-while-running
+  time python test/run_test.py --include test_matmul_cuda test_scaled_matmul_cuda inductor/test_fp8 inductor/test_max_autotune $PYTHON_TEST_EXTRA_OPTION --upload-artifacts-while-running
   assert_git_not_dirty
 }
 

.ci/pytorch/win-test-helpers/arm64/build_pytorch.ps1

@@ -70,7 +70,7 @@ sccache --zero-stats
 sccache --show-stats
 
 # Build the wheel
-python -m build --wheel --no-isolation
+python -m build --wheel --no-build-isolation
 if ($LASTEXITCODE -ne 0) { exit 1 }
 
 # Install the wheel locally

.github/ISSUE_TEMPLATE/release-feature-request.yml

@@ -1,11 +1,11 @@
-name: 🚀 New Feature for Release
+name: 🚀 Release highlight for proposed Feature
 description: Submit a Release highlight for proposed Feature
 labels: ["release-feature-request"]
 
 body:
 - type: textarea
   attributes:
-    label: New Feature for Release
+    label: Release highlight for proposed Feature
     description: >
       Example: “A torch.special module, analogous to SciPy's special module.”
 - type: input

.github/ci_commit_pins/audio.txt

@@ -1 +1 @@
-ad5816f0eee1c873df1b7d371c69f1f811a89387
+3b0e7a6f192ca2715e7e6cbe5db007aea7165fe2

.github/workflows/docker-builds.yml

@@ -77,11 +77,9 @@ jobs:
           pytorch-linux-noble-riscv64-py3.12-gcc14
         ]
         include:
-          - docker-image-name: pytorch-linux-jammy-aarch64-py3.10-gcc13
+          - docker-image-name: pytorch-linux-jammy-aarch64-py3.10-gcc11
             runner: linux.arm64.m7g.4xlarge
-          - docker-image-name: pytorch-linux-jammy-aarch64-py3.10-clang21
-            runner: linux.arm64.m7g.4xlarge
-          - docker-image-name: pytorch-linux-jammy-aarch64-py3.10-gcc13-inductor-benchmarks
+          - docker-image-name: pytorch-linux-jammy-aarch64-py3.10-gcc11-inductor-benchmarks
             runner: linux.arm64.m7g.4xlarge
             timeout-minutes: 600
     # Docker uploads fail from LF runners, see https://github.com/pytorch/pytorch/pull/137358

.github/workflows/inductor-perf-test-nightly-aarch64.yml

@@ -72,7 +72,7 @@ jobs:
       runner_prefix: "${{ needs.get-label-type.outputs.label-type }}"
       runner: linux.arm64.m7g.4xlarge
       build-environment: linux-jammy-aarch64-py3.10
-      docker-image-name: ci-image:pytorch-linux-jammy-aarch64-py3.10-gcc13-inductor-benchmarks
+      docker-image-name: ci-image:pytorch-linux-jammy-aarch64-py3.10-gcc11-inductor-benchmarks
       test-matrix: |
         { include: [
           { config: "inductor_huggingface_perf_cpu_aarch64", shard: 1, num_shards: 9, runner: "linux.arm64.m7g.metal" },

.github/workflows/inductor-rocm.yml

@@ -2,7 +2,7 @@ name: inductor-rocm
 
 on:
   schedule:
-    - cron: 0 */3 * * *
+    - cron: 0 * * * *
   push:
     branches:
       - release/*

.github/workflows/linux-aarch64.yml

@@ -33,7 +33,7 @@ jobs:
     with:
       runner_prefix: ${{ needs.get-label-type.outputs.label-type }}
       build-environment: linux-jammy-aarch64-py3.10
-      docker-image-name: ci-image:pytorch-linux-jammy-aarch64-py3.10-gcc13
+      docker-image-name: ci-image:pytorch-linux-jammy-aarch64-py3.10-gcc11
       runner: linux.arm64.m7g.4xlarge
       test-matrix: |
         { include: [

.github/workflows/operator_benchmark.yml

@@ -60,7 +60,7 @@ jobs:
     with:
       build-environment: linux-jammy-aarch64-py3.10
       runner: linux.arm64.m7g.4xlarge
-      docker-image-name: ci-image:pytorch-linux-jammy-aarch64-py3.10-gcc13
+      docker-image-name: ci-image:pytorch-linux-jammy-aarch64-py3.10-gcc11
       test-matrix: |
         { include: [
           { config: "cpu_operator_benchmark_short", shard: 1, num_shards: 1, runner: "linux.arm64.m8g.4xlarge" },

.github/workflows/rocm.yml

@@ -9,8 +9,7 @@ on:
   workflow_dispatch:
   schedule:
     - cron: 29 8 * * *  # about 1:29am PDT
-    - cron: 0 */3 * * *
-
+    - cron: 0 * * * *
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }}

.gitignore

@@ -127,7 +127,6 @@ torch/test/
 torch/utils/benchmark/utils/valgrind_wrapper/callgrind.h
 torch/utils/benchmark/utils/valgrind_wrapper/valgrind.h
 torch/version.py
-torch/_inductor/kernel/vendored_templates/*
 minifier_launcher.py
 aten/src/ATen/native/transformers/hip/flash_attn/ck/fmha_fwd_d*
 aten/src/ATen/native/transformers/hip/flash_attn/ck/fmha_bwd_d*

.lintrunner.toml

@@ -143,8 +143,7 @@ init_command = [
     'tools/linter/adapters/pip_init.py',
     '--dry-run={{DRYRUN}}',
     'numpy==1.26.4 ; python_version >= "3.10" and python_version <= "3.11"',
-    'numpy==2.1.0 ; python_version >= "3.12" and python_version <= "3.13"',
-    'numpy==2.3.4 ; python_version >= "3.14"',
+    'numpy==2.1.0 ; python_version >= "3.12"',
     'expecttest==0.3.0',
     'pyrefly==0.36.2',
     'sympy==1.13.3',

CONTRIBUTING.md

@@ -18,7 +18,7 @@ aspects of contributing to PyTorch.
   - [Python Unit Testing](#python-unit-testing)
   - [Better local unit tests with `pytest`](#better-local-unit-tests-with-pytest)
   - [Local linting](#local-linting)
-    - [Running `pyrefly`](#running-pyrefly)
+    - [Running `mypy`](#running-mypy)
   - [C++ Unit Testing](#c-unit-testing)
   - [Run Specific CI Jobs](#run-specific-ci-jobs)
 - [Merging your Change](#merging-your-change)
@@ -281,7 +281,7 @@ dependencies as well as the nightly binaries into the repo directory.
 **Prerequisites**:
 The following packages should be installed with `pip`:
 - `expecttest` and `hypothesis` - required to run tests
-- `pyrefly` - recommended for type checking. [Pyrefly](https://pyrefly.org/)
+- `mypy` - recommended for linting
 - `pytest` - recommended to run tests more selectively
 Running
 ```
@@ -350,32 +350,15 @@ make lint
 
 Learn more about the linter on the [lintrunner wiki page](https://github.com/pytorch/pytorch/wiki/lintrunner)
 
-#### Running `pyrefly`
+#### Running `mypy`
 
-[Pyrefly](https://pyrefly.org/) is a high-performance static type checker for Python. It provides fast type checking along with IDE features like autocomplete and instant error feedback.
-
-PyTorch uses Pyrefly for type checking across the codebase. The configuration is managed in `pyrefly.toml` at the root of the repository.
-
-**Getting Started with Pyrefly:**
-
-To run type checking on the PyTorch codebase:
-```bash
-pyrefly check
-```
-
-For more detailed error information with summaries:
-```bash
-pyrefly check --summarize-errors
-```
-
-**Learn More:**
-- [Pyrefly Configuration](https://pyrefly.org/en/docs/configuration/) - Detailed configuration options
-- [Pyrefly IDE Features](https://pyrefly.org/en/docs/IDE-features/) - Set up Pyrefly in your editor for real-time type checking
-- [Python Typing Tutorial](https://pyrefly.org/en/docs/typing-for-python-developers/) - Learn about Python type annotations
+`mypy` is an optional static type checker for Python. We have multiple `mypy`
+configs for the PyTorch codebase that are automatically validated against whenever the linter is run.
 
 See [Guide for adding type annotations to
 PyTorch](https://github.com/pytorch/pytorch/wiki/Guide-for-adding-type-annotations-to-PyTorch)
-for PyTorch-specific guidance on how to set up `pyrefly` and tackle type annotation tasks in this codebase.
+for more information on how to set up `mypy` and tackle type annotation
+tasks.
 
 ### C++ Unit Testing
 

aten/src/ATen/cpu/vec/sve/vec_bfloat16.h

@@ -191,7 +191,7 @@ class Vectorized<BFloat16> {
     auto vals = svreinterpret_u16_bf16(values);
     vals = sveor_u16_x(ptrue, vals, mask);
     return svreinterpret_bf16_u16(vals);
-  }
+  };
   Vectorized<BFloat16> round() const;
   Vectorized<BFloat16> tan() const;
   Vectorized<BFloat16> tanh() const;
@@ -349,47 +349,47 @@ Vectorized<BFloat16> inline Vectorized<BFloat16>::frac() const {
     return convert_float_bfloat16(v1, v2);                     \
   }
 
-DEFINE_BF16_FUNC_VIA_FLOAT(isnan)
-DEFINE_BF16_FUNC_VIA_FLOAT(angle)
-DEFINE_BF16_FUNC_VIA_FLOAT(acos)
-DEFINE_BF16_FUNC_VIA_FLOAT(acosh)
-DEFINE_BF16_FUNC_VIA_FLOAT(asin)
-DEFINE_BF16_FUNC_VIA_FLOAT(atan)
-DEFINE_BF16_FUNC_VIA_FLOAT(atanh)
-DEFINE_BF16_FUNC_VIA_FLOAT_W_ARG(atan2)
-DEFINE_BF16_FUNC_VIA_FLOAT_W_ARG(copysign)
-DEFINE_BF16_FUNC_VIA_FLOAT(erf)
-DEFINE_BF16_FUNC_VIA_FLOAT(erfc)
-DEFINE_BF16_FUNC_VIA_FLOAT(exp)
-DEFINE_BF16_FUNC_VIA_FLOAT(exp2)
-DEFINE_BF16_FUNC_VIA_FLOAT(expm1)
-DEFINE_BF16_FUNC_VIA_FLOAT_W_ARG(fmod)
-DEFINE_BF16_FUNC_VIA_FLOAT_W_ARG(hypot)
-DEFINE_BF16_FUNC_VIA_FLOAT(i0)
-DEFINE_BF16_FUNC_VIA_FLOAT(i0e)
-DEFINE_BF16_FUNC_VIA_FLOAT(digamma)
-DEFINE_BF16_FUNC_VIA_FLOAT_W_ARG(igamma)
-DEFINE_BF16_FUNC_VIA_FLOAT_W_ARG(igammac)
-DEFINE_BF16_FUNC_VIA_FLOAT_W_ARG(nextafter)
-DEFINE_BF16_FUNC_VIA_FLOAT(log)
-DEFINE_BF16_FUNC_VIA_FLOAT(log2)
-DEFINE_BF16_FUNC_VIA_FLOAT(log10)
-DEFINE_BF16_FUNC_VIA_FLOAT(log1p)
-DEFINE_BF16_FUNC_VIA_FLOAT(sin)
-DEFINE_BF16_FUNC_VIA_FLOAT(sinh)
-DEFINE_BF16_FUNC_VIA_FLOAT(cos)
-DEFINE_BF16_FUNC_VIA_FLOAT(cosh)
-DEFINE_BF16_FUNC_VIA_FLOAT(ceil)
-DEFINE_BF16_FUNC_VIA_FLOAT(floor)
-DEFINE_BF16_FUNC_VIA_FLOAT(round)
-DEFINE_BF16_FUNC_VIA_FLOAT(tan)
-DEFINE_BF16_FUNC_VIA_FLOAT(tanh)
-DEFINE_BF16_FUNC_VIA_FLOAT(trunc)
-DEFINE_BF16_FUNC_VIA_FLOAT(lgamma)
-DEFINE_BF16_FUNC_VIA_FLOAT(sqrt)
-DEFINE_BF16_FUNC_VIA_FLOAT(reciprocal)
-DEFINE_BF16_FUNC_VIA_FLOAT(rsqrt)
-DEFINE_BF16_FUNC_VIA_FLOAT_W_ARG(pow)
+DEFINE_BF16_FUNC_VIA_FLOAT(isnan);
+DEFINE_BF16_FUNC_VIA_FLOAT(angle);
+DEFINE_BF16_FUNC_VIA_FLOAT(acos);
+DEFINE_BF16_FUNC_VIA_FLOAT(acosh);
+DEFINE_BF16_FUNC_VIA_FLOAT(asin);
+DEFINE_BF16_FUNC_VIA_FLOAT(atan);
+DEFINE_BF16_FUNC_VIA_FLOAT(atanh);
+DEFINE_BF16_FUNC_VIA_FLOAT_W_ARG(atan2);
+DEFINE_BF16_FUNC_VIA_FLOAT_W_ARG(copysign);
+DEFINE_BF16_FUNC_VIA_FLOAT(erf);
+DEFINE_BF16_FUNC_VIA_FLOAT(erfc);
+DEFINE_BF16_FUNC_VIA_FLOAT(exp);
+DEFINE_BF16_FUNC_VIA_FLOAT(exp2);
+DEFINE_BF16_FUNC_VIA_FLOAT(expm1);
+DEFINE_BF16_FUNC_VIA_FLOAT_W_ARG(fmod);
+DEFINE_BF16_FUNC_VIA_FLOAT_W_ARG(hypot);
+DEFINE_BF16_FUNC_VIA_FLOAT(i0);
+DEFINE_BF16_FUNC_VIA_FLOAT(i0e);
+DEFINE_BF16_FUNC_VIA_FLOAT(digamma);
+DEFINE_BF16_FUNC_VIA_FLOAT_W_ARG(igamma);
+DEFINE_BF16_FUNC_VIA_FLOAT_W_ARG(igammac);
+DEFINE_BF16_FUNC_VIA_FLOAT_W_ARG(nextafter);
+DEFINE_BF16_FUNC_VIA_FLOAT(log);
+DEFINE_BF16_FUNC_VIA_FLOAT(log2);
+DEFINE_BF16_FUNC_VIA_FLOAT(log10);
+DEFINE_BF16_FUNC_VIA_FLOAT(log1p);
+DEFINE_BF16_FUNC_VIA_FLOAT(sin);
+DEFINE_BF16_FUNC_VIA_FLOAT(sinh);
+DEFINE_BF16_FUNC_VIA_FLOAT(cos);
+DEFINE_BF16_FUNC_VIA_FLOAT(cosh);
+DEFINE_BF16_FUNC_VIA_FLOAT(ceil);
+DEFINE_BF16_FUNC_VIA_FLOAT(floor);
+DEFINE_BF16_FUNC_VIA_FLOAT(round);
+DEFINE_BF16_FUNC_VIA_FLOAT(tan);
+DEFINE_BF16_FUNC_VIA_FLOAT(tanh);
+DEFINE_BF16_FUNC_VIA_FLOAT(trunc);
+DEFINE_BF16_FUNC_VIA_FLOAT(lgamma);
+DEFINE_BF16_FUNC_VIA_FLOAT(sqrt);
+DEFINE_BF16_FUNC_VIA_FLOAT(reciprocal);
+DEFINE_BF16_FUNC_VIA_FLOAT(rsqrt);
+DEFINE_BF16_FUNC_VIA_FLOAT_W_ARG(pow);
 
 Vectorized<BFloat16> inline Vectorized<BFloat16>::operator==(
     const Vectorized<BFloat16>& other) const {

aten/src/ATen/cuda/CUDABlas.cpp

@@ -388,7 +388,6 @@ static inline bool bgemm_internal_cublaslt(CUDABLAS_BGEMM_ARGTYPES_AND_C_DTYPE(D
 #ifndef USE_ROCM
   at::Half halpha;
   at::Half hbeta;
-  uint32_t mask = -1;
 #endif
   void * alpha_ptr = α
   void * beta_ptr = β
@@ -428,7 +427,7 @@ static inline bool bgemm_internal_cublaslt(CUDABLAS_BGEMM_ARGTYPES_AND_C_DTYPE(D
     auto fp16_reduction = at::globalContext().allowFP16ReductionCuBLAS();
     if (fp16_reduction !=
         at::CuBLASReductionOption::AllowReducedPrecisionWithSplitK) {
-      mask =
+      uint32_t mask =
           fp16_reduction ==
                   at::CuBLASReductionOption::DisallowReducedPrecisionAllowSplitK
               ? (CUBLASLT_REDUCTION_SCHEME_COMPUTE_TYPE |
@@ -445,7 +444,7 @@ static inline bool bgemm_internal_cublaslt(CUDABLAS_BGEMM_ARGTYPES_AND_C_DTYPE(D
     auto bf16_reduction = at::globalContext().allowBF16ReductionCuBLAS();
     if (bf16_reduction !=
         at::CuBLASReductionOption::AllowReducedPrecisionWithSplitK) {
-      mask =
+      uint32_t mask =
           bf16_reduction ==
                   at::CuBLASReductionOption::DisallowReducedPrecisionAllowSplitK
               ? (CUBLASLT_REDUCTION_SCHEME_COMPUTE_TYPE |
@@ -512,41 +511,17 @@ static inline bool bgemm_internal_cublaslt(CUDABLAS_BGEMM_ARGTYPES_AND_C_DTYPE(D
   cublasStatus_t cublasStatus = CUBLAS_STATUS_SUCCESS;
   cublasLtMatmulHeuristicResult_t heuristicResult = {};
   int returnedResult = 0;
-  // on Blackwell+, we fake a n > 1 matmul when querying heuristics
-  // to prevent cuBLASLt from dispatching to a GEMV kernel for batch-invariance
-#ifndef USE_ROCM
-  const bool lie_to_cublaslt = mask == CUBLASLT_REDUCTION_SCHEME_NONE && n == 1 && at::cuda::getCurrentDeviceProperties()->major >= 10;
-#else
-  const bool lie_to_cublaslt = false;
-#endif
-  if (lie_to_cublaslt) {
-     CuBlasLtMatrixLayout FakeBdesc(abType, k, 2, ldb, opb == CUBLAS_OP_T);
-     CuBlasLtMatrixLayout FakeCdesc(cType, m, 2, ldc);
-
-     TORCH_CUDABLAS_CHECK(cublasLtMatmulAlgoGetHeuristic(
-        ltHandle,
-        computeDesc.descriptor(),
-        Adesc.descriptor(),
-        FakeBdesc.descriptor(),
-        FakeCdesc.descriptor(),
-        FakeCdesc.descriptor(),
-        preference.descriptor(),
-        1,
-        &heuristicResult,
-        &returnedResult));
-  } else {
-    TORCH_CUDABLAS_CHECK(cublasLtMatmulAlgoGetHeuristic(
-        ltHandle,
-        computeDesc.descriptor(),
-        Adesc.descriptor(),
-        Bdesc.descriptor(),
-        Cdesc.descriptor(),
-        Cdesc.descriptor(),
-        preference.descriptor(),
-        1,
-        &heuristicResult,
-        &returnedResult));
-  }
+  TORCH_CUDABLAS_CHECK(cublasLtMatmulAlgoGetHeuristic(
+      ltHandle,
+      computeDesc.descriptor(),
+      Adesc.descriptor(),
+      Bdesc.descriptor(),
+      Cdesc.descriptor(),
+      Cdesc.descriptor(),
+      preference.descriptor(),
+      1,
+      &heuristicResult,
+      &returnedResult));
   if (returnedResult == 0) {
     cublasStatus = CUBLAS_STATUS_NOT_SUPPORTED;
   }
@@ -1597,7 +1572,7 @@ bool gemm_and_bias(
   }
 
   using opmath_t = at::opmath_type<Dtype>;
-  opmath_t beta_val = bias ? 0 : 1; // bias is added in epilogue unless nullptr
+  opmath_t beta_val = 0; // bias is added in epilogue
 
   cudaDataType_t abType = CUDA_R_32F;
   cudaDataType_t cType = CUDA_R_32F;
@@ -1686,22 +1661,15 @@ bool gemm_and_bias(
     _syncCurrentWithCarveoutStream(stream, true);
   }
 #endif
-  const auto epilogue = [&]() -> cublasLtEpilogue_t {
-    // The cuBLAS documentation indicates that
-    // *_<ACTIVATION>_BIAS = *_<ACTIVATION>,
-    // but we keep it verbose here for clarity.
-    switch (activation) {
-      case GEMMAndBiasActivationEpilogue::RELU:
-        return bias ? CUBLASLT_EPILOGUE_RELU_BIAS : CUBLASLT_EPILOGUE_RELU;
-      case GEMMAndBiasActivationEpilogue::GELU:
-        return bias ? CUBLASLT_EPILOGUE_GELU_BIAS : CUBLASLT_EPILOGUE_GELU;
-      default:
-        return bias ? CUBLASLT_EPILOGUE_BIAS : CUBLASLT_EPILOGUE_DEFAULT;
-    }
-  }();
-  computeDesc.setAttribute(CUBLASLT_MATMUL_DESC_EPILOGUE, epilogue);
+  cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_BIAS;
+  if (activation == GEMMAndBiasActivationEpilogue::RELU) {
+    epilogue = CUBLASLT_EPILOGUE_RELU_BIAS;
+  } else if (activation == GEMMAndBiasActivationEpilogue::GELU) {
+    epilogue = CUBLASLT_EPILOGUE_GELU_BIAS;
+  }
 
-  if (bias) {
+  if (bias != nullptr) {
+    computeDesc.setAttribute(CUBLASLT_MATMUL_DESC_EPILOGUE, epilogue);
     computeDesc.setAttribute(CUBLASLT_MATMUL_DESC_BIAS_POINTER, bias);
   }
 

aten/src/ATen/cuda/NumericLimits.cuh

@@ -55,14 +55,6 @@ struct numeric_limits<int8_t> {
   static inline __host__ __device__ int8_t upper_bound() { return INT8_MAX; }
 };
 
-template <>
-struct numeric_limits<uint16_t> {
-  static inline __host__ __device__ uint16_t lowest() { return 0; }
-  static inline __host__ __device__ uint16_t max() { return UINT16_MAX; }
-  static inline __host__ __device__ uint16_t lower_bound() { return 0; }
-  static inline __host__ __device__ uint16_t upper_bound() { return UINT16_MAX; }
-};
-
 template <>
 struct numeric_limits<int16_t> {
   static inline __host__ __device__ int16_t lowest() { return INT16_MIN; }
@@ -71,14 +63,6 @@ struct numeric_limits<int16_t> {
   static inline __host__ __device__ int16_t upper_bound() { return INT16_MAX; }
 };
 
-template <>
-struct numeric_limits<uint32_t> {
-  static inline __host__ __device__ uint32_t lowest() { return 0; }
-  static inline __host__ __device__ uint32_t max() { return UINT32_MAX; }
-  static inline __host__ __device__ uint32_t lower_bound() { return 0; }
-  static inline __host__ __device__ uint32_t upper_bound() { return UINT32_MAX; }
-};
-
 template <>
 struct numeric_limits<int32_t> {
   static inline __host__ __device__ int32_t lowest() { return INT32_MIN; }
@@ -87,21 +71,6 @@ struct numeric_limits<int32_t> {
   static inline __host__ __device__ int32_t upper_bound() { return INT32_MAX; }
 };
 
-template <>
-struct numeric_limits<uint64_t> {
-#ifdef _MSC_VER
-  static inline __host__ __device__ uint64_t lowest() { return 0; }
-  static inline __host__ __device__ uint64_t max() { return _UI64_MAX; }
-  static inline __host__ __device__ uint64_t lower_bound() { return 0; }
-  static inline __host__ __device__ uint64_t upper_bound() { return _UI64_MAX; }
-#else
-  static inline __host__ __device__ uint64_t lowest() { return 0; }
-  static inline __host__ __device__ uint64_t max() { return UINT64_MAX; }
-  static inline __host__ __device__ uint64_t lower_bound() { return 0; }
-  static inline __host__ __device__ uint64_t upper_bound() { return UINT64_MAX; }
-#endif
-};
-
 template <>
 struct numeric_limits<int64_t> {
 #ifdef _MSC_VER

aten/src/ATen/cuda/cub.h

@@ -24,13 +24,7 @@ namespace detail {
 // radix_sort_pairs doesn't interact with value_t other than to copy
 // the data, so we can save template instantiations by reinterpreting
 // it as an opaque type.
-// We use native integer types for 1/2/4/8-byte values to reduce
-// register usage in CUDA kernels. For sizes > 8 fall back to char array.
 template <int N> struct alignas(N) OpaqueType { char data[N]; };
-template <> struct alignas(1) OpaqueType<1> { uint8_t data; };
-template <> struct alignas(2) OpaqueType<2> { uint16_t data; };
-template <> struct alignas(4) OpaqueType<4> { uint32_t data; };
-template <> struct alignas(8) OpaqueType<8> { uint64_t data; };
 
 template<typename key_t, int value_size>
 void radix_sort_pairs_impl(

aten/src/ATen/native/BinaryOps.cpp

@@ -1009,25 +1009,12 @@ static Device correct_out_device(const Tensor& self, const Tensor& other) {
   }
 }
 
-static Tensor send_to_meta(const Tensor& self, const Device& device) {
-  Tensor out_meta;
-  if (self._is_zerotensor() && self.unsafeGetTensorImpl()->is_wrapped_number()) {
-    out_meta = at::_efficientzerotensor(self.sizes(), self.options().device(device));
-    out_meta.unsafeGetTensorImpl()->set_wrapped_number(true);
-  } else {
-    out_meta = self.to(device);
-  }
-  return out_meta;
-}
-
 Tensor mul_zerotensor(const Tensor& self, const Tensor& other) {
   auto out_device = correct_out_device(self, other);
   // hack to use the TensorIterator to get the correct broadcasting and type promotion logic
   auto device_ = Device(DeviceType::Meta);
   constexpr c10::DispatchKeySet meta_dks(at::DispatchKey::Meta);
-  auto self_meta = send_to_meta(self, device_);
-  auto other_meta = send_to_meta(other, device_);
-  auto meta_out = at::_ops::mul_Tensor::redispatch(meta_dks, self_meta, other_meta);
+  auto meta_out = at::_ops::mul_Tensor::redispatch(meta_dks, self.to(device_), other.to(device_));
   return at::_efficientzerotensor(meta_out.sizes(), meta_out.options().device(out_device));
 }
 
@@ -1036,9 +1023,7 @@ Tensor div_zerotensor(const Tensor& self, const Tensor& other) {
   // hack to use the TensorIterator to get the correct broadcasting and type promotion logic
   auto device_ = Device(DeviceType::Meta);
   constexpr c10::DispatchKeySet meta_dks(at::DispatchKey::Meta);
-  auto self_meta = send_to_meta(self, device_);
-  auto other_meta = send_to_meta(other, device_);
-  auto meta_out = at::_ops::div_Tensor::redispatch(meta_dks, self_meta, other_meta);
+  auto meta_out = at::_ops::div_Tensor::redispatch(meta_dks, self.to(device_), other.to(device_));
 
   if (self._is_zerotensor()) {
     if (other._is_zerotensor()) {
@@ -1067,9 +1052,8 @@ static Tensor maybe_add_maybe_sub(const Tensor& self, const Tensor& other, const
   // hack to use the TensorIterator to get the correct broadcasting and type promotion logic
   auto device_ = Device(DeviceType::Meta);
   constexpr c10::DispatchKeySet meta_dks(at::DispatchKey::Meta);
-  auto self_meta = send_to_meta(self, device_);
-  auto other_meta = send_to_meta(other, device_);
-  auto meta_out = at::_ops::add_Tensor::redispatch(meta_dks, self_meta, other_meta, alpha);
+  auto meta_out = at::_ops::add_Tensor::redispatch(
+      meta_dks, self.to(device_), other.to(device_), alpha);
 
   auto get_out_like = [&] (const Tensor& tensor)
   {

aten/src/ATen/native/Linear.cpp

@@ -50,35 +50,18 @@ static inline bool parseLinearFlatten3d() {
 // `_flatten_nd_linear` flattens all but the last dimension of the input tensor
 // before passing it to linear operation
 static inline Tensor _flatten_nd_linear(const Tensor& input, const Tensor& weight, const Tensor& bias) {
-  const auto input_sizes = input.sym_sizes();
-
-  const auto result_flattened = [&]() -> Tensor {
-    const auto input_ncols = input_sizes.back();
-    const auto input_flattened_nrows = [&]() -> c10::SymInt {
-      // can't use -1 in reshape because it errors when a dimension is 0
-      auto flattened_nrows = c10::SymInt{1};
-      for (const auto& size : input_sizes.slice(0, input_sizes.size() - 1)) {
-        flattened_nrows *= size;
-      }
-      return flattened_nrows;
-    }();
-
-    const auto input_flattened = input.view_symint({input_flattened_nrows, input_ncols});
-    if (weight.layout() == c10::kStrided) {
-      return at::addmm(bias, input_flattened, weight.t());
-    } else {
-      // weight is sparse, and addmm for sparse expects matmul lhs to be sparse,
-      // so we transpose the problem.
-      // NOTE: at::matmul handles (dense @ sparse) similarly.
-      const auto bias_t = (bias.dim() >= 2) ? bias.mT() : bias.unsqueeze(-1);
-      return at::addmm(bias_t, weight, input_flattened.t()).t();
+    const auto input_sizes = input.sym_sizes();
+    // can't use -1 in reshape because it errors when a dimension is 0
+    c10::SymInt flattened_dim = 1;
+    for (int64_t i = 0, ndim = input_sizes.size(); i < ndim - 1; ++i) {
+      flattened_dim = flattened_dim * input_sizes[i];
     }
-  }();
-
-  // Unflatten flattened row dims
-  auto result_sizes = c10::SymDimVector{input_sizes.begin(), input_sizes.end()};
-  result_sizes.back() = result_flattened.sym_size(1);
-  return result_flattened.view_symint(result_sizes);
+    auto inp_reshape = input.reshape_symint({flattened_dim, input_sizes.at(input_sizes.size() -1)});
+    const auto result = at::addmm(bias, inp_reshape, weight.t());
+    auto new_size = input_sizes.slice(0, input_sizes.size() - 1);
+    c10::SymDimVector sizes_vec(new_size.begin(), new_size.end());
+    sizes_vec.push_back(result.sym_size(1));
+    return result.view_symint(sizes_vec);
 }
 
 
@@ -107,23 +90,15 @@ Tensor linear(const Tensor& input, const Tensor& weight, const std::optional<Ten
     // Fused op is marginally faster.
     return at::addmm(*bias, input, weight.t());
   }
-
-  const auto is_bias_likely_fusable = (
-      bias->defined() &&
-      // cuBLASLt: will fuse in the epilogue without copies
-      // when input/weight/bias are all strided.
-      // When weight is not strided, bias will not be fused,
-      // but we can still dispatch here to avoid at::matmul
-      // path which will probably use a very similar
-      // flattening optimization.
-      ((bias->dim() == 1 || bias->squeeze().dim() == 1) && bias->is_contiguous_or_false())
-  );
-  if (is_bias_likely_fusable && !input.is_xla()) {
-    // Also hit the fused path for contiguous nD input, if not using xla
+  if (bias->defined() && !input.is_xla()) {
+    // Also hit the fused path for contiguous 3D input, if not using xla
     // backend. Reshaping/flattening has some performance implications on xla.
-    if (input.is_contiguous_or_false()) {
+    bool is_contiguous = input.is_contiguous_or_false();
+    if (is_contiguous && input_dim == 3) {
       return _flatten_nd_linear(input, weight, *bias);
-    } else if (parseLinearFlatten3d()) {
+    } else if (is_contiguous && input.layout() == c10::kStrided && weight.layout() == c10::kStrided && bias->dim() == 1) {
+      return _flatten_nd_linear(input, weight, *bias);
+    } else if (parseLinearFlatten3d() && input_dim == 3) {
       // If user forces flattening via env var
       const Tensor input_cont = input.contiguous();
       return _flatten_nd_linear(input_cont, weight, *bias);

aten/src/ATen/native/TensorShape.cpp

@@ -1,6 +1,5 @@
 #include <ATen/core/ATen_fwd.h>
 #include <c10/core/ScalarType.h>
-#include <c10/core/SymInt.h>
 #define TORCH_ASSERT_ONLY_METHOD_OPERATORS
 #include <ATen/AccumulateType.h>
 #include <ATen/Dispatch.h>
@@ -1711,37 +1710,11 @@ Tensor narrow_symint(
       "], but got ",
       start,
       ")")
-
-  auto cond1 = TORCH_GUARD_OR_FALSE(start.sym_lt(0));
-  auto cond2 = TORCH_GUARD_OR_FALSE(start.sym_ge(0));
-
-  if (cond1 || cond2) {
-    if (cond1) {
-      start = start + cur_size;
-    }
-
-    TORCH_SYM_CHECK(
-        start.sym_le(cur_size - length),
-        "start (",
-        start,
-        ") + length (",
-        length,
-        ") exceeds dimension size (",
-        cur_size,
-        ").");
-    return at::slice_symint(self, dim, start, start + length, 1);
+  if (start < 0) {
+    start = start + cur_size;
   }
-
-  // Unbacked start handling!
-
-  // Bounds check without converting start:
-  // - If start < 0: need (start + cur_size) + length <= cur_size, i.e., start +
-  // length <= 0
-  // - If start >= 0: need start + length <= cur_size
-  auto end = start + length;
   TORCH_SYM_CHECK(
-      (start.sym_lt(0).sym_and((end).sym_le(0)))
-          .sym_or(start.sym_ge(0).sym_and((end).sym_le(cur_size))),
+      start.sym_le(cur_size - length),
       "start (",
       start,
       ") + length (",
@@ -1749,28 +1722,7 @@ Tensor narrow_symint(
       ") exceeds dimension size (",
       cur_size,
       ").");
-
-  if (TORCH_GUARD_OR_FALSE(end.sym_ne(0))) {
-    return at::slice_symint(self, dim, start, end, 1);
-  } else {
-    // Cannot statically determine the condition due to unbacked.
-    // This is an interesting situation; when start is negative and
-    // start + length == 0, slice and narrow do different things.
-    // i.e., x.narrow(0, -2, 2) != x[-2:0]; in that case, we want to
-    // pass curr_size instead of 0. Otherwise, they would do the same thing.
-    // This says at runtime: if start < 0 and end == 0, then pass curr_size
-    // instead of 0.
-
-    auto use_different = start.sym_lt(0).sym_and(end.sym_eq(0)).toSymInt();
-    auto result =
-        at::slice_symint(self, dim, start, end + use_different * cur_size, 1);
-
-    // Ensure slice allocated unbacked size is specialized to length.
-    SymInt new_size = result.sym_size(dim);
-    TORCH_SYM_CHECK(new_size.sym_eq(length), "")
-
-    return result;
-  }
+  return at::slice_symint(self, dim, start, start + length, 1);
 }
 
 // This overload exists purely for XLA, because they wanted to pass in
@@ -1784,8 +1736,8 @@ Tensor narrow_tensor_symint(
       start.dim() == 0 &&
           isIntegralType(start.scalar_type(), /*includeBool=*/false),
       "start must be an 0-dim integral Tensor.");
-  c10::SymInt st = start.item().toSymInt();
-  return at::narrow_symint(self, dim, std::move(st), std::move(length));
+  int64_t st = start.item<int64_t>();
+  return at::narrow_symint(self, dim, c10::SymInt(st), std::move(length));
 }
 
 std::

aten/src/ATen/native/cpu/GridSamplerKernel.cpp

@@ -293,7 +293,7 @@ struct ComputeLocationBase<scalar_t, /*align_corners=*/false> {
     , empty(size <= 0) {}
 
   inline Vec unnormalize(const Vec &in) const {
-    return (in + Vec(static_cast<scalar_t>(1))) * Vec(scaling_factor) - Vec(static_cast<scalar_t>(0.5));
+    return (in + Vec(1)) * Vec(scaling_factor) - Vec(0.5);
   }
 
   inline Vec clip_coordinates(const Vec &in) const {
@@ -831,7 +831,7 @@ struct ApplyGridSample<scalar_t, 2, GridSamplerInterpolation::Bicubic,
 
   // constant used in cubic convolution
   // could be -0.5 or -0.75, use the same value in UpSampleBicubic2d.h
-  const Vec A = Vec(static_cast<scalar_t>(-0.75));
+  const Vec A = Vec(-0.75);
 
   ApplyGridSample(const TensorAccessor<const scalar_t, 4>& input)
     : inp_H(input.size(2))

aten/src/ATen/native/cpu/ReduceAllOpsKernel.cpp

@@ -5,7 +5,6 @@
 #include <ATen/native/ReduceOpsUtils.h>
 
 #include <ATen/Dispatch.h>
-#include <ATen/Dispatch_v2.h>
 #include <ATen/Parallel.h>
 #include <ATen/TensorIterator.h>
 #include <ATen/OpMathType.h>
@@ -79,12 +78,12 @@ void min_all_kernel_impl(Tensor& result, const Tensor& input) {
     reduce_all_impl<int64_t>(result, input, upper_bound<int64_t>(),
       [=](int64_t a, int64_t b) -> int64_t { return min_impl(a, b); });
   } else {
-    AT_DISPATCH_V2(input.scalar_type(), "min_all", AT_WRAP([&] {
+    AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBFloat16, input.scalar_type(), "min_all", [&] {
       using Vec = Vectorized<opmath_type<scalar_t>>;
       reduce_all_impl_vec<scalar_t>(result, input, upper_bound<scalar_t>(),
         [=] (scalar_t a , scalar_t b) -> scalar_t { return min_impl(a, b); },
         [=](Vec a, Vec b) -> Vec { return minimum(a, b); });
-    }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES), kHalf, kBFloat16);
+    });
   }
 }
 
@@ -104,12 +103,12 @@ void max_all_kernel_impl(Tensor& result, const Tensor& input) {
     reduce_all_impl<int64_t>(result, input, lower_bound<int64_t>(),
       [=](int64_t a, int64_t b) -> int64_t { return max_impl(a, b); });
   } else {
-    AT_DISPATCH_V2(input.scalar_type(), "max_all", AT_WRAP([&] {
+    AT_DISPATCH_ALL_TYPES_AND2(kHalf, kBFloat16, input.scalar_type(), "max_all", [&] {
       using Vec = Vectorized<opmath_type<scalar_t>>;
       reduce_all_impl_vec<scalar_t>(result, input, lower_bound<scalar_t>(),
         [=] (scalar_t a , scalar_t b) -> scalar_t { return max_impl(a, b); },
         [=](Vec a, Vec b) -> Vec { return maximum(a, b); });
-    }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES), kHalf, kBFloat16);
+    });
   }
 }
 
@@ -200,7 +199,7 @@ void aminmax_allreduce_kernel(
       }
     );
   } else {
-    AT_DISPATCH_V2(input.scalar_type(), "aminmax_cpu", AT_WRAP([&] {
+    AT_DISPATCH_ALL_TYPES_AND2(kBFloat16, kHalf, input.scalar_type(), "aminmax_cpu", [&] {
       using Vec = Vectorized<opmath_type<scalar_t>>;
       using scalar_t_pair = std::pair<scalar_t, scalar_t>;
       reduce_all_impl_vec_two_outputs<scalar_t>(
@@ -215,7 +214,7 @@ void aminmax_allreduce_kernel(
         [=](Vec a, Vec b) -> Vec { return minimum(a, b); },
         [=](Vec a, Vec b) -> Vec { return maximum(a, b); }
       );
-    }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES), kBFloat16, kHalf);
+    });
   }
 }
 

aten/src/ATen/native/cpu/ReduceOpsKernel.cpp

@@ -3,7 +3,6 @@
 
 #include <ATen/core/Tensor.h>
 #include <ATen/Dispatch.h>
-#include <ATen/Dispatch_v2.h>
 #include <ATen/OpMathType.h>
 #include <ATen/cpu/vec/vec.h>
 #include <ATen/cpu/vec/functional.h>
@@ -348,35 +347,34 @@ struct MinValuesOps: public at::native::MinOps<scalar_t> {
 };
 
 void min_values_kernel_impl(TensorIterator& iter) {
-  // This case is special because of Vectorized<int64_t> does not
-  // handle upper_bound<int64_t>().
-  // See: https://github.com/pytorch/pytorch/issues/43254
-  if (iter.dtype() == kLong || iter.dtype() == kUInt64) {
-    AT_DISPATCH_V2(iter.dtype(), "min_values_cpu", AT_WRAP([&iter] {
-      binary_kernel_reduce(
-        iter,
-        MinValuesOps<scalar_t>{},
-        std::pair<scalar_t, int64_t>(upper_bound<scalar_t>(), -1));
-    }), kLong, kUInt64);
+  if (iter.dtype() == kLong) {
+    // This case is special because of Vectorized<int64_t> does not
+    // handle upper_bound<int64_t>().
+    // See: https://github.com/pytorch/pytorch/issues/43254
+    using scalar_t = int64_t;
+    binary_kernel_reduce(
+      iter,
+      MinValuesOps<scalar_t>{},
+      std::pair<scalar_t, int64_t>(upper_bound<scalar_t>(), -1));
     return;
   }
-  AT_DISPATCH_V2(iter.dtype(), "min_values_cpu", AT_WRAP([&iter] {
+  AT_DISPATCH_ALL_TYPES_AND3(kBFloat16, kHalf, kBool, iter.dtype(), "min_values_cpu", [&iter] {
     binary_kernel_reduce_vec(
       iter,
       [](scalar_t a, scalar_t b) -> scalar_t { return min_impl(a, b); },
       [](Vectorized<scalar_t> a, Vectorized<scalar_t> b) { return minimum(a, b); },
       static_cast<double>(upper_bound<scalar_t>()));
-  }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES), kBFloat16, kHalf, kBool);
+  });
 }
 
 void max_values_kernel_impl(TensorIterator& iter) {
-  AT_DISPATCH_V2(iter.dtype(), "max_values_cpu", AT_WRAP([&iter] {
+  AT_DISPATCH_ALL_TYPES_AND3(kBFloat16, kHalf, kBool, iter.dtype(), "max_values_cpu", [&iter] {
     binary_kernel_reduce_vec(
       iter,
       [](scalar_t a, scalar_t b) -> scalar_t { return max_impl(a, b); },
       [](Vectorized<scalar_t> a, Vectorized<scalar_t> b) { return maximum(a, b); },
       lower_bound<scalar_t>());
-  }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES), kBFloat16, kHalf, kBool);
+  });
 }
 
 void argmax_kernel_impl(TensorIterator &iter) {

aten/src/ATen/native/cpu/TensorCompareKernel.cpp

@@ -11,7 +11,6 @@
 #include <vector>
 
 #include <ATen/Dispatch.h>
-#include <ATen/Dispatch_v2.h>
 #include <ATen/Parallel.h>
 #include <ATen/NumericUtils.h>
 #include <ATen/TensorIterator.h>
@@ -107,7 +106,7 @@ void min_kernel_impl(
     bool keepdim) {
   int64_t self_dim_size = ensure_nonempty_size(self, dim);
 
-  AT_DISPATCH_V2(self.scalar_type(), "min_cpu", AT_WRAP([&] {
+  AT_DISPATCH_ALL_TYPES_AND3(ScalarType::Half, ScalarType::BFloat16, ScalarType::Bool, self.scalar_type(), "min_cpu", [&] {
     compare_base_kernel<scalar_t>(result, indice, self, dim, keepdim, [&] (
       scalar_t* result_data, int64_t* indice_data,
       const scalar_t* self_data, auto self_dim_stride) {
@@ -129,7 +128,7 @@ void min_kernel_impl(
         *indice_data = index;
       }
     );
-  }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES), ScalarType::Half, ScalarType::BFloat16, ScalarType::Bool);
+  });
 }
 
 void max_kernel_impl(
@@ -140,7 +139,7 @@ void max_kernel_impl(
     bool keepdim) {
   int64_t self_dim_size = ensure_nonempty_size(self, dim);
 
-  AT_DISPATCH_V2(self.scalar_type(), "max_cpu", AT_WRAP([&] {
+  AT_DISPATCH_ALL_TYPES_AND3(ScalarType::Half, ScalarType::BFloat16, ScalarType::Bool, self.scalar_type(), "max_cpu", [&] {
     compare_base_kernel<scalar_t>(result, indice, self, dim, keepdim, [&] (
       scalar_t* result_data, int64_t* indice_data,
       const scalar_t* self_data, auto self_dim_stride) {
@@ -162,7 +161,7 @@ void max_kernel_impl(
         *indice_data = index;
       }
     );
-  }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES), ScalarType::Half, ScalarType::BFloat16, ScalarType::Bool);
+  });
 }
 
 void aminmax_kernel(
@@ -187,7 +186,7 @@ void aminmax_kernel(
     return;
   }
 
-  AT_DISPATCH_V2(self.scalar_type(), "aminmax_cpu", AT_WRAP([&] {
+  AT_DISPATCH_ALL_TYPES_AND3(ScalarType::Bool, ScalarType::BFloat16, ScalarType::Half, self.scalar_type(), "aminmax_cpu", [&] {
     compare_base_kernel<scalar_t, scalar_t>(min_result, max_result, self, wrap_dim, keepdim, [&] (
       scalar_t* min_result_data, scalar_t* max_result_data,
       const scalar_t* self_data, auto self_dim_stride) {
@@ -210,7 +209,7 @@ void aminmax_kernel(
         *max_result_data = max_number;
       }
     );
-  }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES), ScalarType::Bool, ScalarType::BFloat16, ScalarType::Half);
+  });
 }
 
 void where_kernel_impl(TensorIterator &iter) {

aten/src/ATen/native/cuda/Blas.cpp

@@ -147,24 +147,14 @@ static bool isGloballyDisabledAddmmCudaLt(const at::Device& device) {
 /*
  * Check whether for the given input we want to enable the Lt interface
  */
-static bool isInputCompliesAddmmCudaLt(
-    Tensor& result,
-    const Tensor& self,
-    const Tensor& mat1,
-    const Tensor& mat2,
-    const Scalar& beta,
-    const Scalar& alpha,
-    Activation activation
-) {
-  #ifdef USE_ROCM
+static bool isInputCompliesAddmmCudaLt(Tensor& result, const Tensor& self, const Tensor& mat1, const Tensor& mat2, const Scalar& beta, const Scalar& alpha) {
   // Implies 2D bias which we currently not send through Lt.
   // TODO: this check is done pre col-major input preparation,
   // so, this condition can be ralexed in cases when a col-major
   // copy of result is needed.
-  if (self.is_same(result) || self.dim() == 2) {
+  if (result.is_same(self)) {
     return false;
   }
-  #endif
 
   #if defined(USE_ROCM) && ROCM_VERSION == 60400
   // hipblaslt TT fp32 regression on ROCm 6.4, cannot use
@@ -179,33 +169,13 @@ static bool isInputCompliesAddmmCudaLt(
   #if defined(CUDA_VERSION) || defined(USE_ROCM)
   const auto scalar_type = mat1.scalar_type();
   return (beta.toComplexDouble() == 1.0
-    // NOTE: row-major result is important when bias is 1D.
-    // This is because Lt broadcasts 1D bias over the columns
-    // while the aten::addmm API broadcasts it over the rows,
-    // and this is in conjuction with the data preparation
-    // procedure that does not transpose arguments with
-    // col-major result. For col-major result we need
-    // to explicitly transpose the problem so that bias is
-    // correctly applied.
-    // TODO: enable col-major result if needed.
-    // TODO: no need to check result's layout when
-    // !result.is_same(self) and self.dim() == 2, because
-    // self needs to be copied into result and the bias ptr
-    // will be ignored.
     && result.dim() == 2 && result.is_contiguous()
+    // Conditions for bias to be fusable
     && (
-      ( // Conditions for bias to be fusable -- implies direct Lt path without copies.
-        self.is_contiguous() &&
-        // NOTE: fine to have 1-len dims to the left from the right-most one
-        (self.dim() == 1 || self.squeeze().dim() == 1) &&
-        self.sizes().back() == mat2_sizes[1]
-      )
-      || ( // 2D bias restrictions. self.is_contiguous() is implicit when result.is_same(self),
-        // and we need to copy self into result otherwise, so the self's layout becomes irrelevant.
-        // See also TODO from above.
-        activation != Activation::None && // Lt is faster when activation is fused
-        (self.dim() == 2 && at::is_expandable_to(self.sizes(), {mat1_sizes[0], mat2_sizes[1]}))
-      )
+      self.is_contiguous() &&
+      // NOTE: fine to have 1-len dims to the left from the right-most one
+      (self.dim() == 1 || self.squeeze().dim() == 1) &&
+      self.sizes().back() == mat2_sizes[1]
     )
     && ( // some dtype restrictions
       #ifndef USE_ROCM
@@ -300,16 +270,7 @@ bool launchGemmAndBiasCublasLt(
     const Scalar& alpha,
     Activation activation = Activation::None
 ) {
-  // We apply bias in the epilogue only when it is 1D,
-  // or when it can be squeezed to 1D.
-  // self_ptr == nullptr implies ignore bias epilogue
-  // and use standard gemm-like API.
-  const auto* self_ptr = [&]() -> auto {
-    if (self.dim() == 1 || self.squeeze().dim() == 1) {
-      return self.const_data_ptr<scalar_t>();
-    }
-    return static_cast<const scalar_t*>(nullptr);
-  }();
+  const auto* self_ptr = self.const_data_ptr<scalar_t>();
 
   const auto tuning_ctx = at::cuda::tunable::getTuningContext();
   if (tuning_ctx->IsTunableOpEnabled()) {
@@ -395,7 +356,7 @@ Tensor& addmm_out_cuda_impl(Tensor& result, const Tensor& self, const Tensor& ma
   disable_addmm_cuda_lt = isGloballyDisabledAddmmCudaLt(self.device()) || disable_addmm_cuda_lt;
   #endif
   // Condition on the input
-  disable_addmm_cuda_lt = !isInputCompliesAddmmCudaLt(result, self, mat1, mat2, beta, alpha, activation) || disable_addmm_cuda_lt;
+  disable_addmm_cuda_lt = !isInputCompliesAddmmCudaLt(result, self, mat1, mat2, beta, alpha) || disable_addmm_cuda_lt;
   // }
 
   at::ScalarType scalar_type = mat1.scalar_type();
@@ -405,20 +366,19 @@ Tensor& addmm_out_cuda_impl(Tensor& result, const Tensor& self, const Tensor& ma
   if (!result.is_same(self)) {
     at::native::resize_output(result, {mat1.sizes()[0], mat2.sizes()[1]});
 
-    // We use bias ptr in the Lt path only when bias is 1D
-    const auto use_bias_ptr_lt = (self.dim() == 1) && !disable_addmm_cuda_lt;
     const auto self_maybe_expanded = [&]() -> c10::MaybeOwned<Tensor> {
-      if (!use_bias_ptr_lt) {
-        // We do expand self even before
+      if (disable_addmm_cuda_lt) {
+        // When in non-Lt path we do expand self even before
         // check for beta != 0.0 to make sure that
         // test_sparse_csr.py::TestSparseCSRCUDA::test_addmm_errors_*
         // runs green.
         return expand_size(self, result.sizes(), "addmm");
       }
+      // copy next, should broadcast
       return c10::MaybeOwned<Tensor>::borrowed(self);
     }();
-    // We do not copy bias only when we need the bias ptr
-    if (beta.toComplexDouble() != 0.0 && !use_bias_ptr_lt) {
+    // We copy bias when in the non-Lt path
+    if (beta.toComplexDouble() != 0.0 && disable_addmm_cuda_lt) {
       // NOTE: self should broadcast over result
       at::native::copy_(result, *self_maybe_expanded);
     }

aten/src/ATen/native/cuda/GroupedBlas.cpp

@@ -22,9 +22,6 @@
 #include <ATen/native/cuda/RowwiseScaledMM.h>
 #include <ATen/native/cuda/ScaledGroupMM.h>
 #include <ATen/native/cuda/GroupMM.h>
-#ifdef USE_ROCM
-#include <ATen/native/hip/ck_group_gemm.h>
-#endif
 #include <ATen/ceil_div.h>
 
 #ifdef USE_FBGEMM_GENAI
@@ -669,19 +666,12 @@ std::optional<c10::ScalarType> out_dtype) {
   // _scaled_mm_allowed_device is used here within _grouped_mm_cuda which seems incorrect since scale is not used.
   // the _grouped_mm_fallback should be safe for any ROCm GPU since it's just calling typical mm/bmm
   bool use_fast_path = false;
-  if (at::detail::getCUDAHooks().isGPUArch({"gfx942", "gfx950"})) {
-    use_fast_path = true;
-  }
 #endif
   const auto out_dtype_ = _resolve_grouped_mm_out_dtype(mat_a, mat_b, out_dtype);
   Tensor out = create_grouped_gemm_output_tensor(mat_a, mat_b, offs, out_dtype_);
   if (use_fast_path) {
     // fast path, no d2h sync needed
-#ifndef USE_ROCM
     at::cuda::detail::bf16bf16_grouped_mm(mat_a, mat_b, offs, bias, out);
-#else
-    at::hip::detail::group_gemm_ck(mat_a, mat_b, offs, bias, out);
-#endif
   } else {
     _grouped_mm_fallback(mat_a, mat_b, offs, bias, out_dtype, out);
   }

aten/src/ATen/native/cuda/ReduceAMinMaxKernel.cu

@@ -1,6 +1,5 @@
 #define TORCH_ASSERT_NO_OPERATORS
 #include <ATen/Dispatch.h>
-#include <ATen/Dispatch_v2.h>
 #include <ATen/NumericUtils.h>
 #include <ATen/native/DispatchStub.h>
 #include <ATen/native/ReduceAllOps.h>
@@ -29,22 +28,22 @@ void _min_max_values_kernel_cuda_impl(TensorIterator& iter) {
 }
 
 void aminmax_allreduce_launch_kernel(TensorIterator& iter) {
-  AT_DISPATCH_V2(
-      iter.input_dtype(), "aminmax_all_cuda", AT_WRAP([&] {
+  AT_DISPATCH_ALL_TYPES_AND3(
+      kBFloat16, kHalf, kBool, iter.input_dtype(), "aminmax_all_cuda", [&] {
         _min_max_values_kernel_cuda_impl<scalar_t>(iter);
-      }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES), kBFloat16, kHalf, kBool);
+      });
 }
 
 void aminmax_launch_kernel(TensorIterator& iter) {
-  AT_DISPATCH_V2(
-      iter.input_dtype(), "aminmax_cuda", AT_WRAP([&]() {
+  AT_DISPATCH_ALL_TYPES_AND3(
+      kBFloat16, kHalf, kBool, iter.input_dtype(), "aminmax_cuda", [&]() {
         gpu_reduce_kernel<scalar_t, scalar_t>(
             iter,
             MinMaxOps<scalar_t, scalar_t, int32_t>{},
             thrust::pair<scalar_t, scalar_t>(
                 at::numeric_limits<scalar_t>::upper_bound(),
                 at::numeric_limits<scalar_t>::lower_bound()));
-      }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES), kBFloat16, kHalf, kBool);
+      });
 }
 
 } // namespace at::native

aten/src/ATen/native/cuda/ReduceMaxValuesKernel.cu

@@ -1,6 +1,5 @@
 #define TORCH_ASSERT_NO_OPERATORS
 #include <ATen/Dispatch.h>
-#include <ATen/Dispatch_v2.h>
 #include <ATen/NumericUtils.h>
 #include <ATen/native/DispatchStub.h>
 #include <ATen/native/ReduceAllOps.h>
@@ -34,27 +33,27 @@ void max_values_kernel_cuda_impl(TensorIterator& iter) {
 }
 
 void max_values_kernel_cuda(TensorIterator& iter) {
-  AT_DISPATCH_V2(
-      iter.dtype(), "max_values_cuda", AT_WRAP([&]() {
+  AT_DISPATCH_ALL_TYPES_AND3(
+      kBFloat16, kHalf, kBool, iter.dtype(), "max_values_cuda", [&]() {
         max_values_kernel_cuda_impl<scalar_t>(iter);
-      }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES), kBFloat16, kHalf, kBool);
+      });
 }
 
 void max_launch_kernel(TensorIterator& iter) {
-  AT_DISPATCH_V2(
-      iter.input_dtype(), "max_cuda", AT_WRAP([&]() {
+  AT_DISPATCH_ALL_TYPES_AND3(
+      kBFloat16, kHalf, kBool, iter.input_dtype(), "max_cuda", [&]() {
         gpu_reduce_kernel<scalar_t, scalar_t>(
             iter,
             MaxOps<scalar_t>{},
             thrust::pair<scalar_t, int64_t>(
                 at::numeric_limits<scalar_t>::lower_bound(), 0));
-      }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES), kBFloat16, kHalf, kBool);
+      });
 }
 
 void max_all_launch_kernel(TensorIterator &iter) {
-  AT_DISPATCH_V2(iter.input_dtype(), "max_all_cuda", AT_WRAP([&] {
+  AT_DISPATCH_ALL_TYPES_AND3(kBFloat16, kHalf, kBool, iter.input_dtype(), "max_all_cuda", [&] {
     max_values_kernel_cuda_impl<scalar_t>(iter);
-  }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES), kBFloat16, kHalf, kBool);
+  });
 }
 
 REGISTER_DISPATCH(max_values_stub, &max_values_kernel_cuda)

aten/src/ATen/native/cuda/ReduceMinValuesKernel.cu

@@ -12,7 +12,6 @@
 #include <ATen/NumericUtils.h>
 
 #include <ATen/Dispatch.h>
-#include <ATen/Dispatch_v2.h>
 #include <ATen/NumericUtils.h>
 #include <ATen/cuda/NumericLimits.cuh>
 
@@ -34,24 +33,24 @@ void min_values_kernel_cuda_impl(TensorIterator& iter) {
 }
 
 void min_values_kernel_cuda(TensorIterator& iter) {
-  AT_DISPATCH_V2(iter.dtype(), "min_values_cuda", AT_WRAP([&]() {
+  AT_DISPATCH_ALL_TYPES_AND3(kBFloat16, kHalf, kBool, iter.dtype(), "min_values_cuda", [&]() {
     min_values_kernel_cuda_impl<scalar_t>(iter);
-  }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES), kBFloat16, kHalf, kBool);
+  });
 }
 
 void min_launch_kernel(TensorIterator &iter) {
-  AT_DISPATCH_V2(iter.input_dtype(), "min_cuda", AT_WRAP([&]() {
+  AT_DISPATCH_ALL_TYPES_AND3(kBFloat16, kHalf, kBool, iter.input_dtype(), "min_cuda", [&]() {
     gpu_reduce_kernel<scalar_t, scalar_t>(
       iter,
       MinOps<scalar_t>{},
       thrust::pair<scalar_t, int64_t>(at::numeric_limits<scalar_t>::upper_bound(), 0));
-  }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES), kBFloat16, kHalf, kBool);
+  });
 }
 
 void min_all_launch_kernel(TensorIterator &iter) {
-  AT_DISPATCH_V2(iter.input_dtype(), "min_all_cuda", AT_WRAP([&] {
+  AT_DISPATCH_ALL_TYPES_AND3(kBFloat16, kHalf, kBool, iter.input_dtype(), "min_all_cuda", [&] {
     min_values_kernel_cuda_impl<scalar_t>(iter);
-  }), AT_EXPAND(AT_ALL_TYPES), AT_EXPAND(AT_BAREBONES_UNSIGNED_TYPES), kBFloat16, kHalf, kBool);
+  });
 }
 
 REGISTER_DISPATCH(min_values_stub, &min_values_kernel_cuda)

aten/src/ATen/native/hip/ck_group_gemm.h

@@ -1,19 +0,0 @@
-#pragma once
-
-#include <ATen/Tensor.h>
-#include <c10/core/ScalarType.h>
-#include <optional>
-
-namespace at {
-namespace hip {
-namespace detail {
-void group_gemm_ck(
-    const at::Tensor& mat_a,
-    const at::Tensor& mat_b,
-    const std::optional<at::Tensor>& offs,
-    const std::optional<at::Tensor>& bias,
-    at::Tensor& out);
-
-} // namespace detail
-} // namespace hip
-} // namespace at

aten/src/ATen/native/hip/ck_group_gemm.hip

@@ -1,462 +0,0 @@
-#undef __HIP_NO_HALF_CONVERSIONS__
-#include <ATen/hip/HIPContext.h>
-#include <ATen/Tensor.h>
-#include <ATen/TensorAccessor.h>
-#include <c10/hip/HIPStream.h>
-#include <iostream>
-#include <vector>
-#include <optional>
-#include <type_traits>
-
-#include <ck/ck.hpp>
-#include <ck/tensor_operation/gpu/device/tensor_layout.hpp>
-#include <ck/tensor_operation/gpu/device/gemm_specialization.hpp>
-#include <ck/tensor_operation/gpu/device/impl/device_grouped_gemm_multiple_d_splitk_xdl_cshuffle_two_stage.hpp>
-#include <ck/tensor_operation/gpu/element/element_wise_operation.hpp>
-#include <ck/utility/tuple.hpp>
-
-template <ck::index_t... Is>
-using S = ck::Sequence<Is...>;
-
-namespace at {
-namespace hip {
-namespace detail {
-
-namespace CkTypes {
-    using BF16 = ck::bhalf_t;
-    using F16 = ck::half_t;
-    using F32 = float;
-    using PassThrough = ck::tensor_operation::element_wise::PassThrough;
-}
-
-template <typename ALayout, typename BLayout, typename DataType>
-using GroupedGemmKernel = ck::tensor_operation::device::DeviceGroupedGemmMultipleDSplitKXdlCShuffleTwoStage<
-    ALayout, BLayout, ck::Tuple<>, ck::tensor_layout::gemm::RowMajor,
-    DataType, DataType, CkTypes::F32, DataType, ck::Tuple<>, DataType,
-    CkTypes::PassThrough, CkTypes::PassThrough, CkTypes::PassThrough,
-    ck::tensor_operation::device::GemmSpecialization::MNKPadding,
-    1, 256, 256, 128, 32, 8, 8, 32, 32, 4, 2,
-    S<1,4,64,1>,  S<0,2,1,3>, S<0,2,1,3>,
-    3, 8, 8, 1,
-    S<1,4,64,1>,  S<0,2,1,3>, S<0,2,1,3>,
-    3, 8, 8, 1,
-    1, 1,
-    S<1,32,1,8>, 4
->;
-
-template <typename ALayout, typename BLayout, typename DataType>
-void launch_grouped_bgemm_ck_impl_dispatch(
-    const at::Tensor& mat_a,
-    const at::Tensor& mat_b,
-    const std::optional<at::Tensor>& offs,
-    at::Tensor& out)
-{
-    using DeviceOp = GroupedGemmKernel<ALayout, BLayout, DataType>;
-    using PassThrough = CkTypes::PassThrough;
-
-    std::vector<ck::tensor_operation::device::GemmDesc> gemm_descs;
-    std::vector<const void*> p_a_ptrs, p_b_ptrs;
-    std::vector<void*> p_e_ptrs;
-    // Note: d_ptrs will be resized after we populate the other vectors
-
-    const int mat_a_dim = mat_a.dim();
-    const int mat_b_dim = mat_b.dim();
-
-    const char* a_ptr_base = reinterpret_cast<const char*>(mat_a.data_ptr());
-    const char* b_ptr_base = reinterpret_cast<const char*>(mat_b.data_ptr());
-    char* out_ptr_base = reinterpret_cast<char*>(out.data_ptr());
-    const size_t a_element_size = mat_a.element_size();
-    const size_t b_element_size = mat_b.element_size();
-    const size_t out_element_size = out.element_size();
-
-    // for each group, calculate m,n,k,lda,ldb,ldc and A,B,out pointer base addresses.
-    if (mat_a_dim == 2 && mat_b_dim == 2) {
-        // 2D*2D case requires offset tensor
-        auto offs_accessor = offs->accessor<int, 1>();
-        int num_groups = offs_accessor.size(0);
-        const int M = mat_a.size(0); // number of rows in A
-        const int N = mat_b.size(1); // number of columns in B
-        const int K = mat_a.size(1); // columns in A == rows in B
-        // for 2d*2d input, output is 3d.
-        // for each group, A columns (K) are sliced. M and N dimensions are not sliced.
-        for (int i = 0; i < num_groups; ++i) {
-            int start_k = (i == 0) ? 0 : offs_accessor[i-1];
-            int end_k = offs_accessor[i];
-            int k = end_k - start_k;
-
-            //K dimension are sliced, hence select stride(1) always.
-            //K dimension is always dimension 1, regardless of memory layout (row/column major)
-            const void* group_a_ptr = a_ptr_base + start_k * mat_a.stride(1) * a_element_size;
-            const void* group_b_ptr;
-            int ldb;
-
-            if (std::is_same<BLayout, ck::tensor_layout::gemm::RowMajor>::value) {
-                // Row-major B [K,N]: K values are horizontally adjacent, use stride(1) for K offset
-                group_b_ptr = b_ptr_base + start_k * mat_b.stride(1) * b_element_size;
-                // Leading dimension = distance between rows = stride(0)
-                ldb = mat_b.stride(0);
-            } else {
-                // Column-major B [K,N]: K values are vertically adjacent, use stride(0) for K offset
-                group_b_ptr = b_ptr_base + start_k * mat_b.stride(0) * b_element_size;
-                // Leading dimension = distance between columns = stride(1)
-                ldb = mat_b.stride(1);
-            }
-
-            // Calculate output pointer for group i in 3D tensor [num_groups, M, N]
-            // stride(0) = M*N elements between groups, so skip i*stride(0) elements to reach group i
-            void* group_e_ptr = out_ptr_base + i * out.stride(0) * out_element_size;
-            int lda, ldc;
-            if (std::is_same<ALayout, ck::tensor_layout::gemm::RowMajor>::value) {
-                // Row-major A [M,K]: leading dimension = distance between rows = stride(0)
-                lda = mat_a.stride(0);
-            } else {
-                // Column-major A [M,K]: leading dimension = distance between columns = stride(1)
-                lda = mat_a.stride(1);
-            }
-            // Output is always row-major in 3D tensor [num_groups, M, N]
-            // Leading dimension for each group's [M,N] slice = stride(1) = N
-            ldc = out.stride(1);
-            size_t output_group_bytes = M * N * out_element_size;
-            void* group_e_ptr_end = (char*)group_e_ptr + output_group_bytes;
-
-            gemm_descs.push_back({
-                static_cast<ck::index_t>(M),
-                static_cast<ck::index_t>(N),
-                static_cast<ck::index_t>(k),
-                static_cast<ck::index_t>(lda),
-                static_cast<ck::index_t>(ldb),
-                static_cast<ck::index_t>(ldc),
-                {} // --> stride_Ds_
-            });
-            p_a_ptrs.push_back(group_a_ptr);
-            p_b_ptrs.push_back(group_b_ptr);
-            p_e_ptrs.push_back(group_e_ptr);
-        }
-    } else if (mat_a_dim == 2 && mat_b_dim == 3) {
-        // 2D*3D case requires offset tensor
-        auto offs_accessor = offs->accessor<int, 1>();
-        int num_groups = offs_accessor.size(0);
-
-        // 2d*3d input, output is 2d.
-        // A: [m * n_groups, k], B: [n_groups, n, k] or [n_groups, k, n], Output: [m * n_groups, n]
-        // Offset divides M dimension (rows of A), each group gets different rows of A and different batch of B
-        const int K = mat_a.size(1); // columns in A
-        // For 2D-3D case: The output determines N (result width)
-        const int N = out.size(1); // N is the width of the output tensor
-
-        for (int i = 0; i < num_groups; ++i) {
-            int start_m = (i == 0) ? 0 : offs_accessor[i - 1];
-            int end_m = offs_accessor[i];
-            int m = end_m - start_m;
-
-            // Skip zero-sized groups but continue processing subsequent groups
-            if (m <= 0) {
-                continue;
-            }
-
-            // Select A rows for group i: skip start_m rows
-            const void* group_a_ptr;
-            int lda;
-            if (std::is_same<ALayout, ck::tensor_layout::gemm::RowMajor>::value) {
-                // Row-major A [total_m, K]: skip start_m rows, each row is stride(0) elements apart
-                group_a_ptr = a_ptr_base + start_m * mat_a.stride(0) * a_element_size;
-                lda = mat_a.stride(0); // distance between rows
-            } else {
-                // Column-major A [total_m, K]: skip start_m elements in the first dimension (stride(0) is between rows)
-                group_a_ptr = a_ptr_base + start_m * mat_a.stride(0) * a_element_size;
-
-                // Detect stride pattern for A tensor to determine appropriate lda calculation
-                bool a_is_strided_tensor = (mat_a.stride(0) > mat_a.size(0));
-
-                if (a_is_strided_tensor) {
-                    // For strided A tensors: stride(0) gives the actual leading dimension
-                    lda = mat_a.stride(0);
-                } else {
-                    // For non-strided A tensors: use the M dimension (total rows)
-                    lda = mat_a.size(0); // Total M dimension for column-major layout
-                }
-            }
-
-            // Select B batch for group i: B[i, :, :]
-            const void* group_b_ptr = b_ptr_base + i * mat_b.stride(0) * b_element_size;
-            int ldb;
-
-            if (std::is_same<BLayout, ck::tensor_layout::gemm::RowMajor>::value) {
-                // Row-major GEMM: expecting B as [K, N] but we have [N, K], so transpose needed
-                ldb = mat_b.stride(2); // Leading dimension for accessing as [K, N]
-            } else {
-                // Detect stride pattern to determine appropriate ldb calculation
-                bool is_strided_tensor = (mat_b.stride(2) > mat_b.size(2));
-
-                if (is_strided_tensor) {
-                    // For strided tensors: stride(2) gives the actual leading dimension
-                    ldb = mat_b.stride(2);
-                } else {
-                    // For non-strided tensors: use the N dimension
-                    ldb = mat_b.size(1);
-                }
-            }
-
-            // Output for this group: rows [start_m:end_m, :] in 2D output [total_m, N]
-            void* group_e_ptr = out_ptr_base + start_m * out.stride(0) * out_element_size;
-            int ldc = out.stride(0); // distance between rows in output (should be N for 2D case)
-
-            gemm_descs.push_back({
-                static_cast<ck::index_t>(m),
-                static_cast<ck::index_t>(N),
-                static_cast<ck::index_t>(K),
-                static_cast<ck::index_t>(lda),
-                static_cast<ck::index_t>(ldb),
-                static_cast<ck::index_t>(ldc),
-                {} // --> stride_Ds_
-            });
-            p_a_ptrs.push_back(group_a_ptr);
-            p_b_ptrs.push_back(group_b_ptr);
-            p_e_ptrs.push_back(group_e_ptr);
-        }
-    } else if (mat_a_dim == 3 && mat_b_dim == 3) {
-        // 3d*3d input, output is 3d - batched matrix multiplication
-        // A: [batch, m, k], B: [batch, k, n] or [batch, n, k] (depending on transpose), Output: [batch, m, n]
-        // Each batch is processed as a separate GEMM operation
-        const int batch_size = mat_a.size(0);
-        const int M = mat_a.size(1); // rows in each A matrix
-        const int K = mat_a.size(2); // columns in A == rows in B (or columns if B is transposed)
-
-        // Determine N from B tensor - it could be B.size(1) or B.size(2) depending on layout
-        int N;
-        if (mat_b.size(1) == K) {
-            // B is [batch, k, n] - normal layout
-            N = mat_b.size(2);
-        } else if (mat_b.size(2) == K) {
-            // B is [batch, n, k] - transposed layout
-            N = mat_b.size(1);
-        } else {
-            TORCH_CHECK(false, "CK Group GEMM 3D-3D: B tensor dimensions incompatible with A. A=[",
-                       batch_size, ",", M, ",", K, "], B=[", mat_b.size(0), ",", mat_b.size(1), ",", mat_b.size(2), "]");
-        }
-
-        for (int i = 0; i < batch_size; ++i) {
-            // Select A batch for group i: A[i, :, :]
-            const void* group_a_ptr = a_ptr_base + i * mat_a.stride(0) * a_element_size;
-
-            // Select B batch for group i: B[i, :, :]
-            const void* group_b_ptr = b_ptr_base + i * mat_b.stride(0) * b_element_size;
-
-            // Select output batch for group i: Output[i, :, :]
-            void* group_e_ptr = out_ptr_base + i * out.stride(0) * out_element_size;
-
-            int lda, ldb, ldc;
-
-            if (std::is_same<ALayout, ck::tensor_layout::gemm::RowMajor>::value) {
-                // Row-major A: leading dimension = distance between rows = stride(1)
-                lda = mat_a.stride(1);
-            } else {
-                // Column-major A: leading dimension = distance between columns = stride(2)
-                lda = mat_a.stride(2);
-            }
-
-            if (std::is_same<BLayout, ck::tensor_layout::gemm::RowMajor>::value) {
-                // Row-major B: leading dimension = distance between rows
-                if (mat_b.size(1) == K) {
-                    // B is [batch, k, n] - normal layout
-                    ldb = mat_b.stride(1); // stride between K rows
-                } else {
-                    // B is [batch, n, k] - transposed layout, treat as [k, n] for GEMM
-                    ldb = mat_b.stride(2); // stride between N rows (since we're accessing as [k,n])
-                }
-            } else {
-                // Column-major B: leading dimension = distance between columns
-                if (mat_b.size(1) == K) {
-                    // B is [batch, k, n] - normal layout
-                    ldb = mat_b.stride(2); // stride between N columns
-                } else {
-                    // B is [batch, n, k] - transposed layout
-                    ldb = mat_b.stride(1); // stride between K columns (since we're accessing as [n,k]→[k,n])
-                }
-            }
-
-            // Output is typically row-major: leading dimension = distance between rows = stride(1)
-            ldc = out.stride(1);
-
-            gemm_descs.push_back({
-                static_cast<ck::index_t>(M),
-                static_cast<ck::index_t>(N),
-                static_cast<ck::index_t>(K),
-                static_cast<ck::index_t>(lda),
-                static_cast<ck::index_t>(ldb),
-                static_cast<ck::index_t>(ldc),
-                {} // --> stride_Ds_
-            });
-            p_a_ptrs.push_back(group_a_ptr);
-            p_b_ptrs.push_back(group_b_ptr);
-            p_e_ptrs.push_back(group_e_ptr);
-        }
-    } else if (mat_a_dim == 3 && mat_b_dim == 2) {
-        // 3D*2D case requires offset tensor
-        auto offs_accessor = offs->accessor<int, 1>();
-        int num_groups = offs_accessor.size(0);
-        // 3d*2d input, output is 3d.
-        // A: [n_groups, m, k], B: [k, total_n] (assuming row-major for both)
-        // Offset divides N dimension of B, each group gets different slice of B and different batch of A
-        const int batch_size = mat_a.size(0); // n_groups
-        const int M = mat_a.size(1); // rows in each A matrix
-        const int K = mat_a.size(2); // columns in A
-
-        // For row-major A and B case: B should be [K, total_N]
-        const int total_N = mat_b.size(1); // B is [K, total_N] for row-major
-
-        for (int i = 0; i < num_groups; ++i) {
-            int start_n = (i == 0) ? 0 : offs_accessor[i - 1];
-            int end_n = offs_accessor[i];
-            int n = end_n - start_n;
-
-            // Skip zero-sized groups but continue processing subsequent groups
-            if (n <= 0) {
-                continue;
-            }
-
-            // Select A batch for group i: A[i, :, :]
-            const void* group_a_ptr = a_ptr_base + i * mat_a.stride(0) * a_element_size;
-
-            // Select B slice for group i: B[:, start_n:end_n] (B[K, total_N])
-            const void* group_b_ptr;
-            int ldb;
-
-            // Check if B is row-major or column-major
-            if (std::is_same<BLayout, ck::tensor_layout::gemm::RowMajor>::value) {
-                // Row-major B [K, total_N]: slice columns [start_n:end_n]
-                group_b_ptr = b_ptr_base + start_n * mat_b.stride(1) * b_element_size;
-                ldb = mat_b.stride(0); // distance between rows (should be total_N)
-            } else {
-                // Column-major B [K, total_N]: slice columns [start_n:end_n]
-                group_b_ptr = b_ptr_base + start_n * mat_b.stride(1) * b_element_size;
-                ldb = mat_b.stride(1); // distance between columns (should be K)
-            }
-
-            // Select output slice for group i: Output[:, start_n:end_n]
-            void* group_e_ptr = out_ptr_base + start_n * out.stride(1) * out_element_size;
-
-            int lda, ldc;
-
-            // Row-major A: leading dimension = distance between rows = stride(1)
-            lda = mat_a.stride(1);
-            // Output is row-major: leading dimension = distance between rows = stride(0)
-            ldc = out.stride(0);
-
-            gemm_descs.push_back({
-                static_cast<ck::index_t>(M),
-                static_cast<ck::index_t>(n),
-                static_cast<ck::index_t>(K),
-                static_cast<ck::index_t>(lda),
-                static_cast<ck::index_t>(ldb),
-                static_cast<ck::index_t>(ldc),
-                {} // --> stride_Ds_
-            });
-            p_a_ptrs.push_back(group_a_ptr);
-            p_b_ptrs.push_back(group_b_ptr);
-            p_e_ptrs.push_back(group_e_ptr);
-        }
-    } else {
-        TORCH_CHECK(false, "CK Group GEMM: Unsupported dimensions, mat A dim is ", mat_a_dim, ", mat B dim is ", mat_b_dim);
-    }
-
-    TORCH_INTERNAL_ASSERT(p_a_ptrs.size() > 0, "CK Group GEMM: No valid groups");
-
-    // Initialize d_ptrs with the correct size
-    std::vector<std::array<const void*, 0>> d_ptrs(p_a_ptrs.size());
-
-    static DeviceOp gemm_instance;
-    auto argument = gemm_instance.MakeArgument(
-        p_a_ptrs, p_b_ptrs, d_ptrs, p_e_ptrs,
-        gemm_descs, PassThrough{}, PassThrough{}, PassThrough{}
-    );
-    TORCH_INTERNAL_ASSERT(gemm_instance.IsSupportedArgument(argument),
-        "CK Group GEMM: argument unsupported (shape/strides/type config)");
-    size_t arg_buf_size = gemm_instance.GetDeviceKernelArgSize(&argument);
-    size_t ws_size = gemm_instance.GetWorkSpaceSize(&argument);
-
-    void* gemm_arg_buf = nullptr;
-    void* ws_buf = nullptr;
-
-    hipMalloc(&gemm_arg_buf, arg_buf_size);
-    hipMalloc(&ws_buf, ws_size);
-
-    gemm_instance.SetDeviceKernelArgs(&argument, gemm_arg_buf);
-    gemm_instance.SetWorkSpacePointer(&argument, ws_buf);
-
-    auto invoker = gemm_instance.MakeInvoker();
-    hipStream_t stream = c10::hip::getCurrentHIPStream();
-    invoker.Run(argument, {stream});
-    hipFree(gemm_arg_buf);
-    hipFree(ws_buf);
-}
-
-void group_gemm_ck(
-    const at::Tensor& input_a,
-    const at::Tensor& input_b_colmajor,
-    const std::optional<at::Tensor>& offs,
-    const std::optional<at::Tensor>& /*bias*/,
-    at::Tensor& out)
-{
-    // Detect if input_a is row-major based on stride pattern
-    bool a_row_major = (input_a.dim() == 3) ? (input_a.stride(2) == 1) : (input_a.stride(1) == 1);
-    bool b_col_major = (input_b_colmajor.dim() == 3) ? (input_b_colmajor.stride(1) == 1) : (input_b_colmajor.stride(0) == 1);
-    // Ensure tensor A is row-major and contiguous if not already
-    at::Tensor mat_a = input_a;
-    if (!a_row_major) {
-        // If A is not row-major, make it contiguous (row-major)
-        mat_a = input_a.contiguous();
-    }
-    // Force tensor B to be column-major using double transpose trick
-    // This guarantees stride(0) == 1 and stride(1) == K for [K, N] shape
-    at::Tensor mat_b = input_b_colmajor;
-    if (!b_col_major) {
-        mat_b = input_b_colmajor.transpose(-2, -1).contiguous().transpose(-2, -1);
-    }
-
-    // For 3D tensors, check the last dimension stride for row-major detection
-    a_row_major = (mat_a.dim() == 3) ? (mat_a.stride(2) == 1) : (mat_a.stride(1) == 1);
-    bool b_row_major = (mat_b.dim() == 3) ? (mat_b.stride(2) == 1) : (mat_b.stride(1) == 1);
-
-    if (mat_a.dtype() == at::kBFloat16) {
-        // bf16 path
-        if (a_row_major && b_row_major) {
-            launch_grouped_bgemm_ck_impl_dispatch<ck::tensor_layout::gemm::RowMajor, ck::tensor_layout::gemm::RowMajor, CkTypes::BF16>(mat_a, mat_b, offs, out);
-        } else if (a_row_major && !b_row_major) {
-            launch_grouped_bgemm_ck_impl_dispatch<ck::tensor_layout::gemm::RowMajor, ck::tensor_layout::gemm::ColumnMajor, CkTypes::BF16>(mat_a, mat_b, offs, out);
-        } else if (!a_row_major && b_row_major) {
-            launch_grouped_bgemm_ck_impl_dispatch<ck::tensor_layout::gemm::ColumnMajor, ck::tensor_layout::gemm::RowMajor, CkTypes::BF16>(mat_a, mat_b, offs, out);
-        } else {
-            launch_grouped_bgemm_ck_impl_dispatch<ck::tensor_layout::gemm::ColumnMajor, ck::tensor_layout::gemm::ColumnMajor, CkTypes::BF16>(mat_a, mat_b, offs, out);
-        }
-    } else if (mat_a.dtype() == at::kHalf) {
-        // fp16 path
-        if (a_row_major && b_row_major) {
-            launch_grouped_bgemm_ck_impl_dispatch<ck::tensor_layout::gemm::RowMajor, ck::tensor_layout::gemm::RowMajor, CkTypes::F16>(mat_a, mat_b, offs, out);
-        } else if (a_row_major && !b_row_major) {
-            launch_grouped_bgemm_ck_impl_dispatch<ck::tensor_layout::gemm::RowMajor, ck::tensor_layout::gemm::ColumnMajor, CkTypes::F16>(mat_a, mat_b, offs, out);
-        } else if (!a_row_major && b_row_major) {
-            launch_grouped_bgemm_ck_impl_dispatch<ck::tensor_layout::gemm::ColumnMajor, ck::tensor_layout::gemm::RowMajor, CkTypes::F16>(mat_a, mat_b, offs, out);
-        } else {
-            launch_grouped_bgemm_ck_impl_dispatch<ck::tensor_layout::gemm::ColumnMajor, ck::tensor_layout::gemm::ColumnMajor, CkTypes::F16>(mat_a, mat_b, offs, out);
-        }
-    } else if (mat_a.dtype() == at::kFloat) {
-        // fp32 path
-        if (a_row_major && b_row_major) {
-            launch_grouped_bgemm_ck_impl_dispatch<ck::tensor_layout::gemm::RowMajor, ck::tensor_layout::gemm::RowMajor, CkTypes::F32>(mat_a, mat_b, offs, out);
-        } else if (a_row_major && !b_row_major) {
-            launch_grouped_bgemm_ck_impl_dispatch<ck::tensor_layout::gemm::RowMajor, ck::tensor_layout::gemm::ColumnMajor, CkTypes::F32>(mat_a, mat_b, offs, out);
-        } else if (!a_row_major && b_row_major) {
-            launch_grouped_bgemm_ck_impl_dispatch<ck::tensor_layout::gemm::ColumnMajor, ck::tensor_layout::gemm::RowMajor, CkTypes::F32>(mat_a, mat_b, offs, out);
-        } else {
-            launch_grouped_bgemm_ck_impl_dispatch<ck::tensor_layout::gemm::ColumnMajor, ck::tensor_layout::gemm::ColumnMajor, CkTypes::F32>(mat_a, mat_b, offs, out);
-        }
-    } else {
-        TORCH_CHECK(false, "CK Group GEMM: Unsupported mat_a dtype");
-    }
-
-}
-
-} // namespace detail
-} // namespace hip
-} // namespace at

aten/src/ATen/native/mps/operations/LossOps.mm

@@ -212,12 +212,17 @@ static Tensor& bce_loss_out_impl(const Tensor& input,
   loss.resize_((reduction == Reduction::None || grad_output.defined()) ? target.sizes() : IntArrayRef({}));
   TORCH_CHECK(loss.is_mps());
 
+  Tensor loss_squeezed = loss.squeeze();
+  Tensor input_squeezed = input.squeeze();
+  Tensor target_squeezed = target.squeeze();
+
   @autoreleasepool {
-    std::string key = op_name + reductionToString(reduction) + getTensorsStringKey({input, target, weight});
+    std::string key =
+        op_name + reductionToString(reduction) + getTensorsStringKey({input_squeezed, target_squeezed, weight});
 
     auto cachedGraph = LookUpOrCreateCachedGraph<CachedGraph>(key, [&](auto mpsGraph, auto newCachedGraph) {
-      newCachedGraph->inputTensor = mpsGraphRankedPlaceHolder(mpsGraph, input);
-      newCachedGraph->targetTensor = mpsGraphRankedPlaceHolder(mpsGraph, target);
+      newCachedGraph->inputTensor = mpsGraphRankedPlaceHolder(mpsGraph, input_squeezed);
+      newCachedGraph->targetTensor = mpsGraphRankedPlaceHolder(mpsGraph, target_squeezed);
 
       MPSGraphTensor* bceLossUnweighted = nil;
       // if grad_output is defined, then it's a backward pass
@@ -247,12 +252,12 @@ static Tensor& bce_loss_out_impl(const Tensor& input,
           newCachedGraph->gradInputTensor = bceLoss;
         }
       } else {
-        newCachedGraph->lossTensor = reduceTensor(bceLoss, reduction, mpsGraph, input.sizes().size());
+        newCachedGraph->lossTensor = reduceTensor(bceLoss, reduction, mpsGraph, input_squeezed.sizes().size());
       }
     });
-    Placeholder inputPlaceholder = Placeholder(cachedGraph->inputTensor, input);
-    Placeholder targetPlaceholder = Placeholder(cachedGraph->targetTensor, target);
-    Placeholder lossPlaceholder = Placeholder(cachedGraph->lossTensor, loss);
+    Placeholder inputPlaceholder = Placeholder(cachedGraph->inputTensor, input_squeezed);
+    Placeholder targetPlaceholder = Placeholder(cachedGraph->targetTensor, target_squeezed);
+    Placeholder lossPlaceholder = Placeholder(cachedGraph->lossTensor, loss_squeezed);
 
     NSMutableDictionary* feeds = [[NSMutableDictionary new] autorelease];
 

benchmarks/operator_benchmark/pt/addmm_test.py

@@ -53,8 +53,10 @@ class AddmmBenchmark(op_bench.TorchBenchmarkBase):
         return torch.addmm(input_one, mat1, mat2)
 
 
-op_bench.generate_pt_test(addmm_short_configs + addmm_long_configs, AddmmBenchmark)
-op_bench.generate_pt_gradient_test(addmm_long_configs, AddmmBenchmark)
+op_bench.generate_pt_test(addmm_long_configs + addmm_long_configs, AddmmBenchmark)
+op_bench.generate_pt_gradient_test(
+    addmm_long_configs + addmm_long_configs, AddmmBenchmark
+)
 
 """Mircobenchmark for addbmm operator."""
 
@@ -105,7 +107,9 @@ addbmm_short_configs = op_bench.cross_product_configs(
 )
 
 op_bench.generate_pt_test(addbmm_long_configs + addbmm_short_configs, AddbmmBenchmark)
-op_bench.generate_pt_gradient_test(addbmm_long_configs, AddbmmBenchmark)
+op_bench.generate_pt_gradient_test(
+    addbmm_long_configs + addbmm_short_configs, AddbmmBenchmark
+)
 
 if __name__ == "__main__":
     op_bench.benchmark_runner.main()

c10/core/SymBool.cpp

@@ -1,5 +1,4 @@
 #include <c10/core/SymBool.h>
-#include <c10/core/SymInt.h>
 #include <c10/core/SymNodeImpl.h>
 
 namespace c10 {
@@ -112,17 +111,4 @@ bool SymBool::has_hint() const {
   return toSymNodeImpl()->has_hint();
 }
 
-SymInt SymBool::toSymInt() const {
-  // If concrete bool, return concrete SymInt
-  if (auto ma = maybe_as_bool()) {
-    return SymInt(*ma ? 1 : 0);
-  }
-
-  // Symbolic case: use sym_ite to convert bool to int (0 or 1)
-  auto node = toSymNodeImpl();
-  auto one_node = node->wrap_int(1);
-  auto zero_node = node->wrap_int(0);
-  return SymInt(node->sym_ite(one_node, zero_node));
-}
-
 } // namespace c10

c10/core/SymBool.h

@@ -12,8 +12,6 @@
 
 namespace c10 {
 
-class SymInt;
-
 class C10_API SymBool {
  public:
   /*implicit*/ SymBool(bool b) : data_(b) {}
@@ -82,10 +80,6 @@ class C10_API SymBool {
     return toSymNodeImplUnowned()->constant_bool();
   }
 
-  // Convert SymBool to SymInt (0 or 1)
-  // This is the C++ equivalent of Python's cast_symbool_to_symint_guardless
-  SymInt toSymInt() const;
-
   bool is_heap_allocated() const {
     return ptr_;
   }

c10/cuda/CUDAAllocatorConfig.cpp

@@ -106,9 +106,6 @@ void CUDAAllocatorConfig::parseArgs(const std::string& env) {
     } else if (key == "graph_capture_record_stream_reuse") {
       i = parseGraphCaptureRecordStreamReuse(tokenizer, i);
       used_native_specific_option = true;
-    } else if (key == "per_process_memory_fraction") {
-      i = parsePerProcessMemoryFraction(tokenizer, i);
-      used_native_specific_option = true;
     } else {
       const auto& keys =
           c10::CachingAllocator::AcceleratorAllocatorConfig::getKeys();
@@ -149,18 +146,6 @@ size_t CUDAAllocatorConfig::parseGraphCaptureRecordStreamReuse(
   return i;
 }
 
-double CUDAAllocatorConfig::parsePerProcessMemoryFraction(
-    const c10::CachingAllocator::ConfigTokenizer& tokenizer,
-    size_t i) {
-  tokenizer.checkToken(++i, ":");
-  double val_env = tokenizer.toDouble(++i);
-  TORCH_CHECK_VALUE(
-      val_env >= 0.0 && val_env <= 1.0,
-      "per_process_memory_fraction is invalid, set it in [0.0, 1.0]");
-  m_per_process_memory_fraction = val_env;
-  return i;
-}
-
 size_t CUDAAllocatorConfig::parsePinnedNumRegisterThreads(
     const c10::CachingAllocator::ConfigTokenizer& tokenizer,
     size_t i) {

c10/cuda/CUDAAllocatorConfig.h

@@ -61,10 +61,6 @@ class C10_CUDA_API CUDAAllocatorConfig {
     return instance().m_graph_capture_record_stream_reuse;
   }
 
-  static double per_process_memory_fraction() {
-    return instance().m_per_process_memory_fraction;
-  }
-
   /** Pinned memory allocator settings */
   static bool pinned_use_cuda_host_register() {
     return instance().m_pinned_use_cuda_host_register;
@@ -156,8 +152,7 @@ class C10_CUDA_API CUDAAllocatorConfig {
         "pinned_use_hip_host_register",
         "graph_capture_record_stream_reuse",
         "pinned_reserve_segment_size_mb",
-        "pinned_num_register_threads",
-        "per_process_memory_fraction"};
+        "pinned_num_register_threads"};
     return keys;
   }
 
@@ -182,9 +177,6 @@ class C10_CUDA_API CUDAAllocatorConfig {
   size_t parseGraphCaptureRecordStreamReuse(
       const c10::CachingAllocator::ConfigTokenizer& tokenizer,
       size_t i);
-  double parsePerProcessMemoryFraction(
-      const c10::CachingAllocator::ConfigTokenizer& tokenizer,
-      size_t i);
 
   std::atomic<size_t> m_pinned_num_register_threads{1};
   std::atomic<size_t> m_pinned_reserve_segment_size_mb{0};
@@ -197,7 +189,6 @@ class C10_CUDA_API CUDAAllocatorConfig {
   std::atomic<bool> m_release_lock_on_cudamalloc{false};
   std::atomic<bool> m_pinned_use_cuda_host_register{false};
   std::atomic<bool> m_graph_capture_record_stream_reuse{false};
-  std::atomic<double> m_per_process_memory_fraction{1.0};
 };
 
 // Keep this for backwards compatibility

c10/cuda/CUDACachingAllocator.cpp

@@ -1100,7 +1100,7 @@ class RingBuffer {
 } // anonymous namespace
 } // namespace Native
 
-static std::string reportProcessMemoryInfo(const cudaDeviceProp& prop) {
+static std::string reportProcessMemoryInfo(c10::DeviceIndex device) {
 #ifdef PYTORCH_C10_DRIVER_API_SUPPORTED
   void* nvml_handle = DriverAPI::get_nvml_handle();
   if (!nvml_handle) {
@@ -1111,6 +1111,9 @@ static std::string reportProcessMemoryInfo(const cudaDeviceProp& prop) {
     return true;
   }();
 
+  cudaDeviceProp prop{};
+  C10_CUDA_CHECK(cudaGetDeviceProperties(&prop, device));
+
   // NOLINTNEXTLINE(*-c-arrays)
   char pci_id[80];
   snprintf(
@@ -1212,16 +1215,14 @@ class DeviceCachingAllocator {
   // record used memory.
   size_t total_allocated_memory = 0;
 
-  cudaDeviceProp device_prop;
-
-  // maximum amount of memory that device is allowed to
-  // allocate. This is set iff memory fraction is less than 1
-  std::optional<size_t> allowed_memory_maximum{std::nullopt};
+  size_t allowed_memory_maximum = 0;
 
   // all live expandable segments
   std::vector<ExpandableSegment*> expandable_segments_;
   std::vector<c10::DeviceIndex> devices_with_peer_access_;
 
+  bool set_fraction = false;
+
   bool record_history = false;
 
   std::atomic<CreateContextFn> context_recorder_;
@@ -1263,9 +1264,6 @@ class DeviceCachingAllocator {
       : device_id(id),
         large_blocks(/*small=*/false),
         small_blocks(/*small=*/true) {
-    C10_CUDA_CHECK(cudaGetDeviceProperties(&device_prop, id));
-
-    setMemoryFraction(CUDAAllocatorConfig::per_process_memory_fraction());
     stats.max_split_size =
         static_cast<int64_t>(AcceleratorAllocatorConfig::max_split_size());
     context_recorder_.store(nullptr);
@@ -1401,7 +1399,7 @@ class DeviceCachingAllocator {
     if (!block_found) {
       // Do garbage collection if the flag is set.
       if (C10_UNLIKELY(
-              allowed_memory_maximum.has_value() &&
+              set_fraction &&
               AcceleratorAllocatorConfig::garbage_collection_threshold() >
                   0.0)) {
         garbage_collect_cached_blocks(context);
@@ -1458,12 +1456,11 @@ class DeviceCachingAllocator {
       C10_CUDA_CHECK(cudaMemGetInfo(&device_free, &device_total));
       std::string allowed_info;
 
-      if (allowed_memory_maximum.has_value()) {
-        allowed_info =
-            format_size(allowed_memory_maximum.value()) + " allowed; ";
+      if (set_fraction) {
+        allowed_info = format_size(allowed_memory_maximum) + " allowed; ";
       }
 
-      std::string proc_info = reportProcessMemoryInfo(device_prop);
+      std::string proc_info = reportProcessMemoryInfo(device_id);
 
       record_trace(
           TraceEntry::OOM,
@@ -1521,7 +1518,7 @@ class DeviceCachingAllocator {
       for (const auto& obs : observers_local) {
         obs(device_id,
             alloc_size,
-            allowed_memory_maximum.value_or(device_total),
+            set_fraction ? allowed_memory_maximum : device_total,
             device_free);
       }
 
@@ -2018,26 +2015,25 @@ class DeviceCachingAllocator {
 
   /** get memory fraction limiting maximum allocated memory **/
   double getMemoryFraction() {
-    if (!allowed_memory_maximum.has_value()) {
+    if (!set_fraction) {
       return 1.0;
     }
 
-    return static_cast<double>(allowed_memory_maximum.value()) /
-        static_cast<double>(device_prop.totalGlobalMem);
+    size_t device_free = 0;
+    size_t device_total = 0;
+    C10_CUDA_CHECK(cudaMemGetInfo(&device_free, &device_total));
+    return static_cast<double>(allowed_memory_maximum) /
+        static_cast<double>(device_total);
   }
 
   /** set memory fraction to limit maximum allocated memory **/
   void setMemoryFraction(double fraction) {
-    TORCH_CHECK(
-        0 <= fraction && fraction <= 1,
-        "invalid fraction:",
-        fraction,
-        ". Please set within [0, 1].");
-    allowed_memory_maximum = std::nullopt;
-    if (fraction < 1.0) {
-      allowed_memory_maximum = static_cast<size_t>(
-          fraction * static_cast<double>(device_prop.totalGlobalMem));
-    }
+    size_t device_free = 0;
+    size_t device_total = 0;
+    C10_CUDA_CHECK(cudaMemGetInfo(&device_free, &device_total));
+    allowed_memory_maximum =
+        static_cast<size_t>(fraction * static_cast<double>(device_total));
+    set_fraction = true;
   }
 
   /** get expandable segment size for all the streams on device **/
@@ -3014,7 +3010,7 @@ class DeviceCachingAllocator {
     BlockPool& pool = *p.pool;
 
     if (C10_UNLIKELY(
-            allowed_memory_maximum.has_value() &&
+            set_fraction &&
             AcceleratorAllocatorConfig::garbage_collection_threshold() > 0.0)) {
       // Track block reuse interval only when garbage collection is enabled.
       ++pool.get_free_blocks_call_count;
@@ -3087,7 +3083,7 @@ class DeviceCachingAllocator {
 
     size_t gc_threshold = static_cast<size_t>(
         AcceleratorAllocatorConfig::garbage_collection_threshold() *
-        static_cast<double>(allowed_memory_maximum.value()));
+        static_cast<double>(allowed_memory_maximum));
     // No need to trigger GC yet
     if (total_allocated_memory <= gc_threshold) {
       return;
@@ -3165,8 +3161,8 @@ class DeviceCachingAllocator {
 
     bool active_pool =
         p.pool->owner_PrivatePool && p.pool->owner_PrivatePool->allocator();
-    if (allowed_memory_maximum.has_value() &&
-        total_allocated_memory + size > allowed_memory_maximum.value()) {
+    if (set_fraction &&
+        total_allocated_memory + size > allowed_memory_maximum) {
       p.err = cudaErrorMemoryAllocation;
       return false;
       // Temporarily disable checkpointing & cudagraphs internally
@@ -3863,6 +3859,7 @@ class NativeCachingAllocator : public CUDAAllocator {
         "Allocator not initialized for device ",
         device,
         ": did you call init?");
+    C10_CUDA_CHECK(c10::cuda::SetDevice(device));
     return device_allocator[device]->getMemoryFraction();
   }
 
@@ -3872,6 +3869,12 @@ class NativeCachingAllocator : public CUDAAllocator {
         "Allocator not initialized for device ",
         device,
         ": did you call init?");
+    TORCH_CHECK(
+        0 <= fraction && fraction <= 1,
+        "invalid fraction:",
+        fraction,
+        ". Please set within [0, 1].");
+    C10_CUDA_CHECK(c10::cuda::SetDevice(device));
     device_allocator[device]->setMemoryFraction(fraction);
   }
 

c10/cuda/CUDACachingAllocator.h

@@ -2,7 +2,6 @@
 
 #include <c10/core/AllocatorConfig.h>
 #include <c10/core/CachingDeviceAllocator.h>
-#include <c10/cuda/CUDAAllocatorConfig.h>
 #include <c10/cuda/CUDAGraphsC10Utils.h>
 #include <c10/cuda/CUDAMacros.h>
 #include <c10/cuda/CUDAStream.h>

c10/cuda/CUDAMallocAsyncAllocator.cpp

@@ -427,6 +427,7 @@ struct CudaMallocAsyncAllocator : public CUDAAllocator {
   // on the current device each later call sees.
   void init(int dev_count) override {
     static bool called = [](int dev_count) {
+      ;
       // Are there external guarantees init will be called before
       // any of the allocator's other functions?
       // std::lock_guard<std::mutex> lk(general_mutex);

c10/util/ArrayRef.h

@@ -18,7 +18,6 @@
 #include <c10/macros/Macros.h>
 #include <c10/util/Exception.h>
 #include <c10/util/SmallVector.h>
-#include <torch/headeronly/util/HeaderOnlyArrayRef.h>
 
 #include <array>
 #include <cstddef>
@@ -41,99 +40,200 @@ namespace c10 {
 ///
 /// This is intended to be trivially copyable, so it should be passed by
 /// value.
-///
-/// NOTE: We have refactored out the headeronly parts of the ArrayRef struct
-/// into HeaderOnlyArrayRef. As adding `virtual` would change the performance of
-/// the underlying constexpr calls, we rely on apparent-type dispatch for
-/// inheritance. This should be fine because their memory format is the same,
-/// and it is never incorrect for ArrayRef to call HeaderOnlyArrayRef methods.
-/// However, you should prefer to use ArrayRef when possible, because its use
-/// of TORCH_CHECK will lead to better user-facing error messages.
 template <typename T>
-class ArrayRef final : public HeaderOnlyArrayRef<T> {
+class ArrayRef final {
  public:
-  /// @name Constructors, all inherited from HeaderOnlyArrayRef except for
-  /// SmallVector. As inherited constructors won't work with class template
-  /// argument deduction (CTAD) until C++23, we add deduction guides after
-  /// the class definition to enable CTAD.
+  using iterator = const T*;
+  using const_iterator = const T*;
+  using size_type = size_t;
+  using value_type = T;
+
+  using reverse_iterator = std::reverse_iterator<iterator>;
+
+ private:
+  /// The start of the array, in an external buffer.
+  const T* Data;
+
+  /// The number of elements.
+  size_type Length;
+
+  void debugCheckNullptrInvariant() {
+    TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
+        Data != nullptr || Length == 0,
+        "created ArrayRef with nullptr and non-zero length! std::optional relies on this being illegal");
+  }
+
+ public:
+  /// @name Constructors
   /// @{
 
-  using HeaderOnlyArrayRef<T>::HeaderOnlyArrayRef;
+  /// Construct an empty ArrayRef.
+  /* implicit */ constexpr ArrayRef() : Data(nullptr), Length(0) {}
+
+  /// Construct an ArrayRef from a single element.
+  // TODO Make this explicit
+  constexpr ArrayRef(const T& OneElt) : Data(&OneElt), Length(1) {}
+
+  /// Construct an ArrayRef from a pointer and length.
+  constexpr ArrayRef(const T* data, size_t length)
+      : Data(data), Length(length) {
+    debugCheckNullptrInvariant();
+  }
+
+  /// Construct an ArrayRef from a range.
+  constexpr ArrayRef(const T* begin, const T* end)
+      : Data(begin), Length(end - begin) {
+    debugCheckNullptrInvariant();
+  }
 
   /// Construct an ArrayRef from a SmallVector. This is templated in order to
   /// avoid instantiating SmallVectorTemplateCommon<T> whenever we
   /// copy-construct an ArrayRef.
-  /// NOTE: this is the only constructor that is not inherited from
-  /// HeaderOnlyArrayRef.
   template <typename U>
   /* implicit */ ArrayRef(const SmallVectorTemplateCommon<T, U>& Vec)
-      : HeaderOnlyArrayRef<T>(Vec.data(), Vec.size()) {}
+      : Data(Vec.data()), Length(Vec.size()) {
+    debugCheckNullptrInvariant();
+  }
+
+  template <
+      typename Container,
+      typename U = decltype(std::declval<Container>().data()),
+      typename = std::enable_if_t<
+          (std::is_same_v<U, T*> || std::is_same_v<U, T const*>)>>
+  /* implicit */ ArrayRef(const Container& container)
+      : Data(container.data()), Length(container.size()) {
+    debugCheckNullptrInvariant();
+  }
+
+  /// Construct an ArrayRef from a std::vector.
+  // The enable_if stuff here makes sure that this isn't used for
+  // std::vector<bool>, because ArrayRef can't work on a std::vector<bool>
+  // bitfield.
+  template <typename A>
+  /* implicit */ ArrayRef(const std::vector<T, A>& Vec)
+      : Data(Vec.data()), Length(Vec.size()) {
+    static_assert(
+        !std::is_same_v<T, bool>,
+        "ArrayRef<bool> cannot be constructed from a std::vector<bool> bitfield.");
+  }
+
+  /// Construct an ArrayRef from a std::array
+  template <size_t N>
+  /* implicit */ constexpr ArrayRef(const std::array<T, N>& Arr)
+      : Data(Arr.data()), Length(N) {}
+
+  /// Construct an ArrayRef from a C array.
+  template <size_t N>
+  // NOLINTNEXTLINE(*c-arrays*)
+  /* implicit */ constexpr ArrayRef(const T (&Arr)[N]) : Data(Arr), Length(N) {}
+
+  /// Construct an ArrayRef from a std::initializer_list.
+  /* implicit */ constexpr ArrayRef(const std::initializer_list<T>& Vec)
+      : Data(
+            std::begin(Vec) == std::end(Vec) ? static_cast<T*>(nullptr)
+                                             : std::begin(Vec)),
+        Length(Vec.size()) {}
 
   /// @}
-  /// @name Simple Operations, mostly inherited from HeaderOnlyArrayRef
+  /// @name Simple Operations
   /// @{
 
+  constexpr iterator begin() const {
+    return Data;
+  }
+  constexpr iterator end() const {
+    return Data + Length;
+  }
+
+  // These are actually the same as iterator, since ArrayRef only
+  // gives you const iterators.
+  constexpr const_iterator cbegin() const {
+    return Data;
+  }
+  constexpr const_iterator cend() const {
+    return Data + Length;
+  }
+
+  constexpr reverse_iterator rbegin() const {
+    return reverse_iterator(end());
+  }
+  constexpr reverse_iterator rend() const {
+    return reverse_iterator(begin());
+  }
+
+  /// Check if all elements in the array satisfy the given expression
+  constexpr bool allMatch(const std::function<bool(const T&)>& pred) const {
+    return std::all_of(cbegin(), cend(), pred);
+  }
+
+  /// empty - Check if the array is empty.
+  constexpr bool empty() const {
+    return Length == 0;
+  }
+
+  constexpr const T* data() const {
+    return Data;
+  }
+
+  /// size - Get the array size.
+  constexpr size_t size() const {
+    return Length;
+  }
+
   /// front - Get the first element.
-  /// We deviate from HeaderOnlyArrayRef by using TORCH_CHECK instead of
-  /// STD_TORCH_CHECK
   constexpr const T& front() const {
     TORCH_CHECK(
-        !this->empty(), "ArrayRef: attempted to access front() of empty list");
-    return this->Data[0];
+        !empty(), "ArrayRef: attempted to access front() of empty list");
+    return Data[0];
   }
 
   /// back - Get the last element.
-  /// We deviate from HeaderOnlyArrayRef by using TORCH_CHECK instead of
-  /// STD_TORCH_CHECK
   constexpr const T& back() const {
-    TORCH_CHECK(
-        !this->empty(), "ArrayRef: attempted to access back() of empty list");
-    return this->Data[this->Length - 1];
+    TORCH_CHECK(!empty(), "ArrayRef: attempted to access back() of empty list");
+    return Data[Length - 1];
+  }
+
+  /// equals - Check for element-wise equality.
+  constexpr bool equals(ArrayRef RHS) const {
+    return Length == RHS.Length && std::equal(begin(), end(), RHS.begin());
   }
 
   /// slice(n, m) - Take M elements of the array starting at element N
-  /// We deviate from HeaderOnlyArrayRef by using TORCH_CHECK instead of
-  /// STD_TORCH_CHECK
   constexpr ArrayRef<T> slice(size_t N, size_t M) const {
     TORCH_CHECK(
-        N + M <= this->size(),
+        N + M <= size(),
         "ArrayRef: invalid slice, N = ",
         N,
         "; M = ",
         M,
         "; size = ",
-        this->size());
-    return ArrayRef<T>(this->data() + N, M);
+        size());
+    return ArrayRef<T>(data() + N, M);
   }
 
   /// slice(n) - Chop off the first N elements of the array.
-  /// We deviate from HeaderOnlyArrayRef by using TORCH_CHECK instead of
-  /// STD_TORCH_CHECK
   constexpr ArrayRef<T> slice(size_t N) const {
     TORCH_CHECK(
-        N <= this->size(),
-        "ArrayRef: invalid slice, N = ",
-        N,
-        "; size = ",
-        this->size());
-    return slice(N, this->size() - N); // should this slice be this->slice?
+        N <= size(), "ArrayRef: invalid slice, N = ", N, "; size = ", size());
+    return slice(N, size() - N);
   }
 
   /// @}
   /// @name Operator Overloads
   /// @{
+  constexpr const T& operator[](size_t Index) const {
+    return Data[Index];
+  }
 
   /// Vector compatibility
-  /// We deviate from HeaderOnlyArrayRef by using TORCH_CHECK instead of
-  /// STD_TORCH_CHECK
   constexpr const T& at(size_t Index) const {
     TORCH_CHECK(
-        Index < this->Length,
+        Index < Length,
         "ArrayRef: invalid index Index = ",
         Index,
         "; Length = ",
-        this->Length);
-    return this->Data[Index];
+        Length);
+    return Data[Index];
   }
 
   /// Disallow accidental assignment from a temporary.
@@ -153,48 +253,16 @@ class ArrayRef final : public HeaderOnlyArrayRef<T> {
   std::enable_if_t<std::is_same_v<U, T>, ArrayRef<T>>& operator=(
       std::initializer_list<U>) = delete;
 
+  /// @}
+  /// @name Expensive Operations
+  /// @{
+  std::vector<T> vec() const {
+    return std::vector<T>(Data, Data + Length);
+  }
+
   /// @}
 };
 
-/// Deduction guides for ArrayRef to support CTAD with inherited constructors
-/// These mirror the constructors inherited from HeaderOnlyArrayRef
-/// @{
-
-// Single element constructor
-template <typename T>
-ArrayRef(const T&) -> ArrayRef<T>;
-
-// Pointer and length constructor
-template <typename T>
-ArrayRef(const T*, size_t) -> ArrayRef<T>;
-
-// Range constructor (begin, end)
-template <typename T>
-ArrayRef(const T*, const T*) -> ArrayRef<T>;
-
-// Generic container constructor (anything with .data() and .size())
-template <typename Container>
-ArrayRef(const Container&) -> ArrayRef<
-    std::remove_pointer_t<decltype(std::declval<Container>().data())>>;
-
-// std::vector constructor
-template <typename T, typename A>
-ArrayRef(const std::vector<T, A>&) -> ArrayRef<T>;
-
-// std::array constructor
-template <typename T, size_t N>
-ArrayRef(const std::array<T, N>&) -> ArrayRef<T>;
-
-// C array constructor
-template <typename T, size_t N>
-ArrayRef(const T (&)[N]) -> ArrayRef<T>;
-
-// std::initializer_list constructor
-template <typename T>
-ArrayRef(const std::initializer_list<T>&) -> ArrayRef<T>;
-
-/// @}
-
 template <typename T>
 std::ostream& operator<<(std::ostream& out, ArrayRef<T> list) {
   int i = 0;

caffe2/CMakeLists.txt

@@ -1307,7 +1307,7 @@ endif()
 
 if(USE_MKLDNN_ACL)
   find_package(ACL REQUIRED)
-  target_include_directories(torch_cpu SYSTEM PRIVATE ${ACL_INCLUDE_DIRS})
+  target_include_directories(torch_cpu PRIVATE ${ACL_INCLUDE_DIRS})
 endif()
 
 target_include_directories(torch_cpu PRIVATE ${ATen_CPU_INCLUDE})

caffe2/perfkernels/batch_box_cox_vec.h

@@ -73,19 +73,6 @@ void box_cox_zero_lambda(
   }
 }
 
-template <typename T>
-at::vec::Vectorized<T> box_cox_nonzero_lambda_impl(
-    at::vec::Vectorized<T> data,
-    at::vec::Vectorized<T> lambda1,
-    at::vec::Vectorized<T> lambda2,
-    at::vec::Vectorized<T> k_eps) {
-  auto sum = data + lambda2;
-  auto max = at::vec::max(sum, k_eps);
-  auto lambda_over_1 = at::vec::fast_recieprocal(lambda1);
-  auto pow = max.pow(lambda1);
-  return at::vec::fmsub(pow, lambda_over_1, lambda_over_1);
-}
-
 template <typename T>
 void box_cox_nonzero_lambda(
     int64_t D,
@@ -101,18 +88,21 @@ void box_cox_nonzero_lambda(
   auto k_eps_vec = Vec(k_eps);
   for(; j + VLEN < D; j += VLEN) {
     auto data = Vec::loadu(data_ptr + j);
-    auto lambda1 = Vec::loadu(lambda1_ptr + j);
     auto lambda2 = Vec::loadu(lambda2_ptr + j);
-    auto res = box_cox_nonzero_lambda_impl(data, lambda1, lambda2, k_eps_vec);
+    auto sum = data + lambda2;
+    auto max = at::vec::max(sum, k_eps_vec);
+    auto lambda1 = Vec::loadu(lambda1_ptr + j);
+    auto lambda_over_1 = at::vec::fast_recieprocal(lambda1);
+    auto pow = max.pow(lambda1);
+    auto res = at::vec::fmsub(pow, lambda_over_1, lambda_over_1);
     res.store(out + j);
   }
-  if (j < D) {
-    auto remaining = D - j;
-    auto data = Vec::loadu(data_ptr + j, remaining);
-    auto lambda1 = Vec::loadu(lambda1_ptr + j, remaining);
-    auto lambda2 = Vec::loadu(lambda2_ptr + j, remaining);
-    auto res = box_cox_nonzero_lambda_impl(data, lambda1, lambda2, k_eps_vec);
-    res.store(out + j, remaining);
+  for ( ;j < D; ++j) {
+    auto sum = data_ptr[j] + lambda2_ptr[j];
+    auto max = std::max(sum, k_eps);
+    auto lambda_over_1 = at::vec::fast_recieprocal(lambda1_ptr[j]);
+    auto pow = std::pow(max, lambda1_ptr[j]);
+    out[j] = pow * lambda_over_1 - lambda_over_1;
   }
 }
 #else

cmake/Modules/FindGloo.cmake

@@ -26,7 +26,7 @@ find_library(Gloo_CUDA_LIBRARY
 # if Gloo + HIP is desired, Gloo_HIP_LIBRARY
 # needs to be linked to desired target
 find_library(Gloo_HIP_LIBRARY
-  NAMES gloo_hip
+  NAMES gloo_hiop
   DOC "Gloo's HIP support/code"
 )
 

docs/source/conf.py

@@ -206,41 +206,6 @@ templates_path = [
     os.path.join(os.path.dirname(pytorch_sphinx_theme2.__file__), "templates"),
 ]
 # TODO: document these and remove them from here.
-# Fixes the duplicated
-autosummary_filename_map = {
-    "torch.nn.utils.prune.identity": "torch.nn.utils.prune.identity_function",
-    "torch.nn.utils.prune.Identity": "torch.nn.utils.prune.Identity_class",
-    "torch.optim.adamw.adamw": "torch.optim.adamw.adamw_function",
-    "torch.optim.adamw.AdamW": "torch.optim.adamw.AdamW_class",
-    "torch.optim.asgd.asgd": "torch.optim.asgd.asgd_function",
-    "torch.optim.asgd.ASGD": "torch.optim.asgd.ASGD_class",
-    "torch.optim.nadam.nadam": "torch.optim.nadam.nadam_function",
-    "torch.optim.nadam.NAdam": "torch.optim.nadam.NAdam_class",
-    "torch.optim.radam.radam": "torch.optim.radam.radam_function",
-    "torch.optim.radam.RAdam": "torch.optim.radam.RAdam_class",
-    "torch.optim.rmsprop.rmsprop": "torch.optim.rmsprop.rmsprop_function",
-    "torch.optim.rmsprop.RMSprop": "torch.optim.rmsprop.RMSprop_class",
-    "torch.optim.rprop.rprop": "torch.optim.rprop.rprop_function",
-    "torch.optim.rprop.Rprop": "torch.optim.rprop.Rprop_class",
-    "torch.optim.sgd.sgd": "torch.optim.sgd.sgd_function",
-    "torch.optim.sgd.SGD": "torch.optim.sgd.SGD_class",
-    "torch.optim.adadelta.adadelta": "torch.optim.adadelta.adadelta_function",
-    "torch.optim.adadelta.Adadelta": "torch.optim.adadelta.Adadelta_class",
-    "torch.optim.adagrad.adagrad": "torch.optim.adagrad.adagrad_function",
-    "torch.optim.adagrad.Adagrad": "torch.optim.adagrad.Adagrad_class",
-    "torch.optim.adam.adam": "torch.optim.adam.adam_function",
-    "torch.optim.adam.Adam": "torch.optim.adam.Adam_class",
-    "torch.optim.adamax.adamax": "torch.optim.adamax.adamax_function",
-    "torch.optim.adamax.Adamax": "torch.optim.adamax.Adamax_class",
-    "torch.mtia.stream": "torch.mtia.stream_function",
-    "torch.mtia.Stream": "torch.mtia.Stream_class",
-    "torch.cpu.stream": "torch.cpu.stream_function",
-    "torch.cpu.Stream": "torch.cpu.Stream_class",
-    "torch.cuda.stream": "torch.cuda.stream_function",
-    "torch.cuda.Stream": "torch.cuda.Stream_class",
-    "torch.xpu.stream": "torch.xpu.stream_function",
-    "torch.xpu.Stream": "torch.xpu.Stream_class",
-}
 
 coverage_ignore_functions = [
     # torch
@@ -3230,11 +3195,6 @@ autodoc_type_aliases = {
 # Enable overriding of function signatures in the first line of the docstring.
 autodoc_docstring_signature = True
 
-# Exclude inherited IntEnum methods that have RST formatting issues in their docstrings
-autodoc_default_options = {
-    "exclude-members": "from_bytes, to_bytes",
-}
-
 # -- katex javascript in header
 #
 #    def setup(app):

docs/source/notes/cuda.rst

@@ -619,10 +619,6 @@ Available options:
   and reallocate buffers across multiple streams, especially when the capture DAG frequently
   reaches joined frontiers.
 
-* ``per_process_memory_fraction`` option limits the amount of memory that can be allocated
-  on all the CUDA devices to a specified fraction of the available memory. This is a value
-  between 0 and 1. Attempting to allocate more memory will raise an out of memory error.
-
 .. note::
 
     Some stats reported by the

docs/source/notes/libtorch_stable_abi.md

@@ -46,108 +46,6 @@ These headers are promised to be ABI stable across releases and adhere to a stro
 Unless absolutely necessary, we recommend the high-level C++ API in `torch/csrc/stable`
 which will handle all the rough edges of the C API for the user.
 
-## Migrating your kernel to the LibTorch stable ABI
-
-If you'd like your kernel to be ABI stable with LibTorch, meaning you'd the ability to build for one version and run on another, your kernel must only use the limited stable ABI. This following section goes through some steps of migrating an existing kernel and APIs we imagine you would need to swap over.
-
-Firstly, instead of registering kernels through `TORCH_LIBRARY`, LibTorch ABI stable kernels must be registered via `STABLE_TORCH_LIBRARY`. Note that, for the time being, implementations registered via `STABLE_TORCH_LIBRARY` must be boxed unlike `TORCH_LIBRARY`. See the simple example below or our docs on [Stack-based APIs](stack-based-apis) for more details. For kernels that are registered via `pybind`, before using the stable ABI, it would be useful to migrate to register them via `TORCH_LIBRARY`.
-
-While previously your kernels might have included APIs from `<torch/*.h>` (for example, `<torch/all.h>`), they are now limited to including from the 3 categories of headers mentioned above (`torch/csrc/stable/*.h`, `torch/headeronly/*.h` and the stable C headers). This means that your extension should no longer use any utilities from the `at::` or `c10::` namespaces but instead use their replacements in `torch::stable` and `torch::headeronly`. To provide a couple examples of the necessary migrations:
-- all uses of `at::Tensor` must be replaced with `torch::stable::Tensor`
-- all uses of `TORCH_CHECK` must be replaced with `STD_TORCH_CHECK`
-- all uses of `at::kCUDA` must be replaced with `torch::headeronly::kCUDA` etc.
-- native functions such as `at::pad` must be replaced with `torch::stable::pad`
-- native functions that are called as Tensor methods (e.g., `Tensor.pad`) must be replaced with the ATen variant through `torch::stable::pad`.
-
-As mentioned above, the LibTorch stable ABI is still under development. If there is any API or feature you would like to see added to the stable ABI/`torch::headeronly`/`torch::stable`, please file a request through a [new issue on the PyTorch repo](https://github.com/pytorch/pytorch/issues).
-
-Below is a simple example of migrating an existing kernel that uses `TORCH_LIBRARY` to the stable ABI (`TORCH_STABLE_LIBRARY`). For a larger end to end example you can take a look at the FA3 repository. Specifically the diff between [`flash_api.cpp`](https://github.com/Dao-AILab/flash-attention/blob/ad70a007e6287d4f7e766f94bcf2f9a813f20f6b/hopper/flash_api.cpp#L1) and the stable variant [`flash_api_stable.cpp`](https://github.com/Dao-AILab/flash-attention/blob/ad70a007e6287d4f7e766f94bcf2f9a813f20f6b/hopper/flash_api_stable.cpp#L1).
-
-
-### Original Version with `TORCH_LIBRARY`
-
-```cpp
-// original_kernel.cpp - Using TORCH_LIBRARY (not stable ABI)
-#include <torch/torch.h>
-#include <ATen/ATen.h>
-
-namespace myops {
-
-// Simple kernel that adds a scalar value to each element of a tensor
-at::Tensor add_scalar(const at::Tensor& input, double scalar) {
-  TORCH_CHECK(input.scalar_type() == at::kFloat, "Input must be float32");
-
-  return input.add(scalar);
-}
-
-// Register the operator
-TORCH_LIBRARY(myops, m) {
-  m.def("add_scalar(Tensor input, float scalar) -> Tensor", &add_scalar);
-}
-
-// Register the implementation
-TORCH_LIBRARY_IMPL(myops, CompositeExplicitAutograd, m) {
-  m.impl("add_scalar", &add_scalar);
-}
-
-} // namespace myops
-```
-
-### Migrated Version with `STABLE_TORCH_LIBRARY`
-
-```cpp
-// stable_kernel.cpp - Using STABLE_TORCH_LIBRARY (stable ABI)
-
-// (1) Don't include <torch/torch.h> <ATen/ATen.h>
-//     only include APIs from torch/csrc/stable, torch/headeronly and C-shims
-#include <torch/csrc/stable/library.h>
-#include <torch/csrc/stable/tensor_struct.h>
-#include <torch/csrc/stable/ops.h>
-#include <torch/csrc/stable/stableivalue_conversions.h>
-#include <torch/headeronly/core/ScalarType.h>
-#include <torch/headeronly/macros/Macros.h>
-
-namespace myops {
-
-// Simple kernel that adds a scalar value to each element of a tensor
-torch::stable::Tensor add_scalar(const torch::stable::Tensor& input, double scalar) {
-  // (2) use STD_TORCH_CHECK instead of TORCH_CHECK
-  STD_TORCH_CHECK(
-      // (3) use torch::headeronly::kFloat instead of at:kFloat
-      input.scalar_type() == torch::headeronly::kFloat,
-      "Input must be float32");
-
-  // (4) Use stable ops namespace instead of input.add
-  return torch::stable::add(input, scalar);
-}
-
-// (5) Add Boxed wrapper required for STABLE_TORCH_LIBRARY
-void boxed_add_scalar(StableIValue* stack, uint64_t num_args, uint64_t num_outputs) {
-  // Extract arguments from stack using `to<T>`
-  auto input = to<torch::stable::Tensor>(stack[0]);
-  auto scalar = to<double>(stack[1]);
-
-  // Call the actual kernel
-  auto result = add_scalar(input, scalar);
-
-  // Put result back on stack using `from()`
-  // Stack slot 0 now holds the return value
-  stack[0] = from(result);
-}
-
-// (6) Register the operator using STABLE_TORCH_LIBRARY
-STABLE_TORCH_LIBRARY(myops, m) {
-  m.def("add_scalar(Tensor input, float scalar) -> Tensor", &boxed_add_scalar);
-}
-
-// (7) Register the implementation using STABLE_TORCH_LIBRARY_IMPL
-STABLE_TORCH_LIBRARY_IMPL(myops, CompositeExplicitAutograd, m) {
-  m.impl("add_scalar", &boxed_add_scalar);
-}
-
-} // namespace myops
-```
-
 
 ## How are objects passed across the ABI boundary when interacting with the dispatcher?
 
@@ -211,7 +109,6 @@ There are two invariants for the stack:
     a. When calling a stack-based API, you must give owning references to the calling stack and steal references from the returned stack.
     b. When registering your function to be called with a stack, you must steal references from your argument stack and push onto the stack new references.
 
-(stack-based-apis)=
 ### Stack-based APIs
 
 The above is relevant in two places:

docs/source/quantization-support.md

@@ -253,6 +253,7 @@ regular full-precision tensor.
 .. autosummary::
     :toctree: generated
     :nosignatures:
+    :template: classtemplate.rst
 
     view
     as_strided

setup.py

@@ -630,37 +630,6 @@ def mirror_files_into_torchgen() -> None:
         raise RuntimeError("Check the file paths in `mirror_files_into_torchgen()`")
 
 
-def mirror_inductor_external_kernels() -> None:
-    """
-    Copy external kernels into Inductor so they are importable.
-    """
-    paths = [
-        (
-            CWD / "torch/_inductor/kernel/vendored_templates/cutedsl_grouped_gemm.py",
-            CWD
-            / "third_party/cutlass/examples/python/CuTeDSL/blackwell/grouped_gemm.py",
-        ),
-    ]
-    for new_path, orig_path in paths:
-        # Create the dirs involved in new_path if they don't exist
-        if not new_path.exists():
-            new_path.parent.mkdir(parents=True, exist_ok=True)
-
-        # Copy the files from the orig location to the new location
-        if orig_path.is_file():
-            shutil.copyfile(orig_path, new_path)
-            continue
-        if orig_path.is_dir():
-            if new_path.exists():
-                # copytree fails if the tree exists already, so remove it.
-                shutil.rmtree(new_path)
-            shutil.copytree(orig_path, new_path)
-            continue
-        raise RuntimeError(
-            "Check the file paths in `mirror_inductor_external_kernels()`"
-        )
-
-
 # ATTENTION: THIS IS AI SLOP
 def extract_variant_from_version(version: str) -> str:
     """Extract variant from version string, defaulting to 'cpu'."""
@@ -1647,8 +1616,6 @@ def main() -> None:
     if RUN_BUILD_DEPS:
         build_deps()
 
-    mirror_inductor_external_kernels()
-
     (
         ext_modules,
         cmdclass,
@@ -1682,7 +1649,6 @@ def main() -> None:
         "_inductor/codegen/aoti_runtime/*.cpp",
         "_inductor/script.ld",
         "_inductor/kernel/flex/templates/*.jinja",
-        "_inductor/kernel/templates/*.jinja",
         "_export/serde/*.yaml",
         "_export/serde/*.thrift",
         "share/cmake/ATen/*.cmake",

test/ao/sparsity/test_scheduler.py

@@ -75,7 +75,6 @@ class TestScheduler(TestCase):
 
 class TestCubicScheduler(TestCase):
     def setUp(self):
-        super().setUp()
         self.model_sparse_config = [
             {"tensor_fqn": "0.weight", "sparsity_level": 0.8},
             {"tensor_fqn": "2.weight", "sparsity_level": 0.4},

test/backends/xeon/test_launch.py

@@ -11,7 +11,6 @@ from torch.testing._internal.common_utils import IS_LINUX, run_tests, TestCase
 @unittest.skipIf(not IS_LINUX, "Only works on linux")
 class TestTorchrun(TestCase):
     def setUp(self):
-        super().setUp()
         self._test_dir = tempfile.mkdtemp(prefix=self.__class__.__name__)
 
     def tearDown(self):

test/cpp/aoti_abi_check/CMakeLists.txt

@@ -12,7 +12,6 @@ set(AOTI_ABI_CHECK_TEST_SRCS
   ${AOTI_ABI_CHECK_TEST_ROOT}/test_devicetype.cpp
   ${AOTI_ABI_CHECK_TEST_ROOT}/test_dtype.cpp
   ${AOTI_ABI_CHECK_TEST_ROOT}/test_exception.cpp
-  ${AOTI_ABI_CHECK_TEST_ROOT}/test_headeronlyarrayref.cpp
   ${AOTI_ABI_CHECK_TEST_ROOT}/test_macros.cpp
   ${AOTI_ABI_CHECK_TEST_ROOT}/test_math.cpp
   ${AOTI_ABI_CHECK_TEST_ROOT}/test_rand.cpp
@@ -45,10 +44,6 @@ endif()
 # Disable unused-variable warnings for variables that are only used to test compilation
 target_compile_options_if_supported(test_aoti_abi_check -Wno-unused-variable)
 target_compile_options_if_supported(test_aoti_abi_check -Wno-unused-but-set-variable)
-# Add -Wno-dangling-pointer for GCC 13
-if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 13)
-  target_compile_options_if_supported(test_aoti_abi_check -Wno-dangling-pointer)
-endif()
 
 foreach(test_src ${AOTI_ABI_CHECK_VEC_TEST_SRCS})
   foreach(i RANGE ${NUM_CPU_CAPABILITY_NAMES})

test/cpp/aoti_abi_check/test_headeronlyarrayref.cpp

@@ -1,52 +0,0 @@
-#include <gtest/gtest.h>
-
-#include <torch/headeronly/util/HeaderOnlyArrayRef.h>
-
-#include <vector>
-
-using torch::headeronly::HeaderOnlyArrayRef;
-
-TEST(TestHeaderOnlyArrayRef, TestEmpty) {
-  HeaderOnlyArrayRef<float> arr;
-  ASSERT_TRUE(arr.empty());
-}
-
-TEST(TestHeaderOnlyArrayRef, TestSingleton) {
-  float val = 5.0f;
-  HeaderOnlyArrayRef<float> arr(val);
-  ASSERT_FALSE(arr.empty());
-  EXPECT_EQ(arr.size(), 1);
-  EXPECT_EQ(arr[0], val);
-}
-
-TEST(TestHeaderOnlyArrayRef, TestAPIs) {
-  std::vector<int> vec = {1, 2, 3, 4, 5, 6, 7};
-  HeaderOnlyArrayRef<int> arr(vec);
-  ASSERT_FALSE(arr.empty());
-  EXPECT_EQ(arr.size(), 7);
-  for (size_t i = 0; i < arr.size(); i++) {
-    EXPECT_EQ(arr[i], i + 1);
-    EXPECT_EQ(arr.at(i), i + 1);
-  }
-  EXPECT_EQ(arr.front(), 1);
-  EXPECT_EQ(arr.back(), 7);
-  ASSERT_TRUE(arr.slice(3, 4).equals(arr.slice(3)));
-}
-
-TEST(TestHeaderOnlyArrayRef, TestFromInitializerList) {
-  std::vector<int> vec = {1, 2, 3, 4, 5, 6, 7};
-  HeaderOnlyArrayRef<int> arr({1, 2, 3, 4, 5, 6, 7});
-  auto res_vec = arr.vec();
-  for (size_t i = 0; i < vec.size(); i++) {
-    EXPECT_EQ(vec[i], res_vec[i]);
-  }
-}
-
-TEST(TestHeaderOnlyArrayRef, TestFromRange) {
-  std::vector<int> vec = {1, 2, 3, 4, 5, 6, 7};
-  HeaderOnlyArrayRef<int> arr(vec.data() + 3, vec.data() + 7);
-  auto res_vec = arr.vec();
-  for (size_t i = 0; i < res_vec.size(); i++) {
-    EXPECT_EQ(vec[i + 3], res_vec[i]);
-  }
-}

test/cpp/api/CMakeLists.txt

@@ -70,13 +70,6 @@ if(NOT MSVC)
   if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12)
     target_compile_options_if_supported(test_api "-Wno-error=nonnull")
   endif()
-
-  # Add -Wno-error=array-bounds for GCC 13+
-  # See: https://gcc.gnu.org/bugzilla/show_bug.cgi?id=113239
-  if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 13)
-    target_compile_options_if_supported(test_api "-Wno-error=array-bounds")
-  endif()
-
 endif()
 
 if(INSTALL_TEST)

test/cpp/api/init_baseline.py

@@ -64,7 +64,7 @@ def run(initializer):
 
 def main():
     initializer_parameter_map = {}
-    for initializer in INITIALIZERS:
+    for initializer in INITIALIZERS.keys():
         sys.stderr.write(f"Evaluating {initializer} ...\n")
         initializer_parameter_map[initializer] = run(initializer)
 

test/cpp/api/optim_baseline.py

@@ -130,7 +130,7 @@ def main():
     options = parser.parse_args()
 
     optimizer_parameter_map = {}
-    for optimizer in OPTIMIZERS:
+    for optimizer in OPTIMIZERS.keys():
         sys.stderr.write(f"Evaluating {optimizer} ...\n")
         optimizer_parameter_map[optimizer] = run(
             optimizer, options.iterations, options.sample_every

test/cpp_extensions/libtorch_agnostic_extension/libtorch_agnostic/csrc/kernel.cpp

@@ -47,10 +47,20 @@ Tensor sgd_out_of_place(
   STD_TORCH_CHECK(param.get_device() == -1, "CPU device index = -1");
   STD_TORCH_CHECK(param.get_device_index() == -1, "CPU device index = -1");
 
-  // testing Tensor strides + stride
-  STD_TORCH_CHECK(param.strides()[0] == param.stride(0));
+  int64_t *param_sizes;
+  int64_t *param_strides;
+  aoti_torch_get_sizes(param.get(), &param_sizes);
+  aoti_torch_get_strides(param.get(), &param_strides);
 
-  auto out = new_empty(param, param.sizes());
+  int32_t param_dtype;
+  aoti_torch_get_dtype(param.get(), &param_dtype);
+
+  int32_t param_device_type;
+  aoti_torch_get_device_type(param.get(), &param_device_type);
+
+  AtenTensorHandle out_ath;
+  aoti_torch_empty_strided(param.dim(), param_sizes, param_strides, param_dtype, param_device_type, param.get_device(), &out_ath);
+  auto out = Tensor(out_ath);
 
   sgd_math(
     reinterpret_cast<float*>(param.data_ptr()),
@@ -301,9 +311,10 @@ void boxed_fill_infinity(
 }
 
 Tensor my_pad(Tensor t) {
+  std::vector<int64_t> padding = {1, 2, 2, 1};
   std::string mode = "constant";
   double value = 0.0;
-  return pad(t, {1, 2, 2, 1}, mode, value);
+  return pad(t, padding, mode, value);
 }
 
 void boxed_my_pad(
@@ -331,11 +342,6 @@ void boxed_my_narrow(
 }
 
 Tensor my_new_empty_dtype_variant(Tensor t) {
-  // Still using a std::vector below even though people can just pass in an
-  // initializer list (which will be implicitly converted to an HeaderOnlyArrayRef)
-  // directly.
-  // This is to test that passing in a std::vector works for BC. (It gets
-  // implicitly converted to HeaderOnlyArrayRef too!)
   std::vector<int64_t> sizes = {2, 5};
   auto dtype = std::make_optional(torch::headeronly::ScalarType::BFloat16);
   return new_empty(t, sizes, dtype);
@@ -347,8 +353,9 @@ void boxed_my_new_empty_dtype_variant(StableIValue* stack, uint64_t num_args, ui
 }
 
 Tensor my_new_zeros_dtype_variant(Tensor t) {
+  std::vector<int64_t> sizes = {2, 5};
   auto dtype = std::make_optional(at::ScalarType::Float);
-  return new_zeros(t, {2, 5}, dtype);
+  return new_zeros(t, sizes, dtype);
 }
 
 void boxed_my_new_zeros_dtype_variant(StableIValue* stack, uint64_t num_args, uint64_t num_outputs) {
@@ -422,7 +429,8 @@ void boxed_my_amax(StableIValue* stack, uint64_t num_args, uint64_t num_outputs)
 }
 
 Tensor my_amax_vec(Tensor t) {
-  return amax(t, {0,1}, false);
+  std::vector<int64_t> v = {0,1};
+  return amax(t, v, false);
 }
 
 void boxed_my_amax_vec(StableIValue* stack, uint64_t num_args, uint64_t num_outputs) {

test/custom_backend/test_custom_backend.py

@@ -11,7 +11,6 @@ from torch.testing._internal.common_utils import run_tests, TestCase
 
 class TestCustomBackend(TestCase):
     def setUp(self):
-        super().setUp()
         # Load the library containing the custom backend.
         self.library_path = get_custom_backend_library_path()
         torch.ops.load_library(self.library_path)

test/custom_operator/test_custom_ops.py

@@ -18,7 +18,6 @@ torch.ops.import_module("pointwise")
 
 class TestCustomOperators(TestCase):
     def setUp(self):
-        super().setUp()
         self.library_path = get_custom_op_library_path()
         ops.load_library(self.library_path)
 

test/distributed/checkpoint/_experimental/test_builder.py

@@ -22,7 +22,6 @@ from torch.testing._internal.common_utils import run_tests, TestCase
 
 class TestMakeCheckpointer(TestCase):
     def setUp(self) -> None:
-        super().setUp()
         # Create a temporary directory for checkpoints
         self.temp_dir = tempfile.mkdtemp()
 

test/distributed/checkpoint/_experimental/test_checkpoint_process.py

@@ -161,7 +161,6 @@ class TestCheckpointProcessConfig(TestCase):
 
 class TestCheckpointProcess(TestCase):
     def setUp(self) -> None:
-        super().setUp()
         """Set up common test fixtures."""
         self.rank_info = RankInfo(
             global_world_size=1,

test/distributed/checkpoint/_experimental/test_checkpoint_reader.py

@@ -14,7 +14,6 @@ from torch.testing._internal.common_utils import run_tests, TestCase
 
 class TestCheckpointReader(TestCase):
     def setUp(self):
-        super().setUp()
         # Create a temporary directory for test checkpoints
         self.temp_dir = tempfile.mkdtemp()
 

test/distributed/checkpoint/_experimental/test_checkpoint_writer.py

@@ -52,7 +52,6 @@ class TestCheckpointWriterConfig(TestCase):
 
 class TestCheckpointWriter(TestCase):
     def setUp(self):
-        super().setUp()
         # Create a temporary directory for test checkpoints
         self.temp_dir = tempfile.mkdtemp()
 

test/distributed/checkpoint/_experimental/test_checkpointer.py

@@ -52,7 +52,6 @@ class TestCheckpointer(TestCase):
     """Parameterized tests that work with both sync and async checkpointers."""
 
     def setUp(self):
-        super().setUp()
         # Create a temporary directory for checkpoints
         self.temp_dir = tempfile.mkdtemp()
 
@@ -398,7 +397,6 @@ class TestAsyncCheckpointerSpecific(TestCase):
     """Tests specific to AsyncCheckpointer functionality."""
 
     def setUp(self):
-        super().setUp()
         # Create a temporary directory for checkpoints
         self.temp_dir = tempfile.mkdtemp()
 

test/distributed/checkpoint/_experimental/test_staging.py

@@ -12,7 +12,6 @@ from torch.testing._internal.common_utils import requires_cuda, run_tests, TestC
 
 class TestDefaultStager(TestCase):
     def setUp(self) -> None:
-        super().setUp()
         # Create a test state dictionary with various data types
         self.state_dict = {
             "model": torch.nn.Linear(10, 5).state_dict(),

test/distributed/checkpoint/test_hf_safetensor_e2e.py

@@ -208,7 +208,7 @@ class TestSingleRankSaveLoad(TestCase):
 
         # Create model.safetensors.index.json with weight mapping
         weight_map = {}
-        for key in quantized_checkpoint:
+        for key in quantized_checkpoint.keys():
             weight_map[key] = "model.safetensors"
 
         index_data = {
@@ -245,7 +245,7 @@ class TestSingleRankSaveLoad(TestCase):
             sorted(original_tensors.keys()), sorted(state_dict_to_load.keys())
         )
 
-        for tensor_name in original_tensors:
+        for tensor_name in original_tensors.keys():
             original = original_tensors[tensor_name]
             loaded = state_dict_to_load[tensor_name]
 

test/distributed/checkpoint/test_quantized_hf_storage.py

@@ -15,7 +15,6 @@ from torch.testing._internal.common_utils import run_tests, TestCase
 
 class TestQuantizedHfStorage(TestCase):
     def setUp(self):
-        super().setUp()
         """Set up common test fixtures."""
         self.temp_dir = tempfile.TemporaryDirectory()
         self.path = self.temp_dir.name

test/distributed/elastic/multiprocessing/test_api.py

@@ -21,7 +21,6 @@ from torch.testing._internal.common_utils import run_tests, TestCase
 
 class SignalHandlingTest(TestCase):
     def setUp(self):
-        super().setUp()
         # Save original environment variable if it exists
         self.original_signals_env = os.environ.get(
             "TORCHELASTIC_SIGNALS_TO_HANDLE", None

test/distributed/fsdp/test_fsdp_mixed_precision.py

@@ -498,7 +498,7 @@ class TestFSDPMixedPrecision(FSDPTest):
                     for name, tensor in state_dict.items():
                         # Parameters and buffers are checkpointed in their
                         # original dtypes, which may be different.
-                        if name in named_buffers:
+                        if name in named_buffers.keys():
                             self.assertEqual(tensor.dtype, _BUFFER_ORIG_DTYPE)
                         else:
                             self.assertEqual(

test/distributed/launcher/test_api.py

@@ -16,7 +16,6 @@ from torch.testing._internal.common_utils import run_tests, TestCase
 
 class LauncherApiTest(TestCase):
     def setUp(self):
-        super().setUp()
         # Save original environment variable if it exists
         self.original_signals_env = os.environ.get(
             "TORCHELASTIC_SIGNALS_TO_HANDLE", None

test/distributed/pipelining/test_schedule.py

@@ -21,7 +21,6 @@ from torch.distributed.pipelining import (
 from torch.distributed.pipelining._utils import generate_stage_to_rank_mapping
 from torch.distributed.pipelining.schedules import (
     _Action,
-    _add_reduce_grad,
     _add_send_recv,
     _add_unshard_reshard,
     _format_pipeline_order,
@@ -575,45 +574,6 @@ class TestScheduleLowering(TestCase):
                 ),
             )
 
-    @parametrize(
-        "test_info",
-        [
-            {
-                "compute": ["0F0", "0F1", "   ", "0B0", "0B1"],
-                "comms": ["0F0", "0F1", "0B0", "0B1", "0REDUCE_GRAD"],
-            },
-            {
-                "compute": ["0F0", "0F1", "1F0", "1F1", "1B0", "1B1", "0B0", "0B1"],
-                "comms": [
-                    "0F0",
-                    "0F1",
-                    "1F0",
-                    "1F1",
-                    "1B0",
-                    "1B1",
-                    "1REDUCE_GRAD",
-                    "0B0",
-                    "0B1",
-                    "0REDUCE_GRAD",
-                ],
-            },
-        ],
-    )
-    def test_reduce_grad(self, test_info):
-        compute_sch = self._parse_actions(test_info["compute"])
-        expected_comms_sch = self._parse_actions(test_info["comms"])
-
-        comms_sch = _add_reduce_grad(compute_sch, 2)
-        for expected, actual in zip(expected_comms_sch, comms_sch, strict=True):
-            self.assertEqual(
-                expected,
-                actual,
-                (
-                    f"Mismatch: expected action {expected} but found {actual}."
-                    f"\nWhole Schedule: {comms_sch}"
-                ),
-            )
-
     @parametrize(
         "test_info",
         [

test/distributed/tensor/debug/test_debug_mode.py

@@ -5,16 +5,8 @@ import contextlib
 import torch
 import torch.distributed as dist
 from torch._subclasses.fake_tensor import FakeTensorMode
-from torch.distributed.tensor import (
-    DeviceMesh,
-    distribute_tensor,
-    DTensor,
-    Partial,
-    Replicate,
-    Shard,
-)
+from torch.distributed.tensor import DeviceMesh, DTensor, Partial, Replicate, Shard
 from torch.distributed.tensor._dtensor_spec import ShardOrderEntry
-from torch.fx.experimental.proxy_tensor import make_fx
 from torch.testing._internal.common_utils import (
     instantiate_parametrized_tests,
     parametrize,
@@ -434,31 +426,6 @@ class TestDTensorDebugMode(TestCase):
         ][-1]
         self.assertTrue("self.l2(self.l1(x))" in sum_op.fwd_stack_trace)
 
-    def test_pretty_print_dtensor_make_fx(self):
-        mesh = DeviceMesh(self.device_type, list(range(self.world_size)))
-
-        A = torch.randn(8, 32)
-        B = torch.randn(32, 32)
-        dA = distribute_tensor(A, mesh, [Shard(0)]).requires_grad_()
-        dB = distribute_tensor(B, mesh, [Replicate()]).requires_grad_()
-
-        def f(dA, dB):
-            dy = dA @ dB
-            loss = dy.sum()
-            loss.backward()
-            return dA.grad, dB.grad
-
-        # We actually need the tracing_mode='fake' here, or to trace under a FakeTensorMode.
-        # make_fx has some logic to ensure we don't accidentally stash real tensors in the graph
-        # so we won't stash our DTensors properly if they don't hold Fake inner tensors
-        gm = make_fx(f, tracing_mode="fake")(dA, dB)
-        # DCE isn't necessary here, there were just a lot of dead detach() nodes that spammed the graph
-        gm.graph.eliminate_dead_code()
-        gm.recompile()
-        # Colored is nice for actual viewing, not using in this test though
-        gm_str = gm.print_readable(colored=False, print_output=False)
-        self.assertTrue('"DTensor(f32[8, 32], S(0))" = torch.ops.aten.mm' in gm_str)
-
 
 instantiate_parametrized_tests(TestDTensorDebugMode)
 

test/distributed/tensor/test_utils.py

@@ -1,18 +1,11 @@
 # Owner(s): ["oncall: distributed"]
 
 import itertools
-from contextlib import nullcontext
 from typing import Any
 
 import torch
-import torch.distributed as dist
-from torch.distributed._local_tensor import (
-    local_tensor_mode,
-    LocalTensor,
-    LocalTensorMode,
-)
 from torch.distributed.device_mesh import init_device_mesh
-from torch.distributed.tensor import DeviceMesh, distribute_tensor, DTensor
+from torch.distributed.tensor import distribute_tensor, DTensor
 from torch.distributed.tensor._dtensor_spec import DTensorSpec, TensorMeta
 from torch.distributed.tensor._utils import (
     _compute_local_shape_and_global_offset,
@@ -21,7 +14,6 @@ from torch.distributed.tensor._utils import (
     compute_global_tensor_shape,
     compute_local_shape_and_global_offset,
     compute_local_tensor_info,
-    ExplicitRedistributionContext,
 )
 from torch.distributed.tensor.debug import CommDebugMode
 from torch.distributed.tensor.placement_types import (
@@ -859,93 +851,5 @@ class Test2DStridedLocalShard(DTensorTestBase):
         self.assertEqual(global_tensor, dtensor_2d.full_tensor())
 
 
-class LocalTensorTestBase(TestCase):
-    def assertEqual(self, lhs, rhs, **kwargs):
-        mode = local_tensor_mode()
-        with nullcontext() if mode is None else mode.disable():
-            if isinstance(lhs, LocalTensor) and isinstance(rhs, LocalTensor):
-                assert isinstance(lhs, LocalTensor) and isinstance(rhs, LocalTensor)
-                super().assertEqual(lhs._ranks, rhs._ranks)
-                for r in lhs._ranks:
-                    super().assertEqual(
-                        lhs._local_tensors[r],
-                        rhs._local_tensors[r],
-                        lambda m: f"rank {r}: {m}",
-                    )
-            elif isinstance(lhs, LocalTensor) or isinstance(rhs, LocalTensor):
-                lhs, rhs = (lhs, rhs) if isinstance(lhs, LocalTensor) else (rhs, lhs)
-                for r in lhs._ranks:
-                    super().assertEqual(
-                        lhs._local_tensors[r], rhs, lambda m: f"rank {r}: {m}"
-                    )
-            else:
-                return super().assertEqual(lhs, rhs, **kwargs)
-
-    @property
-    def world_size(self):
-        raise NotImplementedError("override world-size in your subclass")
-
-    def build_device_mesh(self) -> DeviceMesh:
-        return init_device_mesh("cpu", (self.world_size,))
-
-    def setUp(self):
-        super().setUp()
-        torch.distributed.init_process_group(
-            # TODO: test other ranks too
-            "fake",
-            rank=0,
-            world_size=self.world_size,
-        )
-
-    def tearDown(self):
-        super().tearDown()
-        try:
-            dist.destroy_process_group()
-        except AssertionError:
-            pass
-
-
-class TestExplicitRedistribute(LocalTensorTestBase):
-    @property
-    def world_size(self):
-        return 4
-
-    def test_explicit_matmul(self):
-        with LocalTensorMode(self.world_size):
-            device_mesh = self.build_device_mesh()
-            dim = 128
-            x = torch.randn(8, dim, requires_grad=True)
-            A = torch.randn(dim, dim, requires_grad=True)
-
-            # Prepare DTensors
-            dx = distribute_tensor(x, device_mesh, [Shard(0)])
-            dA = distribute_tensor(A, device_mesh, [Shard(0)])
-
-            # implicit redistribute works as usual by default
-            with CommDebugMode() as comm_mode:
-                torch.matmul(dx, dA)
-            self.assertEqual(comm_mode.get_total_counts(), 1)
-
-            # explicit redistribute works too
-            with ExplicitRedistributionContext():
-                with self.assertRaisesRegex(RuntimeError, "Implicit redistribution"):
-                    torch.matmul(dx, dA)
-
-            # explicit redistribute allows manual redistribute
-            with ExplicitRedistributionContext():
-                dA_repl = dA.redistribute(device_mesh, [Replicate()])
-                torch.matmul(dx, dA_repl)
-
-            dx = distribute_tensor(x, device_mesh, [Shard(0)])
-            dA = distribute_tensor(A, device_mesh, [Replicate()])
-            with ExplicitRedistributionContext():
-                dY = torch.matmul(dx, dA_repl)
-                loss = dY.sum()
-
-                # we now see the error during backwards
-                with self.assertRaisesRegex(RuntimeError, "Implicit redistribution"):
-                    loss.backward()
-
-
 if __name__ == "__main__":
     run_tests()

test/distributed/test_c10d_common.py

@@ -1189,7 +1189,9 @@ class AbstractCommTest:
                 self.assertEqual(len(set(rank_to_seq_num.values())), 2)
                 self.assertEqual(rank_to_seq_num[0], rank_to_seq_num[2])
                 expected_same = {
-                    rank_to_seq_num[i] for i in rank_to_seq_num if i not in [0, 2]
+                    rank_to_seq_num[i]
+                    for i in rank_to_seq_num.keys()
+                    if i not in [0, 2]
                 }
                 self.assertEqual(len(expected_same), 1)
                 self.assertEqual(rank_to_seq_num[0] + 1, rank_to_seq_num[1])
@@ -1556,7 +1558,7 @@ class CommTest(AbstractCommTest, MultiProcessTestCase):
         }
         invalid_debug_modes = ["foo", 0, 1, -1]
 
-        for mode in mapping:
+        for mode in mapping.keys():
             os.environ["TORCH_DISTRIBUTED_DEBUG"] = str(mode)
             dist.set_debug_level_from_env()
             set_debug_mode = dist.get_debug_level()

test/distributed/test_c10d_gloo.py

@@ -2357,7 +2357,6 @@ class ReducerModule(nn.Module):
 
 class ReducerTest(TestCase):
     def setUp(self):
-        super().setUp()
         self.file = tempfile.NamedTemporaryFile(delete=False)
         world_size = 1
         self.store = c10d.FileStore(self.file.name, world_size)

test/distributed/test_c10d_nccl.py

@@ -252,7 +252,6 @@ class ProcessGroupNCCLNoGPUTest(TestCase):
     MAIN_PROCESS_RANK = 0
 
     def setUp(self):
-        super().setUp()
         self.rank = self.MAIN_PROCESS_RANK
         self.world_size = 1
         self.file = tempfile.NamedTemporaryFile(delete=False)
@@ -5790,229 +5789,6 @@ class NCCLTraceTest(NCCLTraceTestBase):
         else:
             self.assertTrue("duration_ms" not in t["entries"][0])
 
-    @requires_nccl()
-    @skip_but_pass_in_sandcastle_if(not TEST_MULTIGPU, "NCCL test requires 2+ GPUs")
-    @parametrize("timing_enabled", [True, False])
-    def test_fr_record_reset_circular_buffer_full(self, timing_enabled):
-        """
-        Test that when the circular buffer in entries_ is full and we call reset,
-        then fill the buffer with new entries, dump_entries returns only the new
-        entries and not the old ones.
-        """
-        if self.rank == self.MAIN_PROCESS_RANK:
-            return
-
-        # Override buffer size to 10 for faster testing
-        os.environ["TORCH_NCCL_TRACE_BUFFER_SIZE"] = "10"
-
-        pg = self._create_process_group_nccl()
-        if timing_enabled:
-            pg._enable_collectives_timing()
-        device = self.local_device
-        self.set_thread_name("fr_test_thread")
-        a = torch.full((3, 4), float(self.rank), device=device)
-
-        # Fill the buffer completely with 10 entries
-        for _ in range(10):
-            f = pg.allreduce(a)
-        f.wait()
-        torch.cuda.synchronize(device=device)
-        time.sleep(1)
-
-        # Verify buffer is full with 10 entries
-        t = pickle.loads(torch._C._distributed_c10d._dump_nccl_trace())
-        self.assertEqual(len(t["entries"]), 10)
-
-        # Now reset the flight recorder
-        torch._C._distributed_c10d._reset_fr_recording_nccl()
-
-        # Add new entries after reset - fill the buffer completely again
-        for _ in range(10):
-            f = pg.allreduce(a)
-        f.wait()
-        torch.cuda.synchronize(device=device)
-        time.sleep(1)
-
-        # Verify we get exactly 10 new entries, not 20
-        t = pickle.loads(torch._C._distributed_c10d._dump_nccl_trace())
-        self.assertEqual(len(t["entries"]), 10)
-
-        # Verify all entries have the expected properties (from after reset)
-        # After reset, record IDs should start from 0 again
-        for i, entry in enumerate(t["entries"]):
-            self.assertIn("profiling_name", entry)
-            self.assertEqual(entry["profiling_name"], "nccl:all_reduce")
-            self.assertIn("record_id", entry)
-            # Record IDs should be sequential starting from 0 after reset
-            self.assertEqual(entry["record_id"], i)
-
-        dist.destroy_process_group()
-
-    @requires_nccl()
-    @skip_but_pass_in_sandcastle_if(not TEST_MULTIGPU, "NCCL test requires 2+ GPUs")
-    @parametrize("timing_enabled", [True, False])
-    def test_fr_record_reset_partial_overwrite(self, timing_enabled):
-        """
-        Test that when the circular buffer is full, we reset, and then add fewer
-        entries than the buffer size, we only get the new entries.
-        This tests that old entries at the end of the circular buffer are properly
-        filtered out based on reset_epoch.
-        """
-        if self.rank == self.MAIN_PROCESS_RANK:
-            return
-
-        # Override buffer size to 10 for faster testing
-        os.environ["TORCH_NCCL_TRACE_BUFFER_SIZE"] = "10"
-
-        pg = self._create_process_group_nccl()
-        if timing_enabled:
-            pg._enable_collectives_timing()
-        device = self.local_device
-        self.set_thread_name("fr_test_thread")
-        a = torch.full((3, 4), float(self.rank), device=device)
-
-        # Fill the buffer completely
-        for _ in range(10):
-            f = pg.allreduce(a)
-        f.wait()
-        torch.cuda.synchronize(device=device)
-        time.sleep(1)
-
-        # Reset the flight recorder
-        torch._C._distributed_c10d._reset_fr_recording_nccl()
-
-        # Add only 3 new entries (much less than buffer size)
-        for _ in range(3):
-            f = pg.allreduce(a)
-        f.wait()
-        torch.cuda.synchronize(device=device)
-        time.sleep(1)
-
-        # Verify we only get the 3 new entries, not 10
-        t = pickle.loads(torch._C._distributed_c10d._dump_nccl_trace())
-        self.assertEqual(len(t["entries"]), 3)
-
-        # Verify record IDs start from 0 after reset
-        for i, entry in enumerate(t["entries"]):
-            self.assertIn("record_id", entry)
-            self.assertEqual(entry["record_id"], i)
-
-        dist.destroy_process_group()
-
-    @requires_nccl()
-    @skip_but_pass_in_sandcastle_if(not TEST_MULTIGPU, "NCCL test requires 2+ GPUs")
-    @parametrize("timing_enabled", [True, False])
-    def test_fr_record_reset_wraparound(self, timing_enabled):
-        """
-        Test that when we reset in the middle of the circular buffer and then
-        wrap around, dump_entries correctly returns only entries from the current
-        epoch in the correct order.
-        """
-        if self.rank == self.MAIN_PROCESS_RANK:
-            return
-
-        # Override buffer size to 10 for faster testing
-        os.environ["TORCH_NCCL_TRACE_BUFFER_SIZE"] = "10"
-
-        pg = self._create_process_group_nccl()
-        if timing_enabled:
-            pg._enable_collectives_timing()
-        device = self.local_device
-        self.set_thread_name("fr_test_thread")
-        a = torch.full((3, 4), float(self.rank), device=device)
-
-        # Fill half the buffer
-        for _ in range(5):
-            f = pg.allreduce(a)
-        f.wait()
-        torch.cuda.synchronize(device=device)
-        time.sleep(1)
-
-        # Reset at this point (reset happens at index 5)
-        torch._C._distributed_c10d._reset_fr_recording_nccl()
-
-        # Now add 8 entries, which will wrap around
-        # (5->9 fills rest of buffer, then 0->2 wraps around)
-        for _ in range(8):
-            f = pg.allreduce(a)
-        f.wait()
-        torch.cuda.synchronize(device=device)
-        time.sleep(1)
-
-        # Should get exactly 8 entries, properly ordered
-        t = pickle.loads(torch._C._distributed_c10d._dump_nccl_trace())
-        self.assertEqual(len(t["entries"]), 8)
-
-        # Entries should be in chronological order
-        # The dump_entries() method returns entries from next_ to end, then 0 to next_
-        # After filtering old entries, we should have 8 entries in order
-        # Verify record IDs start from 0 after reset (id_ is reset in reset_all())
-        for i, entry in enumerate(t["entries"]):
-            self.assertIn("profiling_name", entry)
-            self.assertIn("record_id", entry)
-            self.assertEqual(entry["record_id"], i)
-
-        dist.destroy_process_group()
-
-    @requires_nccl()
-    @skip_but_pass_in_sandcastle_if(not TEST_MULTIGPU, "NCCL test requires 2+ GPUs")
-    @parametrize("timing_enabled", [True, False])
-    def test_fr_record_multiple_resets(self, timing_enabled):
-        """
-        Test multiple consecutive resets to ensure each reset properly increments
-        the epoch and filters out entries from previous epochs.
-        """
-        if self.rank == self.MAIN_PROCESS_RANK:
-            return
-
-        # Override buffer size to 10 for faster testing
-        os.environ["TORCH_NCCL_TRACE_BUFFER_SIZE"] = "10"
-
-        pg = self._create_process_group_nccl()
-        if timing_enabled:
-            pg._enable_collectives_timing()
-        device = self.local_device
-        self.set_thread_name("fr_test_thread")
-        a = torch.full((3, 4), float(self.rank), device=device)
-
-        # First batch: 2 entries
-        for _ in range(2):
-            f = pg.allreduce(a)
-        f.wait()
-        torch.cuda.synchronize(device=device)
-        time.sleep(1)
-
-        # First reset
-        torch._C._distributed_c10d._reset_fr_recording_nccl()
-
-        # Second batch: 3 entries
-        for _ in range(3):
-            f = pg.allreduce(a)
-        f.wait()
-        torch.cuda.synchronize(device=device)
-        time.sleep(1)
-
-        # Second reset
-        torch._C._distributed_c10d._reset_fr_recording_nccl()
-
-        # Third batch: 4 entries
-        for _ in range(4):
-            f = pg.allreduce(a)
-        f.wait()
-        torch.cuda.synchronize(device=device)
-        time.sleep(1)
-
-        # Should only see the last 4 entries
-        t = pickle.loads(torch._C._distributed_c10d._dump_nccl_trace())
-        self.assertEqual(len(t["entries"]), 4)
-
-        # Verify record IDs start from 0 after the last reset
-        for i, entry in enumerate(t["entries"]):
-            self.assertIn("record_id", entry)
-            self.assertEqual(entry["record_id"], i)
-
-        dist.destroy_process_group()
-
 
 def check_if_test_is_skipped(fn):
     def wrapper(self, *args, **kwargs):
@@ -6344,14 +6120,6 @@ class ProcessGroupNCCLLargerScaleTest(MultiProcessTestCase):
         if self.rank == 6 or self.rank == 7:
             dist.broadcast(tensor2, 6, group=ng2)
             self.assertEqual(tensor2, torch.full((1,), 6))
-
-        # Test the case when the split changes the pg option of split group
-        # while the parent pg option is not changed.
-        new_pg = c10d.new_group([0, 1, 2, 3, 4, 5, 6, 7], device_id=device)
-        backend_new_pg = new_pg._get_backend(torch.device(device))
-        self.assertEqual(len(backend_new_pg.options.global_ranks_in_group), 8)
-        c10d.split_group(new_pg, [[0, 2, 4, 6], [1, 3, 5, 7]])
-        self.assertEqual(len(backend_new_pg.options.global_ranks_in_group), 8)
         # a barrier and a cuda sync before destroying all pgs.
         dist.barrier(pg)
         torch.cuda.synchronize()

test/distributed/test_composability.py

@@ -499,7 +499,6 @@ class ComposabilityTest(MultiProcContinuousTest):
             [
                 _ComputationType.UNSHARD,
                 _ComputationType.FORWARD,
-                _ComputationType.REDUCE_GRAD,  # Contains final fsdp post_backward
             ],
             microbatch_index=0,
         )

test/distributed/test_dynamo_distributed.py

@@ -2,7 +2,6 @@
 import contextlib
 import copy
 import functools
-import logging
 import random
 import unittest
 from contextlib import contextmanager
@@ -52,9 +51,6 @@ from torch.testing._internal.inductor_utils import HAS_GPU
 from torch.testing._internal.triton_utils import requires_cuda_and_triton
 
 
-log = logging.getLogger(__name__)
-
-
 def reset_rng_state():
     torch.manual_seed(1337)
     random.seed(1337)
@@ -1204,116 +1200,6 @@ class TestMultiProc(DynamoDistributedMultiProcTestCase):
             for r in res[1:]:
                 self.assertEqual(res[0], r)
 
-    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
-    @patch.object(torch._dynamo.config, "enable_compiler_collectives", True)
-    @patch.object(torch._inductor.config, "max_autotune_gemm", True)
-    @patch.object(torch._inductor.config, "distributed_max_autotune_gemm", True)
-    def test_multiproc_autotune(self):
-        with _dynamo_dist_per_rank_init(self.rank, self.world_size):
-            torch._dynamo.utils.clear_compilation_metrics()
-
-            @torch.compile()
-            def f(a, b, c):
-                res = (
-                    torch.sum((a @ b) + 1.0)
-                    + torch.sum(torch.relu(b @ c))
-                    + torch.sum(c @ a)
-                )
-
-                return res
-
-            a = torch.randn(1024, 1024, device=self.rank, dtype=torch.bfloat16)
-            b = torch.randn(1024, 2048, device=self.rank, dtype=torch.bfloat16)
-            c = torch.randn(2048, 1024, device=self.rank, dtype=torch.bfloat16)
-
-            try:
-                f(a, b, c)
-            except Exception:
-                log.exception("Caught exception running f")
-                raise
-
-            metrics = torch._dynamo.utils.get_compilation_metrics()
-            res = [None] * self.world_size
-            torch.distributed.all_gather_object(res, len(metrics))
-            for r in res[1:]:
-                self.assertEqual(res[0], r)
-
-            print(f"Result from {self.rank} is {f(a, b, c)}")
-
-    @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
-    @patch.object(torch._dynamo.config, "enable_compiler_collectives", True)
-    @patch.object(torch._inductor.config, "max_autotune_gemm", True)
-    @patch.object(torch._inductor.config, "distributed_max_autotune_gemm", True)
-    def test_multiproc_autotune_dynamic_shapes(self):
-        with _dynamo_dist_per_rank_init(self.rank, self.world_size):
-            torch._dynamo.utils.clear_compilation_metrics()
-
-            @torch.compile()
-            def f(a, b, c):
-                res = (
-                    torch.sum((a @ b) + 1.0)
-                    + torch.sum(torch.relu(b @ c))
-                    + torch.sum(c @ a)
-                )
-
-                return res
-
-            a = torch.randn(1024, 1024, device=self.rank, dtype=torch.bfloat16)
-            b = torch.randn(1024, 2048, device=self.rank, dtype=torch.bfloat16)
-            c = torch.randn(2048, 1024, device=self.rank, dtype=torch.bfloat16)
-
-            # Mark tensors as dynamic on dimension 0
-            torch._dynamo.mark_dynamic(a, 0)
-            torch._dynamo.mark_dynamic(a, 1)
-            torch._dynamo.mark_dynamic(b, 0)
-            torch._dynamo.mark_dynamic(b, 1)
-            torch._dynamo.mark_dynamic(c, 0)
-            torch._dynamo.mark_dynamic(c, 1)
-
-            try:
-                f(a, b, c)
-            except Exception:
-                log.exception("Caught exception running f")
-                raise
-
-            metrics = torch._dynamo.utils.get_compilation_metrics()
-            res = [None] * self.world_size
-            torch.distributed.all_gather_object(res, len(metrics))
-            for r in res[1:]:
-                self.assertEqual(res[0], r)
-
-            print(f"Result from {self.rank} is {f(a, b, c)}")
-
-            # Store the initial compilation count
-            initial_compile_count = len(metrics)
-
-            # # Test with different sizes to ensure dynamic shapes work without recompilation
-            a2 = torch.randn(512, 512, device=self.rank, dtype=torch.bfloat16)
-            b2 = torch.randn(512, 2048, device=self.rank, dtype=torch.bfloat16)
-            c2 = torch.randn(2048, 512, device=self.rank, dtype=torch.bfloat16)
-
-            try:
-                result2 = f(a2, b2, c2)
-                print(f"Result2 from {self.rank} is {result2}")
-            except Exception:
-                log.exception("Caught exception running f with different sizes")
-                raise
-
-            # Verify no recompilation occurred
-            metrics_after = torch._dynamo.utils.get_compilation_metrics()
-            final_compile_count = len(metrics_after)
-            self.assertEqual(
-                initial_compile_count,
-                final_compile_count,
-                "Expected no recompilation with dynamic shapes",
-            )
-
-            # Verify all ranks have the same compilation count
-            res_after = [None] * self.world_size
-            torch.distributed.all_gather_object(res_after, final_compile_count)
-            for r in res_after[1:]:
-                self.assertEqual(res_after[0], r)
-
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")
     def test_get_pg_attr(self):
         with _dynamo_dist_per_rank_init(self.rank, self.world_size):

test/distributed/test_inductor_collectives.py

@@ -1985,7 +1985,6 @@ class TestCollectivesInductor(DynamoDistributedSingleProcTestCase):
                     "bucket_reduce_scatters_fx_bucket_size_determinator": lambda _: 2,
                     "reorder_for_compute_comm_overlap": True,
                     "reorder_for_compute_comm_overlap_passes": [
-                        _reorder_communication_preserving_peak_memory,
                         sink_waits_iterative,
                         _reorder_communication_preserving_peak_memory,
                     ],
@@ -2047,6 +2046,11 @@ class TestCollectivesInductor(DynamoDistributedSingleProcTestCase):
         assert node_stats is not None
         self.assertTrue(isinstance(node_stats, dict))
         self.assertEqual(len(node_stats), 4)
+        it = iter(node_stats.values())
+        node_stat0 = next(it)
+        self.assertTrue(node_stat0.limiting_factor == "None")
+        node_stat1 = next(it)
+        self.assertTrue("collective ordering" in node_stat1.limiting_factor)
 
     @skipIfXpu  # https://github.com/intel/torch-xpu-ops/issues/1581
     @unittest.skipIf(not HAS_GPU, "Inductor+gpu needs triton and recent GPU arch")

test/distributed/test_local_tensor.py

@@ -128,14 +128,14 @@ class TestLocalTensorWorld2(LocalTensorTestBase):
         self.assertEqual(len(result_add._local_tensors), 2)
 
         # Verify the operation was applied to each local tensor
-        for rank in identical_local_tensors:
+        for rank in identical_local_tensors.keys():
             expected = identical_local_tensors[rank] + identical_local_tensors[rank]
             self.assertEqual(result_add._local_tensors[rank], expected)
 
         # Test multiplication
         result_mul = lt1 * 2.0
         self.assertIsInstance(result_mul, LocalTensor)
-        for rank in identical_local_tensors:
+        for rank in identical_local_tensors.keys():
             expected = identical_local_tensors[rank] * 2.0
             self.assertEqual(result_mul._local_tensors[rank], expected)
 
@@ -163,7 +163,7 @@ class TestLocalTensorWorld2(LocalTensorTestBase):
         result = lt + regular_tensor
         self.assertIsInstance(result, LocalTensor)
 
-        for rank in identical_local_tensors:
+        for rank in identical_local_tensors.keys():
             expected = identical_local_tensors[rank] + regular_tensor
             self.assertEqual(result._local_tensors[rank], expected)
 
@@ -212,14 +212,14 @@ class TestLocalTensorWorld2(LocalTensorTestBase):
             dist.all_reduce(lt_sum, group=fake_pg)
 
             expected_sum = torch.tensor([[6.0, 8.0], [10.0, 12.0]])
-            for rank in test_tensors:
+            for rank in test_tensors.keys():
                 self.assertEqual(lt_sum._local_tensors[rank], expected_sum)
 
             # Test broadcast within mode
             lt_broadcast = LocalTensor({k: v.clone() for k, v in test_tensors.items()})
             dist.broadcast(lt_broadcast, src=0, group=fake_pg)
 
-            for rank in test_tensors:
+            for rank in test_tensors.keys():
                 self.assertEqual(lt_broadcast._local_tensors[rank], test_tensors[0])
 
             # Test that regular operations still work
@@ -293,21 +293,21 @@ class TestLocalTensorWorld3(LocalTensorTestBase):
         lt_sum = LocalTensor({k: v.clone() for k, v in test_tensors.items()})
         dist.all_reduce(lt_sum, op=dist.ReduceOp.SUM, group=fake_pg)
         expected_sum = torch.tensor([[6.0, 7.0], [6.0, 15.0]])  # Sum of all tensors
-        for rank in test_tensors:
+        for rank in test_tensors.keys():
             self.assertEqual(lt_sum._local_tensors[rank], expected_sum)
 
         # Test MAX reduction
         lt_max = LocalTensor({k: v.clone() for k, v in test_tensors.items()})
         dist.all_reduce(lt_max, op=dist.ReduceOp.MAX, group=fake_pg)
         expected_max = torch.tensor([[3.0, 4.0], [3.0, 6.0]])  # Max across all tensors
-        for rank in test_tensors:
+        for rank in test_tensors.keys():
             self.assertEqual(lt_max._local_tensors[rank], expected_max)
 
         # Test MIN reduction
         lt_min = LocalTensor({k: v.clone() for k, v in test_tensors.items()})
         dist.all_reduce(lt_min, op=dist.ReduceOp.MIN, group=fake_pg)
         expected_min = torch.tensor([[1.0, 1.0], [1.0, 4.0]])  # Min across all tensors
-        for rank in test_tensors:
+        for rank in test_tensors.keys():
             self.assertEqual(lt_min._local_tensors[rank], expected_min)
 
     def test_all_reduce_collective(self):
@@ -328,7 +328,7 @@ class TestLocalTensorWorld3(LocalTensorTestBase):
 
         # Verify all ranks have the sum of all tensors (after adding 1 to each)
         expected_sum = torch.tensor([[114.0, 225.0, 336.0], [447.0, 558.0, 669.0]])
-        for rank in different_tensors:
+        for rank in different_tensors.keys():
             self.assertEqual(lt_sum._local_tensors[rank], expected_sum)
 
     def test_broadcast_collective(self):
@@ -348,7 +348,7 @@ class TestLocalTensorWorld3(LocalTensorTestBase):
 
         # Verify all ranks have rank 1's original tensor
         expected_broadcast = different_tensors[1]
-        for rank in different_tensors:
+        for rank in different_tensors.keys():
             self.assertEqual(lt_broadcast._local_tensors[rank], expected_broadcast)
 
     def test_all_gather_collective(self):

test/distributed/test_run.py

@@ -17,7 +17,6 @@ from torch.testing._internal.common_utils import run_tests, TestCase
 
 class RunTest(TestCase):
     def setUp(self):
-        super().setUp()
         # Save original environment variable if it exists
         self.original_signals_env = os.environ.get(
             "TORCHELASTIC_SIGNALS_TO_HANDLE", None

test/distributed/test_serialization.py

@@ -25,7 +25,6 @@ class MyClass:
 
 class TestSerialization(TestCase):
     def setUp(self) -> None:
-        super().setUp()
         # disable debug asserts
         self._old_debug = os.environ.get(DEBUG_ENV)
         os.environ[DEBUG_ENV] = "0"

test/distributed/test_store.py

@@ -317,7 +317,6 @@ class HashStoreTest(TestCase, StoreTestBase):
 
 class PrefixStoreTest(TestCase):
     def setUp(self):
-        super().setUp()
         # delete is false as FileStore will automatically clean up the file
         self.file = tempfile.NamedTemporaryFile(delete=False)
 

test/dynamo/test_compiler_bisector.py

@@ -283,7 +283,7 @@ class TestCompilerBisector(TestCase):
     )
     def test_bisect_pre_grad_graph(self):
         def f(x):
-            for _ in range(5):
+            for i in range(5):
                 x = x + 1
             return x.relu()
 

test/dynamo/test_dicts.py

@@ -36,15 +36,6 @@ class DummyUserDict(UserDict):
     pass
 
 
-class FakeMapping:
-    def __init__(self, value: Any) -> None:
-        self._value = value
-        self.keys = lambda: ["a", "b", "c"]  # not required to be a method
-
-    def __getitem__(self, key: str) -> Any:
-        return self._value
-
-
 class DictTests(torch._dynamo.test_case.TestCase):
     def test_dict_subclass_instantiation(self):
         def fn(x):
@@ -675,18 +666,6 @@ class DictTests(torch._dynamo.test_case.TestCase):
         for k1, m2 in zip(modules, module_dict.children()):
             self.assertTrue(modules[k1] is m2)
 
-    # FIXME: see comment in torch/_dynamo/polyfills/__init__.py:mutable_mapping_update
-    @unittest.expectedFailure
-    def test_dict_construct_from_mapping_like(self):
-        def fn(x):
-            fm = FakeMapping(x)
-            d = dict(fm, x=x)
-            return d
-
-        x = torch.randn(4)
-        opt_fn = torch.compile(fn, backend="eager", fullgraph=True)
-        self.assertEqual(fn(x), opt_fn(x))
-
     def test_dict_subclass_initialization_in_graph(self):
         for super_class in (
             OrderedDict,
@@ -1108,52 +1087,12 @@ class DictTests(torch._dynamo.test_case.TestCase):
 
         self.assertEqual(ref, res)
 
-    def test_newly_constructed_default_dict_no_default_factory(self):
-        def f1(x):
-            d = defaultdict()
-            try:
-                d[1] += 42
-            except KeyError:
-                d[1] = 1
-            return x + 1, d
-
-        x = torch.ones(2)
-        ref = f1(x)
-        res = torch.compile(f1, backend="eager", fullgraph=True)(x)
-
-        self.assertEqual(ref, res)
-
-        def f2(x):
-            d = defaultdict(None)
-            try:
-                d[1] += 42
-            except KeyError:
-                d[1] = 1
-            return x + 1, d
-
-        ref = f2(x)
-        res = torch.compile(f2, backend="eager", fullgraph=True)(x)
-        self.assertEqual(ref, res)
-
-        def f3(x):
-            d = defaultdict(None, {1: 10})
-            d[1] += 42
-            try:
-                d[2] += 24
-            except KeyError:
-                d[2] = 1
-            return x + 1, d
-
-        ref = f3(x)
-        res = torch.compile(f3, backend="eager", fullgraph=True)(x)
-        self.assertEqual(ref, res)
-
+    @unittest.expectedFailure
     def test_newly_constructed_default_dict_with_dict(self):
         def f(x):
-            d = dict([("a", 1), ("b", 2)], c=3)  # noqa: C406
-            dd = defaultdict(list, d, d=4, e=5)
-            dd["x"].append(42)
-            return x + 1, d, dd
+            d = defaultdict(dict, {2: {"a": 1}})
+            d[0] = {"b": 2}
+            return x + 1, d
 
         x = torch.ones(2)
         ref = f(x)

test/dynamo/test_graph_deduplication.py

@@ -8,11 +8,21 @@ from torch._dynamo.graph_deduplication import apply_graph_deduplication
 from torch._dynamo.graph_utils import _detect_cycles
 from torch._dynamo.output_graph import FakeRootModule
 from torch._dynamo.test_case import TestCase
-from torch._dynamo.testing import extract_graph, extract_graph_and_tracker, normalize_gm
+from torch._dynamo.testing import (
+    AotEagerAndRecordGraphs,
+    extract_graph_and_tracker,
+    normalize_gm,
+)
 from torch.compiler import allow_in_graph
 from torch.utils._ordered_set import OrderedSet
 
 
+def extract_graph(fn, *args, **kwargs):
+    backend = AotEagerAndRecordGraphs()
+    result = torch.compile(backend=backend)(fn)(*args, **kwargs)
+    return result, backend.graphs, backend.fw_graphs
+
+
 def graph_str(gm):
     return normalize_gm(gm.print_readable(print_output=False))
 
@@ -30,7 +40,7 @@ class GraphDededuplicationTests(TestCase):
         super().tearDown()
 
     def run_and_return_graphs(self, fn, *args, **kwargs):
-        return extract_graph(fn, *args, **kwargs)[0:3]
+        return extract_graph(fn, *args, **kwargs)
 
     def run_and_get_simple_graph(self):
         def fn(x, y):

test/dynamo/test_install_free_tensors.py

@@ -1,7 +1,7 @@
 # Owner(s): ["module: dynamo"]
 import unittest
-from collections.abc import Callable, Sequence
-from typing import Any, Union
+from collections.abc import Sequence
+from typing import Any, Callable, Union
 
 import torch
 import torch._dynamo